diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index 695f47dc1..585f089a2 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -107,6 +107,11 @@ steps: - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model + - >- + timeout 20m pytest -s -v + tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py + -m 'core_model and rocm and cards_2' + --run-level "core_model" - label: "Diffusion GPU Worker Test" agent_pool: mi300_2 @@ -122,7 +127,13 @@ steps: mirror_hardwares: [amdproduction] grade: Blocking commands: - - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model" + # Preserve legacy one-GPU tests without cards_1 while excluding every + # registered multi-GPU marker from this one-GPU job. + - >- + timeout 15m pytest -sv tests/diffusion/models/ + -m 'core_model and cuda and not + (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' + --run-level "core_model" - label: "Diffusion · Wan22 Test" agent_pool: mi300_1 diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index 9aa55e5d0..01fb32879 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.buildkite/npu/test-npu-ready.yml @@ -19,6 +19,8 @@ steps: pytest -s -v tests/diffusion/layers/test_activation.py tests/diffusion/layers/test_fused_qk_norm_rope_npu.py + tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py + tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py -m 'core_model and npu and A3' --run-level 'core_model' diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 36cefcf0d..542db3fe9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -2,6 +2,8 @@ You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/). +AI-assisted contributions must follow the [AI Usage Policy](docs/contributing/README.md#ai-usage-policy), including disclosure, human review, and validation requirements. + Local `pre-commit` gates (SPDX, forbidden imports, `torch.cuda`, shellcheck, mypy, test marks, TTS adapter ratchet, markdownlint, Buildkite schema, and how to extend allowlists) are documented in the diff --git a/benchmarks/diffusion/README.md b/benchmarks/diffusion/README.md index 54ed64261..b6ea50f3e 100644 --- a/benchmarks/diffusion/README.md +++ b/benchmarks/diffusion/README.md @@ -17,18 +17,18 @@ The main entrypoint is: vllm serve Qwen/Qwen-Image --omni --port 8099 ``` -2. Run a minimal benchmark: +1. Run a minimal benchmark: ```bash python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - --base-url http://localhost:8099 \ - --model Qwen/Qwen-Image \ - --task t2i \ - --dataset vbench \ - --num-prompts 5 + --base-url http://localhost:8099 \ + --model Qwen/Qwen-Image \ + --task t2i \ + --dataset vbench \ + --num-prompts 5 ``` -**Notes** +### Notes - By default, image tasks talk to `http://:/v1/chat/completions`; video tasks talk to `/v1/videos`. - If you run the server on another host or port, pass `--base-url` accordingly. @@ -52,14 +52,14 @@ Example (`t2v`): ```bash python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ - --base-url http://localhost:8099 \ - --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --task t2v \ - --dataset vbench \ - --num-prompts 50 \ - --width 640 --height 480 \ - --num-frames 81 --fps 16 \ - --num-inference-steps 40 + --base-url http://localhost:8099 \ + --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --task t2v \ + --dataset vbench \ + --num-prompts 50 \ + --width 640 --height 480 \ + --num-frames 81 --fps 16 \ + --num-inference-steps 40 ``` Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths). @@ -140,6 +140,21 @@ Traffic / concurrency flags: - `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately. - `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed. +### Video job timeout and failures + +`--video-job-timeout` controls how long each `/v1/videos` job may be polled, +including time queued on the server (default: `900` seconds). For long video +jobs at high concurrency, use a larger budget, for example +`--video-job-timeout 1800`. This also applies to warmup jobs. The polling budget +starts after job creation; it excludes time waiting for the client concurrency +semaphore. Individual HTTP calls still use the aiohttp session timeout. + +An expired job is counted as failed and deleted, which cancels queued or running +server generation. The benchmark continues processing the remaining requests. +The progress bar counts both successful and failed requests. The final report +shows the failure count and the first 10 errors; `--output-file` saves all errors +with their request IDs in `request_errors`. + ### Batched warmup note For batched serving runs, warm the same in-flight shape you plan to measure. diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index aeff3ebbb..425667d8c 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import base64 import json @@ -43,6 +46,7 @@ class RequestFuncInput: 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 + video_job_timeout: float = 900.0 @dataclass @@ -331,64 +335,64 @@ async def async_request_v1_videos( output = RequestFuncOutput() output.start_time = time.perf_counter() - files = dict(input.extra_body) - if input.prompt: - files.setdefault("prompt", input.prompt) - if input.width and input.height: - files.setdefault("height", input.height) - files.setdefault("width", input.width) - if input.num_frames: - files.setdefault("num_frames", input.num_frames) - if input.num_inference_steps: - files.setdefault("num_inference_steps", input.num_inference_steps) - if input.seed is not None: - files.setdefault("seed", input.seed) - if input.fps: - files.setdefault("fps", input.fps) - - form = aiohttp.FormData() - for k, v in files.items(): - 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 + job_id = None + try: + files = dict(input.extra_body) + if input.prompt: + files.setdefault("prompt", input.prompt) + if input.width and input.height: + files.setdefault("height", input.height) + files.setdefault("width", input.width) + if input.num_frames: + files.setdefault("num_frames", input.num_frames) + if input.num_inference_steps: + files.setdefault("num_inference_steps", input.num_inference_steps) + if input.seed is not None: + files.setdefault("seed", input.seed) + if input.fps: + files.setdefault("fps", input.fps) - if input.image_paths and len(input.image_paths) > 0: - image_path = input.image_paths[0] - image_file = open(image_path, "rb") - form.add_field( - "input_reference", - image_file, - filename=os.path.basename(image_path), - content_type=_guess_mime_type(image_path), - ) - elif input.video_paths and len(input.video_paths) > 0: - video_path = input.video_paths[0] - if _is_minimax_h3_mixed_reference(input.model): - with open(video_path, "rb") as reference_file: - video_b64 = base64.b64encode(reference_file.read()).decode("ascii") - form.add_field( - "video_reference", - json.dumps([{"video_url": f"data:video/mp4;base64,{video_b64}"}]), - ) - else: - image_file = open(video_path, "rb") + form = aiohttp.FormData() + for k, v in files.items(): + form.add_field(k, str(v)) + + 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") form.add_field( "input_reference", image_file, - filename=os.path.basename(video_path), - content_type=_guess_mime_type(video_path), + filename=os.path.basename(image_path), + content_type=_guess_mime_type(image_path), ) + elif input.video_paths and len(input.video_paths) > 0: + video_path = input.video_paths[0] + if _is_minimax_h3_mixed_reference(input.model): + with open(video_path, "rb") as reference_file: + video_b64 = base64.b64encode(reference_file.read()).decode("ascii") + form.add_field( + "video_reference", + json.dumps([{"video_url": f"data:video/mp4;base64,{video_b64}"}]), + ) + else: + 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 - job_status = None - poll_json = {} - resp_json = {} + job_status = None + poll_json = {} + resp_json = {} - try: # invoke a post request (POST /v1/videos) async with session.post(input.api_url, data=form) as response: if response.status == 200: @@ -406,7 +410,7 @@ async def async_request_v1_videos( # invoke a poll request (GET /v1/videos/{video_id}) poll_interval = 2.0 # Unit(s) - timeout_seconds = 600.0 + timeout_seconds = input.video_job_timeout deadline = time.perf_counter() + timeout_seconds job_url = f"{input.api_url}/{job_id}" @@ -422,8 +426,11 @@ async def async_request_v1_videos( poll_json = await poll_response.json() job_status = poll_json.get("status") - if time.perf_counter() >= deadline: - output.error = f"Timed out waiting for video job {job_id} to complete." + if job_status not in {"completed", "failed"} and time.perf_counter() >= deadline: + output.error = ( + f"Timed out after {timeout_seconds:g}s waiting for video job {job_id} to complete. " + "Increase --video-job-timeout for long-running or queued video jobs." + ) output.success = False return output @@ -452,7 +459,7 @@ async def async_request_v1_videos( elif "peak_memory_mb" in resp_json: output.peak_memory_mb = resp_json["peak_memory_mb"] except Exception as e: - output.error = str(e) + output.error = f"{type(e).__name__}: {e}" output.success = False finally: if image_file is not None: @@ -465,13 +472,13 @@ async def async_request_v1_videos( except Exception as e: print(f"Failed to clean up video job {job_id}: {e}") - output.latency = time.perf_counter() - output.start_time + output.latency = time.perf_counter() - output.start_time + if pbar is not None: + pbar.update(1) if output.success and input.slo_ms is not None: output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms) - if pbar: - pbar.update(1) return output diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 6daac9df4..1e5aaf29c 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -1,6 +1,6 @@ # adapted from fastvideo # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Benchmark online serving for diffusion models (Image/Video Generation). @@ -87,6 +87,7 @@ import base64 import glob import json import logging +import math import os import random import tempfile @@ -1116,6 +1117,11 @@ def calculate_metrics( "duration": total_duration, "completed_requests": num_success, "failed_requests": len(error_outputs), + "request_errors": [ + {"request_id": req.request_id, "error": out.error} + for req, out in zip(requests_list, outputs) + if not out.success + ], "throughput_qps": num_success / total_duration if total_duration > 0 else 0, "latency_mean": np.mean(latencies) if latencies else 0, "latency_median": np.median(latencies) if latencies else 0, @@ -1261,6 +1267,9 @@ def _default_endpoint_for_task(task: str) -> str: async def benchmark(args): + if not math.isfinite(args.video_job_timeout) or args.video_job_timeout <= 0: + raise ValueError("--video-job-timeout must be a positive finite number of seconds.") + # Construct base_url if not provided if args.base_url is None: args.base_url = f"http://{args.host}:{args.port}" @@ -1312,6 +1321,8 @@ async def benchmark(args): print("Loading requests...") requests_list = dataset.get_requests() print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.") + for req in requests_list: + req.video_job_timeout = args.video_job_timeout if args.return_stage_metrics and args.endpoint in _STAGE_METRICS_ENDPOINTS: for req in requests_list: @@ -1376,6 +1387,8 @@ async def benchmark(args): metrics["model"] = args.model metrics["dataset"] = args.dataset metrics["task"] = args.task + if args.endpoint == "/v1/videos": + metrics["video_job_timeout"] = args.video_job_timeout if args.endpoint == "/v1/images/edits": metrics["bot_task"] = args.bot_task @@ -1398,6 +1411,11 @@ async def benchmark(args): ) ) print("{:<40} {}/{:<15}".format("Successful requests:", metrics["completed_requests"], len(requests_list))) + print("{:<40} {:<15}".format("Failed requests:", metrics["failed_requests"])) + for failure in metrics["request_errors"][:10]: + print(f"Request {failure['request_id']} failed: {failure['error']}") + if metrics["failed_requests"] > 10: + print("Showing first 10 failures; use --output-file to save all request errors.") # Section 3: Performance Metrics print(f"{'-' * 50}") @@ -1481,6 +1499,13 @@ if __name__ == "__main__": help="Path to local dataset file (optional).", ) parser.add_argument("--num-prompts", type=int, default=10, help="Number of prompts to benchmark.") + parser.add_argument( + "--video-job-timeout", + type=float, + default=900.0, + help="Maximum polling time in seconds per video job, including server queue time (default: 900). " + "Increase this for long-running video generation or high concurrency.", + ) parser.add_argument( "--max-concurrency", type=int, diff --git a/benchmarks/tts/README.md b/benchmarks/tts/README.md index c403e5dd0..918e0727a 100644 --- a/benchmarks/tts/README.md +++ b/benchmarks/tts/README.md @@ -85,6 +85,25 @@ vllm bench serve --omni \ --save-result --result-dir ./results ``` +#### Set the sample rate for models that are not 24 kHz + +Streamed PCM carries no header, so the harness derives `audio_duration` from the +byte count and a **default 24000 Hz** assumption +(`vllm_omni/metrics/definitions.py:DEFAULT_AUDIO_SAMPLE_RATE`). For a model that +emits another rate, export the real one or every duration-derived metric is +wrong by `real_sr / 24000`: + +```bash +# Audio8 TTS, Fish Speech S2 Pro, Ming-omni-tts, ... all decode at 44.1 kHz +export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100 +``` + +Without it, a 44.1 kHz model reports `audio_duration` 1.84x too long, so +`audio_rtf` looks 1.84x *better* than reality and `audio_underrun` / +`Streaming continuity OK rate` are measured against inflated chunk budgets. +A quick sanity check: compare `Mean AUDIO_DURATION (s)` against the actual +length of a saved WAV for the same prompt. + #### Streaming continuity (`audio_underrun`) `audio_underrun` (seconds) is the per-request worst-case buffer deficit diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 69fe09b50..53c84e66a 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -17,6 +17,17 @@ models: extra_params: lang: en text_normalization: true + # Audio8 TTS Preview: DualAR, 44.1 kHz. trust_remote_code stays false. + # Bench at 44.1 kHz: export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100, otherwise + # the streamed-PCM duration defaults to 24 kHz and RTF looks 1.84x too good. + Audio8/Audio8-TTS-Preview-0.6b: + supported_tasks: [default_voice, voice_clone] + backend: openai-audio-speech + endpoint: /v1/audio/speech + trust_remote_code: false + task_extra_body: + default_voice: {} + voice_clone: {} # -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is # NOT supported (an attempt raises ValueError from _extract_speaker_embedding diff --git a/docs/cli/bench/serve.md b/docs/cli/bench/serve.md index b9c590095..a162c001a 100644 --- a/docs/cli/bench/serve.md +++ b/docs/cli/bench/serve.md @@ -436,6 +436,76 @@ This should be seen as an edge case, and if this behavior can be avoided by sett +### Image / video generation + +For `/v1/images/generations`, `/v1/images/edits`, and `/v1/videos`, set `--endpoint` to that path and omit `--backend`. The endpoint path is registered as the bench request adapter, so a separate named backend is not required. + +Example (`/v1/images/generations`): + +```bash +vllm bench serve --omni \ + --endpoint /v1/images/generations \ + --dataset-name random \ + --model ~/models/Qwen/Qwen-Image \ + --tokenizer ~/models/Qwen/Qwen-Image/tokenizer \ + --max-concurrency 1 \ + --num-warmups 2 \ + --num-prompts 10 \ + --random-input-len 64 \ + --random-output-len 1 \ + --ignore-eos \ + --percentile-metrics e2el \ + --extra-body '{ + "num_inference_steps": 20, + "seed": 42, + "true_cfg_scale": 4.0 + }' +``` + +If successful, following output is like: + +```text +============ Serving Benchmark Result ============ +Successful requests: 3 +Failed requests: 0 +Maximum request concurrency: 1 +Benchmark duration (s): 8.42 +Request throughput (req/s): 0.36 +Peak concurrent requests: 2.00 +-------------------Peak Memory-------------------- +Mean PEAK_MEMORY_MB (MB): 58832.00 +Median PEAK_MEMORY_MB (MB): 58832.00 +P99 PEAK_MEMORY_MB (MB): 58832.00 +----------------End-to-end Latency---------------- +Mean E2EL (ms): 2798.16 +Median E2EL (ms): 2799.83 +P99 E2EL (ms): 2800.03 +================== Image Result ================== +Total images generated: 3 +Image throughput (img/s): 0.36 +Average pixels per image: 1048576.00 +Mean denoise step latency (ms): 136.59 +---------------- Image Generation ---------------- +Mean IMAGE_GENERATION (ms): 2731.84 +Median IMAGE_GENERATION (ms): 2732.02 +P99 IMAGE_GENERATION (ms): 2734.65 +================================================== +``` + +Use the same pattern with `--endpoint /v1/images/edits` or `--endpoint /v1/videos` (and model-specific `--extra-body` as needed). Pure image/video runs omit the Text Result section when there is no generated text. + +`/v1/images/edits` defaults to **non-streaming JSON** (`stream=false`) so single-stage edit models are not rejected by the server. For multi-stage pipelines that need SSE (AR TTFT / image chunks), pass `"stream": true` in `--extra-body`. + +`/v1/videos` is an async job API: the client creates a job, then polls until `completed`/`failed`. Measured **e2el** therefore includes client poll sleep and any overshoot after the job actually finishes. Tune polling via `--extra-body`: + +- `poll_interval_s` (default `2.0`): sleep between status polls +- `poll_timeout_s` (default `21600`, i.e. 6 hours): give up waiting for the job + +**VIDEO_RTF** prefers server-reported generation time when available: + +- `video_rtf = video_generation_time_ms / 1000 / video_duration`, where `video_generation_time_ms` comes from response `stage_durations` (or `inference_time_s` as fallback) +- If generation time is missing, falls back to `e2el / video_duration` (this fallback *does* include poll overhead) + ### Multi-Stage Benchmark
diff --git a/docs/contributing/README.md b/docs/contributing/README.md index dbfc917bf..a45854b14 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -262,6 +262,28 @@ or serving, explain the required L2/L3 environment and provide the exact command Repository skills guide the agent, but they do not replace contributor judgment, the accepted issue or RFC, required test evidence, or maintainer review. Review all generated changes before submitting them. +## AI Usage Policy + +AI tools are welcome in vLLM-Omni contributions. Contributors, including maintainers, remain responsible for everything they submit. These requirements apply to pull requests, issues, discussions, and review comments: + +- **Disclose AI assistance.** Name the tools used and describe their role in the submission, including assistance with code, tests, documentation, or writing. +- **Understand your changes.** Review every generated change and be able to explain its behavior, design choices, and interaction with vLLM-Omni without relying on AI to answer for you. +- **Validate before submitting.** Run relevant checks and report actual results. For model, kernel, or performance changes, include the hardware, model, configuration, and commands needed to reproduce accuracy or benchmark evidence. Clearly identify checks you could not run; generated claims are not test results. +- **Keep communication useful.** Verify technical claims, search for existing reports, and edit AI-assisted text for accuracy and brevity. Submit only content you have personally reviewed, and take responsibility for responding to review feedback. +- **Label generated media.** Model-generated images, video, and audio are valid examples or test artifacts for this multimodal project. Identify them as generated output and provide the model, inputs, and reproduction settings. Do not present them as ground-truth references or independent evidence of correctness. + +For example, a PR description could include: + +```text +AI assistance: Used Codex to draft the regression test and PR description. +I reviewed the changes and verified the test reproduces the reported failure. +Validation: . +``` + +Maintainer review time is limited. Contributions that are unverified, misleading, or beyond the author's understanding may be closed. Repeated violations may lead to contribution restrictions. New contributors are welcome to ask for help and submit small, well-understood changes. + +This policy is adapted from [Ghostty's AI Usage Policy](https://github.com/ghostty-org/ghostty/blob/main/AI_POLICY.md) for vLLM-Omni's contribution and multimodal validation workflows. The [coding-agent workflow](#using-repository-skills-with-coding-agents) provides practical guidance for applying it. + ## Issues If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible. diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md index 5b14dd8ec..58722ffa0 100644 --- a/docs/design/architecture_overview.md +++ b/docs/design/architecture_overview.md @@ -323,18 +323,24 @@ inputs, fields, and ownership rules are described below. ```mermaid flowchart TB layer1["Layer 1 · Authoring inputs
PipelineConfig + DeployConfig"] - layer2["Layer 2 · Resolve once
StageConfigFactory.create_from_model()
VllmOmniConfig.from_pipeline_config()"] - layer3["Layer 3 · Transport-safe control plane
VllmOmniConfig"] + layer2["Layer 2 · Production resolution boundary
resolve_omni_config()
(StageConfigFactory + VllmOmniConfig internally)"] + layer3["Layer 3 · Startup hand-off
OmniConfigResolution
PipelineConfig + temporary OmegaConf stage bridge"] layer4["Layer 4 · Runtime launch planning
StageRuntime"] layer5["Layer 5 · Engine materialization
VllmConfig / OmniDiffusionConfig"] layer1 --> layer2 --> layer3 --> layer4 --> layer5 ``` -The single resolution boundary is implemented by -`StageConfigFactory.create_from_model()` and -`VllmOmniConfig.from_pipeline_config()` in +The production resolution boundary is `resolve_omni_config()` in [`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config). +It delegates typed construction to `StageConfigFactory.create_from_model()` and +`VllmOmniConfig.from_pipeline_config()`, then returns an +`OmniConfigResolution` consumed by both `AsyncOmniEngine` and headless startup. +Until `StageRuntime` consumes typed stage configs directly, this envelope carries +the effective `PipelineConfig` alongside OmegaConf-compatible `stage_configs` as +a temporary runtime bridge; both views describe the same resolved topology, +including injected stages. + The legacy `stage_args` YAML path has been removed. Model topology now resolves through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`. @@ -354,9 +360,12 @@ The important ownership rules are: deploy-time defaults; it does not redefine the model graph. 3. CLI and Python overrides are applied at the resolution boundary, with per-stage overrides taking precedence over global values where supported. -4. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan` +4. `OmniConfigResolution` is the sole startup hand-off. Its `pipeline_config` + and temporary `stage_configs` compatibility view must describe the same + topology. +5. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan` is runtime-private state, not a user configuration object. -5. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the +6. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the process that owns the corresponding engine. ## Main features diff --git a/docs/design/feature/offloader/README.md b/docs/design/feature/offloader/README.md index 7b767f41f..96cba8f60 100644 --- a/docs/design/feature/offloader/README.md +++ b/docs/design/feature/offloader/README.md @@ -53,6 +53,11 @@ must use distinct hook names and remove only hooks they own. ## Discovery and topology +`resolve_offload_plan(pipeline, config)` is the single topology entry point. +Every backend calls it once when `enable()` starts and then consumes the +resolved components, block stacks, resident heads, and staged-residency flags +it returns. Backends do not re-read model declarations. + Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: - `_dit_modules`; @@ -61,16 +66,26 @@ Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: - `_resident_modules`. Dotted paths are supported. Legacy pipelines may use the fallback scan of -well-known attribute names, but new integrations should declare components -explicitly. +well-known attribute names; it now warns once per pipeline class, and new +integrations must declare components explicitly. -Both layerwise backends first consume pipeline `OffloadPlan` metadata. A -plan's `block_attrs` maps each DiT path to its ordered block containers, while -`encoder_block_attrs` declares streamable encoder stacks. DiTs absent from the -plan fall back to `_layerwise_offload_blocks_attrs` (including the deprecated -singular-name compatibility path). Discovery metadata describes structure -only; the backend remains responsible for transfer, synchronization, and -storage ownership. +Block topology comes from the pipeline `OffloadPlan` first: `block_attrs` maps +each DiT path to its ordered block containers, `encoder_block_attrs` declares +streamable encoder stacks, `on_demand_component_paths` marks pipeline-managed +residency, `resident_dit_paths` marks the DiTs that may hold resident layers, +and `encoder_dlo_weight_replication` marks the encoders whose loader-produced +weights are safe for AllGather. DiTs absent from the plan fall back to +`_layerwise_offload_blocks_attrs` (including the deprecated singular-name +compatibility path). + +The resolver is pure: it moves no tensor, installs no hook, writes no module +attribute, and reads no process group — multi-rank facts come from +`OffloadConfig`. It owns topology validation, so a configuration the model +cannot serve fails before the first component is placed. Explicit component +selection turns every declaration problem into an error, while the +compatibility topology (no component selector) keeps the historical +warn-and-skip behavior. Metadata still describes structure only; the backend +remains responsible for transfer, synchronization, and storage ownership. ## Cross-strategy invariants diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md index 2311ee1cf..87fb1f79f 100644 --- a/docs/design/fullduplex.md +++ b/docs/design/fullduplex.md @@ -66,7 +66,8 @@ The checkpoint keeps these verified contracts: - scheduler data-plane append over a resumable request; - stale epoch/turn/response fencing; - `/v1/duplex` and OpenAI Realtime projection; -- optional OpenAI `server_vad` turn detection backed by per-session Silero VAD; +- optional OpenAI `server_vad` turn detection backed by a process-shared Silero + model and per-session VAD state; - existing JoyVL behavior. The checkpoint does not claim: @@ -652,7 +653,7 @@ EOS decisions advance the model conversation. Explicit Realtime clients may still use `input_audio_buffer.commit` to create a conversation item. The translator validates that wire input before producing a -commit carrying `realtime_item_id`. Native auto-response may already have +commit carrying `item_id`. Native auto-response may already have streamed those PCM samples into the runtime, so the runner accepts a validated commit even when no runtime-side chunk remains. A truly empty explicit buffer continues to return `input_audio_buffer_empty`. @@ -677,8 +678,7 @@ multi-turn runtime evidence. During auto-response overlap, `preserve_realtime_input` distinguishes "do not append this silent chunk to the native buffer" from "clear the open Realtime -item". Silent overlap no longer discards earlier user PCM. A separate explicit -`server_vad` policy is documented in the Realtime contract above. +item". Silent overlap no longer discards earlier user PCM. The first chunk of one overlapping input item also reserves its target model turn. A later Realtime commit uses that reserved identity even if response EOS diff --git a/docs/design/module/diffusion/diffusion_runtime.md b/docs/design/module/diffusion/diffusion_runtime.md index 61224cd5b..0d36d08c2 100644 --- a/docs/design/module/diffusion/diffusion_runtime.md +++ b/docs/design/module/diffusion/diffusion_runtime.md @@ -8,22 +8,44 @@ owners: - "@SamitHuang" - "@fhfuih" primary_code_paths: + - vllm_omni/diffusion/diffusion_engine.py + - vllm_omni/diffusion/request.py + - vllm_omni/diffusion/data.py - vllm_omni/diffusion/sched/** - vllm_omni/diffusion/executor/** - vllm_omni/diffusion/worker/** related_code_paths: + - vllm_omni/diffusion/stage_diffusion_client.py + - vllm_omni/diffusion/stage_diffusion_proc.py + - vllm_omni/diffusion/inline_stage_diffusion_client.py + - vllm_omni/diffusion/ipc.py - vllm_omni/diffusion/hooks/** - vllm_omni/diffusion/postprocess/** depends_on: - ../engine_orchestration.md + - ../stage_runtime.md - ../input_output_modality_contracts.md - diffusion_model_integration.md validation_paths: - - tests/diffusion/diffusion_backend/** - - tests/diffusion/hooks/** + - tests/diffusion/test_diffusion_engine.py + - tests/diffusion/test_diffusion_engine_cleanup.py + - tests/diffusion/test_diffusion_engine_dummy_run.py + - tests/diffusion/test_diffusion_engine_rpc_routing.py + - tests/diffusion/test_diffusion_scheduler.py + - tests/diffusion/test_diffusion_model_runner.py + - tests/diffusion/test_diffusion_worker.py + - tests/diffusion/test_multiproc_engine_concurrency.py + - tests/diffusion/test_uniproc_executor.py + - tests/diffusion/test_result_pump.py + - tests/diffusion/test_async_output_worker.py + - tests/diffusion/test_diffusion_ipc.py + - tests/diffusion/test_diffusion_streaming_output.py + - tests/diffusion/test_stage_diffusion_proc.py + - tests/diffusion/test_inline_stage_diffusion_client.py upstream_refs: - diffusers.DiffusionPipeline -last_reviewed: 2026-07-16 +last_reviewed: 2026-09-02 +last_verified_commit: e51fe6ec1b9a9a0e14bb1fdb296d61b6593b93c6 --- # Diffusion runtime @@ -31,6 +53,380 @@ last_reviewed: 2026-07-16 The diffusion runtime owns request admission, scheduling, execution, progress, output, cancellation, and cleanup inside a diffusion stage. +It is easiest to think of the runtime as a small control loop: + +1. the **scheduler** chooses ready requests; +2. the **executor** sends that choice to the worker layer; +3. each **worker** manages its device and delegates model work to a **runner**; +4. the runner calls the model pipeline and returns per-request results; and +5. the engine feeds those results back to the scheduler and output streams. + +!!! note + + This page describes the runtime plumbing. Model loading and pipeline + contracts belong to [Diffusion model integration](diffusion_model_integration.md). + Batch compatibility and step-state details belong to + [Diffusion continuous batching](continuous_batching.md). + +## Goals and non-goals + +The runtime aims to keep request policy separate from device execution. It +also keeps request identity stable from admission through the final output, +including cancellation and failures. + +This module does **not** choose the next Omni stage, place stage replicas, or +define a model's denoising algorithm: + +- cross-stage routing belongs to [Engine orchestration](../engine_orchestration.md); +- process placement and replica lifecycle belong to + [Stage runtime](../stage_runtime.md); and +- model-specific work belongs to + [Diffusion model integration](diffusion_model_integration.md). + +## Runtime at a glance + +The stage owns `DiffusionEngine` and a `DiffusionExecutor` backend. The inline +stage client keeps that stack in the caller process; a process-backed stage +runs it inside `StageDiffusionProc`. + +The executor then chooses how the worker runs: + +- **`uni`** (default when `num_gpus == 1`): `UniProcDiffusionExecutor` builds + one in-process worker. There is no worker subprocess and no shared-memory + RPC. +- **`mp`** (default when `num_gpus > 1`, or when set explicitly): + `MultiprocDiffusionExecutor` starts one `WorkerProc` per device. + +Either way, each worker owns its device, distributed state, +`DiffusionWorker`, `DiffusionModelRunner`, and pipeline instance. + +```mermaid +flowchart TB + client["Stage client"] + + subgraph stage["Diffusion stage process"] + stageProc["StageDiffusionProc
stage transport and lifecycle"] + engine["DiffusionEngine
admission and control loop"] + scheduler["RequestScheduler or StepScheduler
request state and policy"] + executor["DiffusionExecutor
uni or mp backend"] + streams["Per-request output streams"] + + stageProc --> engine + engine --> scheduler + scheduler -->|"DiffusionSchedulerOutput"| engine + engine --> executor + engine --> streams + end + + subgraph workers["Worker layer"] + workerProc["WorkerProc
mp IPC loop only"] + worker["DiffusionWorker
device and distributed setup"] + runner["DiffusionModelRunner
model state and execution"] + pipeline["Diffusion pipeline
model-specific computation"] + + workerProc --> worker + worker --> runner --> pipeline + end + + client -->|"process-backed: ZMQ"| stageProc + client -.->|"inline: direct call"| engine + executor -->|"mp: broadcast / result queues"| workerProc + executor -.->|"uni: in-process call"| worker + streams -->|"process-backed"| stageProc + stageProc -->|"ZMQ results"| client + streams -.->|"inline"| client +``` + +!!! info "Stage process vs worker process" + + A deployed stage may already run in a `StageDiffusionProc` child process. + That is separate from the optional `mp` worker processes. With `uni`, the + worker lives in the same process as the engine. + +## Component responsibilities + +| Component | Owns | Does not own | +| --- | --- | --- | +| `DiffusionEngine` | Admission, the busy loop, scheduler/executor coordination, RPC ordering, output streams, warmup, abort routing | Batch policy details, device setup, model computation | +| `BaseScheduler` | Request states, waiting/running sets, capacity, compatibility, terminal transitions | IPC, worker calls, output formatting | +| `RequestScheduler` | Complete-request waves and optional admission delay for request-level batching | Denoising-step progress | +| `StepScheduler` | Per-request denoising progress and step-wise completion | Pipeline tensors and model state | +| `DiffusionExecutor` | Execution backend contract, worker RPC, health, shutdown | Admission and request-state transitions | +| `UniProcDiffusionExecutor` | Single in-process worker for `num_gpus == 1`; no IPC or async output pump | Multi-GPU execution | +| `MultiprocDiffusionExecutor` | Worker processes, shared-memory message queues, result dispatch, worker monitoring | In-process single-GPU path | +| `WorkerProc` | One `mp` worker process's IPC loop and reply rules | Scheduling policy; unused by `uni` | +| `DiffusionWorker` | Device/distributed setup, LoRA activation, profiling, sleep/wake, runner delegation | Request admission | +| `DiffusionModelRunner` | Pipeline loading, request-local model state, cache/compile setup, request or step execution | Queueing and cross-stage routing | + +The boundary between policy and execution is +`DiffusionSchedulerOutput`. It contains newly admitted request payloads, +IDs for requests whose runner state is already cached, finished IDs, +optional KV-prefetch work, and per-request `diffusion_kv_metadata` on +`NewRequestData` envelopes when paged Diffusion KV is enabled. Executors and +workers consume this output; they do not decide what should run next. + +`BaseScheduler.initialize()` sets `max_num_running_reqs` from +`od_config.max_num_seqs` (default 1). The engine may override that capacity +for distributed layerwise offload with AllGather DP concurrency. + +## One scheduling tick + +Async requests share one engine busy loop. Worker calls and control RPCs pass +through that loop, so they cannot race on the executor transport. + +```mermaid +sequenceDiagram + participant Caller + participant Engine as DiffusionEngine + participant Scheduler + participant Executor + participant Worker + participant Runner + + Caller->>Engine: add request + Engine->>Scheduler: add_request(request) + Engine->>Scheduler: schedule() + Scheduler-->>Engine: DiffusionSchedulerOutput + Engine->>Executor: execute_batch() or execute_step() + Executor->>Worker: run worker method + Worker->>Runner: execute model or one step + Runner-->>Worker: RunnerOutput(s) + Worker-->>Executor: result + Executor-->>Engine: BaseRunnerOutput + Engine->>Scheduler: update_from_output(...) + Scheduler-->>Engine: finished request IDs + Engine-->>Caller: chunk or final output +``` + +The engine catches request execution errors and turns them into per-request +error outputs. A dead worker group is different: the executor marks itself +failed, health checks raise `EngineDeadError`, and the owning stage client +handles the stage-level failure. + +## Execution modes + +The engine resolves one mode at startup and binds the matching scheduler and +executor call. + +| Mode | Scheduler | Executor call | Runner path | +| --- | --- | --- | --- | +| Request batch | `RequestScheduler` | `execute_batch()` | `execute_model()` for one request, or `execute_model_batch()` for a fused batch | +| Step batch | `StepScheduler` | `execute_step()` | `execute_stepwise()` | + +**Request mode** runs a complete pipeline forward for each scheduled wave. A +single-request wave is the conservative path. A multi-request wave usually +uses fused `execute_model_batch()` and requires the pipeline to declare +request-batch support. + +**Distributed layerwise offload with AllGather (DLO DP concurrency)** is a +separate multiproc dispatch path. It activates when +`data_parallel_size > 1`, `enable_distributed_layerwise_offload` is set, and +`dlo_use_allgather` is true. The engine then sets `dp_concurrent = True` and +raises `scheduler.max_num_running_reqs` to `dp_size`, overriding the +`max_num_seqs` value from `initialize()`. The scheduler still emits a +multi-request `DiffusionSchedulerOutput`; the multiproc executor routes that +wave through `execute_request()` instead of fused `execute_model_batch()`. + +Optional admission coalescing can fill a `dp_size` wave: when +`request_batch_max_wait_ms > 0`, `RequestScheduler` may wait briefly before +the first schedule of a wave (under `dp_concurrent`, the stable window is +`min(0.3s, wait/2)`). With the default `request_batch_max_wait_ms == 0`, +there is no wait. + +Before dispatch, `MultiprocDiffusionExecutor.execute_request()` rejects waves +whose requests differ in sampling-parameter compatibility or `extra_args` +(AllGather requires every DP rank to follow the same forward schedule). Each +worker rank then picks one envelope from the wave: +`req[dp_rank % len(req)]`, so every rank enters the same collective while +computing different requests. + +**Step mode** keeps `StepRequestState` in the runner. New requests run +`prepare_encode()` once; every tick runs `denoise_step()` and +`step_scheduler()`; completed requests run `post_decode()`. The scheduler keeps +only lifecycle and progress metadata—it does not hold model tensors. + +**Step mode is also the path for streaming diffusion output**. Enabling +`streaming_output` turns on step execution if needed. Chunk-capable pipelines +can emit intermediate outputs through the same request stream; final-only step +pipelines still emit only the final result. If the pipeline does not implement +step execution, initialization fails. + +!!! tip + + For user-facing flags and examples, see + [Diffusion execution modes](../../../user_guide/diffusion/execution_modes.md). + That guide explains how to select a mode; this page explains what happens + after the choice is made. + +## Request state and identity + +The scheduler is the source of truth for an admitted request. Its normal state +flow is: + +```mermaid +stateDiagram-v2 + [*] --> WAITING: add_request + WAITING --> RUNNING: schedule + WAITING --> FINISHED_ABORTED: abort + RUNNING --> FINISHED_ABORTED: abort + RUNNING --> FINISHED_COMPLETED: successful output + RUNNING --> FINISHED_ERROR: failed or missing output + FINISHED_COMPLETED --> [*]: deliver and remove state + FINISHED_ABORTED --> [*]: deliver and remove state + FINISHED_ERROR --> [*]: deliver and remove state +``` + +`PREEMPTED` and `BaseScheduler.preempt_request()` exist on the scheduler API, +but the engine loop does not call them today—only scheduler unit tests do. + +`request_id` ties together scheduler state, runner state, executor results, and +the output queue. Batch positions are temporary and must never replace request +identity when results are mapped back. + +Only compatible requests enter the same running wave. The scheduler compares a +mode-specific sampling-parameter key and stops at the first incompatible +waiting request. This is intentionally conservative and can cause +head-of-line blocking. + +## Executor and IPC + +`DiffusionExecutor` is the backend interface. Built-in backends are selected by +`distributed_executor_backend`: + +| Backend | When chosen | Worker layout | +| --- | --- | --- | +| `uni` | Default for `num_gpus == 1`, or set explicitly | One in-process worker | +| `mp` | Default for `num_gpus > 1`, or set explicitly | One worker process per device | + +Ray and external-launcher diffusion backends are not implemented. A custom +`DiffusionExecutor` subclass or import path is also accepted. + +### Uniproc backend + +`UniProcDiffusionExecutor` constructs `WorkerWrapperBase` in the engine process +and calls worker methods directly. It avoids a second model load, MessageQueue +rings, ZMQ IPC sockets, and `/dev/shm` tensor packing. RPC timeouts are +accepted for interface parity but are not enforced: a hung worker blocks the +calling thread. Sticky accelerator faults mark the executor dead through +failure callbacks; ordinary request errors stay per-request. + +### Multiproc backend + +`MultiprocDiffusionExecutor`: + +1. creates a broadcast queue shared by all workers; +2. spawns one worker process per configured device; +3. waits until every worker reports ready; +4. sends generation and control calls as RPC messages; +5. applies rank-aware reply rules so only expected ranks respond; and +6. monitors worker process sentinels and fails the executor if a worker dies. + +In request mode on this backend, each async output produces two correlated +messages keyed by `async_output_id`: + +- `COMPUTE_DONE` — the forward finished and the device can start the next + request. +- `OUTPUT_READY` — background D2H/SHM packing finished and the final output + is ready. + +The worker queues background packing work before it enqueues `COMPUTE_DONE`, +but either message may reach the executor first. The result pump and +`execute_batch()` must handle both orderings. Step mode keeps the synchronous +result path and does not start those pumps. + +See [Async diffusion output](../../feature/async_diffusion_output.md) for the +multiproc request-mode timeline. + +## Worker and runner boundary + +On the `mp` path, `WorkerProc` receives messages and calls methods through +`WorkerWrapperBase`. On the `uni` path, the executor owns that wrapper +directly. In both cases, `DiffusionWorker` handles the parts tied to a device: +distributed initialization, model-runner construction, LoRA activation, +profiling, and memory sleep/wake. It delegates actual model work to +`DiffusionModelRunner`. + +The runner owns long-lived model-side state: + +- the loaded pipeline and compilation setup; +- cache and offload integration; +- random-generator setup; +- request batches for complete forwards; and +- `StepRequestState` plus `InputBatch` for step execution. + +This split keeps infrastructure out of model pipelines. A pipeline receives a +request batch or step state and performs model work; it should not inspect +engine queues or change scheduler state. + +## Lifecycle and cleanup + +### Startup + +`DiffusionEngine.make_engine()` resolves the engine class, constructs the +executor and scheduler, then normally runs a small dummy request to warm up the +model. + +Warmup is skipped only when distributed layerwise offload with AllGather +is enabled (via `enable_distributed_layerwise_offload` and `dlo_use_allgather`) +and `max(data_parallel_size, sequence_parallel_size) > 1`, because +the dummy run sends one request while AllGather needs every shard rank to +enter the same collective. If initialization or warmup fails, the runtime +closes scheduler state and worker resources before returning the error. + +### Cancellation + +`abort()` places request IDs on an abort queue. The busy loop marks those +requests `FINISHED_ABORTED`; finalization removes scheduler state and emits an +aborted output when a consumer still exists. The runner removes cached +step-mode state when it receives a scheduler output that reports the finished +request ID. + +On the multiproc request-mode async path, abort or consumer drop while output +is still materializing should release the associated async-output bookkeeping +so late results are not retained after the request ends +*(pending #6253/#6439/#6580)*. Today, an unconsumed `OUTPUT_READY` is cached +by design until `wait_output_ready()` or executor teardown. Dropping an +output-stream consumer removes that consumer's queue. It does not take +ownership of scheduler cleanup. + +### Shutdown and worker failure + +`DiffusionEngine.close()` is idempotent. It stops the busy loop, wakes pending +streams with an error, closes scheduler state, and shuts down the executor. + +- **`mp`**: the executor asks workers to stop, waits for them, then terminates + workers that miss the grace period. Shutdown marks unfinished RPC and + async-output futures as failed, then clears them. Completed async outputs + cached for later `wait_output_ready()` calls are dropped with the executor + object. +- **`uni`**: the executor shuts down the in-process worker, drops the final + model reference, and empties the accelerator cache so a later engine can reuse + the device. + +!!! warning + + Do not continue using a worker group after a fatal collective timeout or + unexpected worker exit. Distributed state may be incomplete. The + multiprocess executor deliberately fails closed so the stage can restart. + The uniproc executor fails closed when the accelerator context itself is + poisoned. + +## Extension boundaries + +- A custom engine may set `default_diffusion_model_runner_cls`; an explicit + `diffusion_model_runner_cls` configuration still wins. +- Tests and custom engine integrations may inject a `BaseScheduler` subclass. + `SchedulerInterface` remains only as a deprecated compatibility name. +- `distributed_executor_backend` may be `"uni"`, `"mp"`, a custom `DiffusionExecutor` + subclass, or an import path. A backend must preserve scheduler output, + request identity, health, and cleanup contracts. Use `"mp"` when you need + process isolation or multi-GPU; `"uni"` is the single-GPU default. +- Worker extensions go through `WorkerWrapperBase`. They add worker methods; + they do not become a second scheduler or request lifecycle. + +These are advanced Python integration points, not stable end-user CLI +customization surfaces. + ## Candidate invariants ### DIFF-RUNTIME-INV-001: One lifecycle owner @@ -41,19 +437,50 @@ until completion, cancellation, or failure. ### DIFF-RUNTIME-INV-002: Execution follows scheduler output **Rule:** Executors and workers MUST execute scheduler decisions without -admitting, reordering, or forwarding requests independently. +admitting, reordering, or forwarding requests independently. Distributed +layerwise offload with AllGather is an exception at dispatch time: the +scheduler may schedule a multi-request wave, but each DP rank executes one +envelope from that wave so every rank enters the same collective schedule. ### DIFF-RUNTIME-INV-003: Terminal cleanup is complete **Rule:** Every terminal path MUST release request state, temporary tensors, -hooks, and runtime-owned resources. +hooks, and runtime-owned resources. Async-output bookkeeping on abort or +consumer drop before consumption is also required *(pending #6253/#6439/#6580)*. ### DIFF-RUNTIME-INV-004: Optional features use runtime hooks **Rule:** Cache, profiling, offload, and parallel features SHOULD integrate at defined hooks instead of creating another request lifecycle. +### DIFF-RUNTIME-INV-005: Results keep request identity + +**Rule:** Batched and asynchronous results MUST be mapped by stable request ID, +not by assumed completion order. + ## Safe-change guide -Test success, cancellation, failure, shutdown, multiple requests, and resource -cleanup. +Test the smallest affected slice, then cover its neighboring boundary: + +| Change area | Minimum evidence | +| --- | --- | +| Admission or state transitions | `test_diffusion_scheduler.py`, including duplicate IDs, compatibility, abort, and missing output | +| Engine loop or output delivery | `test_diffusion_engine.py`, `test_diffusion_engine_cleanup.py`, `test_diffusion_engine_rpc_routing.py` | +| Executor or IPC | `test_multiproc_engine_concurrency.py`, `test_uniproc_executor.py`, `test_result_pump.py`, `test_diffusion_ipc.py`, `test_async_output_worker.py` | +| Worker or runner | `test_diffusion_worker.py`, `test_diffusion_model_runner.py` | +| Stage boundary | `test_stage_diffusion_proc.py`, `test_inline_stage_diffusion_client.py` | +| Warmup or streaming | `test_diffusion_engine_dummy_run.py`, `test_diffusion_streaming_output.py` | + +Always exercise success, cancellation, per-request failure, fatal worker +failure, repeated shutdown, one request, and multiple compatible requests. +For step-mode changes, also test partial progress and cleanup of runner state. + +## Related documents + +- [Diffusion module overview](index.md) +- [Diffusion model integration](diffusion_model_integration.md) +- [Diffusion continuous batching](continuous_batching.md) +- [Diffusion execution modes](../../../user_guide/diffusion/execution_modes.md) +- [Async diffusion output](../../feature/async_diffusion_output.md) +- [Engine orchestration](../engine_orchestration.md) +- [Stage runtime](../stage_runtime.md) diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index b695185d1..4c3ad3373 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/hooks/generate_argparse.py @@ -124,6 +124,7 @@ def extract_omni_serve_subparser_init(): "make_arg_parser": lambda parser: parser, # no-op for doc "_ensure_vllm_platform": lambda: None, # no-op for doc "_nonneg_finite_float": float, + "_parse_stage_overrides": __import__("json").loads, "VLLM_SUBCMD_PARSER_EPILOG": "", "logger": logger, "DummySubparsers": DummySubparsers, diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 9e3a5a9f7..6b0952eac 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -64,6 +64,7 @@ th { | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) | | `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) | +| `BooguImageTurboPipeline` | Boogu-Image Turbo distilled T2I | `Boogu/Boogu-Image-0.1-Turbo` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) | | `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` | ✅︎ | ✅︎ | | ✅︎ | — | @@ -96,6 +97,7 @@ th { | `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` | ✅︎ | ✅︎ | | ✅︎ | — | +| `Audio8TTSSlowARForConditionalGeneration` | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✅︎ | | | | — | | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) | | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — | diff --git a/docs/serving/openpi_api.md b/docs/serving/openpi_api.md index 27f092d87..52d94cbbf 100644 --- a/docs/serving/openpi_api.md +++ b/docs/serving/openpi_api.md @@ -64,7 +64,7 @@ connect <- msgpack(policy_server_config) infer - -> msgpack({"endpoint": "infer", "session_id": "...", ...observation}) + -> msgpack({"endpoint": "infer", "session_id": "...", "seed": , ...observation}) <- msgpack(ndarray | dict[str, ndarray]) reset @@ -92,6 +92,12 @@ frames are not observation messages. - The policy pipeline owns observation transforms and persistent model state, normally keyed by `session_id`; the API layer forwards the raw observation dictionary. +- An optional integer `seed` in an inference message becomes the engine + request's `sampling_params.seed`. A policy whose sampling consumes the + request generator (GR00T-N1.7) returns the same action chunk for the same + observation and seed. If omitted, the engine assigns a random per-request + seed, so unseeded inferences are independent of each other. The API layer + consumes `seed`; it is not forwarded inside the observation. ## Limits and Errors diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index 30021e107..d37d02291 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -1,17 +1,17 @@ # Realtime Duplex API vLLM-Omni serves full-duplex speech models over a WebSocket endpoint, -`/v1/realtime?duplex=1`, that speaks the OpenAI Realtime vocabulary plus a +`/v1/realtime`, that speaks the OpenAI Realtime vocabulary plus a set of duplex extensions: the model decides when to listen and when to speak, the user can talk over the assistant, playback progress is reported back so history stays honest, and a dropped connection can resume the same session. This page covers how to run a duplex deployment, how to drive it from Python with `vllm_omni.clients.duplex.DuplexClient`, and the complete wire contract. -The endpoint is served for models that ship a duplex plugin — currently -MiniCPM-o 4.5, PersonaPlex, and Nemotron VoiceChat — and only for deploy -configurations that declare `session_mode: duplex`. The runtime architecture -is described in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md). +The handler is enabled by either a model-native `session_mode: duplex` +deployment or an explicit `duplex_session` configuration for turn-based +Server VAD. The runtime architecture is described in +[Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md). ## Quick Start @@ -27,7 +27,7 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \ `vllm_omni/deploy/minicpmo_4_5.yaml` declares `session_mode: duplex` and `duplex_session.max_sessions: 4`, which mounts the WebSocket routes -`ws://:8099/v1/realtime?duplex=1` (this page) and `ws://:8099/v1/duplex` +`ws://:8099/v1/realtime` (this page) and `ws://:8099/v1/duplex` (the native dialect for raw-protocol clients) next to the usual HTTP API. PersonaPlex and Nemotron VoiceChat use `vllm_omni/deploy/personaplex.yaml` and `vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml` in the same way. @@ -240,7 +240,7 @@ on the same helpers. ## API Reference -This section is the complete wire contract of `/v1/realtime?duplex=1`: how +This section is the complete wire contract of `/v1/realtime`: how the dialect relates to the OpenAI Realtime protocol, which surfaces each model's capability flags gate, a catalogue of every client-to-server and server-to-client event, and a JSON example for each. `DuplexClient` speaks @@ -253,10 +253,11 @@ its protocol tests. ### Endpoint and transport -The endpoint is `ws(s):///v1/realtime?duplex=1`; the `duplex` query -parameter (`1`, `true`, or `on`) selects the duplex session handler, and the -route is mounted only for deployments whose deploy configuration declares -`session_mode: duplex`. Optional query parameters are `model`, `session_id`, +The endpoint is `ws(s):///v1/realtime`. Bare connections select the +session handler when the deployment declares either `session_mode: duplex` +or `duplex_session`. `?duplex=1` remains a compatibility alias, while +`?duplex=0` selects the legacy Realtime handler. Optional query parameters +are `model`, `session_id`, `autostart` (`0` means resume-only), `resume`, and `native_duplex` (`minicpmo45_native_duplex` is accepted as a deprecated alias and folded into the canonical name). Every message is one JSON object per WebSocket @@ -342,9 +343,14 @@ Semantic divergences hidden behind shared names: before any commit; a commit may end in `response.listen` and no response, and the model may open a response with no commit at all. OpenAI: commit ⇒ item, `response.create` ⇒ exactly one response. -- `turn_detection.interrupt_response=false` is rejected; - `create_response` is ignored; `semantic_vad` is unsupported; the VAD - runs per session (Silero) and implies `overlap_policy=barge_in_on_speech`. +- With `turn_detection={"type":"server_vad"}`, omitted `interrupt_response` + defaults to `false` on turn-based models such as Qwen3-Omni, enabling + endpointing and automatic responses without interruption. Explicit `true` + is rejected there. Model-native duplex keeps the `true` default and rejects + `false`; its VAD implies `overlap_policy=barge_in_on_speech`. + This capability-dependent default differs from OpenAI's `true` default; + the effective configuration is returned in `session.created` / `session.updated`. + `semantic_vad` is unsupported. - `rate_limits.updated` is always an empty list (compatibility only). - Cancellation never reuses a `response_id`; truncation is driven by `playback.ack` as well as `conversation.item.truncate`. @@ -1201,9 +1207,10 @@ Realtime projector renames and fans out these events before they reach a barge-in, or audio truncation; Nemotron VoiceChat does not support barge-in or audio truncation; camera frames are consumed only by MiniCPM-o 4.5; tool calls are produced only by Nemotron VoiceChat. -- `turn_detection` supports only `server_vad` with `interrupt_response=true`; - `semantic_vad`, `interrupt_response=false`, and `create_response` are not - supported. Server VAD requires a model-native session. +- `turn_detection` supports `server_vad` or `null`, not `semantic_vad`. + Turn-based models support endpointing without interruption, including + `create_response=false` for explicit response creation. Model-native duplex + requires `interrupt_response=true` and `create_response=true`. - `input_audio_transcription` and `input_audio_noise_reduction` are accepted and echoed but no separate transcription or noise-reduction stage runs; transcripts come from the model. diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index d4871ce15..f0658f278 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -279,10 +279,15 @@ curl -X POST http://localhost:8091/v1/audio/voices \ ## Streaming Text Input (WebSocket) -The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and generates audio per sentence as boundaries are detected. +The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally. +By default (`split_granularity=none`) it buffers until `input.done` and +synthesizes that flush as **one** TTS request, which keeps long-form timbre +stable. Set `split_granularity` to `sentence` or `clause` to emit a request +at each detected boundary (lower time-to-first-audio for STT/LLM pipelines). -> Note: text input is always streamed incrementally. Audio output remains sentence-scoped: -> use `stream_audio=false` for one binary frame per sentence, or `stream_audio=true` for one or more PCM chunks per sentence. +> Note: `stream_audio` only changes how **audio bytes** are framed (one WAV/PCM +> payload vs chunked PCM). Text segmentation is controlled separately by +> `split_granularity`. ### WebSocket Protocol @@ -314,14 +319,16 @@ upstream LLM) pays the WebSocket handshake once instead of once per utterance. - The session config is sticky. Send `input.text` again straight after `session.done` to reuse it, or send another `session.config` first to change - voice, format, or reference audio. A `session.config` sent while text is - still buffered is rejected so no pending input is silently dropped. -- An utterance is the flush unit, not a linguistic one: it is whatever text was - buffered when `input.done` arrived, of any length, synthesized as one request. - `utterance_index` counts those flushes across the connection, so it tells you - which `input.done` a frame belongs to. `sentence_index` counts within one - flush and so pairs with `total_sentences`, which means every utterance reports - `sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer). + voice, format, or reference audio. A `session.config` sent in the middle of + an utterance is rejected so no pending input is silently dropped and so a + split utterance cannot end up half in one voice and half in another. +- An utterance is the flush unit, not a linguistic one: it is one `input.done` + cycle. `utterance_index` counts those flushes across the connection, so it + tells you which `input.done` a frame belongs to. `sentence_index` counts the + TTS requests inside one flush and so pairs with `total_sentences`: with the + default `split_granularity=none` that is always `sentence_index: 0` of + `total_sentences: 1` (or `0` for an empty buffer), while `sentence` or + `clause` counts the linguistic units actually synthesized. - End the connection with `session.close`, or by closing the socket. An idle connection is still closed after the server's idle timeout, which now also applies to the gap between utterances. @@ -332,7 +339,15 @@ All REST API parameters are supported, plus: | Parameter | Type | Default | Description | | ----------- | ------ | --------- | ------------- | -| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket | +| `stream_audio` | bool | false | Stream one or more PCM chunks for each TTS request over WebSocket | +| `split_granularity` | string | `"none"` | `"none"`: one request per `input.done`. `"sentence"`: split on `.!?` plus CJK `。！？…`, Indic danda `।॥`, and Arabic `؟`. `"clause"`: also split on `,;，；،؛`. | +| `seed` | integer | null | Forwarded to the speech engine for this session | + +ASCII punctuation only ends a unit when whitespace or `input.done` follows it, +and decimals (`3.14`), thousands separators (`1,000`), abbreviations (`Dr.`, +`e.g.`) and initials (`J. R.`) are not treated as boundaries. A punctuation run +and any closing quote or bracket stay with the unit they close, so `Wait...` +and `He said "Hello."` are one request each. ```bash DELETE /v1/audio/voices/{name} @@ -752,6 +767,10 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne | ------- | ------------- | | `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. | +OmniVoice uses packed variable-length attention for batched generator execution. The attention operator accepts FP16 and BF16 inputs, so its query, +key, and value tensors are evaluated in BF16 even when the stage is configured with `dtype: float32`; the attention output is converted back to the model's +hidden-state dtype before the output projection. Consequently, a float32 stage configuration does not imply FP32 attention arithmetic. + ### VoxCPM2 | Model | Description | diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 5a1d8ecf2..09f023d27 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -68,6 +68,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4 | Parameter | Type | Default | Description | | ----------- | ------ | --------- | ------------- | | `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests | +| `control_reference` | file | null | Optional uploaded image/video control, up to 512 MiB, for models that declare control-upload support | +| `control_type` | string | null | Model control name associated with `control_reference`; currently Cosmos3 supports `edge`, `blur`, `depth`, `seg`, and `wsm` | | `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` | | `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` | | `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs | @@ -177,6 +179,23 @@ corresponding addition. `negative_metadata_mode` accepts `same`, `inverse`, or `none` and defaults to `same` for transfer. See the Cosmos3 recipe for complete examples. +A client that cannot place the control on the server filesystem can upload one +control with `control_reference` and identify it with `control_type`. The API +streams the upload to request-scoped storage, supplies its path to the model, +and removes it after synchronous or asynchronous generation completes. Other +options for the selected control can remain in `extra_params`; do not also set +`control` or `control_path` there. Uploads larger than 512 MiB are rejected. + +```bash +curl -s http://localhost:8091/v1/videos/sync \ + -F "prompt=Preserve the scene while following the world-state control" \ + -F "input_reference=@input.mp4;type=video/mp4" \ + -F "control_reference=@wsm.mp4;type=video/mp4" \ + -F "control_type=wsm" \ + -F 'extra_params={"wsm":{"control_weight":1.0}}' \ + -o output.mp4 +``` + HTTP redirects for `image_reference.image_url` follow vLLM's `VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to `1` (the default) to allow redirects or `0` to reject them. A redirect target diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 00c245489..a8fb6d972 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -162,9 +162,9 @@ The following tables show which models support each feature: | **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ | | **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ | | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ | +| **LTX-2** | ❌ | ✅ | ✅ (Ulysses only) | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **LTX-2.3** | ❌ | ✅ | ✅ (Ulysses only) | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **LTX-2.5** | ❌ | ❓ (one-stage) | ✅* (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ | | **SANA-Video-2B T2V I2V** | ❌ | ❌ | ✅ (`--usp`, frame-sharded) | ✅ | ✅ (TP=2 only) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md index a6544c1e5..95ffc5543 100644 --- a/docs/user_guide/examples/offline_inference/text_to_speech.md +++ b/docs/user_guide/examples/offline_inference/text_to_speech.md @@ -126,11 +126,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ ## Fish Speech S2 Pro 4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz. - -### Prerequisites -```bash -pip install fish-speech -``` +No extra packages are required; the DAC codec is vendored in vLLM-Omni. ### Quick start ```bash 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 f1b643ac7..9a33dc858 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -155,13 +155,8 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex ## Fish Speech S2 Pro -4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec. - -### Prerequisites - -```bash -pip install fish-speech -``` +4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec, which is vendored in +vLLM-Omni (no extra packages required). ### Launch diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 6dbafdfc1..bbeef32b0 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -5,10 +5,12 @@ This example shows how to use vLLM-Omni for running offline inference with the correct prompt format on Qwen2.5-Omni """ +import argparse +import functools import json import os import time -from typing import NamedTuple +from typing import Any, NamedTuple import numpy as np import soundfile as sf @@ -26,6 +28,20 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser SEED = 42 +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"{flag_name} must be valid JSON: {e}") from e + if not isinstance(config, dict): + 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") + + class QueryResult(NamedTuple): inputs: dict limit_mm_per_prompt: dict[str, int] @@ -378,11 +394,10 @@ def main(args): for i, prompt in enumerate(prompts): prompt["modalities"] = output_modalities - profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) - if profiler_enabled and hasattr(omni, "start_profile"): + profiler_enabled = args.profiler_config is not None + if profiler_enabled: + print("[Profiler] Starting profiling...") omni.start_profile(stages=[0]) - elif profiler_enabled: - print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.") omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator) # Determine output directory: prefer --output-dir; fallback to --output-wav @@ -418,7 +433,7 @@ def main(args): print(f"Request ID: {request_id}, Saved audio to {output_wav}") processed_count += 1 - if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests: + if profiler_enabled and processed_count >= total_requests: print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...") # Stop the profiler while workers are still alive omni.stop_profile() @@ -476,6 +491,12 @@ def parse_args(): default=300, help="Timeout for initializing stages in seconds (default: 300)", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) parser.add_argument( "--shm-threshold-bytes", type=int, diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index c21bff5fe..a8a273026 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -14,6 +14,7 @@ list of supported architectures across all modalities, see | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate | |---|---|---|---|---|---|---| +| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | | Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz | @@ -44,6 +45,57 @@ python examples/offline_inference/text_to_speech//end2end.py \ --- +## Audio8 TTS Preview + +0.6B DualAR text-to-speech model (Audio8) with its own 44.1 kHz neural audio +codec: a slow AR transformer predicts one semantic token per codec frame, and a +4-layer fast AR predicts that frame's remaining 9 codebooks. + +### Prerequisites +No extra packages: the codec architecture ships in tree and its weights +(`codec.pth`) come with the checkpoint. + +### Quick start +```bash +python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ + --text "Welcome to Audio8 TTS, a compact multilingual text to speech model." +``` + +### Voice cloning +```bash +python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ + --text "Welcome to Audio8 TTS." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript of the reference recording." +``` +The reference transcript must match what is actually spoken in the clip; +a mismatch degrades both stability and speaker similarity. + +### Streaming +```bash +python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ + --text "Welcome to Audio8 TTS." --streaming +``` + +### Notes +- Output: 44.1 kHz mono WAV; the codec emits ~21.5 frames/s (2048 samples per frame). +- Context is 2048 packed text+audio positions, so `max_model_len` is 2048. +- Greedy decoding (`temperature: 0`) is degenerate for this checkpoint — it + ends the utterance almost immediately. Keep the deploy defaults + (`temperature 0.7`, `top_k 50`, `top_p 0.9`), which also match + `generation_config.json`. +- Repetition-Aware Sampling is applied to the semantic token: a repeat inside + the last 10 frames is re-drawn from a flatter distribution instead of being + masked out. +- `--num-prompts 4` smoke-tests concurrency (stage 0 runs `max_num_seqs: 4`). + +### Measured behaviour (1x H20, shared GPU) +`--streaming` on a 58-character English sentence: TTFA ~0.74 s (4-frame first +chunk), 7 chunks, 1.55 s wall for 3.16 s of audio (RTF ~0.49). Treat these as +smoke-test magnitudes, not a benchmark. + +--- + ## CosyVoice3 2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz. @@ -130,11 +182,7 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 ## Fish Speech S2 Pro 4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz. - -### Prerequisites -```bash -pip install fish-speech -``` +No extra packages are required; the DAC codec is vendored in vLLM-Omni. ### Quick start ```bash diff --git a/examples/offline_inference/text_to_speech/audio8_tts/end2end.py b/examples/offline_inference/text_to_speech/audio8_tts/end2end.py new file mode 100644 index 000000000..8ce94f385 --- /dev/null +++ b/examples/offline_inference/text_to_speech/audio8_tts/end2end.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Offline inference demo for Audio8 TTS Preview 0.6B via vLLM Omni. + +Usage: + # Text-only synthesis + python end2end.py --text "Welcome to Audio8 TTS." + + # Zero-shot voice cloning (the transcript must match the reference audio) + python end2end.py --text "Welcome to Audio8 TTS." \ + --ref-audio reference.wav --ref-text "The exact transcript of the reference recording." + + # Streaming (per-chunk TTFA / inter-chunk timings) + python end2end.py --text "Welcome to Audio8 TTS." --streaming +""" + +import asyncio +import os +import time + +import numpy as np +import soundfile as sf +import torch + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +from vllm_omni import AsyncOmni, Omni +from vllm_omni.model_executor.models.audio8_tts.codec_utils import estimate_reference_code_frames +from vllm_omni.model_executor.models.audio8_tts.prompt_utils import ( + build_text_only_prompt_ids, + estimate_voice_clone_prompt_len, + normalize_text, +) +from vllm_omni.utils.tracking_parser import TrackingArgumentParser + +DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b" + + +def build_prompt( + text: str, + ref_audio_path: str | None = None, + ref_text: str | None = None, + model_name: str = DEFAULT_MODEL, +) -> dict: + """Build an engine prompt, using the same protocol as online serving.""" + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_name) + if ref_audio_path is None and ref_text is None: + prompt_ids, normalized_text = build_text_only_prompt_ids(tokenizer, text) + return { + "prompt_token_ids": prompt_ids, + "additional_information": {"text": [normalized_text]}, + } + + if not ref_audio_path or not ref_text: + raise ValueError("Audio8 TTS voice cloning requires both --ref-audio and --ref-text") + + normalized_text = normalize_text(text) + normalized_ref_text = normalize_text(ref_text, add_default_speaker=True) + ref_audio_wav, ref_audio_sr = sf.read(ref_audio_path, dtype="float32", always_2d=False) + if ref_audio_wav.ndim > 1: + ref_audio_wav = ref_audio_wav.mean(axis=-1) + ref_frames = estimate_reference_code_frames(int(ref_audio_wav.shape[0]), int(ref_audio_sr)) + placeholder_len = estimate_voice_clone_prompt_len(tokenizer, normalized_text, normalized_ref_text, ref_frames) + + # The clone prompt embeds the reference audio's own codec codes, so the + # model builds it in preprocess(); we only reserve the exact length here. + return { + "prompt_token_ids": [1] * placeholder_len, + "additional_information": { + "text": normalized_text, + "ref_text": normalized_ref_text, + "ref_audio_wav": torch.from_numpy(np.asarray(ref_audio_wav, dtype=np.float32)), + "ref_audio_sr": int(ref_audio_sr), + "audio8_structured_voice_clone": True, + }, + } + + +def extract_audio(multimodal_output: dict) -> tuple[torch.Tensor, int]: + """Pull the waveform and sample rate out of a multimodal output. + + The output processor concatenates the per-step delta tensors under + ``model_outputs``; ``audio`` is the older alias. Never use ``a or b`` on + these values -- truthiness on a tensor raises. + """ + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is None: + raise ValueError(f"No audio in multimodal_output: {sorted(multimodal_output)}") + if isinstance(audio, list): + audio = torch.cat([torch.as_tensor(chunk).reshape(-1) for chunk in audio], dim=0) + audio = torch.as_tensor(audio).reshape(-1) + + sr_raw = multimodal_output.get("sr") + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else None + if sr_raw is None: + raise ValueError("Missing sample rate in multimodal_output") + sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw) + return audio, sample_rate + + +def _save_wav(output_dir: str, request_id: str, audio: torch.Tensor, sample_rate: int) -> None: + out_wav = os.path.join(output_dir, f"output_{request_id}.wav") + sf.write(out_wav, audio.float().cpu().numpy(), samplerate=sample_rate, format="WAV") + duration = audio.numel() / sample_rate + # print(), not logging: vLLM reconfigures logging on engine start and + # disables loggers created before that, which would swallow these lines. + print(f"Request {request_id}: saved {out_wav} (sr={sample_rate}, {duration:.2f}s)") + + +def _omni_kwargs(args, model_name: str) -> dict: + kwargs = { + "model": model_name, + "log_stats": args.log_stats, + "stage_init_timeout": args.stage_init_timeout, + } + if args.deploy_config: + kwargs["deploy_config"] = args.deploy_config + return kwargs + + +def main(args) -> None: + os.makedirs(args.output_dir, exist_ok=True) + model_name = args.model or DEFAULT_MODEL + prompts = [ + build_prompt(text, args.ref_audio, args.ref_text, model_name) for text in ([args.text] * args.num_prompts) + ] + + omni = Omni(**_omni_kwargs(args, model_name)) + t_start = time.perf_counter() + total_audio_seconds = 0.0 + for stage_outputs in omni.generate(prompts): + request_output = stage_outputs + if request_output is None or not request_output.outputs: + continue + audio, sample_rate = extract_audio(request_output.outputs[0].multimodal_output) + assert audio.numel() > 0, "Audio8 TTS produced an empty waveform" + total_audio_seconds += audio.numel() / sample_rate + _save_wav(args.output_dir, request_output.request_id, audio, sample_rate) + elapsed = time.perf_counter() - t_start + rtf = elapsed / max(total_audio_seconds, 1e-6) + print(f"Total {elapsed * 1000:.1f} ms for {total_audio_seconds:.2f} s of audio (RTF={rtf:.3f})") + + +async def main_streaming(args) -> None: + os.makedirs(args.output_dir, exist_ok=True) + model_name = args.model or DEFAULT_MODEL + prompt = build_prompt(args.text, args.ref_audio, args.ref_text, model_name) + + omni = AsyncOmni(**_omni_kwargs(args, model_name)) + request_id = "0" + chunks: list[torch.Tensor] = [] + sample_rate = None + t_start = time.perf_counter() + t_prev = t_start + chunk_index = 0 + + async for stage_output in omni.generate(prompt, request_id=request_id): + multimodal_output = stage_output.outputs[0].multimodal_output + if not multimodal_output: + continue + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is not None: + if isinstance(audio, list): + chunks.extend(torch.as_tensor(chunk).reshape(-1) for chunk in audio) + else: + chunks.append(torch.as_tensor(audio).reshape(-1)) + sr_raw = multimodal_output.get("sr") + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else None + if sr_raw is not None: + sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw) + + if stage_output.finished: + break + now = time.perf_counter() + if chunk_index == 0: + print(f"chunk 0: TTFA={(now - t_start) * 1000:.1f} ms") + else: + print(f"chunk {chunk_index}: inter_chunk={(now - t_prev) * 1000:.1f} ms") + t_prev = now + chunk_index += 1 + + if not chunks or sample_rate is None: + raise RuntimeError("Audio8 TTS streaming produced no audio") + _save_wav(args.output_dir, request_id, torch.cat(chunks, dim=0), sample_rate) + print(f"Streaming done: total={(time.perf_counter() - t_start) * 1000:.1f} ms chunks={chunk_index}") + + +def parse_args(): + parser = TrackingArgumentParser(description="Audio8 TTS Preview offline inference with vLLM Omni") + parser.add_argument("--model", type=str, default=DEFAULT_MODEL, help="Model path or HuggingFace repo ID.") + parser.add_argument( + "--text", + type=str, + default="Welcome to Audio8 TTS, a compact multilingual text to speech model.", + help="Text to synthesize.", + ) + parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.") + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Transcript of --ref-audio; required for voice cloning and must match the recording.", + ) + parser.add_argument("--output-dir", type=str, default="output_audio", help="Directory for output WAV files.") + parser.add_argument( + "--num-prompts", + type=int, + default=1, + help="Duplicate the prompt N times to smoke-test concurrent requests.", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override the deploy config path; defaults to vllm_omni/deploy/audio8_tts.yaml via model_type.", + ) + parser.add_argument("--streaming", action="store_true", help="Use AsyncOmni and report per-chunk latency.") + parser.add_argument("--log-stats", action="store_true", help="Enable engine stats logging.") + parser.add_argument("--stage-init-timeout", type=int, default=600, help="Per-stage init timeout in seconds.") + return parser.parse_args() + + +if __name__ == "__main__": + parsed_args = parse_args() + if parsed_args.streaming: + asyncio.run(main_streaming(parsed_args)) + else: + main(parsed_args) diff --git a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py index 62917de07..40e194120 100644 --- a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py +++ b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import functools +import json import os import urllib.request from pathlib import Path +from typing import Any import numpy as np import soundfile as sf @@ -29,6 +32,20 @@ def _default_ref_audio() -> str: return str(dest) +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"{flag_name} must be valid JSON: {e}") from e + if not isinstance(config, dict): + 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 run_e2e(): parser = argparse.ArgumentParser() # ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512 @@ -64,6 +81,12 @@ def run_e2e(): required=True, help="Path to tokenizer directory (e.g., /CosyVoice-BlankEN).", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) args = parser.parse_args() # Ensure tokenizer directory exists if not os.path.exists(args.tokenizer): @@ -79,6 +102,7 @@ def run_e2e(): deploy_config=args.deploy_config, tokenizer=args.tokenizer, log_stats=True, + profiler_config=args.profiler_config, ) sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1} @@ -149,8 +173,8 @@ def run_e2e(): sampling_params_list = [gpt_sampling, s2mel_sampling] - # Start profiling (requires VLLM_TORCH_PROFILER_DIR env var) - if os.environ.get("VLLM_TORCH_PROFILER_DIR"): + profiler_enabled = args.profiler_config is not None + if profiler_enabled: print("Starting profiler...") omni.start_profile() @@ -158,7 +182,7 @@ def run_e2e(): outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2])) # Stop profiling and get results - if os.environ.get("VLLM_TORCH_PROFILER_DIR"): + if profiler_enabled: print("Stopping profiler...") profile_results = omni.stop_profile() print(f"Profile traces saved to: {profile_results}") diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index 926a58682..ef2367412 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -264,7 +264,7 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene #### Realtime WebSocket client (`openai_realtime_client.py`) -[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and receives **`response.audio.delta`** (incremental PCM for the reply) plus **`transcription.*`** events. By default it concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging. +[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks, and receives **`response.audio.*`**, **`response.output_text.*`**, and transcript events. By default it uses explicit `input_audio_buffer.commit`; pass **`--server-vad`** to let the server detect speech endpoints and commit the turn automatically. The client concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging. Streaming input works well for translation-style use cases; if the Thinker runs while input is still incomplete, consider limiting **`max_tokens`** in your session / server defaults to avoid over-generation. @@ -282,6 +282,7 @@ python openai_realtime_client.py \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct \ --input-wav /path/to/input_16k_mono.wav \ --output-wav realtime_output.wav \ + --server-vad \ --delta-dump-dir ./rt_delta_wavs ``` @@ -297,15 +298,32 @@ python openai_realtime_client.py \ | `--chunk-ms` | `200` | Size of each uploaded audio chunk (milliseconds of audio) | | `--send-delay-ms` | `0` | Delay between chunk sends (simulate realtime upload) | | `--delta-dump-dir` | *(optional)* | Directory to write per-`response.audio.delta` WAV files | +| `--server-vad` | disabled | Enable server-side Silero VAD instead of sending explicit commits | | `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) | | `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 | -Ensure the server is running, for example: +Server VAD uses the pinned Silero v6.2 ONNX artifact from `istupakov/silero-vad-onnx` and never downloads model files while processing a session. Make the artifact available in the Hugging Face cache before startup, or configure a local artifact in the deployment YAML: + +```yaml +duplex_session: + server_vad_model_path: /models/silero_vad.onnx +``` + +Start from the bundled Qwen3-Omni deployment YAML, add the fields above at the top level, and serve that complete configuration: ```bash -vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 +cp vllm_omni/deploy/qwen3_omni_moe.yaml /path/to/qwen3_omni_server_vad.yaml +# Edit /path/to/qwen3_omni_server_vad.yaml and add duplex_session. +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \ + --omni \ + --port 8091 \ + --deploy-config /path/to/qwen3_omni_server_vad.yaml ``` +Keep Qwen's default `session_mode: turn`; `duplex_session` enables the Realtime handler without changing scheduler +semantics. Bare `/v1/realtime` selects that handler. The client uses `?duplex=0` only for the existing non-Server-VAD +wire flow; `?duplex=1` remains a supported compatibility alias. + The Python client supports the following command-line arguments: - `--query-type` (or `-q`): Query type (default: `use_video`). Options: `text`, `use_audio`, `use_image`, `use_video` diff --git a/examples/online_serving/qwen3_omni/openai_realtime_client.py b/examples/online_serving/qwen3_omni/openai_realtime_client.py index 79e30a3f5..8a6871580 100644 --- a/examples/online_serving/qwen3_omni/openai_realtime_client.py +++ b/examples/online_serving/qwen3_omni/openai_realtime_client.py @@ -1,9 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Realtime client for vLLM-Omni /v1/realtime (audio + text events). This client: 1) Reads a local WAV file (must be mono, 16-bit PCM, 16kHz), 2) Streams PCM16 chunks to /v1/realtime with OpenAI-style events, -3) Receives response.audio.* and transcription.* events, +3) Receives response.audio.*, response.output_text.*, and transcript events, 4) Saves synthesized audio to an output WAV file and optional text file. By default each ``response.audio.delta`` is treated as an **incremental PCM** @@ -18,6 +21,7 @@ Usage: --model Qwen/Qwen3-Omni-30B-A3B-Instruct \ --input-wav input_16k_mono.wav \ --output-wav realtime_output.wav \ + --server-vad \ --delta-dump-dir ./rt_delta_wavs Dependencies: @@ -32,6 +36,7 @@ import base64 import json import wave from pathlib import Path +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit try: import websockets @@ -70,6 +75,14 @@ def _write_wav_pcm16(path: Path, pcm16_bytes: bytes, sample_rate_hz: int) -> Non wf.writeframes(pcm16_bytes) +def _with_realtime_route(url: str, *, server_vad: bool) -> str: + parts = urlsplit(url) + query = [(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key != "duplex"] + if not server_vad: + query.append(("duplex", "0")) + return urlunsplit(parts._replace(query=urlencode(query))) + + async def run_client( url: str, model: str, @@ -79,6 +92,7 @@ async def run_client( chunk_ms: int, send_delay_ms: int, delta_dump_dir: Path | None, + server_vad: bool, request_idx: int = 1, total_requests: int = 1, ) -> None: @@ -92,23 +106,44 @@ async def run_client( delta_index = 0 text_chunks: list[str] = [] final_text: str = "" + final_audio_transcript: str = "" if delta_dump_dir is not None: delta_dump_dir.mkdir(parents=True, exist_ok=True) + # Keep the CLI mode authoritative even if the input URL already contains + # an explicit duplex query parameter. + url = _with_realtime_route(url, server_vad=server_vad) async with websockets.connect(url, max_size=64 * 1024 * 1024) as ws: # 1) Validate model. - await ws.send( - json.dumps( - { - "type": "session.update", + if server_vad: + session_update = { + "type": "session.update", + "session": { "model": model, - } - ) - ) + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 16_000}, + "turn_detection": { + "type": "server_vad", + "threshold": 0.5, + "prefix_padding_ms": 300, + "silence_duration_ms": 500, + "create_response": True, + "interrupt_response": False, + }, + } + }, + }, + } + else: + # Preserve the existing non-duplex wire format. + session_update = {"type": "session.update", "model": model} + await ws.send(json.dumps(session_update)) - # 2) Start generation once (non-final commit). - await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False})) + # 2) Preserve the legacy explicit-commit flow unless Server VAD owns endpointing. + if not server_vad: + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False})) # 3) Stream audio chunks. for i in range(0, len(pcm16), chunk_bytes): @@ -124,8 +159,20 @@ async def run_client( if send_delay_ms > 0: await asyncio.sleep(send_delay_ms / 1000.0) - # 4) Final commit closes input stream. - await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) + # 4) Finalize the turn explicitly, or append silence for Server VAD endpointing. + if server_vad: + trailing_silence = bytes(16_000 * 2) + for i in range(0, len(trailing_silence), chunk_bytes): + await ws.send( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(trailing_silence[i : i + chunk_bytes]).decode("utf-8"), + } + ) + ) + else: + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) # 5) Receive server events until audio done. while True: @@ -144,7 +191,7 @@ async def run_client( sr = event.get("sample_rate_hz") if isinstance(sr, int) and sr > 0: output_sample_rate = sr - audio_b64 = event.get("audio", "") + audio_b64 = event.get("delta" if server_vad else "audio", "") if audio_b64: pcm_delta = base64.b64decode(audio_b64) incremental_pcm_parts.append(pcm_delta) @@ -158,14 +205,14 @@ async def run_client( ) continue - if event_type == "transcription.delta": + if event_type in {"response.output_text.delta", "transcription.delta"}: delta = event.get("delta", "") if delta: text_chunks.append(delta) print(delta, end="", flush=True) continue - if event_type == "transcription.done": + if event_type in {"response.output_text.done", "transcription.done"}: final_text = event.get("text", "") or "".join(text_chunks) usage = event.get("usage") final_text_with_tag = f"Final transcription: {final_text}" @@ -176,7 +223,16 @@ async def run_client( print(f"{log_prefix}text usage: {usage}") continue + if event_type == "response.audio_transcript.done": + final_audio_transcript = event.get("transcript", "") + continue + if event_type == "response.audio.done": + if not server_vad: + break + continue + + if event_type == "response.done" and server_vad: break if event_type == "error": @@ -191,7 +247,7 @@ async def run_client( print(f"{log_prefix}Saved realtime audio to: {output_wav} (incremental chunks joined)") if output_text is not None: - text_to_save = final_text if final_text else "".join(text_chunks) + text_to_save = final_text or "".join(text_chunks) or final_audio_transcript output_text.parent.mkdir(parents=True, exist_ok=True) output_text.write_text(text_to_save, encoding="utf-8") print(f"{log_prefix}Saved realtime text to: {output_text}") @@ -213,6 +269,7 @@ async def run_clients_concurrent( chunk_ms: int, send_delay_ms: int, delta_dump_dir: Path | None, + server_vad: bool, num_requests: int, concurrency: int, ) -> None: @@ -235,6 +292,7 @@ async def run_clients_concurrent( chunk_ms=chunk_ms, send_delay_ms=send_delay_ms, delta_dump_dir=per_delta_dir, + server_vad=server_vad, request_idx=index, total_requests=num_requests, ) @@ -283,6 +341,11 @@ def main() -> None: default=None, help="If set, each response.audio.delta is saved as delta_NNNNNN.wav under this directory", ) + parser.add_argument( + "--server-vad", + action="store_true", + help="Let the server detect speech endpoints instead of sending input_audio_buffer.commit", + ) parser.add_argument("--num-requests", type=int, default=1, help="Total number of requests to send") parser.add_argument( "--concurrency", @@ -309,6 +372,7 @@ def main() -> None: chunk_ms=args.chunk_ms, send_delay_ms=args.send_delay_ms, delta_dump_dir=args.delta_dump_dir, + server_vad=args.server_vad, ) ) else: @@ -322,6 +386,7 @@ def main() -> None: chunk_ms=args.chunk_ms, send_delay_ms=args.send_delay_ms, delta_dump_dir=args.delta_dump_dir, + server_vad=args.server_vad, num_requests=args.num_requests, concurrency=concurrency, ) diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 4d0715a7e..5141a37c2 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -14,6 +14,7 @@ For the full list of supported architectures across all modalities, see | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | | --- | --- | --- | --- | --- | --- | +| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ | | 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`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — | @@ -253,16 +254,152 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \ --- -## Fish Speech S2 Pro +## Audio8 TTS Preview -4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec. +0.6B DualAR TTS at 44.1 kHz, 11 languages, zero-shot voice cloning. ### Prerequisites +None beyond the base install: the neural audio codec is implemented in tree and +its weights (`codec.pth`) ship with the checkpoint. + +### Launch + ```bash -pip install fish-speech +vllm serve Audio8/Audio8-TTS-Preview-0.6b --omni --port 8092 +# or: +./audio8_tts/run_server.sh ``` +The deploy config auto-loads from `vllm_omni/deploy/audio8_tts.yaml` (HF +`model_type` is `arktts`). Do **not** pass `--trust-remote-code`: vllm-omni +registers its own `arktts` config, and transformers would otherwise prefer the +checkpoint's remote code and bypass it. + +### Text-only synthesis + +```bash +curl -X POST http://localhost:8092/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Welcome to Audio8 TTS.", + "response_format": "wav" + }' --output output.wav +``` + +### Voice cloning + +```bash +curl -X POST http://localhost:8092/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Welcome to Audio8 TTS.", + "ref_audio": "https://example.com/reference.wav", + "ref_text": "The exact transcript of the reference recording." + }' --output cloned.wav +``` + +`ref_audio` also accepts a `data:audio/wav;base64,...` URL. `ref_text` is +mandatory whenever `ref_audio` is present and must match the recording. +Uploading a voice via `POST /v1/audio/voices` lets the server reuse the encoded +reference codes across requests (`voice: ""`). + +### Streaming + +```bash +python audio8_tts/speech_client.py --text "Welcome to Audio8 TTS." --stream --output out.pcm +ffplay -f s16le -ar 44100 -ac 1 out.pcm +``` + +### Gradio demo + +```bash +./audio8_tts/run_gradio_demo.sh # server + demo +python audio8_tts/gradio_demo.py --api-base http://localhost:8092 # demo only +``` + +### Benchmark + +```bash +export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100 # required: harness defaults to 24 kHz +vllm bench serve --omni --host 127.0.0.1 --port 8092 \ + --model Audio8/Audio8-TTS-Preview-0.6b \ + --backend openai-audio-speech --endpoint /v1/audio/speech \ + --dataset-name seed-tts-text --dataset-path benchmarks/build_dataset/seed_tts_smoke \ + --seed-tts-locale en --num-prompts 20 --num-warmups 2 \ + --max-concurrency 1 --request-rate inf \ + --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun +``` + +Measured on 1x H20 **shared with an unrelated job holding ~80% SM**, so treat +these as a lower bound. 20 prompts from `seed_tts_smoke/en`, mean audio 4.5 s, +deploy defaults: + +| | HF reference (batch=1) | vllm-omni c=1 | c=4 | c=8 | +| --- | --- | --- | --- | --- | +| RTF (mean) | 1.008 | **0.19** | 0.30 | 0.54 | +| Throughput (req/s) | 0.23 | 1.18 | 2.84 | **2.93** | +| E2E latency (mean, ms) | 4375 | 850 | 1349 | 2361 | +| Time to first packet (mean, ms) | n/a (no streaming) | 63 | 105 | 1124 | +| Underrun p99 / continuity OK | n/a | 0 s / 100% | 0 s / 100% | 0 s / 100% | + +~5.3x lower per-request RTF than the upstream reference implementation at +concurrency 1, ~12.8x aggregate throughput at concurrency 8. + +#### Stage-0 `max_num_seqs`: first-packet vs throughput + +This subsection is operator tuning guidance, not a claim this PR optimizes +anything: the shipped deploy default is unchanged (`max_num_seqs: 4`). The +numbers below are a single run on the same contended H20 as above (one unrelated +job holding ~80% SM, no repeats), so read them as directional, not as measured +speedups. Same setup, varying stage 0's `max_num_seqs` (stage 1 stays at 1): + +| stage-0 `max_num_seqs` | client concurrency | req/s | RTF | TTFP mean / p99 (ms) | underrun p99 (s) | +| --- | --- | --- | --- | --- | --- | +| 4 (default) | 1 | 1.18 | 0.19 | 63 / 67 | 0.00 | +| 4 | 4 | 2.84 | 0.30 | 105 / 148 | 0.00 | +| 4 | 8 | 2.93 | 0.54 | 1124 / 1642 | 0.00 | +| 8 | 1 | 1.16 | 0.19 | 64 / 70 | 0.00 | +| 8 | 4 | 2.88 | 0.30 | 107 / 148 | 0.00 | +| 8 | 8 | **3.72** | 0.40 | **247** / 733 | 0.22 | +| 8 | 16 | 3.91 | 0.65 | 1366 / 2530 | 0.29 | + +Reading of this: + +- Below the limit (`concurrency <= max_num_seqs`) the setting is irrelevant, as + expected — c=1 and c=4 are identical for both. +- At c=8, raising the limit to 8 admits every request instead of queueing half + of them: in this single run, roughly +27% throughput and ~4.5x lower + first-packet latency. +- Throughput saturates around 3.7-3.9 req/s (~17 s of audio per second). Past + c=8 you only buy queueing delay. +- The cost of admitting more streams is a 0.22-0.29 s worst-case buffer deficit + in 1 of 20 requests (`Streaming continuity OK rate` 95%), because stage 1 runs + `max_num_seqs: 1` and the codec window round-robins across streams. Client-side + prebuffering (the bundled Gradio demo holds 2 chunks) absorbs that; raising + stage 1's `max_num_seqs` is the untested knob if you need both. + +The shipped default stays at 4 because every measured point there had zero +underrun. Raise it to 8 if first-packet latency under load matters more than a +rare sub-300 ms gap. + +### Notes + +- Output: 44.1 kHz mono; ~21.5 codec frames per second. +- No built-in speaker presets. Omit `voice` for a random timbre, or clone one. +- `max_new_tokens` caps generated codec frames (1 frame ~= 46 ms). +- Context is 2048 packed text+audio positions (`max_model_len: 2048`). +- Do not send `temperature: 0`: greedy decoding is degenerate for this + checkpoint and truncates the utterance. The deploy defaults + (0.7 / top_k 50 / top_p 0.9) mirror `generation_config.json`. + +--- + +## Fish Speech S2 Pro + +4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec, which is vendored in +vLLM-Omni (no extra packages required). + ### Kvcache attention fast path Fish Speech S2 Pro uses a Triton decode-only kvcache attention fast path by @@ -708,13 +845,21 @@ Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and ### Streaming WebSocket -The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`: +The `/v1/audio/speech/stream` endpoint accepts text incrementally and, by default, synthesizes the buffered text as one continuous request on `input.done`: ```bash python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine." python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1 ``` +For per-sentence audio (including Indic danda `।`) before `input.done`: + +```bash +python qwen3_tts/streaming_speech_client.py \ + --text "नमस्ते। कैसे हो?" \ + --split-granularity sentence +``` + `input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection: ```bash diff --git a/examples/online_serving/text_to_speech/audio8_tts/gradio_demo.py b/examples/online_serving/text_to_speech/audio8_tts/gradio_demo.py new file mode 100644 index 000000000..368ddfaa0 --- /dev/null +++ b/examples/online_serving/text_to_speech/audio8_tts/gradio_demo.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Gradio demo for Audio8 TTS Preview via the /v1/audio/speech API. + +Supports: + - text-to-speech synthesis in 11 languages + - zero-shot voice cloning from reference audio (upload, microphone, or URL) + - streaming (progressive PCM chunks) and non-streaming modes + +Usage: + # Start the server first (see run_server.sh), then: + python gradio_demo.py --api-base http://localhost:8092 +""" + +import argparse +import base64 +import io + +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 httpx +import numpy as np +import soundfile as sf + +#: Audio8 TTS decodes at 44.1 kHz; raw PCM must be interpreted at that rate. +PCM_SAMPLE_RATE = 44100 +DEFAULT_API_BASE = "http://localhost:8092" + + +def encode_audio_to_base64(audio_data: tuple) -> str: + """Encode a Gradio ``(sample_rate, numpy_array)`` input as a data URL.""" + sample_rate, audio_np = audio_data + if audio_np.dtype != np.int16: + if audio_np.dtype in (np.float32, np.float64): + audio_np = (np.clip(audio_np, -1.0, 1.0) * 32767).astype(np.int16) + else: + audio_np = audio_np.astype(np.int16) + buf = io.BytesIO() + sf.write(buf, audio_np, sample_rate, format="WAV") + return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}" + + +def build_payload( + text: str, + ref_audio: tuple | None, + ref_audio_url: str, + ref_text: str, + response_format: str = "wav", + stream: bool = False, +) -> dict: + """Build the ``/v1/audio/speech`` payload. + + Audio8 TTS has no built-in speakers: omit ``voice`` for a random timbre, or + supply ``ref_audio`` + ``ref_text`` to clone a voice. + """ + if not text or not text.strip(): + raise gr.Error("Please enter text to synthesize.") + + payload: dict = { + "input": text.strip(), + "response_format": "pcm" if stream else response_format, + "stream": stream, + } + if stream: + payload["stream_format"] = "audio" + + ref_url = ref_audio_url.strip() if ref_audio_url else "" + if ref_audio is not None: + payload["ref_audio"] = encode_audio_to_base64(ref_audio) + elif ref_url: + payload["ref_audio"] = ref_url + + if "ref_audio" in payload: + if not ref_text or not ref_text.strip(): + raise gr.Error("Voice cloning requires the transcript of the reference audio.") + payload["ref_text"] = ref_text.strip() + + return payload + + +def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format): + """Non-streaming request; returns the whole waveform.""" + payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False) + try: + with httpx.Client(timeout=300.0) as client: + resp = client.post( + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") from None + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None + + if resp.status_code != 200: + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + if "application/json" in resp.headers.get("content-type", ""): + raise gr.Error(f"Server error: {resp.text}") + + try: + if response_format == "pcm": + audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0 + return (PCM_SAMPLE_RATE, audio_np) + audio_np, sample_rate = sf.read(io.BytesIO(resp.content)) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + return (sample_rate, audio_np.astype(np.float32)) + except Exception as e: + raise gr.Error(f"Failed to decode audio: {e}") from None + + +def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5): + """Streaming request; yields PCM chunks for progressive playback.""" + payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True) + min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds) + prebuffer_chunks = 2 + pending: list[np.ndarray] = [] + pending_len = 0 + chunks_yielded = 0 + prebuffer: list[np.ndarray] = [] + leftover = b"" + + try: + with ( + httpx.Client(timeout=300.0) as client, + client.stream( + "POST", + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) as resp, + ): + if resp.status_code != 200: + resp.read() + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + for chunk in resp.iter_bytes(): + if not chunk: + continue + raw = leftover + chunk + # s16le: never split a sample across chunk boundaries. + usable = len(raw) - (len(raw) % 2) + leftover = raw[usable:] + if usable == 0: + continue + samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0 + pending.append(samples) + pending_len += len(samples) + if pending_len < min_chunk_samples: + continue + audio_chunk = np.concatenate(pending) + pending.clear() + pending_len = 0 + if chunks_yielded < prebuffer_chunks: + prebuffer.append(audio_chunk) + chunks_yielded += 1 + if chunks_yielded == prebuffer_chunks: + yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer)) + prebuffer.clear() + else: + yield (PCM_SAMPLE_RATE, audio_chunk) + remaining = prebuffer + pending + if remaining: + yield (PCM_SAMPLE_RATE, np.concatenate(remaining)) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") from None + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None + + +def on_stream_change(stream: bool): + """Streaming only makes sense with raw PCM.""" + if stream: + return gr.update(value="pcm", interactive=False) + return gr.update(value="wav", interactive=True) + + +def build_interface(api_base: str, stream_chunk_seconds: float = 0.5): + with gr.Blocks(title="Audio8 TTS Preview Demo") as demo: + gr.Markdown("# Audio8 TTS Preview 0.6B - Text to Speech") + gr.Markdown(f"**Server:** `{api_base}` | **Model:** Audio8/Audio8-TTS-Preview-0.6b | **Output:** 44.1 kHz") + + with gr.Row(): + with gr.Column(scale=3): + text_input = gr.Textbox(label="Text to Synthesize", placeholder="Enter text here...", lines=4) + + with gr.Accordion("Voice Cloning (optional)", open=False): + gr.Markdown( + "Upload or link a short reference recording (5-30 s) and paste its **exact** transcript. " + "A transcript that does not match the audio degrades both stability and speaker similarity." + ) + ref_audio = gr.Audio(label="Reference Audio", type="numpy", sources=["upload", "microphone"]) + ref_audio_url = gr.Textbox( + label="Reference Audio URL (alternative to upload)", + placeholder="https://example.com/reference.wav", + lines=1, + ) + ref_text = gr.Textbox( + label="Reference Audio Transcript (required for cloning)", + placeholder="Exact transcript of the reference audio...", + lines=2, + ) + + with gr.Row(): + response_format = gr.Dropdown( + choices=["wav", "mp3", "flac", "pcm"], value="wav", label="Audio Format", scale=1 + ) + stream_checkbox = gr.Checkbox( + label="Stream output", value=False, info="Progressive PCM streaming", scale=1 + ) + + generate_btn = gr.Button("Generate Speech", variant="primary", size="lg") + + with gr.Column(scale=2): + audio_output = gr.Audio(label="Generated Audio", interactive=False, streaming=True, autoplay=True) + gr.Markdown( + "### About\n" + "- **Audio8 TTS Preview** — 0.6B DualAR model (slow AR + fast AR codebook decoder)\n" + "- **Languages**: Cantonese, Chinese, Dutch, English, French, German, " + "Italian, Japanese, Korean, Polish, Spanish\n" + "- **Voice cloning**: zero-shot from a reference clip + transcript\n" + "- **44.1 kHz** output, ~21.5 codec frames per second" + ) + + stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format]) + all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format] + + def dispatch(stream_enabled, *args): + if stream_enabled: + yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds) + else: + yield generate_speech(api_base, *args) + + generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output]) + demo.queue() + return demo + + +def main() -> None: + parser = argparse.ArgumentParser(description="Gradio demo for Audio8 TTS Preview") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})") + parser.add_argument("--host", default="0.0.0.0", help="Gradio host") + parser.add_argument("--port", type=int, default=7861, help="Gradio port") + parser.add_argument("--share", action="store_true", help="Create a public share link") + parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds of audio per yielded chunk") + args = parser.parse_args() + + print(f"Connecting to vLLM server at: {args.api_base}") + build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds).launch( + server_name=args.host, server_port=args.port, share=args.share + ) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/text_to_speech/audio8_tts/run_gradio_demo.sh b/examples/online_serving/text_to_speech/audio8_tts/run_gradio_demo.sh new file mode 100644 index 000000000..5588b921e --- /dev/null +++ b/examples/online_serving/text_to_speech/audio8_tts/run_gradio_demo.sh @@ -0,0 +1,44 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# Launch the Audio8 TTS Preview server + Gradio demo together. +# +# Usage: +# ./run_gradio_demo.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8092 GRADIO_PORT=7861 ./run_gradio_demo.sh + +set -e + +MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}" +PORT="${PORT:-8092}" +GRADIO_PORT="${GRADIO_PORT:-7861}" +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +echo "Starting Audio8 TTS Preview server (port $PORT)..." +vllm serve "$MODEL" \ + --omni \ + --host 0.0.0.0 \ + --port "$PORT" & +SERVER_PID=$! + +cleanup() { + echo "Stopping server (PID $SERVER_PID)..." + kill $SERVER_PID 2>/dev/null + wait $SERVER_PID 2>/dev/null +} +trap cleanup EXIT + +echo "Waiting for server to start..." +for _ in $(seq 1 120); do + if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then + echo "Server ready." + break + fi + sleep 2 +done + +echo "Starting Gradio demo (port $GRADIO_PORT)..." +python "$SCRIPT_DIR/gradio_demo.py" \ + --api-base "http://localhost:$PORT" \ + --port "$GRADIO_PORT" diff --git a/examples/online_serving/text_to_speech/audio8_tts/run_server.sh b/examples/online_serving/text_to_speech/audio8_tts/run_server.sh new file mode 100644 index 000000000..a208f0a5c --- /dev/null +++ b/examples/online_serving/text_to_speech/audio8_tts/run_server.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# Launch a vLLM-Omni server for Audio8 TTS Preview 0.6B. +# +# Usage: +# ./run_server.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8092 ./run_server.sh + +set -e + +MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}" +PORT="${PORT:-8092}" + +echo "Starting Audio8 TTS Preview server with model: $MODEL" + +# --trust-remote-code is deliberately NOT passed: vllm-omni registers its own +# `arktts` config, and transformers would otherwise prefer the checkpoint's +# remote code and bypass it. +vllm serve "$MODEL" \ + --omni \ + --host 0.0.0.0 \ + --port "$PORT" diff --git a/examples/online_serving/text_to_speech/audio8_tts/speech_client.py b/examples/online_serving/text_to_speech/audio8_tts/speech_client.py new file mode 100644 index 000000000..d0a2c6f6a --- /dev/null +++ b/examples/online_serving/text_to_speech/audio8_tts/speech_client.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Client for Audio8 TTS Preview via the /v1/audio/speech endpoint. + +Examples: + # Basic TTS + python speech_client.py --text "Welcome to Audio8 TTS." + + # Zero-shot voice cloning (the transcript must match the reference audio) + python speech_client.py --text "Welcome to Audio8 TTS." \ + --ref-audio ref.wav --ref-text "The exact transcript of the reference recording." + + # Streaming PCM output + python speech_client.py --text "Welcome to Audio8 TTS." --stream --output output.pcm +""" + +import argparse +import base64 +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8092" +DEFAULT_API_KEY = "EMPTY" +DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b" +#: The codec runs at 44.1 kHz, so raw PCM must be played back at that rate. +PCM_SAMPLE_RATE = 44100 + + +def encode_audio_to_base64(audio_path: str) -> str: + """Encode a local audio file as a base64 data URL.""" + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + ext = audio_path.lower().rsplit(".", 1)[-1] + mime_type = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}.get( + ext, "audio/wav" + ) + with open(audio_path, "rb") as handle: + audio_b64 = base64.b64encode(handle.read()).decode("utf-8") + return f"data:{mime_type};base64,{audio_b64}" + + +def build_payload(args) -> dict: + payload = { + "model": args.model, + "input": args.text, + "response_format": args.response_format, + } + if args.ref_audio: + if args.ref_audio.startswith(("http://", "https://")): + payload["ref_audio"] = args.ref_audio + else: + payload["ref_audio"] = encode_audio_to_base64(args.ref_audio) + if args.ref_text: + payload["ref_text"] = args.ref_text + if args.max_new_tokens is not None: + payload["max_new_tokens"] = args.max_new_tokens + if args.stream: + payload["stream"] = True + payload["stream_format"] = "audio" + payload["response_format"] = "pcm" + return payload + + +def run_tts(args) -> None: + payload = build_payload(args) + api_url = f"{args.api_base}/v1/audio/speech" + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {args.api_key}"} + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + if args.ref_audio: + print(f"Voice cloning: ref_audio={args.ref_audio!r} ref_text={args.ref_text!r}") + + if args.stream: + output_path = args.output or "output.pcm" + with ( + httpx.Client(timeout=300.0) as client, + client.stream("POST", api_url, json=payload, headers=headers) as resp, + ): + if resp.status_code != 200: + print(f"Error {resp.status_code}: {resp.read().decode()}") + return + total_bytes = 0 + with open(output_path, "wb") as handle: + for chunk in resp.iter_bytes(): + handle.write(chunk) + total_bytes += len(chunk) + seconds = total_bytes / 2 / PCM_SAMPLE_RATE + print(f"Streamed {total_bytes} bytes (~{seconds:.2f}s of s16le @ {PCM_SAMPLE_RATE} Hz) to {output_path}") + print(f"Play with: ffplay -f s16le -ar {PCM_SAMPLE_RATE} -ac 1 {output_path}") + return + + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + if response.status_code != 200: + print(f"Error {response.status_code}: {response.text}") + return + if response.headers.get("content-type", "").startswith("application/json"): + print(f"Error: {response.text}") + return + + output_path = args.output or f"output.{args.response_format}" + with open(output_path, "wb") as handle: + handle.write(response.content) + print(f"Audio saved to: {output_path} ({len(response.content)} bytes)") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Audio8 TTS Preview client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)") + parser.add_argument("--ref-text", default=None, help="Transcript of the reference audio") + parser.add_argument("--max-new-tokens", type=int, default=None, help="Cap the number of generated codec frames") + parser.add_argument("--stream", action="store_true", help="Stream PCM instead of returning a whole file") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio format (default: wav)", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + run_tts(parser.parse_args()) + + +if __name__ == "__main__": + main() 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 e141a4c48..793a5a6d5 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 @@ -24,6 +24,11 @@ Usage: --text "Hello world. How are you? I am fine." \ --simulate-stt --stt-delay 0.1 + # Opt into per-sentence synthesis (Indic danda, CJK, Latin) + python streaming_speech_client.py \ + --text "नमस्ते। कैसे हो?" \ + --split-granularity sentence + # Receive JSON sidecar chunks with word-level timestamps python streaming_speech_client.py \ --text "Hello world. How are you?" \ @@ -73,8 +78,9 @@ except ImportError: def frame_basename(msg: dict) -> str: """Name a file after the utterance and sentence a frame belongs to. - Every utterance is sentence 0, so the utterance index is what keeps the - files of one connection from overwriting each other. + Every utterance uses `sentence_index` for units inside that flush. + With default `split_granularity=none` that index stays 0; sentence mode + uses 0, 1, ... so files from one connection do not overwrite each other. """ return f"utterance_{msg['utterance_index']:03d}_sentence_{msg['sentence_index']:03d}" @@ -368,6 +374,13 @@ def main(): ) 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") + parser.add_argument("--seed", type=int, default=None, help="Sampling seed forwarded to the server") + parser.add_argument( + "--split-granularity", + default=None, + choices=["none", "sentence", "clause"], + help="Text split mode (default: server none = one request per input.done)", + ) # Base task options parser.add_argument("--ref-audio", default=None, help="Reference audio") @@ -409,6 +422,8 @@ def main(): "max_new_tokens", "ref_audio", "ref_text", + "seed", + "split_granularity", ]: val = getattr(args, key.replace("-", "_"), None) if val is not None: diff --git a/pyproject.toml b/pyproject.toml index cd7544af8..3ba902977 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -77,7 +77,6 @@ dev = [ # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo, # whose top-level stepaudio2 package conflicts with this distribution. "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml) - "silero-vad==6.1", # Realtime server-VAD duplex E2E ] demo = [ @@ -87,9 +86,6 @@ demo = [ "cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次") ] -server-vad = [ - "silero-vad==6.1", -] forced-aligner = [ "qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner @@ -378,3 +374,7 @@ ba = "ba" SER = "SER" # NumPy ndarray.flags.writeable (canonical flag name, not "writable"). writeable = "writeable" +# "convnet" abbreviates convolutional network. It appears in the vendored Fish +# Speech DAC codec identifiers (convnet_type, transconvnet_type) and is not a +# typo of "convent". +convnet = "convnet" diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index ecc2c1d0f..0d832b0ae 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -1,12 +1,13 @@ # Boogu-Image > Text-to-image and image-editing online serving -(Boogu-Image-0.1-Base / -Edit / -Edit-Turbo) +(Boogu-Image-0.1-Base / -Turbo / -Edit / -Edit-Turbo) ## Summary - Vendor: Boogu - Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image), + `Boogu/Boogu-Image-0.1-Turbo` (four-step text-to-image), `Boogu/Boogu-Image-0.1-Edit` (image editing), and `Boogu/Boogu-Image-0.1-Edit-Turbo` (four-step image editing) - Task: Text-to-image generation and text-guided image editing (TI2I) @@ -16,8 +17,9 @@ ## When to use this recipe Use this recipe when you want a known-good starting point for serving -`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, or -`Boogu/Boogu-Image-0.1-Edit-Turbo` with vLLM-Omni's native pipeline (no +`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Turbo`, +`Boogu/Boogu-Image-0.1-Edit`, or `Boogu/Boogu-Image-0.1-Edit-Turbo` with +vLLM-Omni's native pipeline (no `--diffusion-load-format diffusers`, and the upstream `boogu` package is not required). @@ -25,12 +27,14 @@ Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model family. The Base text-to-image checkpoint pairs a Qwen3-VL multimodal encoder with a Diffusion Transformer (DiT) and a flow-match Euler scheduler with time-shift. It handles photorealistic generation and Chinese/English text -rendering. The Edit and Edit-Turbo checkpoints use the same native -`BooguImagePipeline` for image editing. +rendering. The Edit checkpoint uses the native `BooguImagePipeline`, while the +distilled Turbo and Edit-Turbo checkpoints run the four-step DMD path through +`BooguImageTurboPipeline`. ## References - Upstream model card: +- Turbo model card: - Edit model card: - Edit-Turbo 1K hotfix: - Project page: @@ -186,6 +190,19 @@ ID passed to `vllm serve`. For example, use corresponding Edit model. Base text-to-image: +## Fast text-to-image (Boogu-Image-0.1-Turbo) + +Turbo is the distilled text-to-image checkpoint. Its `model_index.json` +declares `BooguImageTurboPipeline`, which runs the four-step DMD path instead +of the scheduler-driven path used by Base. + +### Command + +```bash +vllm serve Boogu/Boogu-Image-0.1-Turbo --omni --port 8091 +``` + +### Verification ```bash curl -X POST http://localhost:8091/v1/images/generations \ @@ -230,6 +247,23 @@ curl -X POST http://localhost:8091/v1/images/edits \ checkpoint-declared FP8 configuration. With the regular checkpoints used for online FP8, the MLLM remains in BF16. The VAE and scheduler are outside this quantization routing. + "model": "Boogu/Boogu-Image-0.1-Turbo", + "prompt": "A mountain lake at sunset, photorealistic, cinematic lighting", + "size": "1024x1024", + "seed": 42 + }' | jq -r '.data[0].b64_json' | base64 -d > output-turbo.png +``` + +### Notes + +- **Required settings:** `guidance_scale=1.0` and `guidance_scale_2=1.0`; any + other value is rejected. `num_inference_steps` defaults to 4. All three are + pipeline defaults, so the request above omits them. +- **Resolution:** upstream constrains Turbo to 1K, so `1024x1024` is the + supported working resolution. The shared pipeline still clamps at 2K; larger + sizes are not validated for this checkpoint. +- **Known limitations:** the same single-GPU limitations as Base apply; CPU + offload, Cache-DiT, and multi-GPU parallelism are not yet validated. ## Image editing (Boogu-Image-0.1-Edit) @@ -354,16 +388,18 @@ curl -s http://localhost:8091/v1/chat/completions \ ## Fast image editing (Boogu-Image-0.1-Edit-Turbo) Edit-Turbo is the distilled image-editing checkpoint. Its `model_index.json` -declares `BooguImagePipeline`, so it uses the same native TI2I path as Edit. -The upstream 1K hotfix is recommended for stable results and is pinned below -to avoid silently loading the older checkpoint from the repository's default -revision. +declares `BooguImagePipeline`, so `--model-class-name BooguImageTurboPipeline` +is required to reach the four-step DMD path; without it the checkpoint is +served on the dense Edit path. The upstream 1K hotfix is recommended for +stable results and is pinned below to avoid silently loading the older +checkpoint from the repository's default revision. ### Command ```bash vllm serve Boogu/Boogu-Image-0.1-Edit-Turbo \ --omni \ + --model-class-name BooguImageTurboPipeline \ --revision hotfix-1k-20260708 \ --port 8091 ``` @@ -407,11 +443,10 @@ curl -s http://localhost:8091/v1/chat/completions \ - **Pinned revision:** use `hotfix-1k-20260708`; upstream recommends the 1K hotfix over the 1.5K variant for more stable results. -- **Recommended settings:** `num_inference_steps=4` and - `guidance_scale=1.0`. Edit-Turbo is guidance-distilled, meaning the guidance - behavior is already baked into the distilled checkpoint. A scale of `1.0` - avoids applying additional classifier-free guidance; unlike the regular Edit - checkpoint, it should not be increased to `5.0`. +- **Required settings:** `guidance_scale=1.0` and `guidance_scale_2=1.0`; the + Turbo pipeline rejects any other value, so the `guidance_scale_2=2.0` shown + for the regular Edit checkpoint does not apply here. `num_inference_steps` + defaults to 4. - **Reference images:** the same single-reference and `align_res` behavior as the regular Edit checkpoint applies. - **Known limitations:** the same single-GPU limitations as Base and Edit apply; diff --git a/recipes/JD/JoyAI-VL-Interaction.md b/recipes/JD/JoyAI-VL-Interaction.md index e09d5b16c..7efc45d38 100644 --- a/recipes/JD/JoyAI-VL-Interaction.md +++ b/recipes/JD/JoyAI-VL-Interaction.md @@ -219,3 +219,51 @@ the audio-track caveat. 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. + +## Native multi-stage pipeline (experimental) + +An experimental native speech path is available alongside the existing Day-0 serving path: + +```text +one request with image/video and text + -> JoyAI complete action text + -> Qwen3-TTS Talker + -> Code2Wav + -> 24 kHz audio +``` + +It runs behind one `vllm serve --omni` endpoint: + +```bash +vllm serve jdopensource/JoyAI-VL-Interaction-Preview --omni \ + --deploy-config vllm_omni/deploy/joyai_vl_interaction.yaml \ + --port 8091 +``` + +With `modalities: ["text", "audio"]`, its behavior is: + +| JoyAI action | Native output | +|---|---| +| ` ` | action text + speech for `` | +| `` | action text + empty audio output (zero samples) | +| ` ` | action text + speech for `` only | + +Native speech supports the Qwen3-TTS `CustomVoice` task. The request-level `voice` and +`language` fields select the speaker and language; when omitted, they default to +`Vivian` and `Auto`. Values supplied as `tts_speaker` and `tts_language` inside +`additional_information` take priority. + +This experimental native multi-stage pipeline is request-scoped, stateless, and +all-sync. It currently does not include: + +- continuous frame sessions, persistent standing instructions, or session memory; +- audio input / ASR, full-duplex interaction, or barge-in; +- background Agent execution for delegated questions; +- async-chunk support from Talker to Code2Wav; +- integration with the existing WebUI. + +For actions with spoken text, the Talker starts only after the complete JoyAI action is +available, and Code2Wav starts only after the Talker completes. Silence actions skip both +TTS stages and return an empty audio output with zero samples. All three stages are placed +on GPU 0 by the provided +[`deploy config`](../../vllm_omni/deploy/joyai_vl_interaction.yaml). diff --git a/recipes/LTX/LTX-2.5.md b/recipes/LTX/LTX-2.5.md index f0be6948b..538eed1ce 100644 --- a/recipes/LTX/LTX-2.5.md +++ b/recipes/LTX/LTX-2.5.md @@ -110,6 +110,29 @@ vllm serve Lightricks/LTX-2.5-Diffusers \ Use the T2V or I2V request below after the server is ready. +### Multi-GPU low-latency path + +For low-latency 1920x1088 distilled two-stage serving, match the Ulysses and +VAE patch-parallel degrees to the GPU count. For example, on four GPUs: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +vllm serve Lightricks/LTX-2.5-Diffusers \ + --omni \ + --model-class-name LTX2DistilledTwoStagePipeline \ + --usp 4 \ + --vae-patch-parallel-size 4 \ + --vae-use-tiling \ + --host 0.0.0.0 \ + --port 8000 \ + --stage-init-timeout 900 +``` + +Strict Ulysses only: `((F - 1) / 8 + 1) * (H / 32) * (W / 32)` and TP-local +SP attention head counts must be divisible by `ulysses_degree`. Use each +phase's `H, W` (half/full resolution for two-stage); non-divisible shapes fail. +Each GPU must still fit the resident weights. Warm up with the production shape. + ## Offline inference Choose values from the pipeline table. For example, the distilled two-stage diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md index ec4e8fb3b..5a6a049ee 100644 --- a/recipes/LTX/LTX-2.md +++ b/recipes/LTX/LTX-2.md @@ -101,6 +101,12 @@ sigmas and input latents. Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer and automatically switches to dynamic LoRA when quantization is enabled. +## Sequence Parallelism + +Strict Ulysses only: `((F - 1) / 8 + 1) * (H / 32) * (W / 32)` and TP-local +SP attention head counts must be divisible by `ulysses_degree`. Use each +phase's `H, W` (half/full resolution for two-stage); non-divisible shapes fail. + ## Serving Start either one-stage checkpoint: @@ -335,9 +341,8 @@ bundled offline CLI do not currently expose `sigmas`. `--max-num-seqs 1`. - `--cfg-parallel-size` shards the complete LTX guidance plan, including STG, modality guidance, and rescale-compatible prediction gathering. -- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding; - Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a - divisible request shape. +- Ulysses masks padded audio latents; video sequences must satisfy the + [sequence-parallel divisibility constraints](#sequence-parallelism). - Cache-DiT is one-stage only; multi-stage configurations are rejected. ## Operational Notes diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 4aaebf0be..490f5b1f4 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -207,6 +207,18 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ -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 +# The same transfer control can be uploaded by a remote client instead of +# being placed on the server filesystem. Uploads are limited to 512 MiB. +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 world-state control." \ + -F "input_reference=@/path/to/input.mp4;type=video/mp4" \ + -F "control_reference=@/path/to/wsm.mp4;type=video/mp4" \ + -F "control_type=wsm" \ + -F 'extra_params={"wsm":{"control_weight":1.0},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \ + -o cosmos3_transfer_wsm.mp4 + # Text-to-video-with-sound curl -sS -X POST http://localhost:8000/v1/videos/sync \ -H "Accept: video/mp4" \ diff --git a/recipes/fishaudio/Fish-Speech-S2-Pro.md b/recipes/fishaudio/Fish-Speech-S2-Pro.md index dca9f537e..167bedb04 100644 --- a/recipes/fishaudio/Fish-Speech-S2-Pro.md +++ b/recipes/fishaudio/Fish-Speech-S2-Pro.md @@ -36,19 +36,13 @@ Other hardware and configurations are welcome as community validation lands. - OS: Linux - Python: 3.10+ -- Required package: `fish-speech` (for DAC codec) - CUDA 12.8 - vLLM version: 0.19.0 - vLLM-Omni version or commit: c93359bb354a6aa5c14d062430cb85b2c4db251e -```bash -# Install PortAudio dependency if missing (required by pyaudio which is a dependency of fish-speech) -# For Ubuntu/Debian: -if ! dpkg -l | grep -q libportaudio2; then - sudo apt-get update && sudo apt-get install -y libportaudio2 portaudio19-dev -fi -pip install fish-speech -``` +No extra packages are required: the DAC codec architecture is vendored in +vLLM-Omni (`vllm_omni/model_executor/models/fish_speech/dac_modules/`), so the +model runs out of the box, including in the official Docker image. #### Kvcache attention fast path @@ -129,7 +123,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \ - OS: Linux - Python: 3.10+ -- Required package: `fish-speech` (for DAC codec) - CUDA 12.8 - vLLM version: 0.19.0+ diff --git a/tests/benchmarks/metrics/test_metrics.py b/tests/benchmarks/metrics/test_metrics.py index 315685258..d7dd6c14a 100644 --- a/tests/benchmarks/metrics/test_metrics.py +++ b/tests/benchmarks/metrics/test_metrics.py @@ -457,5 +457,136 @@ def test_text_benchmark_still_reports_ttft(capsys): assert "Time to First Token" in out, "text benchmarks must keep TTFT" +# ============================================================================ +# Suppress text metrics for pure image / video endpoints +# ============================================================================ + + +def _make_image_output(prompt_len: int) -> MixRequestFuncOutput: + output = MixRequestFuncOutput() + output.success = True + output.prompt_len = prompt_len + output.output_tokens = 0 + output.generated_text = "" + output.ttft = 0.0 + output.text_latency = 0.0 + output.latency = 2.5 + output.start_time = 0.0 + output.itl = [] + output.image_count = 1 + output.image_generation_time_ms = 2000.0 + output.image_pixels = 1024 * 1024 + output.input_audio_duration = 0.0 + output.error = "" + return output + + +def _make_video_output(prompt_len: int) -> MixRequestFuncOutput: + output = MixRequestFuncOutput() + output.success = True + output.prompt_len = prompt_len + output.output_tokens = 0 + output.generated_text = "" + output.ttft = 0.0 + output.text_latency = 0.0 + output.latency = 5.0 + output.start_time = 0.0 + output.itl = [] + output.video_duration = 2.0 + output.video_frames = 48 + output.video_generation_time_ms = 4000.0 + output.video_rtf = 2.0 + output.input_audio_duration = 0.0 + output.error = "" + return output + + +_MEDIA_PERCENTILE_METRICS = ["e2el", "ttft"] + + +def test_pure_image_benchmark_omits_text_result(capsys): + """Pure image generation must not print Text Result or invent peak tok/s.""" + outputs = [_make_image_output(64), _make_image_output(64)] + + metrics, actual_output_lens = calculate_metrics( + input_requests=[], + outputs=outputs, + dur_s=5.0, + tokenizer=None, + selected_percentiles=[99.0], + goodput_config_dict={}, + task_type=TaskType.GENERATION, + selected_percentile_metrics=_MEDIA_PERCENTILE_METRICS, + max_concurrency=None, + request_rate=float("inf"), + benchmark_duration=5.0, + ) + + out = capsys.readouterr().out + assert actual_output_lens == [0, 0] + assert metrics.total_output == 0 + assert metrics.max_output_tokens_per_s == 0.0 + assert " Text Result " not in out + assert "Peak output token throughput" not in out + assert "Time to First Token" not in out + assert " Image Result " in out + + +def test_pure_video_benchmark_omits_text_result(capsys): + """Pure video generation must not print Text Result or invent peak tok/s.""" + outputs = [_make_video_output(63), _make_video_output(63)] + + metrics, actual_output_lens = calculate_metrics( + input_requests=[], + outputs=outputs, + dur_s=10.0, + tokenizer=None, + selected_percentiles=[99.0], + goodput_config_dict={}, + task_type=TaskType.GENERATION, + selected_percentile_metrics=_MEDIA_PERCENTILE_METRICS, + max_concurrency=None, + request_rate=float("inf"), + benchmark_duration=10.0, + ) + + out = capsys.readouterr().out + assert actual_output_lens == [0, 0] + assert metrics.total_output == 0 + assert metrics.max_output_tokens_per_s == 0.0 + assert " Text Result " not in out + assert "Peak output token throughput" not in out + assert "Time to First Token" not in out + assert " Video Result " in out + + +def test_image_with_generated_text_still_reports_text_result(capsys): + """Image edits / AR text alongside images must keep Text Result.""" + output = _make_image_output(64) + output.output_tokens = 8 + output.generated_text = "revised caption" + output.ttft = 0.05 + output.itl = [0.01] * 7 + + calculate_metrics( + input_requests=[], + outputs=[output], + dur_s=3.0, + tokenizer=_EmptyAwareTokenizer(), + selected_percentiles=[99.0], + goodput_config_dict={}, + task_type=TaskType.GENERATION, + selected_percentile_metrics=_MEDIA_PERCENTILE_METRICS, + max_concurrency=None, + request_rate=float("inf"), + benchmark_duration=3.0, + ) + + out = capsys.readouterr().out + assert " Text Result " in out + assert "Time to First Token" in out + assert " Image Result " in out + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index e72d5a263..accb30100 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -22,9 +22,16 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( ) from vllm_omni.benchmarks.patch.patch import ( MixRequestFuncOutput, + _add_video_extra_body_to_form, + _add_video_reference_to_form, _apply_stage0_token_timings, + _apply_video_metrics_from_payload, _attach_seed_tts_to_request_func_input, + _build_benchmark_session, + _omni_request_timeout_s, async_request_openai_chat_omni_completions, + async_request_openai_image_edits_omni, + async_request_openai_image_generations_omni, async_request_openai_realtime_duplex, should_request_stage_metrics, ) @@ -1196,5 +1203,337 @@ async def test_prompt_len_assigned_from_usage(mocker: MockerFixture): ) +class TestOmniRequestTimeout: + """``--omni-request-timeout-s`` precedence: explicit value > 900 s default.""" + + _OVERRIDE = "vllm_omni.benchmarks.patch.patch._REQUEST_TIMEOUT_OVERRIDE_S" + + def test_default_timeout_is_900s_when_not_set(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(self._OVERRIDE, None) + assert _omni_request_timeout_s() == 900.0 + + def test_explicit_value_wins_over_default(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(self._OVERRIDE, 123.5) + assert _omni_request_timeout_s() == 123.5 + + def test_non_positive_restores_legacy_6h_cap(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(self._OVERRIDE, 0) + assert _omni_request_timeout_s() == 6 * 60 * 60.0 + + async def test_benchmark_session_uses_configured_timeout(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(self._OVERRIDE, 42.0) + session = _build_benchmark_session(max_concurrency=8, ssl_setting=False) + try: + assert session.timeout.total == 42.0 + assert session.connector.limit == 8 + assert session.connector.limit_per_host == 8 + finally: + await session.close() + + async def test_hung_server_request_times_out_as_failed(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A server that accepts the request but never responds must surface as ``failed``.""" + + async def handler(reader, writer): + # Drain the request but never respond — the original hang failure mode. + try: + while await reader.read(4096): + pass + except ConnectionResetError: + pass + finally: + # Close our side of the socket: the client aborts after the + # timeout and only sends FIN, so without writer.close() the + # half-closed connection keeps Server.wait_closed() (which on + # Python 3.12+ waits for every accepted connection to drop) + # blocked forever. + writer.close() + try: + await writer.wait_closed() + except (ConnectionResetError, BrokenPipeError): + pass + + monkeypatch.setattr(self._OVERRIDE, 1.0) + server = await asyncio.start_server(handler, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + session = _build_benchmark_session(max_concurrency=1, ssl_setting=False) + try: + request_input = _seed_tts_request_func_input() + request_input.api_url = f"http://127.0.0.1:{port}/v1/chat/completions" + output = await async_request_openai_chat_omni_completions(request_input, session) + assert output.success is False + assert output.error + finally: + await session.close() + server.close() + await server.wait_closed() + + +def test_video_rtf_prefers_generation_time_over_poll_latency(): + """RTF must not grow with client poll_interval overshoot in output.latency.""" + payload = { + "duration_s": 2.0, + "num_frames": 48, + "fps": 24.0, + "stage_durations": {"stage_0_gen_ms": 4000.0}, + } + request_body: dict[str, object] = {} + + short_poll = MixRequestFuncOutput() + short_poll.latency = 4.2 # ~generation + small poll overshoot + _apply_video_metrics_from_payload(short_poll, payload, request_body) + + long_poll = MixRequestFuncOutput() + long_poll.latency = 8.0 # same job, larger poll_interval_s overshoot + _apply_video_metrics_from_payload(long_poll, payload, request_body) + + assert short_poll.video_generation_time_ms == pytest.approx(4000.0) + assert long_poll.video_generation_time_ms == pytest.approx(4000.0) + assert short_poll.video_rtf == pytest.approx(2.0) + assert long_poll.video_rtf == pytest.approx(2.0) + assert short_poll.video_rtf == long_poll.video_rtf + + +def test_video_rtf_falls_back_to_e2e_latency_without_generation_time(): + output = MixRequestFuncOutput() + output.latency = 6.0 + _apply_video_metrics_from_payload( + output, + {"duration_s": 2.0, "num_frames": 48, "fps": 24.0}, + {}, + ) + assert output.video_generation_time_ms == 0.0 + assert output.video_rtf == pytest.approx(3.0) + + +# 1x1 PNG used as a valid b64_json image payload in edit-client tests. +_MIN_PNG_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" + + +@pytest.mark.asyncio +async def test_image_edits_defaults_to_non_streaming_json(mocker: MockerFixture) -> None: + """Single-stage servers reject stream=true; default path must send stream=false + JSON.""" + + class MockJsonResponse: + status = 200 + + def __init__(self): + self._payload = { + "created": 1, + "data": [ + { + "b64_json": _MIN_PNG_B64, + "stage_durations": {"stage_0_gen_ms": 12.5}, + } + ], + } + + async def json(self): + return self._payload + + async def text(self): + return json.dumps(self._payload) + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + captured_stream: list[str] = [] + real_add_field = None + + def tracking_add_field(self, name, value=None, **kwargs): + if name == "stream": + captured_stream.append(str(value)) + return real_add_field(self, name, value, **kwargs) + + import aiohttp + + real_add_field = aiohttp.FormData.add_field + mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field) + + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=MockJsonResponse()) + + request = RequestFuncInput( + model="single-stage-edit", + model_name="single-stage-edit", + prompt="make it sunny", + api_url="http://test.com/v1/images/edits", + prompt_len=4, + output_len=1, + multi_modal_content=[{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"}}], + extra_body={"num_inference_steps": 4}, + ) + output = await async_request_openai_image_edits_omni(request, mock_session, pbar=None) + + assert captured_stream == ["false"] + assert output.success is True + assert not output.error + assert output.image_count == 1 + assert output.image_generation_time_ms == pytest.approx(12.5) + + +@pytest.mark.asyncio +async def test_image_edits_stream_true_uses_sse_path(mocker: MockerFixture) -> None: + """Explicit stream=true keeps the multi-stage SSE client path.""" + sse_chunk = ( + b'data: {"type":"image","data":[{"b64_json":"' + _MIN_PNG_B64.encode() + b'"}]}\n\n' + b"data: [DONE]\n\n" + ) + mock_response = MockResponse(200, [sse_chunk]) + captured_stream: list[str] = [] + + import aiohttp + + real_add_field = aiohttp.FormData.add_field + + def tracking_add_field(self, name, value=None, **kwargs): + if name == "stream": + captured_stream.append(str(value)) + return real_add_field(self, name, value, **kwargs) + + mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field) + + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) + + request = RequestFuncInput( + model="multi-stage-edit", + model_name="multi-stage-edit", + prompt="edit", + api_url="http://test.com/v1/images/edits", + prompt_len=2, + output_len=1, + multi_modal_content=[{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"}}], + extra_body={"stream": True}, + ) + output = await async_request_openai_image_edits_omni(request, mock_session, pbar=None) + + assert captured_stream == ["true"] + assert output.success is True + assert output.image_count == 1 + + +@pytest.mark.asyncio +async def test_image_generations_e2el_includes_json_body_consume(mocker: MockerFixture) -> None: + """E2EL must include body transfer/decode, not stop at HTTP headers.""" + + class SlowJsonResponse: + status = 200 + + def __init__(self): + self._payload = { + "created": 1, + "data": [{"b64_json": _MIN_PNG_B64, "stage_durations": {"stage_0_gen_ms": 1.0}}], + } + + async def json(self): + await asyncio.sleep(0.05) + return self._payload + + async def text(self): + return json.dumps(self._payload) + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=SlowJsonResponse()) + request = RequestFuncInput( + model="img-gen", + model_name="img-gen", + prompt="a cat", + api_url="http://test.com/v1/images/generations", + prompt_len=2, + output_len=1, + extra_body={}, + ) + output = await async_request_openai_image_generations_omni(request, mock_session, pbar=None) + assert output.success is True + assert output.latency >= 0.05 + + +def test_video_local_image_reference_not_forwarded_as_raw_extra_field(tmp_path, mocker: MockerFixture) -> None: + """Local image_reference must only be uploaded via dedicated serializer.""" + import aiohttp + + ref_path = tmp_path / "ref.png" + ref_path.write_bytes(base64.b64decode(_MIN_PNG_B64)) + + captured: list[tuple[str, object]] = [] + real_add_field = aiohttp.FormData.add_field + + def tracking_add_field(self, name, value=None, **kwargs): + captured.append((str(name), value)) + return real_add_field(self, name, value, **kwargs) + + mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field) + + form = aiohttp.FormData() + extra_body = { + "num_inference_steps": 2, + "image_reference": str(ref_path), + } + request_body = {"model": "vid", "prompt": "p", **extra_body} + _add_video_extra_body_to_form(form, extra_body, request_body) + assert _add_video_reference_to_form(form, extra_body["image_reference"]) is True + + field_names = [name for name, _ in captured] + assert "image_reference" not in field_names + assert field_names.count("input_reference") == 1 + # Dedicated uploader sends file bytes, not the raw local path string. + uploaded = next(value for name, value in captured if name == "input_reference") + assert uploaded == base64.b64decode(_MIN_PNG_B64) + + +@pytest.mark.parametrize( + "reference", + [ + {"image_url": "https://example.com/ref.png"}, + {"file_id": "file-abc"}, + [{"image_url": "https://example.com/a.png"}, {"file_id": "file-xyz"}], + ], +) +def test_video_structured_image_reference_serialized_to_form(reference: object, mocker: MockerFixture) -> None: + """Object-form image_reference must be JSON-serialized, not silently dropped.""" + import aiohttp + + captured: list[tuple[str, object]] = [] + real_add_field = aiohttp.FormData.add_field + + def tracking_add_field(self, name, value=None, **kwargs): + captured.append((str(name), value)) + return real_add_field(self, name, value, **kwargs) + + mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field) + + form = aiohttp.FormData() + extra_body = {"image_reference": reference} + request_body = {"model": "vid", "prompt": "p", **extra_body} + _add_video_extra_body_to_form(form, extra_body, request_body) + assert _add_video_reference_to_form(form, reference) is True + + field_names = [name for name, _ in captured] + # Reserved key must not be double-forwarded as a generic extra field; + # exactly one dedicated image_reference field. + assert field_names.count("image_reference") == 1 + assert "input_reference" not in field_names + payload = next(value for name, value in captured if name == "image_reference") + assert json.loads(payload) == reference + + +def test_video_unsupported_image_reference_raises() -> None: + import aiohttp + + form = aiohttp.FormData() + with pytest.raises(ValueError, match="Unsupported image_reference"): + _add_video_reference_to_form(form, {"not_a_supported_key": "x"}) + with pytest.raises(ValueError, match="Unsupported image_reference"): + _add_video_reference_to_form(form, "/tmp/does-not-exist-ref.png") + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/benchmarks/test_accuracy_bench_utils.py b/tests/benchmarks/test_accuracy_bench_utils.py index ab5392346..8d78da2b2 100644 --- a/tests/benchmarks/test_accuracy_bench_utils.py +++ b/tests/benchmarks/test_accuracy_bench_utils.py @@ -1,9 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # ruff: noqa: E402, I001 import argparse import math import os import sys -import types from pathlib import Path import pytest @@ -106,11 +108,9 @@ def test_resolve_seed_tts_root_downloads_only_requested_locale(monkeypatch, tmp_ captured["allow_patterns"] = allow_patterns return str(downloaded_root) - monkeypatch.setitem( - sys.modules, - "huggingface_hub", - types.SimpleNamespace(snapshot_download=fake_snapshot_download), - ) + from vllm_omni.transformers_utils import repo_utils + + monkeypatch.setattr(repo_utils.hf_api(), "snapshot_download", fake_snapshot_download) resolved = resolve_seed_tts_root( "zhaochenyang20/seed-tts-eval", diff --git a/tests/benchmarks/test_serve_cli.py b/tests/benchmarks/test_serve_cli.py index 5365d56ac..dae70a6c4 100644 --- a/tests/benchmarks/test_serve_cli.py +++ b/tests/benchmarks/test_serve_cli.py @@ -167,6 +167,26 @@ def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default( {"bot_task": "recaption"}, {"backend", "print_stage", "bot_task"}, ), + ( + [ + "--endpoint", + "/v1/images/edits", + "--image-edits-bot-task", + "think", + ], + {"bot_task": "think"}, + {"endpoint", "bot_task"}, + ), + ( + [ + "--backend", + "/v1/images/edits", + "--image-edits-bot-task", + "recaption", + ], + {"bot_task": "recaption"}, + {"backend", "bot_task"}, + ), ( ["--extra-body", '{"bot_task":"vanilla"}'], {"bot_task": "vanilla"}, @@ -182,6 +202,7 @@ def test_omni_args_parse_and_preprocess( parser = TrackingArgumentParser() parser.add_argument("--extra-body", type=json.loads, default=None) parser.add_argument("--backend", default="openai-chat-omni") + parser.add_argument("--endpoint", default="/v1/chat/completions") add_omni_args(parser) args = parser.parse_args(argv) @@ -358,3 +379,17 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path): ) assert bench_requests assert all(url == expected_url for url in bench_requests), f"Unexpected target URLs: {bench_requests}" + + +def test_omni_request_timeout_s_flag_defaults_and_parses() -> None: + parser = TrackingArgumentParser() + add_omni_args(parser) + + args = parser.parse_args([]) + assert args.omni_request_timeout_s is None + + args = parser.parse_args(["--omni-request-timeout-s", "0"]) + assert args.omni_request_timeout_s == 0.0 + + args = parser.parse_args(["--omni-request-timeout-s", "3600"]) + assert args.omni_request_timeout_s == 3600.0 diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 3723e5523..51c7d549e 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -6,7 +6,6 @@ Unit tests for StageConfigFactory and related classes. import importlib import warnings -from dataclasses import dataclass from pathlib import Path from unittest.mock import patch @@ -37,6 +36,7 @@ from vllm_omni.config.stage_config import ( normalize_pipeline_cli_overrides, pipeline_cfg_resolver, ) +from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -451,13 +451,8 @@ class TestStageConfigFactory: def test_default_diffusion_with_parallel_config(self): """Test diffusion config calculates devices from parallel_config.""" - - @dataclass - class MockParallelConfig: - world_size: int = 4 - kwargs = { - "parallel_config": MockParallelConfig(), + "parallel_config": DiffusionParallelConfig(tensor_parallel_size=4), "cache_backend": "tea_cache", } configs = StageConfigFactory.create_default_diffusion(kwargs) @@ -937,6 +932,44 @@ class TestPipelineRegistration: assert isinstance(omni_config, VllmOmniConfig) assert omni_config.stage_by_id(0).diffusion_config.model == "fake/model" + def test_cosyvoice3_deploy_resolves_hash_snapshot_without_hf_metadata(self, tmp_path): + snapshot = tmp_path / "models--FunAudioLLM--Fun-CosyVoice3-0.5B-2512" / "snapshots" / ("a" * 40) + snapshot.mkdir(parents=True) + (snapshot / "config.json").write_text("{}\n", encoding="utf-8") + model = str(snapshot) + deploy_path = get_deploy_config_path("cosyvoice3.yaml") + + assert StageConfigFactory.try_infer_model_type(model, trust_remote_code=False) is None + assert StageConfigFactory.get_pipeline_config(model, trust_remote_code=False) is None + + omni_config = StageConfigFactory.create_from_model( + model, + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=deploy_path, + ) + legacy_configs, _ = StageConfigFactory.create_legacy_stage_configs_from_model( + model, + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=deploy_path, + ) + + assert omni_config is not None + assert omni_config.pipeline_config.model_type == "cosyvoice3" + assert [ + (stage.stage_id, stage.stage_type, stage.worker_type, stage.model_stage) + for stage in omni_config.stage_configs + ] == [ + (0, StageType.LLM, "ar", "cosyvoice3_talker"), + (1, StageType.LLM, "generation", "cosyvoice3_code2wav"), + ] + assert all(stage.model_config.model == model for stage in omni_config.stage_configs) + assert all(stage.stage_type is not StageType.DIFFUSION for stage in omni_config.stage_configs) + + assert legacy_configs is not None + assert [stage.stage_type for stage in legacy_configs] == [StageType.LLM, StageType.LLM] + 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" @@ -1602,7 +1635,7 @@ stages: assert stage.max_model_len == 1024 assert stage.max_num_batched_tokens == 1024 assert stage.max_num_seqs == 1 - assert stage.gpu_memory_utilization == 0.7 + assert stage.gpu_memory_utilization == 0.8 assert stage.skip_mm_profiling is True assert stage.enforce_eager is True assert stage.async_scheduling is False diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 52ff1e274..7603945a0 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 24, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 62, + EnvironmentVariableCategory.MODEL_SPECIFIC: 61, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -179,7 +179,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { ModelEnvironmentVariableDisposition.PROMOTE: 37, - ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 7, + ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 13, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5, diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index f91e6af42..3e499dc38 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -261,6 +261,57 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy) +@pytest.mark.parametrize("stage_id", [0, 2], ids=["ar", "generation"]) +def test_llm_additional_config_roundtrip_and_isolation(stage_id, monkeypatch): + from vllm_omni.engine import stage_init_utils + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + + # Worker discovery requires hardware support; this test covers config transport. + monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None) + additional_config = {"backend_options": {"enabled": True}} + pipeline = _resolve_pipeline_or_skip("minicpmo_4_5") + deploy = DeployConfig( + stages=[StageDeployConfig(stage_id=i, engine_extras={"additional_config": additional_config}) for i in (0, 2)] + ) + config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy) + stage = config.stage_by_id(stage_id) + assert stage.runtime_config.additional_config == additional_config + engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model") + assert engine_args["additional_config"] == additional_config + + engine_args["additional_config"]["backend_options"]["enabled"] = False + assert stage.runtime_config.additional_config["backend_options"]["enabled"] is True + stage.runtime_config.additional_config["backend_options"]["enabled"] = False + assert config.stage_by_id(2 if stage_id == 0 else 0).runtime_config.additional_config == additional_config + assert additional_config["backend_options"]["enabled"] is True + + +@pytest.mark.parametrize("deploy_name", ["minicpmo_4_5", "minicpmo_4_5_2gpu", "minicpmo_4_5_3gpu"]) +def test_minicpmo_npu_additional_config_reaches_engine_args(monkeypatch, deploy_name): + from vllm_omni.engine import stage_init_utils + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "device_name", "npu") + monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None) + stage = _from_pipeline_key("minicpmo_4_5", deploy_config_path=deploy_name).stage_by_id(2) + expected = {"code2wav_enable_npu_graph": True, "code2wav_max_npu_graphs": 32} + assert stage.runtime_config.additional_config == expected + engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model") + assert engine_args["additional_config"] == expected + + +def test_diffusion_additional_config_keeps_diffusion_owner(): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + + additional_config = {"torchair_graph_config": {"enabled": True}} + stage = _from_pipeline_key("dreamzero", cli_overrides={"additional_config": additional_config}).stage_by_id(0) + assert stage.runtime_config.additional_config is None + assert stage.diffusion_config.additional_config == additional_config + engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model") + assert engine_args["additional_config"] == additional_config + + @pytest.mark.parametrize( ("engine_extras", "cli_overrides"), [ @@ -532,6 +583,17 @@ def test_from_pipeline_config_dispatches_async_chunk_processors_without_mutating assert pipeline.get_stage(1).custom_process_input_func is None +def test_joyai_code2wav_waits_for_full_payload(): + config = _from_pipeline_key("joyai_vl_interaction") + talker = config.stage_by_id(1) + code2wav = config.stage_by_id(2) + + assert talker.custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload") + assert code2wav.custom_process_input_func.endswith("talker2code2wav_token_only") + assert code2wav.connector_config.async_chunk is False + assert code2wav.model_config.requires_full_payload_input is True + + 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") @@ -559,6 +621,7 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations(): def test_runtime_config_fields_match_structured_runtime_scope(): assert {f.name for f in fields(OmniStageRuntimeConfig)} == { + "additional_config", "distributed_executor_backend", "worker_cls", "devices", @@ -1253,6 +1316,70 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): assert cfg.diffusers_call_kwargs == {} +def test_diffusion_config_none_values_preserve_dataclass_defaults(): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + normalized = OmniDiffusionConfig.normalize_init_kwargs( + { + "lora_scale": None, + "enable_sleep_mode": None, + "diffusers_load_kwargs": None, + } + ) + + assert "lora_scale" not in normalized + assert "enable_sleep_mode" not in normalized + assert normalized["diffusers_load_kwargs"] == {} + + config = OmniDiffusionConfig.from_kwargs( + lora_scale=None, + enable_sleep_mode=None, + ) + assert config.lora_scale == 1.0 + assert config.enable_sleep_mode is False + + +@pytest.mark.parametrize( + ("canonical_key", "alias_key", "canonical_value", "alias_value"), + [ + ("lora_scale", "static_lora_scale", 0.75, 0.25), + ( + "quantization_config", + "diffusion_quantization_config", + {"method": "canonical"}, + {"method": "diffusion-alias"}, + ), + ( + "quantization_config", + "quantization", + {"method": "canonical"}, + "legacy-alias", + ), + ("diffusion_kv_cache_dtype", "kv_cache_dtype", "fp8", "fp16"), + ("diffusion_kv_cache_skip_steps", "kv_cache_skip_steps", "0-1", "2-3"), + ("diffusion_kv_cache_skip_layers", "kv_cache_skip_layers", "1-2", "3-4"), + ("streaming_output", "diffusion_streaming_output", False, True), + ], +) +def test_diffusion_alias_conflicts_prefer_canonical_key( + canonical_key, + alias_key, + canonical_value, + alias_value, +): + from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs + + normalized = normalize_omni_diffusion_kwargs( + { + canonical_key: canonical_value, + alias_key: alias_value, + } + ) + + assert normalized[canonical_key] == canonical_value + assert alias_key not in normalized + + def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch): from vllm_omni.platforms import current_omni_platform diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index bb69391fd..ffcfc45b7 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Performance benchmark CI runner for diffusion models. @@ -356,6 +359,8 @@ def _resolve_offline_model(model: str) -> str: """ import huggingface_hub + from vllm_omni.transformers_utils.repo_utils import hf_api + if not model or os.path.isdir(model): return model @@ -380,9 +385,7 @@ def _resolve_offline_model(model: str) -> str: if len(parts) >= 3: repo_id = "/".join(parts[:2]) subfolder = "/".join(parts[2:]) - from huggingface_hub import snapshot_download - - snapshot_root = snapshot_download( + snapshot_root = hf_api().snapshot_download( repo_id, allow_patterns=[f"{subfolder}/**"], local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, @@ -391,9 +394,7 @@ def _resolve_offline_model(model: str) -> str: if not huggingface_hub.constants.HF_HUB_OFFLINE: return model - from huggingface_hub import snapshot_download - - return snapshot_download(model, local_files_only=True) + return hf_api().snapshot_download(model, local_files_only=True) class DiffusionServer: diff --git a/tests/diffusion/distributed/test_wan_chunked_decode.py b/tests/diffusion/distributed/test_wan_chunked_decode.py new file mode 100644 index 000000000..87f3ddff7 --- /dev/null +++ b/tests/diffusion/distributed/test_wan_chunked_decode.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import weakref +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from diffusers.models.autoencoders import AutoencoderKLWan + +from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + DistributedAutoencoderKLWan, + OmniAutoencoderKLWan, +) +from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture +def vae(): + # Real causal convolutions and caches, with small random weights for CPU CI. + return OmniAutoencoderKLWan( + base_dim=4, + z_dim=2, + dim_mult=[1, 2], + num_res_blocks=1, + temperal_downsample=[True], + scale_factor_spatial=2, + scale_factor_temporal=2, + ).eval() + + +@pytest.mark.parametrize("patch_size", [None, 2]) +@pytest.mark.parametrize("return_dict", [False, True]) +@torch.no_grad() +def test_chunked_decode_matches_diffusers(vae, patch_size, return_dict): + if patch_size is not None: + vae = OmniAutoencoderKLWan.from_config(vae.config, patch_size=patch_size, out_channels=12).eval() + z = torch.randn(1, 2, 3, 2, 2) + expected = AutoencoderKLWan._decode(vae, z).sample + actual = vae.decode(z, return_dict=return_dict) + torch.testing.assert_close(actual.sample if return_dict else actual[0], expected, rtol=0, atol=0) + chunks: list[torch.Tensor] = [] + assert supports_chunked_vae_decode(vae) + assert not supports_chunked_vae_decode(object()) + assert vae.decode_with_chunks(z, on_chunk=chunks.append) is None + assert [chunk.shape[2] for chunk in chunks] == [1, 2, 2] + torch.testing.assert_close(torch.cat(chunks, dim=2), expected, rtol=0, atol=0) + assert all(entry is None for entry in vae._feat_map) + + +@torch.no_grad() +def test_slicing_and_forward_hooks(vae): + z = torch.randn(2, 2, 2, 2, 2) + vae.enable_slicing() + batches = [] + handle = vae.post_quant_conv.register_forward_pre_hook(lambda module, args: batches.append(args[0].shape[0])) + vae._hf_hook = SimpleNamespace(pre_forward=Mock()) + try: + assert vae.decode(z).sample.shape[0] == 2 + assert batches == [1, 1] + vae._hf_hook.pre_forward.assert_called_once_with(vae) + batches.clear() + with pytest.raises(ValueError, match="use_slicing"): + vae.decode_with_chunks(z, on_chunk=Mock()) + assert batches == [] + vae._hf_hook.pre_forward.reset_mock() + vae.decode_with_chunks(z[:1], on_chunk=Mock()) + assert batches == [1] + vae._hf_hook.pre_forward.assert_called_once_with(vae) + finally: + handle.remove() + + +@torch.no_grad() +def test_callback_failure_drains_and_clears_cache(vae): + z = torch.randn(1, 2, 4, 2, 2) + decoded = Mock(return_value=None) + handle = vae.decoder.register_forward_hook(lambda *args: decoded()) + error = RuntimeError("consumer failed") + callback = Mock(side_effect=error) + try: + with pytest.raises(RuntimeError, match="consumer failed") as exc: + vae.decode_with_chunks(z, on_chunk=callback) + assert exc.value is error + assert callback.call_count == 1 + assert decoded.call_count == z.shape[2] + assert all(entry is None for entry in vae._feat_map) + finally: + handle.remove() + # The same instance can decode again after the failure. + chunks: list[torch.Tensor] = [] + vae.decode_with_chunks(z, on_chunk=chunks.append) + torch.testing.assert_close(torch.cat(chunks, 2), vae.decode(z).sample, rtol=0, atol=0) + + +def test_chunked_decode_rejects_spatial_tiling(vae): + vae.enable_tiling(tile_sample_min_height=2, tile_sample_min_width=2) + with pytest.raises(ValueError, match="spatial tiling"): + vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=Mock()) + + +@pytest.mark.parametrize("rank", [0, 1]) +@pytest.mark.parametrize("fail_callback", [False, True]) +def test_spatial_callback_ownership_and_bounded_retention(monkeypatch, rank, fail_callback): + # Isolate the temporal loop from halo kernels; weakrefs detect retention of + # decoded GPU-sized tensors even after the consumer has failed. + monkeypatch.setattr(wan_spatial_shard, "install_wan_spatial_shard_decode", lambda *args, **kwargs: None) + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (rank, 2)) + references: list[weakref.ReferenceType[torch.Tensor]] = [] + + def decode(x, **kwargs): + assert sum(ref() is not None for ref in references) <= 1 + chunk = x.clone() + references.append(weakref.ref(chunk)) + return chunk + + vae = SimpleNamespace( + config=SimpleNamespace(patch_size=None), + clear_cache=Mock(), + post_quant_conv=lambda z: z, + decoder=decode, + _feat_map=[], + ) + callback = Mock(side_effect=RuntimeError("consumer failed") if fail_callback else None) + z = torch.zeros(1, 3, 6, 2, 2) + if rank == 0 and fail_callback: + with pytest.raises(RuntimeError, match="consumer failed"): + wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback) + else: + assert wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback) is None + assert len(references) == z.shape[2] + assert callback.call_count == (0 if rank else 1 if fail_callback else z.shape[2]) + assert vae.clear_cache.call_count == 2 + + +@pytest.mark.parametrize("rank", [0, 1]) +@torch.no_grad() +def test_distributed_small_latent_callback_ownership(vae, monkeypatch, rank): + vae.is_distributed_enabled = lambda: True + vae.distributed_executor = SimpleNamespace(group=None) + monkeypatch.setattr(torch.distributed, "get_rank", lambda **kwargs: rank) + callback = Mock() + vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=callback) + assert callback.call_count == (2 if rank == 0 else 0) + + +@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width", "tile"]) +def test_distributed_chunked_routing(monkeypatch, mode): + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + torch.nn.Module.__init__(vae) + vae.use_tiling = True + vae.use_slicing = False + vae.tile_sample_min_height = vae.tile_sample_min_width = 2 + vae.spatial_compression_ratio = 2 + vae.distributed_executor = SimpleNamespace(group=None, parallel_mode=mode, parallel_size=2) + monkeypatch.setattr(vae, "is_distributed_enabled", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda **kwargs: 2) + decode = Mock() + monkeypatch.setattr(wan_spatial_shard, "spatial_shard_decode", decode) + callback = Mock() + z = torch.zeros(1, 2, 2, 2, 2) + if mode == "tile": + with pytest.raises(ValueError, match="require spatial-shard"): + vae.decode_with_chunks(z, on_chunk=callback) + decode.assert_not_called() + else: + assert vae.decode_with_chunks(z, on_chunk=callback) is None + decode.assert_called_once_with( + vae, z, group=None, return_dict=True, split_dim=mode.removeprefix("spatial_shard_"), on_chunk=callback + ) diff --git a/tests/diffusion/lora/helpers.py b/tests/diffusion/lora/helpers.py index 8b9b1ef4d..4f1dfe176 100644 --- a/tests/diffusion/lora/helpers.py +++ b/tests/diffusion/lora/helpers.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared test helpers for diffusion LoRA tests.""" from __future__ import annotations @@ -26,15 +26,39 @@ class DummyBaseLayerWithLoRA(torch.nn.Module): tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor] ] = [] self.reset_calls: int = 0 + self.suspend_calls: int = 0 + self.resume_calls: int = 0 + self.active_slices: tuple[bool, ...] = () + self.suspended_slices: tuple[bool, ...] | None = None self.create_calls: int = 0 def set_lora(self, index: int, lora_a, lora_b): assert index == 0 self.set_calls.append((lora_a, lora_b)) + if isinstance(lora_b, list): + self.active_slices = tuple(b is not None for b in lora_b) + else: + self.active_slices = (True,) + self.suspended_slices = None def reset_lora(self, index: int): assert index == 0 self.reset_calls += 1 + self.active_slices = () + self.suspended_slices = None + + def suspend_lora(self) -> None: + if self.suspended_slices is not None: + return + self.suspended_slices = self.active_slices + self.active_slices = (False,) * len(self.active_slices) + self.suspend_calls += 1 + + def resume_lora(self) -> None: + if self.suspended_slices is not None: + self.active_slices = self.suspended_slices + self.suspended_slices = None + self.resume_calls += 1 def create_lora_weights(self, max_loras, lora_config, model_config): self.create_calls += 1 diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 8edc6ce4b..0814bf15f 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -30,14 +30,38 @@ class _DummyLoRALayer: tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor] ] = [] self.reset_calls: int = 0 + self.suspend_calls: int = 0 + self.resume_calls: int = 0 + self.active_slices: tuple[bool, ...] = () + self.suspended_slices: tuple[bool, ...] | None = None def set_lora(self, index: int, lora_a, lora_b): assert index == 0 self.set_calls.append((lora_a, lora_b)) + if isinstance(lora_b, list): + self.active_slices = tuple(b is not None for b in lora_b) + else: + self.active_slices = (True,) + self.suspended_slices = None def reset_lora(self, index: int): assert index == 0 self.reset_calls += 1 + self.active_slices = () + self.suspended_slices = None + + def suspend_lora(self) -> None: + if self.suspended_slices is not None: + return + self.suspended_slices = self.active_slices + self.active_slices = (False,) * len(self.active_slices) + self.suspend_calls += 1 + + def resume_lora(self) -> None: + if self.suspended_slices is not None: + self.active_slices = self.suspended_slices + self.suspended_slices = None + self.resume_calls += 1 # Aliases for backward compatibility within this file @@ -723,3 +747,104 @@ def test_lora_manager_discovers_unet_component(monkeypatch): 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) + + +def _suspend_harness(monkeypatch, adapter_id: int = 7, rank: int = 2): + """Manager wired to a single dummy layer, ready to activate `adapter_id`.""" + import vllm_omni.diffusion.lora.manager as manager_mod + + monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA) + + lora_model = _DummyLM(rank=rank) + manager = DiffusionLoRAManager( + pipeline=_DummyPipeline(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + monkeypatch.setattr(manager, "_load_adapter", lambda _req: (lora_model, type("PH", (), {"r": rank})())) + manager._registered_adapters = {adapter_id: lora_model} + manager._lora_modules = {"transformer.foo": lora_model.transformer.foo} + return manager, lora_model.transformer.foo + + +def test_reactivating_same_adapter_resumes_without_rebinding(monkeypatch): + """activate -> deactivate -> activate at the same scale must re-arm the + saved mask instead of re-uploading every layer.""" + adapter_id = 7 + req = _dummy_lora_request(adapter_id) + manager, layer = _suspend_harness(monkeypatch, adapter_id) + + manager.set_active_adapter(req, lora_scale=0.5) + assert len(layer.set_calls) == 1 + assert layer.active_slices == (True,) + + manager.set_active_adapter(None) + assert manager._suspended_adapter_id == adapter_id + assert layer.suspend_calls == 1 + assert layer.active_slices == (False,) + # Suspending must not tear the upload down. + assert layer.reset_calls == 0 + + manager.set_active_adapter(req, lora_scale=0.5) + assert layer.resume_calls == 1 + assert len(layer.set_calls) == 1, "resume must not rebind" + assert layer.active_slices == (True,) + assert manager._active_adapter_id == adapter_id + assert manager._suspended_adapter_id is None + + +@pytest.mark.parametrize("second_scale", [0.5, 0.25]) +def test_scale_or_adapter_change_forces_rebind(monkeypatch, second_scale): + """A different scale, or a different adapter, must take the full bind + path rather than re-arming a stale mask.""" + adapter_id = 7 + manager, layer = _suspend_harness(monkeypatch, adapter_id) + + manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5) + manager.set_active_adapter(None) + assert manager._suspended_adapter_id == adapter_id + + # Same id at a new scale, or a different id entirely. + other_id = adapter_id if second_scale != 0.5 else adapter_id + 1 + if other_id != adapter_id: + manager._registered_adapters[other_id] = manager._registered_adapters[adapter_id] + + manager.set_active_adapter(_dummy_lora_request(other_id), lora_scale=second_scale) + assert len(layer.set_calls) == 2, "must rebind" + assert manager._suspended_adapter_id is None + assert manager._active_adapter_id == other_id + + +def test_partial_packed_mask_survives_suspend_resume(monkeypatch): + """A packed layer where only some slices carry LoRA must come back with + the same mask, not an all-True one.""" + adapter_id = 7 + manager, layer = _suspend_harness(monkeypatch, adapter_id) + + # Emulate a packed bind that left slice 1 empty. + layer.set_lora(0, [torch.ones(2, 2), None], [torch.ones(2, 2), None]) + manager._active_adapter_id = adapter_id + manager._update_adapter_scale(adapter_id, 0.5) + assert layer.active_slices == (True, False) + + manager.set_active_adapter(None) + assert layer.active_slices == (False, False) + + manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5) + assert layer.active_slices == (True, False), "mask must survive unchanged" + + +def test_removing_suspended_adapter_drops_the_upload(monkeypatch): + """A removed adapter can never be resumed, so its weights must not stay + in the stacked buffers.""" + adapter_id = 7 + manager, layer = _suspend_harness(monkeypatch, adapter_id) + + manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5) + manager.set_active_adapter(None) + assert manager._suspended_adapter_id == adapter_id + + manager.remove_adapter(adapter_id) + assert manager._suspended_adapter_id is None + assert layer.reset_calls == 1, "the upload must be torn down" + assert layer.suspended_slices is None diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 555cb19bb..565949ff6 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -13,7 +13,6 @@ from unittest.mock import MagicMock import pytest import torch import torch.nn as nn -from huggingface_hub import snapshot_download from safetensors.torch import save_file from vllm.config.load import LoadConfig @@ -32,6 +31,7 @@ from vllm_omni.diffusion.models.helios import HeliosPipeline from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract from vllm_omni.diffusion.registry import initialize_model from vllm_omni.quantization.component_config import ComponentQuantizationConfig +from vllm_omni.transformers_utils.repo_utils import hf_api pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -41,7 +41,7 @@ model_path = "hf-internal-testing/tiny-helios-modular-pipe" @pytest.fixture(scope="module") def prefetch_helios_model(): """Downloads the tiny helios model prior to running a test.""" - snapshot_download(model_path) + hf_api().snapshot_download(model_path) @pytest.fixture(scope="function") diff --git a/tests/diffusion/model_loader/test_hub_prefetch.py b/tests/diffusion/model_loader/test_hub_prefetch.py index 39eb2b277..94170edcc 100644 --- a/tests/diffusion/model_loader/test_hub_prefetch.py +++ b/tests/diffusion/model_loader/test_hub_prefetch.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import contextlib @@ -14,10 +14,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] def test_prefetch_subfolders_propagates_revision(monkeypatch): calls = [] - def fake_snapshot_download(**kwargs): + def fake_snapshot_download(self, **kwargs): calls.append(kwargs) - monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", fake_snapshot_download) monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext()) hub_prefetch.prefetch_subfolders( diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_attention_mask.py b/tests/diffusion/models/boogu_image/test_boogu_image_attention_mask.py new file mode 100644 index 000000000..dbaf3b71a --- /dev/null +++ b/tests/diffusion/models/boogu_image/test_boogu_image_attention_mask.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.diffusion.models.utils import make_attention_mask + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.mark.parametrize("seq_lengths", [[16], [16, 16]]) +def test_attention_mask_is_skipped_for_dense_batch(seq_lengths): + hidden_states = torch.empty(len(seq_lengths), 16, 64) + + assert make_attention_mask(hidden_states, seq_lengths) is None + + +def test_attention_mask_preserves_variable_length_padding(): + hidden_states = torch.empty(2, 16, 64) + + attention_mask = make_attention_mask(hidden_states, [16, 12]) + + expected = torch.tensor( + [ + [True] * 16, + [True] * 12 + [False] * 4, + ] + ) + torch.testing.assert_close(attention_mask, expected) diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_registry.py b/tests/diffusion/models/boogu_image/test_boogu_image_registry.py new file mode 100644 index 000000000..8048540dc --- /dev/null +++ b/tests/diffusion/models/boogu_image/test_boogu_image_registry.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""L1 tests for checkpoint-to-pipeline routing of Boogu-Image Turbo.""" + +import json + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_turbo_registry_entry_loads_the_turbo_class(): + from vllm_omni.diffusion.models.boogu_image import BooguImageTurboPipeline + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + assert DiffusionModelRegistry._try_load_model_cls("BooguImageTurboPipeline") is BooguImageTurboPipeline + + +def test_turbo_entry_disables_request_batching(): + from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline, BooguImageTurboPipeline + + assert BooguImagePipeline.supports_request_batch + assert not BooguImageTurboPipeline.supports_request_batch + + +def test_turbo_process_funcs_resolve_to_the_shared_boogu_helpers(): + from vllm_omni.diffusion.models import boogu_image + from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS + + module = boogu_image.pipeline_boogu_image + for funcs in (_DIFFUSION_PRE_PROCESS_FUNCS, _DIFFUSION_POST_PROCESS_FUNCS): + name = funcs["BooguImageTurboPipeline"] + assert name == funcs["BooguImagePipeline"] + assert callable(getattr(module, name)) + + +def test_turbo_inherits_the_base_multimodal_metadata(): + from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata + + metadata = get_diffusion_model_metadata("BooguImageTurboPipeline") + + assert metadata == get_diffusion_model_metadata("BooguImagePipeline") + assert metadata.supports_multimodal_inputs + assert metadata.max_multimodal_image_inputs == 1 + + +def test_turbo_model_index_resolves_and_enriches(tmp_path): + from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name + + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "BooguImageTurboPipeline"}), + encoding="utf-8", + ) + + assert resolve_model_class_name(str(tmp_path)) == "BooguImageTurboPipeline" + + config = OmniDiffusionConfig(model=str(tmp_path)) + config.enrich_config() + + assert config.model_class_name == "BooguImageTurboPipeline" + assert config.supports_multimodal_inputs + assert config.max_multimodal_image_inputs == 1 + + +def test_edit_turbo_override_selects_the_turbo_class(tmp_path): + """Edit-Turbo publishes ``BooguImagePipeline``; ``--model-class-name`` must win.""" + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "BooguImagePipeline"}), + encoding="utf-8", + ) + + config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="BooguImageTurboPipeline") + config.enrich_config() + + assert config.model_class_name == "BooguImageTurboPipeline" + assert DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo + + +def test_edit_turbo_without_override_keeps_the_base_class(tmp_path): + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "BooguImagePipeline"}), + encoding="utf-8", + ) + + config = OmniDiffusionConfig(model=str(tmp_path)) + config.enrich_config() + + assert config.model_class_name == "BooguImagePipeline" + assert not DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 0b2ed3a57..72e4ab962 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -2401,6 +2401,42 @@ def test_forward_transfer_runs_multichunk_overlap_path( torch.testing.assert_close(captured["targets"][1][:, :, 0], torch.full((1, 3, 16, 16), -0.2)) +def test_forward_transfer_non_output_rank_uses_canonical_envelope( + make_cosmos3_pipeline, + sequential_cfg_parallel, +) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.vae.distributed_executor = SimpleNamespace(rank=1) + pipeline.vae.is_distributed_enabled = lambda: True + pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1") + pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask()) + pipeline._set_flow_shift = lambda *_args, **_kwargs: None + decoded = torch.zeros(1, 3, 1, 16, 16) + pipeline._decode_latents = lambda latents: decoded + + request = SimpleNamespace( + prompts=[{"prompt": "transfer", "modalities": ["video"]}], + sampling_params=make_sampling_params( + height=16, + width=16, + num_inference_steps=1, + guidance_scale=1.0, + extra_args={ + "edge": {"control": torch.zeros(3, 1, 16, 16, dtype=torch.uint8)}, + "max_frames": 1, + "num_video_frames_per_chunk": 1, + }, + ), + ) + + output = pipeline.forward(request) + + assert set(output.output) == {"payload", "metadata"} + assert set(output.output["payload"]) == {"video"} + torch.testing.assert_close(output.output["payload"]["video"], decoded) + assert output.output["metadata"] == {"video": {"fps": 24.0}} + + 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 diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.py b/tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.py index 7ab04cc21..1aa653d4a 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -81,9 +81,8 @@ def test_from_config_loads_local_diffusers_component(tmp_path, monkeypatch: pyte def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: - import huggingface_hub - from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer + from vllm_omni.transformers_utils import repo_utils cache_dir = tmp_path / "hf" _write_component(cache_dir, checkpoint_name=DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME) @@ -95,7 +94,7 @@ def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pyt calls.append((repo_id, revision, allow_patterns)) return str(cache_dir) - monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(repo_utils.hf_api(), "snapshot_download", fake_snapshot_download) sound_tokenizer.Cosmos3SoundTokenizer.from_config( SimpleNamespace( diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 2bfa7b558..44bd8e9a7 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -383,6 +383,117 @@ def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None: ) +def test_gen_sp_plan_auto_pads_hidden_states_and_rope_together() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer + + prepare_plan = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"] + assert set(prepare_plan) == {0, 1, 2} + assert all(spec.auto_pad for spec in prepare_plan.values()) + assert [prepare_plan[index].split_dim for index in range(3)] == [1, 1, 1] + + +def test_gen_sp_prepare_resets_branch_local_padding_metadata() -> None: + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare + + ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=12121) + hidden = torch.zeros(1, 4, 8) + rope = torch.zeros(1, 4, 1, 2) + with override_forward_context(ctx): + output = Cosmos3GenSPPrepare()(hidden, rope, rope) + + assert output[0] is hidden + assert output[1] is rope + assert output[2] is rope + assert ctx.sp_padding_size == 0 + assert ctx.sp_original_seq_len is None + + +def test_sp_attention_masks_padded_gen_and_joint_und_keys() -> None: + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention + + class RecordingAttention(nn.Module): + def __init__(self) -> None: + super().__init__() + self.metadata: Any | None = None + + def forward(self, query, key, value, metadata): + del key, value + self.metadata = metadata + return query + + module = object.__new__(Cosmos3CrossAttention) + nn.Module.__init__(module) + module.num_heads_local = 2 + module.head_dim = 4 + recording_attention = RecordingAttention() + module.attn = recording_attention + + q = torch.zeros(1, 2, 2, 4) + k = torch.zeros_like(q) + v = torch.zeros_like(q) + k_und = torch.zeros(1, 3, 2, 4) + v_und = torch.zeros_like(k_und) + ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=3) + with override_forward_context(ctx): + output = module._forward_sp(q, k, v, k_und, v_und) + + assert output.shape == (1, 2, 8) + assert recording_attention.metadata is not None + assert recording_attention.metadata.attn_mask.tolist() == [[True, True, True, False]] + assert recording_attention.metadata.joint_attn_mask.tolist() == [[True, True, True]] + + +def test_ulysses_2d_mask_uses_key_length_with_kv_only_joint_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + from vllm_omni.diffusion.attention.parallel import ulysses + + def fake_all_to_all( + _group: object, + tensor: torch.Tensor, + _scatter_idx: int, + _gather_idx: int, + _use_sync: bool, + ) -> torch.Tensor: + return torch.cat([tensor, tensor], dim=1)[:, :, :1] + + monkeypatch.setattr(ulysses.SeqAllToAll4D, "apply", staticmethod(fake_all_to_all)) + strategy = ulysses.UlyssesParallelAttention( + sp_group=SimpleNamespace( + ulysses_group=object(), + ulysses_world_size=2, + ulysses_rank=0, + ring_world_size=1, + ), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + query = torch.zeros(1, 2, 2, 4) + key = torch.zeros_like(query) + value = torch.zeros_like(query) + joint_query = torch.zeros(1, 0, 2, 4) + joint_key = torch.zeros(1, 3, 2, 4) + metadata = AttentionMetadata( + attn_mask=torch.tensor([[True, True, True, False]]), + joint_attn_mask=torch.ones(1, 3, dtype=torch.bool), + joint_query=joint_query, + joint_key=joint_key, + joint_value=torch.zeros_like(joint_key), + joint_strategy="front", + ) + + query_out, key_out, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata) + + assert query_out.shape[1] == 4 + assert key_out.shape[1] == 7 + assert metadata_out is not None + assert metadata_out.attn_mask.tolist() == [[True, True, True, True, True, True, False]] + + def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 @@ -721,7 +832,7 @@ def test_forward_returns_video_plus_optional_modality_predictions( assert [tuple(tensor.shape) for tensor in output] == expected_shapes -def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None: +def test_forward_with_sound_odd_sequence_uses_auto_padding_contract(monkeypatch: pytest.MonkeyPatch) -> None: import vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 as cosmos3_module model = cosmos3_module.Cosmos3VFMTransformer( @@ -732,16 +843,17 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch ) monkeypatch.setattr(cosmos3_module, "_get_ulysses_state", lambda: (2, 0, None)) - with pytest.raises(ValueError, match=r"GEN sequence length \(3 = video tokens 2 \+ sound tokens 1\)"): - model( - hidden_states=torch.zeros(1, 2, 1, 1, 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, 1, 2), - fps=24.0, - sound_latents=torch.zeros(1, 3, 1), - ) + output = model( + hidden_states=torch.zeros(1, 2, 1, 1, 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, 1, 2), + fps=24.0, + sound_latents=torch.zeros(1, 3, 1), + ) + + assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 1, 2), (1, 3, 1)] def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/diffusion/models/gr00t/test_noise_seed.py b/tests/diffusion/models/gr00t/test_noise_seed.py new file mode 100644 index 000000000..a865d9a05 --- /dev/null +++ b/tests/diffusion/models/gr00t/test_noise_seed.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The request's seed must reach GR00T's flow-matching noise. + +The runner builds ``sampling_params.generator`` from ``sampling_params.seed``, the pipeline hands +that generator to the policy, and the action head draws its initial noise from it, so the same +seed reproduces the same action chunk and different seeds give independent ones. These tests pin +both ends of that plumbing. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + +from vllm_omni.diffusion.models.gr00t import pipeline_gr00t +from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config +from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7ActionHead +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] + +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", +} + + +class _OptionsRecordingPolicy: + """Stand-in for ``Gr00tPolicy``; the pipeline only needs ``language_key`` and ``get_action``.""" + + def __init__(self, *, model_path: str, embodiment_tag: str, device: str, strict: bool) -> None: + del model_path, embodiment_tag, device, strict # the pipeline's constructor contract; unused here + self.language_key = "annotation.language.language_instruction" + self.seen_options: dict[str, object] | None = None + + def get_action(self, obs: dict[str, object], options: dict[str, object] | None = None): + self.seen_options = options + return {"arm": np.zeros((1, 2, 2), dtype=np.float32)}, {} + + +@pytest.fixture +def pipeline(monkeypatch) -> Gr00tN1d7Pipeline: + monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", _OptionsRecordingPolicy) + 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 _request(**sampling_kwargs) -> OmniDiffusionRequest: + return OmniDiffusionRequest( + prompt="pick", + request_id="req", + sampling_params=OmniDiffusionSamplingParams(extra_args={"robot_obs": ROBOT_OBS}, **sampling_kwargs), + ) + + +def test_runner_generator_is_passed_through_untouched(pipeline): + """In serving the runner has already built ``sampling_params.generator``; the head must get that object.""" + generator = torch.Generator().manual_seed(99) + + assert pipeline.forward(_request(generator=generator)).error is None + + assert pipeline.policy.seen_options["generator"] is generator + + +def test_seed_without_generator_builds_one_on_the_pipeline_device(pipeline): + """Callers that bypass the runner (direct use, tests) still get the seed honoured.""" + pipeline.forward(_request(seed=1234)) + + generator = pipeline.policy.seen_options["generator"] + assert isinstance(generator, torch.Generator) + assert generator.initial_seed() == 1234 + assert generator.device.type == pipeline.device + + +def test_auto_assigned_seed_reaches_the_head(pipeline): + """``OmniDiffusionRequest`` seeds every request, so the head is never left on the global RNG.""" + req = _request() + assert req.sampling_params.seed is not None + + pipeline.forward(req) + + assert pipeline.policy.seen_options["generator"].initial_seed() == req.sampling_params.seed + + +def _captured_warnings(monkeypatch) -> list[tuple]: + """vLLM loggers do not propagate to the root logger, so ``caplog`` cannot see them; record calls instead.""" + captured: list[tuple] = [] + monkeypatch.setattr(pipeline_gr00t.logger, "warning", lambda *args, **kwargs: captured.append(args)) + return captured + + +def test_unexpected_generator_form_warns_and_falls_back_to_the_seed(pipeline, monkeypatch): + """GR00T draws one noise tensor per request, so a list of several generators cannot be honoured.""" + warnings = _captured_warnings(monkeypatch) + generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)] + + pipeline.forward(_request(generator=generators, seed=7)) + + assert len(warnings) == 1 + assert "torch.Generator" in warnings[0][0] + assert pipeline.policy.seen_options["generator"].initial_seed() == 7 + + +def test_unexpected_generator_form_without_a_seed_warns_and_leaves_the_global_rng(pipeline, monkeypatch): + warnings = _captured_warnings(monkeypatch) + generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)] + + pipeline.forward(_request(generator=generators)) + + assert len(warnings) == 1 + assert pipeline.policy.seen_options is None + + +def _noise_only_head(action_horizon: int, action_dim: int) -> Gr00tN1d7ActionHead: + """An action head whose single denoising step is the identity, so ``action_pred`` is the initial noise.""" + config = Gr00tN1d7Config( + max_action_dim=action_dim, + action_horizon=action_horizon, + num_inference_timesteps=1, + add_pos_embed=False, + use_alternate_vl_dit=False, + ) + head = Gr00tN1d7ActionHead.__new__(Gr00tN1d7ActionHead) + nn.Module.__init__(head) + head.config = config + head.action_dim = config.max_action_dim + head.action_horizon = config.action_horizon + head.num_inference_timesteps = config.num_inference_timesteps + head.num_timestep_buckets = config.num_timestep_buckets + head.action_encoder = lambda actions, timesteps, embodiment_id: actions + head.model = lambda hidden_states, encoder_hidden_states, timestep: hidden_states + head.action_decoder = lambda hidden_states, embodiment_id: torch.zeros_like(hidden_states) + return head + + +def _predict(head: Gr00tN1d7ActionHead, options) -> torch.Tensor: + batch, seq, dim = 2, 3, head.action_dim + return head.get_action_with_features( + backbone_features=torch.zeros(batch, seq, dim), + state_features=torch.zeros(batch, 1, dim), + embodiment_id=torch.zeros(batch, dtype=torch.long), + backbone_output=BatchFeature(data={}), + action_input=BatchFeature(data={}), + options=options, + )["action_pred"] + + +def test_head_draws_its_noise_from_the_request_generator(): + head = _noise_only_head(action_horizon=4, action_dim=3) + expected = torch.randn(2, 4, 3, generator=torch.Generator().manual_seed(5)) + + assert torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(5)}), expected) + assert not torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(6)}), expected) + + +def test_head_without_generator_stays_on_the_global_rng(): + """Upstream Isaac-GR00T parity: no generator means a bare ``torch.randn`` on the global RNG.""" + head = _noise_only_head(action_horizon=4, action_dim=3) + + torch.manual_seed(5) + expected = torch.randn(2, 4, 3) + torch.manual_seed(5) + + assert torch.equal(_predict(head, None), expected) diff --git a/tests/diffusion/models/gr00t/test_pipeline.py b/tests/diffusion/models/gr00t/test_pipeline.py index 322112107..3310d049e 100644 --- a/tests/diffusion/models/gr00t/test_pipeline.py +++ b/tests/diffusion/models/gr00t/test_pipeline.py @@ -1,4 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace +from typing import Any import numpy as np import pytest @@ -12,12 +16,12 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class FakeGr00tPolicy: - instances = [] + instances: list["FakeGr00tPolicy"] = [] def __init__(self, **kwargs): self.kwargs = kwargs self.reset_calls = 0 - self.seen_obs = None + self.seen_obs: dict[str, Any] | None = None self.embodiment_tag = SimpleNamespace(value="fake_embodiment") self.language_key = "annotation.language.language_instruction" self.modality_configs = { @@ -40,7 +44,7 @@ class FakeGr00tPolicy: ) FakeGr00tPolicy.instances.append(self) - def get_action(self, obs): + def get_action(self, obs, options=None): self.seen_obs = obs return { "arm": np.array([[[1.0, 2.0]]], dtype=np.float64), @@ -107,6 +111,7 @@ def test_forward_returns_dict_actions_in_output(): 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 policy.seen_obs is not None 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 diff --git a/tests/diffusion/models/gr00t/test_processing_eval_image_transform.py b/tests/diffusion/models/gr00t/test_processing_eval_image_transform.py new file mode 100644 index 000000000..452273185 --- /dev/null +++ b/tests/diffusion/models/gr00t/test_processing_eval_image_transform.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""GR00T-N1.7 eval image transform: chosen from ``processor_config.json`` and bit-exact with Isaac-GR00T. + +``GOLDEN`` holds sha256 digests of the ``(C, H, W)`` uint8 output of Isaac-GR00T's own eval transform +(``build_image_transformations_albumentations(...)[1]`` applied through ``apply_with_replay``) for the +synthetic frames built below, at Isaac-GR00T 51d4c89 / albumentations 1.4.18 / opencv 5.0.0.93. The frame +generators are integer-only so they are identical on every platform. + +Regenerating a digest needs Isaac-GR00T installed, whose transformers pin conflicts with this repo's, so it +cannot run here: in an Isaac-GR00T environment push each ``synthetic_frames()`` frame through that transform +with ``letter_box_transform`` False and True and hash the output with ``digest()``. +""" + +from __future__ import annotations + +import hashlib +import json + +import numpy as np +import pytest +import torch +import torchvision.transforms.v2 as transforms + +from vllm_omni.diffusion.models.gr00t.modeling import processing_gr00t_n1d7 as processing +from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import ( + AlbumentationsEvalImageTransform, + Gr00tN1d7Processor, + LetterBoxTransform, + _build_eval_image_transform, + _fractional_center_crop, + _smallest_edge_resize, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# nvidia/GR00T-N1.7-3B and nvidia/GR00T-N1.7-DROID both ship exactly these processor_kwargs. +TARGET, CROP, EDGE, FRACTION = [256, 256], [230, 230], 256, 0.95 +MODALITY_CONFIGS = { + "new_embodiment": { + "video": {"delta_indices": [0], "modality_keys": ["cam"]}, + "state": {"delta_indices": [0], "modality_keys": ["joint"]}, + "action": {"delta_indices": [0, 1], "modality_keys": ["joint"]}, + } +} + +# Isaac-GR00T eval transform output for synthetic_frames(); key = "|letterbox=". +GOLDEN: dict[str, tuple[tuple[int, int, int], str]] = { + "gradient_256x256|letterbox=False": ( + (3, 256, 256), + "4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715", + ), + "checker_256x256|letterbox=False": ( + (3, 256, 256), + "d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c", + ), + "noise_256x256|letterbox=False": ( + (3, 256, 256), + "a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6", + ), + "noise_512x768|letterbox=False": ( + (3, 256, 383), + "45bc3a309f3d06c3f31051911bc8a544ff3f067b023b999dd9c2b7e6029ecc93", + ), + "gradient_256x256|letterbox=True": ( + (3, 256, 256), + "4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715", + ), + "checker_256x256|letterbox=True": ( + (3, 256, 256), + "d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c", + ), + "noise_256x256|letterbox=True": ((3, 256, 256), "a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6"), + "noise_512x768|letterbox=True": ((3, 256, 256), "52a3db185e6ee4f00617b87711dd198b8f2e2ad53865713d34e56225a4ec61a2"), +} + + +def _hash_noise(height: int, width: int, seed: int) -> np.ndarray: + """Integer-only pseudo-noise (murmur-style mixing on uint32); no RNG or float, so identical everywhere.""" + ys, xs = np.mgrid[0:height, 0:width].astype(np.uint32) + channels = [] + for channel in range(3): + v = xs * np.uint32(73856093) ^ ys * np.uint32(19349663) ^ np.uint32((seed * 3 + channel) * 83492791) + v ^= v >> np.uint32(13) + v *= np.uint32(0x5BD1E995) + v ^= v >> np.uint32(15) + channels.append((v & np.uint32(0xFF)).astype(np.uint8)) + return np.stack(channels, axis=-1) + + +def synthetic_frames() -> dict[str, np.ndarray]: + """Deterministic HWC uint8 RGB frames: smooth, hard-edged, noisy, and a non-square noisy one. + + The non-square frame is 512x768 so both INTER_AREA downscales (2x without letterbox, 3x with it) + are integer factors, which OpenCV computes with exact integer sums: the goldens do not depend on + the CPU architecture. The 243->256 upscale is a fixed-point bilinear on uint8, also exact. + """ + ys, xs = np.mgrid[0:256, 0:256].astype(np.int64) + frames = { + "gradient_256x256": np.stack([xs % 256, ys % 256, (xs + ys) // 2 % 256], axis=-1).astype(np.uint8), + "checker_256x256": np.stack( + [((ys // 16 + xs // 16) % 2) * 255, ((ys // 8 + xs // 8) % 2) * 255, ((ys // 32) % 2) * 255], + axis=-1, + ).astype(np.uint8), + "noise_256x256": _hash_noise(256, 256, seed=1), + "noise_512x768": _hash_noise(512, 768, seed=2), + } + return frames + + +def digest(tensor: torch.Tensor) -> str: + return hashlib.sha256(np.ascontiguousarray(tensor.numpy()).tobytes()).hexdigest() + + +class _FakeTokenizer: + padding_side = "right" + + +class _FakeVlmProcessor: + """The two attributes ``Gr00tN1d7Processor`` and its collator touch on the Qwen3-VL processor.""" + + def __init__(self) -> None: + self.tokenizer = _FakeTokenizer() + + def apply_chat_template(self, *_args: object, **_kwargs: object) -> str: + return "" + + +@pytest.fixture(autouse=True) +def fake_vlm_processor(monkeypatch: pytest.MonkeyPatch) -> None: + """The real Qwen3-VL processor is a gated download and irrelevant to the image pipeline.""" + monkeypatch.setattr(processing, "build_processor", lambda _model_name, _kwargs: _FakeVlmProcessor()) + + +def _processor(**kwargs) -> Gr00tN1d7Processor: + return Gr00tN1d7Processor( + modality_configs=MODALITY_CONFIGS, image_crop_size=CROP, image_target_size=TARGET, **kwargs + ) + + +def _has_letterbox(compose: transforms.Compose) -> bool: + return any(isinstance(step, LetterBoxTransform) for step in compose.transforms) + + +def test_processor_builds_the_pipeline_named_by_processor_config(): + albumentations = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION) + transform = albumentations.eval_image_transform + assert isinstance(transform, AlbumentationsEvalImageTransform) + assert (transform.max_size, transform.fraction, transform.letter_box_transform) == (EDGE, FRACTION, False) + assert albumentations.use_albumentations is True + + torchvision = _processor(use_albumentations=False) + assert isinstance(torchvision.eval_image_transform, transforms.Compose) + assert not _has_letterbox(torchvision.eval_image_transform) + + letterboxed = _processor(use_albumentations=False, letter_box_transform=True) + assert _has_letterbox(letterboxed.eval_image_transform) + assert _processor(use_albumentations=True, letter_box_transform=True).eval_image_transform.letter_box_transform + + +def test_legacy_config_without_the_flags_keeps_the_torchvision_recipe(): + """Checkpoints predating the flags behave like Isaac-GR00T's defaults; training-only keys stay ignored.""" + processor = _processor(state_dropout_prob=0.2, extra_augmentation_config=None) + assert processor.use_albumentations is False + assert processor.letter_box_transform is False + assert isinstance(processor.eval_image_transform, transforms.Compose) + assert not _has_letterbox(processor.eval_image_transform) + + +def test_fraction_and_max_size_fall_back_like_isaac_gr00t(): + transform = _build_eval_image_transform( + TARGET, CROP, shortest_image_edge=None, crop_fraction=None, use_albumentations=True + ) + assert isinstance(transform, AlbumentationsEvalImageTransform) + assert transform.fraction == CROP[0] / TARGET[0] + assert transform.max_size == TARGET[0] + with pytest.raises(ValueError, match="crop_fraction is None"): + _build_eval_image_transform(TARGET, None, crop_fraction=None, use_albumentations=True) + + +def test_geometry_follows_isaac_rounding_rules(): + frame = _hash_noise(540, 640, seed=3) + resized = _smallest_edge_resize(frame, EDGE) + assert resized.shape == (256, 303, 3) # round() per dimension, shortest edge -> 256 + cropped = _fractional_center_crop(resized, FRACTION) + assert cropped.shape == (243, 287, 3) # int() truncation: 243.2 -> 243, 287.85 -> 287 + assert np.array_equal(cropped, resized[6:249, 8:295]) # (256-243)//2, (303-287)//2 + square = frame[:256, :256] + assert _smallest_edge_resize(square, EDGE) is square # scale == 1.0 is a no-op, not a resample + + +@pytest.mark.parametrize("case", sorted(GOLDEN)) +def test_albumentations_recipe_is_bit_exact_with_isaac_gr00t(case: str): + name, flag = case.split("|letterbox=") + transform = _build_eval_image_transform( + TARGET, + CROP, + shortest_image_edge=EDGE, + crop_fraction=FRACTION, + use_albumentations=True, + letter_box_transform=(flag == "True"), + ) + out = transform(synthetic_frames()[name]) + expected_shape, expected_digest = GOLDEN[case] + assert out.dtype == torch.uint8 + assert tuple(out.shape) == expected_shape + assert digest(out) == expected_digest + + +def test_letterbox_is_gated_on_both_recipes(): + frame = _hash_noise(540, 640, seed=3) + + def albumentations(letterbox: bool) -> torch.Tensor: + return _build_eval_image_transform( + TARGET, + CROP, + shortest_image_edge=EDGE, + crop_fraction=FRACTION, + use_albumentations=True, + letter_box_transform=letterbox, + )(frame) + + padded, unpadded = albumentations(True), albumentations(False) + assert tuple(padded.shape) == (3, 256, 256) # 540x640 -> 640x640 -> 256x256 -> 243x243 -> 256x256 + assert tuple(unpadded.shape) == (3, 256, 302) # aspect preserved: 256x303 -> 243x287 -> 256x302 + + tv_padded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=True)(frame) + tv_unpadded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=False)(frame) + assert tuple(tv_padded.shape) == tuple(tv_unpadded.shape) == (3, 256, 256) + assert not torch.equal(tv_padded, tv_unpadded) + assert (tv_padded[:, :2, :] == 0).all() and not (tv_unpadded[:, :2, :] == 0).all() # black bars only when gated + + +def test_albumentations_recipe_rejects_frames_outside_the_contract(): + transform = _build_eval_image_transform(TARGET, CROP, use_albumentations=True) + with pytest.raises(ValueError, match="HWC uint8"): + transform(np.zeros((256, 256, 3), dtype=np.float32)) + with pytest.raises(ValueError, match="HWC uint8"): + transform(np.zeros((3, 256, 256), dtype=np.uint8)) + with pytest.raises(ValueError, match="HWC uint8"): + transform(np.zeros((256, 256), dtype=np.uint8)) + + +def test_get_vlm_inputs_runs_the_selected_transform_per_frame(): + processor = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION) + frames = np.stack([_hash_noise(540, 640, seed=3)] * 2) # (T, H, W, C) like policy.py hands over + vlm_inputs = processor._get_vlm_inputs( + image_keys=["cam"], + images={"cam": frames}, + masks=None, + image_transform=processor.eval_image_transform, + language="x", + ) + images = vlm_inputs["vlm_content"]["images"] + assert len(images) == 2 + assert all(img.size == (302, 256) for img in images) # PIL reports (W, H); no letterbox for this checkpoint + + +def test_flags_round_trip_through_save_pretrained_and_from_pretrained(tmp_path): + saved = _processor( + use_albumentations=True, letter_box_transform=True, shortest_image_edge=EDGE, crop_fraction=FRACTION + ) + saved.save_pretrained(tmp_path) + kwargs = json.loads((tmp_path / "processor_config.json").read_text())["processor_kwargs"] + assert kwargs["use_albumentations"] is True + assert kwargs["letter_box_transform"] is True + + loaded = Gr00tN1d7Processor.from_pretrained(tmp_path) + assert loaded.use_albumentations is True + assert loaded.letter_box_transform is True + assert isinstance(loaded.eval_image_transform, AlbumentationsEvalImageTransform) + assert loaded.eval_image_transform.letter_box_transform is True + + +def test_from_pretrained_defaults_missing_flags_to_isaac_gr00t_defaults(tmp_path): + (tmp_path / "processor_config.json").write_text( + json.dumps( + { + "processor_class": "Gr00tN1d7Processor", + "processor_kwargs": { + "modality_configs": MODALITY_CONFIGS, + "image_crop_size": CROP, + "image_target_size": TARGET, + }, + } + ) + ) + (tmp_path / "statistics.json").write_text("{}") + loaded = Gr00tN1d7Processor.from_pretrained(tmp_path) + assert loaded.use_albumentations is False + assert loaded.letter_box_transform is False + assert isinstance(loaded.eval_image_transform, transforms.Compose) diff --git a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py index 950906190..ae90ea1a3 100644 --- a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py @@ -5,6 +5,7 @@ import json from dataclasses import replace +from threading import Lock from types import SimpleNamespace from typing import Any @@ -500,6 +501,32 @@ def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypa assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE +def test_ltx_runtime_profiles_phase_methods_separately(monkeypatch): + pipe = object.__new__(LTX2DistilledTwoStagePipeline) + pipe.enable_diffusion_pipeline_profiler = True + pipe._profiler_lock = Lock() + pipe._stage_durations = {} + monkeypatch.setattr(ltx2_runtime.current_omni_platform, "is_available", lambda: False) + + with pipe._profile_phase_method("run_phase", "generate_lowres"): + pass + with pipe._profile_phase_method("run_phase", "refine"): + pass + with pipe._profile_phase_method("decode_phase"): + pass + output = DiffusionOutput() + assert pipe._attach_profiled_stage_durations(output) is output + + expected_keys = { + "LTX2DistilledTwoStagePipeline.generate_lowres.run_phase", + "LTX2DistilledTwoStagePipeline.refine.run_phase", + "LTX2DistilledTwoStagePipeline.decode_phase", + } + assert set(pipe.stage_durations) == expected_keys + assert set(output.stage_durations) == expected_keys + assert all(duration >= 0.0 for duration in pipe.stage_durations.values()) + + def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path): from vllm_omni.diffusion.data import OmniDiffusionConfig diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py index fbac0bfe5..724498f7c 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_latents.py +++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py @@ -103,7 +103,7 @@ def test_prepare_audio_latents_samples_directly_in_packed_token_space(): torch.testing.assert_close(actual, expected, rtol=0, atol=0) -def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(): +def test_prepare_audio_latents_pads_short_audio_for_sp(): pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2) latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py index 0573dfb37..272cdc7d8 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py +++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for LTX guidance and forward-parallel behavior.""" @@ -85,14 +85,14 @@ class TestCFGParallelHelpers: LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size) - def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch): + def test_guidance_parallel_preserves_full_guidance_plan(self, monkeypatch): from vllm_omni.diffusion.models.ltx2 import ltx2_runtime from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2) pipe = object.__new__(LTX2Pipeline) - req = SimpleNamespace(is_dummy_run=lambda: True) + req = SimpleNamespace() request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance) guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None) diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index c2fe7cbfa..a7f6b95df 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -10,6 +10,7 @@ from dataclasses import replace from types import SimpleNamespace from typing import Any +import huggingface_hub import numpy as np import pytest import torch @@ -175,7 +176,19 @@ def test_ltx_device_uint8_output_rejects_invalid_contract(video, error): _prepare_ltx2_video_output(video, do_normalize=True) -def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch): +@pytest.mark.parametrize( + "parallel_config", + [ + None, + SimpleNamespace(ulysses_degree=4), + SimpleNamespace(tensor_parallel_size=2), + SimpleNamespace(use_hsdp=True, hsdp_shard_size=2), + SimpleNamespace(ulysses_degree=2, tensor_parallel_size=2), + SimpleNamespace(ulysses_degree=2, use_hsdp=True, hsdp_shard_size=2), + ], + ids=["single", "ulysses", "tp", "hsdp", "tp-ulysses", "hsdp-ulysses"], +) +def test_ltx_supported_parallel_configs_keep_output_profiling(tmp_path, monkeypatch, parallel_config): from vllm_omni.diffusion.models.ltx2 import ltx2_runtime (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]})) @@ -193,7 +206,13 @@ def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch): monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", capture_profiler_setup) - LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False)) + LTX2Pipeline( + od_config=SimpleNamespace( + model=str(tmp_path), + parallel_config=parallel_config, + enable_diffusion_pipeline_profiler=False, + ) + ) assert "video_processor.postprocess_video" in profiler_setup["profiler_targets"] assert "_prepare_video_output_for_transport" in profiler_setup["profiler_targets"] @@ -284,11 +303,11 @@ def test_ltx_artifact_uses_source_revision_and_hub_fallback( filename = "ltx-sidecar.safetensors" calls = [] - def fake_download(**kwargs): + def fake_download(self, **kwargs): calls.append(kwargs) return "/cache/ltx-sidecar.safetensors" - monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download) + monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download) assert ( resolve_ltx_artifact( @@ -313,7 +332,9 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch): filename = "ltx-sidecar.safetensors" expected = tmp_path / filename expected.write_bytes(b"sidecar") - monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup")) + monkeypatch.setattr( + huggingface_hub.HfApi, "hf_hub_download", lambda *_args, **_kwargs: pytest.fail("unexpected Hub lookup") + ) assert resolve_ltx_artifact( str(tmp_path), @@ -327,11 +348,11 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch): def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch): calls = [] - def fake_download(**kwargs): + def fake_download(self, **kwargs): calls.append(kwargs) return "/cache/ltx-sidecar.safetensors" - monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download) + monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download) assert ( resolve_ltx_artifact( @@ -383,14 +404,42 @@ def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expecte ) -def test_ltx_rejects_advanced_uaa_before_component_initialization(): +@pytest.mark.parametrize( + "pipeline_cls", + [ + LTX2Pipeline, + LTX2DistilledOneStagePipeline, + LTX2TwoStagePipeline, + LTX2DistilledTwoStagePipeline, + LTX2T2VDMD2Pipeline, + LTX2I2VDMD2Pipeline, + ], +) +@pytest.mark.parametrize( + ("parallel_config", "error"), + [ + (SimpleNamespace(ulysses_mode="advanced_uaa"), "does not support ulysses_mode='advanced_uaa'"), + (SimpleNamespace(ring_degree=2), "pure Ulysses sequence parallelism only"), + (SimpleNamespace(ulysses_degree=2, ring_degree=2), "pure Ulysses sequence parallelism only"), + (SimpleNamespace(allgather_degree=2), "pure Ulysses sequence parallelism only"), + ], + ids=["advanced-uaa", "ring", "hybrid", "allgather-kv"], +) +def test_ltx_rejects_unsupported_sp_before_component_initialization(pipeline_cls, parallel_config, error, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + def unexpected_model_access(*_args, **_kwargs): + pytest.fail("Unsupported SP must be rejected before accessing model files or initializing components") + + monkeypatch.setattr(ltx2_runtime, "detect_ltx_model_version", unexpected_model_access) + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", unexpected_model_access) od_config = SimpleNamespace( model="unused", - parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"), + parallel_config=parallel_config, ) - with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"): - LTX2Pipeline(od_config=od_config) + with pytest.raises(ValueError, match=error): + pipeline_cls(od_config=od_config) def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch): @@ -1409,39 +1458,19 @@ def test_ltx_distilled_forward_rejects_fixed_recipe_overrides(direct_kwargs, sam pipeline.forward(req, **direct_kwargs) -def test_ltx_distilled_dummy_run_uses_fixed_recipe_values(): - from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest - from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - req = DiffusionRequestBatch( - [ - OmniDiffusionRequest( - prompt="dummy run", - sampling_params=OmniDiffusionSamplingParams( - height=512, - width=512, - num_frames=LTX2DistilledPipeline.dummy_run_num_frames, - num_inference_steps=1, - guidance_scale=0.0, - ), - request_id=DUMMY_DIFFUSION_REQUEST_ID, - ) - ] - ) - pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline) - seen = {} - - def fake_run_recipe(_req, request_inputs, **_kwargs): - seen["request_inputs"] = request_inputs - return "dummy-output" - - object.__setattr__(pipeline, "_run_recipe", fake_run_recipe) - - assert pipeline.forward(req) == "dummy-output" - assert seen["request_inputs"].num_frames == 1 - assert seen["request_inputs"].num_inference_steps == 8 - assert seen["request_inputs"].guidance == LTXGuidanceSpec.positive_only() +@pytest.mark.parametrize( + "pipeline_cls", + [ + LTX2Pipeline, + LTX2DistilledOneStagePipeline, + LTX2TwoStagePipeline, + LTX2DistilledTwoStagePipeline, + LTX2T2VDMD2Pipeline, + LTX2I2VDMD2Pipeline, + ], +) +def test_ltx_pipelines_disable_startup_dummy_run(pipeline_cls): + assert pipeline_cls.dummy_run_num_frames == 0 def test_ltx_custom_sigmas_bypass_scheduler_shifting(): diff --git a/tests/diffusion/models/ltx2/test_ltx2_sequence_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_sequence_parallel.py new file mode 100644 index 000000000..395b53ba2 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_sequence_parallel.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.models.ltx2.ltx2_sequence_parallel import ( + LTX2VideoToAudioParallelAttention, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _run_video_to_audio_parity(rank: int, world_size: int, master_port: int) -> None: + dist.init_process_group( + "gloo", + init_method=f"tcp://127.0.0.1:{master_port}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=60), + ) + try: + # Non-contiguous SP subgroups match TP-fastest rank ordering. Different + # seeds across groups catch collectives that accidentally use WORLD. + group_count = world_size // 2 + ulysses_group = dist.group.WORLD + if group_count > 1: + for offset in range(group_count): + ranks = list(range(offset, world_size, group_count)) + group = dist.new_group(ranks) + if rank in ranks: + ulysses_group = group + ulysses_rank = dist.get_rank(ulysses_group) + sp_group = SimpleNamespace( + ring_world_size=1, + ulysses_group=ulysses_group, + ulysses_world_size=2, + ulysses_rank=ulysses_rank, + ) + strategy = LTX2VideoToAudioParallelAttention(sp_group) + + for batch_size, masked in [(1, False), (2, True)]: + torch.manual_seed(17 + rank % group_count) + query = torch.randn(batch_size, 3, 4, 8) + global_key = torch.randn(batch_size, 6, 4, 8) + global_value = torch.randn(batch_size, 6, 4, 8) + key = global_key.chunk(2, dim=1)[ulysses_rank].contiguous() + value = global_value.chunk(2, dim=1)[ulysses_rank].contiguous() + mask = None + if masked: + mask = torch.ones(batch_size, 1, 1, 6, dtype=torch.bool) + mask[0, ..., -1] = False + mask[1, ..., -2:] = False + metadata = AttentionMetadata(attn_mask=mask) if masked else None + + local_query, local_key, local_value, local_metadata, ctx = strategy.pre_attention( + query, key, value, metadata + ) + assert local_metadata is metadata + local_output = F.scaled_dot_product_attention( + local_query.transpose(1, 2), + local_key.transpose(1, 2), + local_value.transpose(1, 2), + attn_mask=local_metadata.attn_mask if local_metadata is not None else None, + ).transpose(1, 2) + actual = strategy.post_attention(local_output, ctx) + expected = F.scaled_dot_product_attention( + query.transpose(1, 2), + global_key.transpose(1, 2), + global_value.transpose(1, 2), + attn_mask=mask, + ).transpose(1, 2) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world_size", [2, 4], ids=["world", "subgroups"]) +def test_ltx_video_to_audio_sp_matches_replicated_attention(unused_tcp_port, world_size): + torch.multiprocessing.spawn( + _run_video_to_audio_parity, + args=(world_size, unused_tcp_port), + nprocs=world_size, + ) + + +def test_ltx_video_to_audio_sp_rejects_ring_parallelism(): + sp_group = SimpleNamespace(ring_world_size=2) + + with pytest.raises(NotImplementedError, match="ring_degree must be 1"): + LTX2VideoToAudioParallelAttention(sp_group) diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index 937d487e6..f12fbc466 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -141,16 +141,54 @@ def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks(): @pytest.mark.parametrize("rope_type", ["interleaved", "split"]) -def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type): - root_plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)[""] +def test_ltx_sp_plan_only_shards_video_stream(rope_type): + plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type) + root_plan = plan[""] - video_timestep = root_plan["timestep"] - audio_timestep = root_plan["audio_timestep"] + assert set(root_plan) == {"hidden_states", "keyframes_mask", "timestep"} + assert root_plan["hidden_states"].split_dim == 1 + assert root_plan["keyframes_mask"].split_dim == 1 + assert root_plan["timestep"].split_dim == 1 + assert root_plan["timestep"].expected_dims == 2 + assert not root_plan["timestep"].split_output + assert set(plan) == {"", "rope", "cross_attn_rope", "proj_out"} - assert video_timestep.split_dim == audio_timestep.split_dim == 1 - assert video_timestep.expected_dims == audio_timestep.expected_dims == 2 - assert not video_timestep.split_output - assert not audio_timestep.split_output + +def test_ltx_attention_assigns_video_only_sp_modes(monkeypatch): + modes = {} + + class FakeAttention(nn.Module): + def __init__(self, *, prefix, sequence_parallel_mode="local", **_kwargs): + super().__init__() + modes[prefix.rsplit(".", 1)[-1]] = sequence_parallel_mode + + class FakeFeedForward(nn.Module): + def __init__(self, *_args, **_kwargs): + super().__init__() + + monkeypatch.setattr(ltx2_transformer, "LTX2Attention", FakeAttention) + monkeypatch.setattr(ltx2_transformer, "LTX2FeedForward", FakeFeedForward) + + LTX2VideoTransformerBlock( + dim=8, + num_attention_heads=2, + attention_head_dim=4, + cross_attention_dim=8, + audio_dim=8, + audio_num_attention_heads=2, + audio_attention_head_dim=4, + audio_cross_attention_dim=8, + prefix="transformer_blocks.0", + ) + + assert modes == { + "attn1": "video_self", + "audio_attn1": "local", + "attn2": "local", + "audio_attn2": "local", + "audio_to_video_attn": "local", + "video_to_audio_attn": "video_to_audio", + } def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch): diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py b/tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py new file mode 100644 index 000000000..0e080e015 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Full-transformer numeric parity for LTX Ulysses sequence parallelism. + +This FP32/SDPA check verifies that Ulysses preserves the complete transformer +mathematical path. It does not guarantee pixel-level E2E parity: production +BF16 and compiled attention kernels can change rounding order, and diffusion +can amplify that drift across denoising steps. +""" + +from __future__ import annotations + +import os + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + init_distributed_environment, + initialize_model_parallel, +) +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig +from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context +from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel +from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.parallel] + +_FP32_REL_L2 = 1e-5 +_TINY_CONFIG = { + "in_channels": 4, + "out_channels": 4, + "num_attention_heads": 2, + "attention_head_dim": 4, + "cross_attention_dim": 8, + "vae_scale_factors": (1, 1, 1), + "pos_embed_max_pos": 4, + "base_height": 2, + "base_width": 2, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 4, + "audio_cross_attention_dim": 8, + "audio_scale_factor": 1, + "audio_pos_embed_max_pos": 4, + "audio_sampling_rate": 16, + "audio_hop_length": 4, + "num_layers": 1, + "caption_channels": 8, + "rope_type": "split", + "use_prompt_embeddings": False, + "use_keyframes_abs_pos_embedding": True, +} + + +def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float: + return ((actual - expected).double().norm() / expected.double().norm()).item() + + +def _omni_config(sp_size: int) -> OmniDiffusionConfig: + return OmniDiffusionConfig.from_kwargs( + model="test", + dtype=torch.float32, + parallel_config=DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + sequence_parallel_size=sp_size, + ulysses_degree=sp_size, + ring_degree=1, + cfg_parallel_size=1, + ), + diffusion_attention_backend="TORCH_SDPA", + ) + + +def _inputs(device: torch.device) -> dict[str, torch.Tensor | int | bool]: + generator = torch.Generator().manual_seed(123) + return { + "hidden_states": torch.randn(1, 8, 4, generator=generator, device="cpu").to(device), + "audio_hidden_states": torch.randn(1, 4, 4, generator=generator, device="cpu").to(device), + "encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device), + "audio_encoder_hidden_states": torch.randn(1, 3, 8, generator=generator, device="cpu").to(device), + "timestep": torch.linspace(0.1, 0.8, 8, device=device).view(1, 8), + "audio_timestep": torch.linspace(0.2, 0.5, 4, device=device).view(1, 4), + "keyframes_mask": torch.tensor([[[1.0], [0.0], [1.0], [0.0], [0.0], [1.0], [0.0], [0.0]]], device=device), + "sigma": torch.tensor([0.1], device=device), + "audio_sigma": torch.tensor([0.2], device=device), + "num_frames": 2, + "height": 2, + "width": 2, + "audio_num_frames": 4, + "return_dict": False, + } + + +def _run_transformer_parity(rank: int, world_size: int, master_port: int) -> None: + device = torch.device(f"{current_omni_platform.device_type}:{rank}") + current_omni_platform.set_device(device) + os.environ.update( + { + "RANK": str(rank), + "LOCAL_RANK": str(rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": str(master_port), + } + ) + + try: + init_distributed_environment() + initialize_model_parallel(sequence_parallel_size=world_size, ulysses_degree=world_size) + + sp_config = _omni_config(world_size) + torch.manual_seed(42) + with set_forward_context(omni_diffusion_config=sp_config), set_current_diffusion_config(sp_config): + model = LTX2VideoTransformer3DModel(**_TINY_CONFIG).to(device).eval() + with torch.no_grad(): + for _, parameter in sorted(model.named_parameters()): + torch.nn.init.normal_(parameter, mean=0.0, std=0.02) + + inputs = _inputs(device) + reference_config = _omni_config(1) + with torch.no_grad(), set_forward_context(omni_diffusion_config=reference_config): + expected_video, expected_audio = model(**inputs) + + apply_sequence_parallel( + model, + SequenceParallelConfig(ulysses_degree=world_size, ring_degree=1), + model._sp_plan, + ) + with torch.no_grad(), set_forward_context(omni_diffusion_config=sp_config): + get_forward_context().sp_plan_hooks_applied = True + actual_video, actual_audio = model(**inputs) + + v2a_attention = model.transformer_blocks[0].video_to_audio_attn.attn + assert v2a_attention._video_to_audio_strategy is not None + assert torch.isfinite(expected_video).all() + assert torch.isfinite(expected_audio).all() + for modality, actual, expected in ( + ("video", actual_video, expected_video), + ("audio", actual_audio, expected_audio), + ): + drift = _relative_l2(actual, expected) + assert drift <= _FP32_REL_L2, ( + f"LTX SP{world_size} {modality} output drifted rel_l2={drift:.3e} from SP1 (bound {_FP32_REL_L2:.1e})" + ) + finally: + destroy_distributed_env() + + +@pytest.mark.full_model +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2) +def test_ltx_ulysses_transformer_matches_sp1(unused_tcp_port) -> None: + """Exercise real plan hooks, RoPE/timestep sharding, and attention routing.""" + world_size = 2 + device_count = current_omni_platform.device_count() + assert device_count >= world_size, ( + f"LTX Ulysses parity requires {world_size} accelerator devices; only {device_count} are visible" + ) + torch.multiprocessing.spawn( + _run_transformer_parity, + args=(world_size, unused_tcp_port), + nprocs=world_size, + ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index cb12ce9e5..a94f842c5 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -3,6 +3,7 @@ import json import sys +from contextlib import contextmanager from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace from typing import Any, TypeVar @@ -176,9 +177,9 @@ def _write_partition_index(path, *, partition, tasks): def test_modular_diffusers_index_is_resolved_generically(tmp_path): + from vllm_omni.config.resolver import resolve_omni_config from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model - from vllm_omni.entrypoints.utils import resolve_model_config_path (tmp_path / "modular_model_index.json").write_text( json.dumps( @@ -192,7 +193,16 @@ def test_modular_diffusers_index_is_resolved_generically(tmp_path): assert is_diffusion_model(str(tmp_path)) assert resolve_model_class_name(str(tmp_path)) == "MiniMaxH3ModularPipeline" - assert resolve_model_config_path(str(tmp_path)) is None + resolved = resolve_omni_config( + str(tmp_path), + trust_remote_code=False, + deploy_config_path=None, + cli_overrides=None, + stage_overrides=None, + strategy_config_path=None, + ) + assert resolved.config_path is None + assert resolved.stage_by_id(0).engine_args.model_class_name == "MiniMaxH3ModularPipeline" config = OmniDiffusionConfig(model=str(tmp_path)) config.enrich_config() @@ -1550,6 +1560,97 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch): assert next(video_vae.parameters()).dtype == torch.float32 +def test_keyframe_encode_pins_and_restores_cudnn_settings(): + from vllm_omni.diffusion.models.minimax_h3.vae import ( + _minimax_h3_keyframe_encode_context, + ) + + cudnn = torch.backends.cudnn + original = ( + cudnn.enabled, + cudnn.benchmark, + cudnn.deterministic, + cudnn.allow_tf32, + ) + try: + cudnn.enabled = False + cudnn.benchmark = True + cudnn.deterministic = False + cudnn.allow_tf32 = False + + with _minimax_h3_keyframe_encode_context(torch.device("cuda")): + assert cudnn.enabled + assert not cudnn.benchmark + assert cudnn.deterministic + assert cudnn.allow_tf32 + + assert not cudnn.enabled + assert cudnn.benchmark + assert not cudnn.deterministic + assert not cudnn.allow_tf32 + finally: + ( + cudnn.enabled, + cudnn.benchmark, + cudnn.deterministic, + cudnn.allow_tf32, + ) = original + + +@pytest.mark.parametrize("fail_encode", [False, True]) +def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode): + from PIL import Image + + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + events = [] + input_image = Image.new("RGB", (32, 32)) + + @contextmanager + def track_context(device): + assert device == torch.device("cpu") + events.append("enter") + try: + yield + finally: + events.append("exit") + + class FakeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.parallel_tiling = True + + def encode_images(self, image, *, use_fp16_latent): + assert events == ["enter"] + assert image is input_image + assert use_fp16_latent + events.append("encode") + if fail_encode: + raise RuntimeError("encode failed") + return [torch.ones(1, 1, 2, 2)] + + video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE) + torch.nn.Module.__init__(video_vae) + video_vae.model = FakeModel() + video_vae.config_dict = { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + } + monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context) + + if fail_encode: + with pytest.raises(RuntimeError, match="encode failed"): + video_vae.encode_image(input_image) + else: + rows = video_vae.encode_image(input_image) + torch.testing.assert_close(rows, torch.ones(1, 4)) + + assert events == ["enter", "encode", "exit"] + assert video_vae.model.parallel_tiling + + def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py index dbeeaa66a..e10fbe56a 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py @@ -5,14 +5,154 @@ import pytest import torch from vllm.triton_utils import HAS_TRITON +from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( + _MAX_1D_GRID_SIZE, + _iter_row_chunks, + _launch_row_chunks, + indexed_gate, indexed_gate_rms_norm_scale_shift, + indexed_scale_shift_, rms_norm_indexed_scale_shift, ) +from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] +@pytest.mark.cpu +@pytest.mark.parametrize( + ("rows", "expected"), + [ + (0, []), + (_MAX_1D_GRID_SIZE, [(0, _MAX_1D_GRID_SIZE)]), + (_MAX_1D_GRID_SIZE + 1, [(0, _MAX_1D_GRID_SIZE), (_MAX_1D_GRID_SIZE, 1)]), + (69888, [(0, _MAX_1D_GRID_SIZE), (_MAX_1D_GRID_SIZE, 4353)]), + ], +) +def test_row_chunks_respect_ascend_grid_limit(rows, expected): + assert list(_iter_row_chunks(rows)) == expected + + +@pytest.mark.cpu +@pytest.mark.parametrize( + ("device_type", "expected_grids_and_offsets"), + [ + ("npu", [(_MAX_1D_GRID_SIZE, 0), (4353, _MAX_1D_GRID_SIZE)]), + ("cuda", [(69888, 0)]), + ], +) +def test_launch_row_chunks_is_npu_only(device_type, expected_grids_and_offsets): + class RecordingKernel: + def __init__(self): + self.calls = [] + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self.calls.append((grid, args, kwargs)) + + return launch + + kernel = RecordingKernel() + first_arg = object() + second_arg = object() + + _launch_row_chunks(kernel, 69888, device_type, first_arg, second_arg, block_n=16) + + assert kernel.calls == [ + ((grid,), (first_arg, second_arg, row_offset), {"block_n": 16}) + for grid, row_offset in expected_grids_and_offsets + ] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize( + "device", + [pytest.param("npu", marks=hardware_marks(res={"npu": "A3"}, num_cards=1))], +) +def test_modulation_wrappers_match_reference_above_grid_limit(device): + rows = 69888 + hidden_size = 16 + num_conditions = 7 + eps = 1e-6 + dtype = torch.bfloat16 + + torch.manual_seed(0) + x = torch.randn(rows, hidden_size, dtype=dtype) + residual = torch.randn_like(x) + branch = torch.randn_like(x) + weight = torch.randn(hidden_size, dtype=dtype) + shift = torch.randn(num_conditions, hidden_size, dtype=dtype) + scale = torch.randn_like(shift) + gate = torch.randn_like(shift) + indices = torch.arange(rows) % num_conditions + + indexed_shift = shift.index_select(0, indices) + indexed_scale = scale.index_select(0, indices) + indexed_gate_values = gate.index_select(0, indices) + + x_float = x.float() + branch_float = branch.float() + indexed_shift_float = indexed_shift.float() + indexed_scale_float = indexed_scale.float() + indexed_gate_float = indexed_gate_values.float() + + expected_scale_shift = (x_float * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype) + expected_gate = (x_float + indexed_gate_float * branch_float).to(dtype) + + normalized = x_float + normalized = normalized * torch.rsqrt(normalized.pow(2).mean(-1, keepdim=True) + eps) + normalized = normalized * weight.float() + expected_rms = (normalized * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype) + + updated_residual = residual.float() + indexed_gate_float * branch_float + expected_residual = updated_residual.to(dtype) + normalized_residual = updated_residual + normalized_residual = normalized_residual * torch.rsqrt(normalized_residual.pow(2).mean(-1, keepdim=True) + eps) + normalized_residual = normalized_residual * weight.float() + expected_modulated = (normalized_residual * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype) + + npu_x = x.to(device) + npu_residual = residual.to(device) + npu_branch = branch.to(device) + npu_weight = weight.to(device) + npu_shift = shift.to(device) + npu_scale = scale.to(device) + npu_gate = gate.to(device) + npu_indices = indices.to(device) + + actual_scale_shift = indexed_scale_shift_(npu_x.clone(), npu_shift, npu_scale, npu_indices) + actual_gate = indexed_gate(npu_x, npu_gate, npu_branch, npu_indices) + actual_rms = rms_norm_indexed_scale_shift( + npu_x, + npu_weight, + npu_shift, + npu_scale, + npu_indices, + eps, + ) + actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift( + npu_residual, + npu_gate, + npu_branch, + npu_weight, + npu_shift, + npu_scale, + npu_indices, + eps, + ) + + for actual, expected in ( + (actual_scale_shift, expected_scale_shift), + (actual_gate, expected_gate), + (actual_rms, expected_rms), + (actual_residual, expected_residual), + (actual_modulated, expected_modulated), + ): + torch.testing.assert_close(actual.cpu(), expected, atol=2e-2, rtol=2e-2) + + +@pytest.mark.cuda @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.skipif(not HAS_TRITON, reason="Triton required") def test_fused_modulation_preserves_bf16_residual_boundary() -> None: diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index 0fb043a1e..a45631f7d 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -42,10 +42,10 @@ _QUALITY_CONFIG = QualityTestConfig( def _resolve_fl2va_model_ref() -> str: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api repo_root = Path( - snapshot_download( + hf_api().snapshot_download( repo_id=_MINIMAX_H3_REPO, revision=_MINIMAX_H3_REVISION, allow_patterns=["FL2VA/**"], @@ -253,13 +253,13 @@ def test_resolve_fl2va_model_ref(tmp_path, monkeypatch): fl2va_root.mkdir() (fl2va_root / "model_index.json").write_text("{}", encoding="utf-8") - def fake_snapshot_download(*, repo_id, revision, allow_patterns): + def fake_snapshot_download(self, *, repo_id, revision, allow_patterns): assert repo_id == _MINIMAX_H3_REPO assert revision == _MINIMAX_H3_REVISION assert allow_patterns == ["FL2VA/**"] return str(tmp_path) - monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot_download) + monkeypatch.setattr("huggingface_hub.HfApi.snapshot_download", fake_snapshot_download) assert _resolve_fl2va_model_ref() == str(fl2va_root) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py index 3ed7649c6..699dd47d6 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py @@ -49,3 +49,29 @@ def test_qwen3_vl_npu_swiglu_matches_two_projection_reference() -> None: atol=2e-2, rtol=2e-2, ) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_qwen3_vl_npu_swiglu_forward_uses_projection_module(mocker) -> None: + from vllm_omni.platforms.npu.models import minimax_h3 + from vllm_omni.platforms.npu.models.minimax_h3 import ( + _forward_minimax_h3_qwen3vl_text_mlp_npu, + ) + + x = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16) + packed = torch.randn(1, 16, 256, device="npu", dtype=torch.bfloat16) + activated = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16) + gate_up_proj = mocker.Mock(return_value=packed) + down_proj = mocker.Mock(side_effect=lambda value: value) + mlp = mocker.Mock(gate_up_proj=gate_up_proj, down_proj=down_proj) + swiglu = mocker.patch.object(minimax_h3, "npu_swiglu_from_packed", return_value=activated) + + output = _forward_minimax_h3_qwen3vl_text_mlp_npu(mlp, x) + + gate_up_proj.assert_called_once() + assert gate_up_proj.call_args.args[0] is x + swiglu.assert_called_once() + assert swiglu.call_args.args[0] is packed + down_proj.assert_called_once() + assert down_proj.call_args.args[0] is activated + assert output is activated diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 05d7d93b5..c24639434 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -52,6 +52,7 @@ from vllm_omni.diffusion.offloader.offload_plan import ( OffloadPlan, get_offload_plan, ) +from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state from vllm_omni.host_weight_runtime import MappedHostRegion from vllm_omni.platforms import current_omni_platform @@ -2813,18 +2814,21 @@ class TestDistributedComponentSelection: def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self): """Every rank must reject an unsafe encoder group, including stub ranks.""" - backend = DistributedLayerwiseOffloadBackend( - OffloadConfig( - strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, - pin_cpu_memory=False, - dp_size=2, - components=frozenset({"text_encoder"}), - dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, - ), - torch.device("cpu"), + + class StubEncoderPipeline(nn.Module): + _offload_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)}) + + def __init__(self): + super().__init__() + self.text_encoder = nn.Module() + + config = OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"text_encoder"}), + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, ) - backend.dp_group = object() - stub_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)}) with pytest.raises(ValueError, match="not declared replicated"): - backend._try_layerwise_offload_encoder(nn.Module(), "text_encoder", stub_plan) + resolve_offload_plan(StubEncoderPipeline(), config) diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 2dd8804cf..099e93f79 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -488,7 +488,8 @@ class TestLayerwiseComponentSelection: with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"): backend.enable(pipeline) - assert not pipeline.text_encoder._omni_layerwise_enabled + # The plan resolver rejects the topology before any encoder is hooked. + assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False) def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime): pipeline = _LegacyComponentPipeline() diff --git a/tests/diffusion/offloader/test_plan_resolver.py b/tests/diffusion/offloader/test_plan_resolver.py new file mode 100644 index 000000000..483ce4060 --- /dev/null +++ b/tests/diffusion/offloader/test_plan_resolver.py @@ -0,0 +1,458 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Unit tests for the shared offload plan resolver.""" + +from typing import ClassVar + +import pytest +import torch +from torch import nn + +from tests.diffusion.offloader.helpers import ( + _DummyBlock, + _PlainEncoder, + _SingleBlockModel, + _StagedEncoder, + _StagedVAE, + patch_offload_runtime, +) +from vllm_omni.diffusion.offloader import layerwise_backend as layerwise_backend_module +from vllm_omni.diffusion.offloader import plan_resolver as plan_resolver_module +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan +from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + +CPU = torch.device("cpu") + + +@pytest.fixture +def patched_offload_runtime(monkeypatch): + patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform) + + +def _config( + components: set[str] | None = None, + strategy: OffloadStrategy = OffloadStrategy.LAYER_WISE, + **kwargs, +) -> OffloadConfig: + return OffloadConfig( + strategy=strategy, + pin_cpu_memory=False, + components=None if components is None else frozenset(components), + **kwargs, + ) + + +def _stack_of(component): + assert len(component.stacks) == 1 + return component.stacks[0] + + +class _WanStylePipeline(nn.Module): + """Streamable DiT plus a plan-declared, replicated text encoder.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + encoder_dlo_weight_replication=frozenset({"text_encoder"}), + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=4) + self.text_encoder = _PlainEncoder() + self.vae = _StagedVAE() + + +class _StagedComponentPipeline(nn.Module): + """MiniMax-H3 shape: two encoder stacks plus pipeline-staged components.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, + on_demand_component_paths=frozenset({"text_encoder", "vae"}), + resident_dit_paths=frozenset({"transformer"}), + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=4) + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + +class _LegacyScanPipeline(nn.Module): + """No SupportsComponentDiscovery declaration and no OffloadPlan.""" + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=3) + self.text_encoder = _PlainEncoder() + + +class TestProducerPrecedence: + def test_plan_block_attrs_win_over_class_attributes(self): + class TwoStackDiT(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([_DummyBlock() for _ in range(2)]) + self.alt_blocks = nn.ModuleList([_DummyBlock() for _ in range(3)]) + + class Pipeline(nn.Module): + _offload_plan = OffloadPlan(block_attrs={"transformer": ("alt_blocks",)}) + + def __init__(self): + super().__init__() + self.transformer = TwoStackDiT() + + pipeline = Pipeline() + resolved = resolve_offload_plan(pipeline, _config({"dit"})) + + stack = _stack_of(resolved.dits[0]) + assert stack.attrs == ("alt_blocks",) + assert list(stack.blocks) == list(pipeline.transformer.alt_blocks) + + def test_class_attributes_are_used_without_a_plan(self): + pipeline = _LegacyScanPipeline() + resolved = resolve_offload_plan(pipeline, _config({"dit"})) + + assert list(_stack_of(resolved.dits[0]).blocks) == list(pipeline.transformer.blocks) + + def test_declared_component_type_wins_over_leaf_name(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["mllm"] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={"mllm": "text_encoder"}, + encoder_block_attrs={"mllm": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.mllm = _PlainEncoder() + + resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"})) + + encoder = resolved.encoders[0] + assert encoder.path == "mllm" + assert encoder.selected + assert len(_stack_of(encoder).blocks) == 2 + + def test_component_declaration_wins_over_attribute_scan(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["custom_dit"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + + def __init__(self): + super().__init__() + self.custom_dit = _SingleBlockModel(num_blocks=2) + # Would be picked up by the deprecated attribute scan. + self.transformer = _SingleBlockModel(num_blocks=2) + + resolved = resolve_offload_plan(Pipeline(), _config({"dit"})) + + assert [component.path for component in resolved.dits] == ["custom_dit"] + + def test_attribute_scan_warns_once_per_pipeline_class(self, monkeypatch): + warnings: list[str] = [] + + class _Recorder: + def warning(self, message, *args): + warnings.append(message % args if args else message) + + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + monkeypatch.setattr(plan_resolver_module, "logger", _Recorder()) + + class UndeclaredPipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + + resolve_offload_plan(UndeclaredPipeline(), _config({"dit"})) + resolve_offload_plan(UndeclaredPipeline(), _config({"dit"})) + + scan_warnings = [message for message in warnings if "SupportsComponentDiscovery" in message] + assert len(scan_warnings) == 1 + assert "deprecated" in scan_warnings[0] + + +class TestResolvedTopology: + def test_encoder_plan_resolves_one_stack_per_declared_path(self): + pipeline = _WanStylePipeline() + resolved = resolve_offload_plan(pipeline, _config({"dit", "text_encoder"})) + + encoder = resolved.encoders[0] + assert encoder.selected and not encoder.on_demand + assert list(_stack_of(encoder).blocks) == list(pipeline.text_encoder.encoder.block) + assert [component.path for component in resolved.vaes] == ["vae"] + + def test_staged_plan_resolves_every_declared_stack(self): + pipeline = _StagedComponentPipeline() + resolved = resolve_offload_plan(pipeline, _config({"text_encoder"})) + + encoder = resolved.encoders[0] + assert encoder.on_demand + assert [len(stack.blocks) for stack in encoder.stacks] == [2, 2] + # An explicit selection never stages a VAE through the legacy plan. + assert not resolved.vaes[0].on_demand + # DiT is unselected, so it keeps no streaming topology. + assert resolved.dits[0].selected is False + assert resolved.dits[0].stacks == () + + def test_omitted_selector_keeps_the_legacy_staged_vae(self): + resolved = resolve_offload_plan(_StagedComponentPipeline(), _config()) + + assert resolved.vaes[0].on_demand + assert resolved.encoders[0].on_demand + assert resolved.dits[0].selected + + def test_on_demand_only_plan_resolves_without_stacks(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder", "vae"})) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"})) + + assert resolved.encoders[0].on_demand + assert resolved.encoders[0].stacks == () + + def test_module_strategy_resolves_selection_without_stacks(self): + resolved = resolve_offload_plan( + _WanStylePipeline(), + _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL), + ) + + assert [component.path for component in resolved.components] == ["transformer", "text_encoder", "vae"] + assert all(component.stacks == () for component in resolved.components) + assert resolved.encoders[0].selected + + def test_module_strategy_leaves_staged_residency_to_the_pipeline(self): + resolved = resolve_offload_plan( + _StagedComponentPipeline(), + _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL), + ) + + assert not any(component.on_demand for component in resolved.components) + + +class TestResidentLayers: + def _dlo_config(self, resident_layers: int, components: set[str] | None = None) -> OffloadConfig: + return _config( + components, + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + dlo_transfers={"dit": "rank-local", "text_encoder": "rank-local"}, + dlo_resident_layers=resident_layers, + ) + + def test_resident_head_splits_the_declared_dit(self): + pipeline = _StagedComponentPipeline() + resolved = resolve_offload_plan(pipeline, self._dlo_config(2, {"dit"})) + + stack = _stack_of(resolved.dits[0]) + assert stack.resident_head == 2 + assert list(stack.resident) == list(pipeline.transformer.blocks)[:2] + assert list(stack.streaming) == list(pipeline.transformer.blocks)[2:] + + def test_resident_count_beyond_the_stack_keeps_every_block(self): + resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(9, {"dit"})) + + stack = _stack_of(resolved.dits[0]) + assert stack.resident_head == 4 + assert stack.streaming == () + + def test_resident_head_leaving_one_streaming_block_is_rejected(self): + with pytest.raises(ValueError, match="leaves only one streaming block"): + resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3, {"dit"})) + + def test_legacy_single_streaming_block_stays_resident(self): + resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3)) + + stack = _stack_of(resolved.dits[0]) + assert stack.resident_head == 4 + assert stack.streaming == () + + def test_resident_layers_without_a_declared_path_is_rejected(self): + with pytest.raises(ValueError, match="no matching resident_dit_paths"): + resolve_offload_plan(_WanStylePipeline(), self._dlo_config(2, {"dit"})) + + +class TestValidation: + def test_selected_dit_requires_a_discovered_module(self): + pipeline = nn.Module() + pipeline.text_encoder = _PlainEncoder() + + with pytest.raises(ValueError, match="No DiT/transformer modules found"): + resolve_offload_plan(pipeline, _config({"dit"})) + + def test_model_level_selection_requires_both_swap_sides(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + + with pytest.raises(ValueError, match="requires an encoder execution stage"): + resolve_offload_plan(pipeline, _config({"dit"}, strategy=OffloadStrategy.MODEL_LEVEL)) + + def test_selected_text_encoder_requires_a_matching_module(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + + with pytest.raises(ValueError, match="No text encoder modules found"): + resolve_offload_plan(pipeline, _config({"text_encoder"})) + + def test_selected_encoder_requires_a_streamable_or_staged_plan(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + pipeline.text_encoder = _PlainEncoder() + + with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"): + resolve_offload_plan(pipeline, _config({"text_encoder"})) + + def test_staged_component_requires_the_explicit_lifecycle(self): + class Pipeline(nn.Module): + _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder"})) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _PlainEncoder() + + with pytest.raises(ValueError, match="must implement load_to_device"): + resolve_offload_plan(Pipeline(), _config({"text_encoder"})) + + def test_single_block_dit_is_rejected_for_layer_offload(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=1) + + with pytest.raises(ValueError, match="requires at least two streamable"): + resolve_offload_plan(pipeline, _config({"dit"})) + + def test_encoder_allgather_requires_declared_replication(self): + config = _config( + {"text_encoder"}, + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + dp_size=2, + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, + ) + + with pytest.raises(ValueError, match="not declared replicated"): + resolve_offload_plan(_StagedComponentPipeline(), config) + + def test_two_dits_cannot_own_the_same_blocks(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.transformer_2 = _SingleBlockModel(num_blocks=2) + self.transformer_2.blocks = self.transformer.blocks + + with pytest.raises(ValueError, match="claimed by both"): + resolve_offload_plan(Pipeline(), _config({"dit"})) + + def test_legacy_selection_keeps_unstreamable_components(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=1) + pipeline.text_encoder = _PlainEncoder() + + resolved = resolve_offload_plan(pipeline, _config()) + + assert resolved.dits[0].stacks == () + assert resolved.encoders[0].stacks == () + + +class TestPurity: + def _snapshot(self, pipeline: nn.Module): + return { + "devices": [(name, tensor.device) for name, tensor in pipeline.named_parameters()], + "attrs": {name: sorted(vars(module)) for name, module in pipeline.named_modules()}, + } + + def test_resolution_does_not_touch_the_model(self): + pipeline = _StagedComponentPipeline() + before = self._snapshot(pipeline) + + resolve_offload_plan(pipeline, _config({"dit", "text_encoder"})) + + assert self._snapshot(pipeline) == before + assert pipeline.text_encoder.offload_calls == 0 + assert pipeline.text_encoder.to_calls == 0 + assert pipeline.vae.to_calls == 0 + + def test_resolution_is_repeatable(self): + pipeline = _WanStylePipeline() + config = _config({"dit", "text_encoder"}) + + first = resolve_offload_plan(pipeline, config) + second = resolve_offload_plan(pipeline, config) + + assert [component.path for component in first.components] == [component.path for component in second.components] + for left, right in zip(first.components, second.components, strict=True): + assert left.selected == right.selected + assert left.on_demand == right.on_demand + assert [list(stack.blocks) for stack in left.stacks] == [list(stack.blocks) for stack in right.stacks] + + def test_backend_enable_rejects_before_touching_earlier_components(self, patched_offload_runtime): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={ + "text_encoder": "text_encoder", + "text_encoder_2": "text_encoder", + }, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _PlainEncoder() + self.text_encoder_2 = _PlainEncoder() + self.vae = _StagedVAE() + + pipeline = Pipeline() + backend = LayerWiseOffloadBackend(_config({"text_encoder"}), CPU) + + with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"): + backend.enable(pipeline) + + assert not backend.enabled + assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False) + for block in pipeline.text_encoder.encoder.block: + assert getattr(block, "_hook_registry", None) is None + assert pipeline.vae.to_calls == 0 diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 164fcfbeb..394f7c04f 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -131,7 +131,7 @@ class TestCreateDefaultDiffusion: def test_dtype_serialized_as_string(self): stages = StageConfigFactory.create_default_diffusion({"dtype": torch.float16, "model": "x"}) - assert stages[0]["engine_args"]["dtype"] == "torch.float16" + assert stages[0]["engine_args"]["dtype"] == "float16" def test_cache_backend_defaults_to_none(self): stages = StageConfigFactory.create_default_diffusion({"model": "x"}) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 94c3eb48f..28bba6e4d 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -15,6 +15,7 @@ from pytest_mock import MockerFixture import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test +from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode @@ -46,7 +47,6 @@ from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunn from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState -from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.errors import OmniClientError from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -1182,7 +1182,7 @@ class TestSupportedPipelines: """Step-execution protocol checks for supported pipelines.""" def test_default_stage_config_includes_step_execution(self): - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "step_execution": True, } diff --git a/tests/diffusion/utils/test_chunked_mp4.py b/tests/diffusion/utils/test_chunked_mp4.py new file mode 100644 index 000000000..e7ef465f3 --- /dev/null +++ b/tests/diffusion/utils/test_chunked_mp4.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import hashlib + +import av +import numpy as np +import pytest + +from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder, mux_av_video_audio_bytes + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _frames() -> np.ndarray: + rng = np.random.default_rng(6872) + return rng.integers(0, 256, size=(8, 16, 24, 3), dtype=np.uint8) + + +def test_chunked_mp4_is_byte_identical_to_whole_mux() -> None: + frames = _frames() + baseline = mux_av_video_audio_bytes( + (av.VideoFrame.from_ndarray(frame, format="rgb24") for frame in frames), + width=24, + height=16, + fps=24, + video_codec_options={"threads": "1"}, + ) + encoder = ChunkedMP4Encoder( + width=24, + height=16, + fps=24, + max_pending=2, + video_codec_options={"threads": "1"}, + ) + for start in (0, 1, 5): + stop = {0: 1, 1: 5, 5: len(frames)}[start] + encoder.push(np.ascontiguousarray(frames[start:stop])) + chunked = encoder.finish() + assert chunked == baseline + assert hashlib.sha256(chunked).digest() == hashlib.sha256(baseline).digest() + + audio = np.zeros((2, 320), dtype=np.float32) + baseline_audio = mux_av_video_audio_bytes( + (av.VideoFrame.from_ndarray(frame, format="rgb24") for frame in frames), + width=24, + height=16, + fps=24, + audio_waveform=audio, + audio_sample_rate=32000, + video_codec_options={"threads": "1"}, + ) + encoder_audio = ChunkedMP4Encoder( + width=24, + height=16, + fps=24, + audio_waveform=audio, + audio_sample_rate=32000, + video_codec_options={"threads": "1"}, + ) + encoder_audio.push(frames) + assert encoder_audio.finish() == baseline_audio + + +def test_chunked_mp4_abort_joins_worker() -> None: + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + encoder.push(_frames()[:1]) + encoder.abort() + assert not encoder._thread.is_alive() + with pytest.raises(RuntimeError, match="already closed"): + encoder.push(_frames()[:1]) + + +def test_chunked_mp4_validates_shape_and_dtype() -> None: + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + with pytest.raises(ValueError, match="shape"): + encoder.push(np.zeros((16, 24, 3), dtype=np.uint8)) + with pytest.raises(ValueError, match="dtype"): + encoder.push(np.zeros((1, 16, 24, 3), dtype=np.float32)) + encoder.abort() + + +def test_chunked_mp4_close_returns_encoded_bytes() -> None: + import io + + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + encoder.push(_frames()) + result = encoder.close() + assert result == encoder.finish() == encoder.close() + with av.open(io.BytesIO(result)) as container: + assert len(list(container.decode(video=0))) == len(_frames()) + + +def test_chunked_mp4_failure_during_push(monkeypatch) -> None: + from threading import Event + + from vllm_omni.diffusion.utils import media_utils + + fail = Event() + draining = Event() + error = RuntimeError("encoding failed") + + def mux(*args, **kwargs): + assert fail.wait(5) + raise error + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", mux) + drain = ChunkedMP4Encoder._drain_until_done + + def notified_drain(self): + draining.set() + drain(self) + + monkeypatch.setattr(ChunkedMP4Encoder, "_drain_until_done", notified_drain) + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + put = encoder._queue.put + + def fail_before_put(item): + fail.set() + assert draining.wait(5) + put(item) + + monkeypatch.setattr(encoder._queue, "put", fail_before_put) + try: + with pytest.raises(RuntimeError, match="encoding failed") as exc: + encoder.push(_frames()) + assert exc.value is error + with pytest.raises(RuntimeError, match="encoding failed"): + encoder.finish() + finally: + fail.set() + encoder.abort() + assert not encoder._thread.is_alive() + + +def test_chunked_mp4_flush_failure_does_not_hang(monkeypatch) -> None: + from threading import Thread + + from vllm_omni.diffusion.utils import media_utils + + error = RuntimeError("flush failed") + + def mux(frames, **kwargs): + list(frames) + raise error + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", mux) + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24) + encoder.push(_frames()) + errors = [] + + def finish(): + try: + encoder.finish() + except RuntimeError as exc: + errors.append(exc) + + closer = Thread(target=finish, daemon=True) + closer.start() + closer.join(5) + try: + assert not closer.is_alive(), "finish hung after the worker consumed its sentinel" + assert errors == [error] + finally: + if closer.is_alive(): + encoder._send_done() + closer.join(5) + + +def test_chunked_mp4_abort_discards_pending_chunks(monkeypatch) -> None: + from threading import Event, Thread + + from vllm_omni.diffusion.utils import media_utils + + entered = Event() + release = Event() + done_sent = Event() + encoded = [] + + def mux(frames, **kwargs): + encoded.append(next(frames)) + entered.set() + assert release.wait(5) + encoded.extend(frames) + return b"unused" + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", mux) + encoder = ChunkedMP4Encoder(width=24, height=16, fps=24, max_pending=1) + send_done = encoder._send_done + + def notified_done(): + send_done() + done_sent.set() + + monkeypatch.setattr(encoder, "_send_done", notified_done) + encoder.push(_frames()[:1]) + assert entered.wait(5) + encoder.push(_frames()[1:]) + closer = Thread(target=encoder.abort, daemon=True) + closer.start() + try: + assert done_sent.wait(5), "abort waited for encoding before discarding the pending queue" + finally: + release.set() + closer.join(5) + assert not closer.is_alive() + assert len(encoded) == 1 + with pytest.raises(RuntimeError, match="aborted"): + encoder.finish() diff --git a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py index 63ce73923..c1b083495 100644 --- a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py +++ b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E accuracy guard against a pinned Lightricks LTX pipeline revision. @@ -22,11 +22,11 @@ from typing import Literal import numpy as np import pytest import torch -from huggingface_hub import hf_hub_download, snapshot_download from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure from tests.e2e.accuracy.helpers import reset_artifact_dir from tests.helpers.mark import hardware_test +from vllm_omni.transformers_utils.repo_utils import hf_api OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git" OFFICIAL_REVISION = "9377758131b1ffde4b7f766804590a6617bf2ab9" @@ -326,7 +326,7 @@ def _resolve_model(case: LTXAccuracyCase) -> Path: if revision is None and model_id == case.model_id: revision = case.model_revision return Path( - snapshot_download( + hf_api().snapshot_download( repo_id=model_id, revision=revision, allow_patterns=[ @@ -359,7 +359,7 @@ def _resolve_gemma_root(case: LTXAccuracyCase, model: Path) -> Path: if configured_model and Path(configured_model).is_dir(): return Path(configured_model) return Path( - snapshot_download( + hf_api().snapshot_download( repo_id=case.gemma_model_id, revision=case.gemma_model_revision, allow_patterns=[ @@ -383,7 +383,7 @@ def _resolve_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path: if model_path.is_file(): return model_path return Path( - hf_hub_download( + hf_api().hf_hub_download( repo_id=artifact.repo_id, repo_type=None if artifact.repo_type == "model" else artifact.repo_type, filename=artifact.filename, diff --git a/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py index dab5feef1..8d8062092 100644 --- a/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py +++ b/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E test that follows the EXACT same flow as ``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py`` @@ -35,7 +35,6 @@ from uuid import uuid4 import pytest import ray import torch -from huggingface_hub import snapshot_download from omegaconf import OmegaConf from transformers import AutoTokenizer @@ -46,6 +45,7 @@ from tests.e2e.features.helpers.verl_omni_server import ( vLLMOmniHttpServerLocal, ) from tests.helpers.mark import hardware_test +from vllm_omni.transformers_utils.repo_utils import hf_api MODEL = "tiny-random/Qwen-Image" TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct" @@ -55,7 +55,7 @@ _MIN_PROMPT_TOKENS = 35 def _resolve_model_path(repo_id: str) -> str: if os.path.isdir(repo_id): return repo_id - return snapshot_download(repo_id=repo_id) + return hf_api().snapshot_download(repo_id=repo_id) @lru_cache(maxsize=1) diff --git a/tests/e2e/offline_inference/test_audio8_tts.py b/tests/e2e/offline_inference/test_audio8_tts.py new file mode 100644 index 000000000..687fff4ae --- /dev/null +++ b/tests/e2e/offline_inference/test_audio8_tts.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""E2E offline tests for Audio8 TTS Preview (DualAR, 44.1 kHz codec). + +The prompt is pre-tokenized here (not built from ``mm_processor_kwargs``) +because Audio8 TTS shares Fish Speech's protocol: text-only prompts are plain +token ids, and voice-clone prompts are a placeholder of the exact final length +whose real embeddings are built model-side from the encoded reference audio. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner, OmniRunnerHandler +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = os.environ.get("AUDIO8_TTS_MODEL_PATH", "Audio8/Audio8-TTS-Preview-0.6b") +DEPLOY_CONFIG = get_deploy_config_path("audio8_tts.yaml") +SAMPLE_RATE = 44100 +TEXT = "The weather is nice today, perfect for a walk in the park." + +# Four distinct inputs for the concurrency test: distinct text is required so a +# leaked per-request codec buffer (which would make two requests decode to the +# same audio) actually fails the pairwise-difference assertion. +CONCURRENT_TEXTS = [ + "The weather is nice today, perfect for a walk in the park.", + "She sold seashells by the seashore all through the summer.", + "Our train departs at a quarter past nine tomorrow morning.", + "He planted rows of tomatoes and basil behind the old house.", +] + +# The checkpoint registers its own `arktts` config in vllm-omni, so remote code +# must stay disabled (transformers would otherwise prefer the checkpoint's +# auto_map and bypass the Qwen2-shaped view Qwen2Model needs). +_OMNI_RUNNER_PARAM = (MODEL, DEPLOY_CONFIG, {"trust_remote_code": False}) + +pytestmark = pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True) + + +def _text_only_prompt(text: str = TEXT) -> dict: + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.audio8_tts.prompt_utils import build_text_only_prompt_ids + + tokenizer = AutoTokenizer.from_pretrained(MODEL) + prompt_ids, normalized_text = build_text_only_prompt_ids(tokenizer, text) + return {"prompt_token_ids": prompt_ids, "additional_information": {"text": [normalized_text]}} + + +# At ``core_model`` (L1/L2) the harness patches every stage to +# ``load_format: dummy``, so the AR runs on random weights: it never emits EOS +# and always generates to ``max_tokens`` (1024 frames ~= 47.5 s). Duration +# bounds are therefore only meaningful once real weights load, at +# ``advanced_model`` / ``full_model``. +_REAL_WEIGHT_RUN_LEVELS = frozenset({"advanced_model", "full_model"}) + + +def _duration_bounds(run_level: str) -> dict[str, float]: + if run_level not in _REAL_WEIGHT_RUN_LEVELS: + return {} + return {"min_duration_s": 0.5, "max_duration_s": 30.0} + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_text_to_audio_001(omni_runner: OmniRunner, run_level: str) -> None: + """Default deploy, single request, text-only synthesis. + + Deploy Setting: audio8_tts.yaml + Input Modal: text + Output Modal: audio + """ + prompt = _text_only_prompt() + request_config = { + "input": TEXT, + "prompt_token_ids": prompt["prompt_token_ids"], + "additional_information": prompt["additional_information"], + "response_format": "wav", + "expected_sample_rate": SAMPLE_RATE, + **_duration_bounds(run_level), + } + OmniRunnerHandler(omni_runner).send_tokenized_tts_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_text_to_audio_concurrent_002(omni_runner: OmniRunner, run_level: str) -> None: + """Four concurrent requests with distinct text: stage 0 runs ``max_num_seqs: 4``. + + Per-request codec state is keyed by request id. The four inputs are + deliberately distinct and the decoded outputs are asserted pairwise + different, so leaked codec state (which would make two requests decode to + identical audio) fails here, and only here. + """ + prompts = [_text_only_prompt(text) for text in CONCURRENT_TEXTS] + request_config = { + "input": CONCURRENT_TEXTS[0], + "prompt_token_ids": [p["prompt_token_ids"] for p in prompts], + "additional_information": [p["additional_information"] for p in prompts], + "response_format": "wav", + "expected_sample_rate": SAMPLE_RATE, + "assert_distinct_outputs": True, + **_duration_bounds(run_level), + } + OmniRunnerHandler(omni_runner).send_tokenized_tts_request(request_config) diff --git a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py index 7e57433ef..dd7c8ec85 100644 --- a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py +++ b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Offline E2E smoke test for CosyVoice3 zero-shot reference inference. @@ -20,7 +20,6 @@ from pathlib import Path import numpy as np import pytest import soundfile as sf -from huggingface_hub import snapshot_download from vllm.sampling_params import SamplingParams from tests.helpers.mark import hardware_test @@ -29,6 +28,7 @@ from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config +from vllm_omni.transformers_utils.repo_utils import hf_api MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512" MODEL_DIR_ENV = "VLLM_OMNI_COSYVOICE3_MODEL_DIR" @@ -58,7 +58,7 @@ def _resolve_model_dir() -> Path: override = os.environ.get(MODEL_DIR_ENV) if override: return Path(override).expanduser().resolve() - return Path(snapshot_download(MODEL, allow_patterns=["*"])) + return Path(hf_api().snapshot_download(MODEL, allow_patterns=["*"])) def _reference_zero_shot_stage0_sampling(*, text: str) -> SamplingParams: diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py index f570a6ea7..4f21dd1c4 100644 --- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py +++ b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ End-to-end test for MammothModa2 text-to-image generation. @@ -19,12 +22,12 @@ from pathlib import Path import pytest import torch -from huggingface_hub import snapshot_download from vllm.sampling_params 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 +from vllm_omni.transformers_utils.repo_utils import hf_api # --------------------------------------------------------------------------- # Constants @@ -64,7 +67,7 @@ _PIXEL_SAMPLE_COORDS = [ # Helpers # --------------------------------------------------------------------------- def _load_t2i_gen_config(repo_id: str) -> dict: - weights_dir = Path(snapshot_download(repo_id)) + weights_dir = Path(hf_api().snapshot_download(repo_id)) cfg_path = weights_dir / "t2i_generation_config.json" if not cfg_path.exists(): pytest.skip(f"t2i_generation_config.json not found at {cfg_path}") diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py index 34d36118d..81d97cd74 100644 --- a/tests/e2e/offline_inference/test_moss_tts_realtime.py +++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B). Uses the standard omni_runner + pytestmark pattern (one module-scoped engine @@ -160,12 +160,13 @@ def _build_request(ref_audio_path: str, text: str) -> dict: 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 + from vllm_omni.transformers_utils.repo_utils import hf_api + # Step 1: locate the realtime processor module in the snapshot. try: - snap_dir = Path(snapshot_download(repo_id=MODEL)) + snap_dir = Path(hf_api().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"): diff --git a/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py b/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py index 432124e76..66f8be616 100644 --- a/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py +++ b/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py @@ -12,7 +12,6 @@ from pathlib import Path from types import SimpleNamespace import pytest -from huggingface_hub import snapshot_download from tests.helpers.runtime import OmniServerParams, get_model_prefix from tests.helpers.stage_config import ( @@ -20,6 +19,7 @@ from tests.helpers.stage_config import ( get_deploy_duplex_max_sessions, modify_stage_config, ) +from vllm_omni.transformers_utils.repo_utils import hf_api MODEL = "openbmb/MiniCPM-o-4_5" DEPLOY_CONFIG_REL = "minicpmo_4_5.yaml" @@ -111,7 +111,7 @@ def resolve_ref_audio() -> Path: model_prefix = get_model_prefix() model_root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL) if not model_root.is_dir(): - model_root = Path(snapshot_download(MODEL, local_files_only=True)) + model_root = Path(hf_api().snapshot_download(MODEL, local_files_only=True)) ref_audio = model_root / REF_AUDIO_RELATIVE_PATH if not ref_audio.is_file(): raise FileNotFoundError(f"MiniCPM-o checkpoint ref audio is missing: {ref_audio}") diff --git a/tests/e2e/online_serving/test_audio8_tts.py b/tests/e2e/online_serving/test_audio8_tts.py new file mode 100644 index 000000000..1ef725d68 --- /dev/null +++ b/tests/e2e/online_serving/test_audio8_tts.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""E2E online tests for Audio8 TTS Preview via /v1/audio/speech. + +Real model inference (no mocks): text-only synthesis, streaming, and zero-shot +voice cloning. The reference audio is inlined as a base64 data URL because CI +hosts cannot fetch external URLs. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = os.environ.get("AUDIO8_TTS_MODEL_PATH", "Audio8/Audio8-TTS-Preview-0.6b") +DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0 +# The codec decodes at 44.1 kHz; raw PCM carries no header, so assertions that +# depend on the rate (HNR pitch search) have to be told explicitly. +SAMPLE_RATE = 44100 +MAX_CONCURRENT = 4 + +# ~0.5 s of 44.1 kHz mono PCM_16 WAV. +_MIN_AUDIO_BYTES = 40_000 + +# Reused vendored asset (3.5 s, Chinese) plus its exact transcript: Audio8 TTS +# requires the reference transcript to match the recording. +REF_AUDIO_URL = get_asset_path("cosyvoice3/zero_shot_prompt.wav", as_data_url=True) +REF_TEXT = "希望你以后能够做的比我还好呦。" + + +def get_prompt(prompt_type: str = "en") -> str: + prompts = { + "en": "The weather is nice today, perfect for a walk in the park.", + "zh": "今天天气很好，非常适合去公园散步。", + } + return prompts.get(prompt_type, prompts["en"]) + + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("audio8_tts.yaml"), + # No --trust-remote-code: vllm-omni owns the `arktts` config. + server_args=["--disable-log-stats"], + ), + id="audio8_tts", + ) +] + + +@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_text_to_audio_001(omni_server, openai_client) -> None: + """Baseline smoke: default deploy, non-streaming WAV. + + Deploy Setting: audio8_tts.yaml + Input Modal: text + Output Modal: audio + Input Setting: stream=False + Datasets: few requests (max_num_seqs=4) + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) + + +@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_text_to_audio_streaming_002(omni_server, openai_client) -> None: + """Streaming PCM: exercises the async_chunk path end to end. + + Deploy Setting: audio8_tts.yaml (async_chunk: true) + Input Setting: stream=True, response_format=pcm + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": True, + "stream_format": "audio", + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "pcm", + "min_audio_bytes": _MIN_AUDIO_BYTES, + # Raw PCM has no header: without this the harness assumes 24 kHz and the + # HNR pitch search lands on 147-735 Hz instead of 80-400 Hz, scoring + # clean 44.1 kHz speech ~1.9 dB too low. + "expected_sample_rate": SAMPLE_RATE, + } + openai_client.send_audio_speech_request(request_config) + + +@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_003(omni_server, openai_client) -> None: + """Zero-shot voice cloning: the reference audio is encoded model-side and + its codec codes are spliced into the prompt. + + Input Modal: text + ref_audio + ref_text + """ + request_config = { + "model": omni_server.model, + "input": get_prompt("zh"), + "ref_audio": REF_AUDIO_URL, + "ref_text": REF_TEXT, + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + openai_client.send_audio_speech_request(request_config) + + +@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_requires_ref_text_004(omni_server, openai_client) -> None: + """``ref_audio`` without ``ref_text`` must be rejected, not silently + synthesized with an untethered voice.""" + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "ref_audio": REF_AUDIO_URL, + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "status_code": (400, 422), + "err_message": "ref_text", + } + openai_client.send_audio_speech_request(request_config) + + +@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_empty_input_is_rejected_005(omni_server, openai_client) -> None: + request_config = { + "model": omni_server.model, + "input": " ", + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "status_code": (400, 422), + "err_message": "Input text cannot be empty", + } + openai_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_gr00t_openpi_expansion.py b/tests/e2e/online_serving/test_gr00t_openpi_expansion.py index 70e11088d..2f81ab3fa 100644 --- a/tests/e2e/online_serving/test_gr00t_openpi_expansion.py +++ b/tests/e2e/online_serving/test_gr00t_openpi_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """End-to-end online serving test for GR00T N1.7 through the OpenPI robot endpoint.""" import os @@ -29,8 +29,7 @@ test_params = [ 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"}, + env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1"}, init_timeout=1200, stage_init_timeout=900, ), @@ -162,14 +161,14 @@ def test_gr00t_n1d7_openpi_online(omni_server, openai_client) -> None: @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, openai_client) -> None: - """Assert actions match Isaac-GR00T reference (GR00T_NOISE_SEED=42, zero inputs).""" + """Assert actions match Isaac-GR00T reference (request seed=42, zero inputs).""" response = openai_client.send_robot_openpi_ws_request( { "operations": [ {"endpoint": "reset", "payload": {}}, { "endpoint": "infer", - "payload": build_openpi_droid_observation(session_id="gr00t-precision-e2e"), + "payload": build_openpi_droid_observation(session_id="gr00t-precision-e2e", seed=42), }, ], } diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py index 5b35c4b53..c3c6fd90a 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py @@ -30,10 +30,27 @@ from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import run_soft_interrupt, ) from tests.helpers.mark import hardware_test +from vllm_omni.entrypoints.duplex.server_vad import ( + SILERO_VAD_FILENAME, + SILERO_VAD_REPO_ID, + SILERO_VAD_REVISION, +) pytestmark = [pytest.mark.full_model, pytest.mark.omni] +@pytest.fixture(scope="module") +def _cached_server_vad_artifact() -> None: + """Cache the pinned ONNX artifact before the Realtime server starts.""" + from huggingface_hub import hf_hub_download + + hf_hub_download( + repo_id=SILERO_VAD_REPO_ID, + filename=SILERO_VAD_FILENAME, + revision=SILERO_VAD_REVISION, + ) + + @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) def test_duplex_admission_and_expiry_reaper(omni_server, tmp_path: Path) -> None: @@ -97,7 +114,7 @@ def test_duplex_soft_interrupt(omni_server, tmp_path: Path) -> None: @hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) -def test_duplex_server_vad_hard_interrupt(omni_server) -> None: +def test_duplex_server_vad_hard_interrupt(_cached_server_vad_artifact, omni_server) -> None: result = asyncio.run( run_server_vad_interrupt( SimpleNamespace( diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py index 3a042082e..d77405a73 100644 --- a/tests/e2e/online_serving/test_omnivoice_expansion.py +++ b/tests/e2e/online_serving/test_omnivoice_expansion.py @@ -7,7 +7,10 @@ Tests verify that the OmniVoice model generates valid audio when accessed through the standard OpenAI-compatible speech API. """ +import io import os +import wave +from collections.abc import Sequence os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -15,7 +18,7 @@ import pytest from tests.helpers.mark import hardware_test from tests.helpers.media import get_asset_path -from tests.helpers.runtime import OmniServerParams +from tests.helpers.runtime import OmniResponse, OmniServerParams from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.entrypoints.openai.serving_speech import _DEFAULT_VOICE_NAME @@ -34,6 +37,10 @@ STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml") EXTRA_ARGS = [ "--trust-remote-code", "--disable-log-stats", + "--max-num-seqs", + "8", + "--request-batch-max-wait-ms", + "50", ] TEST_PARAMS = [ OmniServerParams( @@ -42,7 +49,14 @@ TEST_PARAMS = [ server_args=EXTRA_ARGS, ) ] - +STEP_EXECUTION_ARGS = EXTRA_ARGS + ["--step-execution"] +STEP_EXECUTION_PARAMS = [ + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIG, + server_args=STEP_EXECUTION_ARGS, + ) +] # Lower this in ``request_config`` via ``min_audio_bytes`` if a run produces legitimately short WAVs. _DEFAULT_MIN_AUDIO_BYTES = 5000 @@ -58,6 +72,33 @@ def get_prompt(prompt_type="text"): return prompts.get(prompt_type, prompts["text"]) +def _assert_valid_omnivoice_wav_responses(responses: Sequence[OmniResponse]) -> None: + """Validate that every concurrent response contains plausible OmniVoice audio.""" + audio_shapes: list[tuple[int, int, int, int]] = [] + for response in responses: + audio_bytes = response.audio_bytes + assert audio_bytes is not None + assert audio_bytes.startswith(b"RIFF") + with wave.open(io.BytesIO(audio_bytes), "rb") as wav_file: + num_channels = wav_file.getnchannels() + sample_width = wav_file.getsampwidth() + sample_rate = wav_file.getframerate() + num_frames = wav_file.getnframes() + assert num_channels == 1 + assert sample_width == 2 + assert sample_rate == 24000 + assert num_frames > 0 + audio_shapes.append((num_channels, sample_width, sample_rate, num_frames)) + + duration_s = num_frames / sample_rate + assert 0.1 <= duration_s <= 30.0 + + # Identical prompts use the same estimated target length. Their samples may + # differ numerically across batch layouts, but the decoded audio shape must + # remain consistent across concurrent requests. + assert len(set(audio_shapes)) == 1 + + @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) class TestOmniVoiceTTS: """E2E tests for OmniVoice TTS model.""" @@ -74,6 +115,57 @@ class TestOmniVoiceTTS: } online_client.send_audio_speech_request(request_config) + @hardware_test(res={"cuda": "L4"}, num_cards=1) + def test_speech_auto_voice_batch(self, omni_server, openai_client) -> None: + """Test concurrent request-batch TTS generation.""" + batch_request_config = { + "model": omni_server.model, + "input": get_prompt("text"), + "response_format": "wav", + "seed": 42, + "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, + } + batch_r = openai_client.send_audio_speech_request(batch_request_config, request_num=2) + assert len(batch_r) == 2 + _assert_valid_omnivoice_wav_responses(batch_r) + + +@pytest.mark.parametrize("omni_server", STEP_EXECUTION_PARAMS, indirect=True) +class TestOmniVoiceStepExecution: + """E2E tests for OmniVoice TTS model.""" + + @hardware_test(res={"cuda": "L4"}, num_cards=1) + def test_speech_auto_voice_step_execution(self, omni_server, openai_client) -> None: + """Test auto voice TTS generation (text only, no reference audio).""" + request_config = { + "model": omni_server.model, + "input": get_prompt("text"), + "response_format": "wav", + "timeout": 180.0, + "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, + "extra_params": { + "num_inference_steps": 32, + }, + } + openai_client.send_audio_speech_request(request_config) + + @hardware_test(res={"cuda": "L4"}, num_cards=1) + def test_speech_auto_voice_batch_step_execution(self, omni_server, openai_client) -> None: + """Test concurrent step-execution TTS generation.""" + batch_request_config = { + "model": omni_server.model, + "input": get_prompt("text"), + "response_format": "wav", + "seed": 42, + "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, + "extra_params": { + "num_inference_steps": 32, + }, + } + batch_r = openai_client.send_audio_speech_request(batch_request_config, request_num=2) + assert len(batch_r) == 2 + _assert_valid_omnivoice_wav_responses(batch_r) + @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) class TestOmniVoiceSeed: diff --git a/tests/e2e/online_serving/test_omnivoice_parity.py b/tests/e2e/online_serving/test_omnivoice_parity.py new file mode 100644 index 000000000..e9b267f74 --- /dev/null +++ b/tests/e2e/online_serving/test_omnivoice_parity.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cross-mode full-model parity tests for OmniVoice.""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest +import requests + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer +from tests.helpers.stage_config import get_deploy_config_path + +pytestmark = [pytest.mark.slow, pytest.mark.tts] + +MODEL = "k2-fsa/OmniVoice" +STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml") +PROMPT = "The weather is nice today, perfect for a walk in the park." + +payload = { + "model": MODEL, + "input": PROMPT, + "language": "English", + "seed": 42, + "response_format": "wav", + "extra_params": {"num_inference_steps": 32}, +} + + +def _generate_without_graph(server_args: list[str]) -> bytes: + with OmniServer( + MODEL, + server_args, + use_omni=True, + env_dict={"OMNIVOICE_CUDA_GRAPH": "0"}, + ) as server: + response = requests.post( + f"http://{server.host}:{server.port}/v1/audio/speech", + json=payload, + timeout=600, + ) + response.raise_for_status() + assert response.content.startswith(b"RIFF") + return response.content + + +def _generate_with_graph(server_args: list[str]) -> bytes: + with OmniServer( + MODEL, + server_args, + use_omni=True, + ) as server: + response = requests.post( + f"http://{server.host}:{server.port}/v1/audio/speech", + json=payload, + timeout=600, + ) + response.raise_for_status() + assert response.content.startswith(b"RIFF") + return response.content + + +def _common_args() -> list[str]: + """Return the shared B=1 server arguments for parity tests.""" + return [ + "--trust-remote-code", + "--disable-log-stats", + "--deploy-config", + STAGE_CONFIG, + "--max-num-seqs", + "1", + ] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_request_mode_and_step_execution_b1_parity_without_graph() -> None: + """B=1 eager request and step modes must produce identical seeded WAV bytes.""" + common_args = _common_args() + + request_audio = _generate_without_graph(common_args) + step_audio = _generate_without_graph([*common_args, "--step-execution", "--enforce-eager"]) + + assert request_audio == step_audio + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_request_mode_and_step_execution_b1_parity_with_graph() -> None: + """B=1 Graph request and step modes must produce identical seeded WAV bytes.""" + common_args = _common_args() + + request_audio = _generate_with_graph(common_args) + step_audio = _generate_with_graph([*common_args, "--step-execution", "--enforce-eager"]) + + assert request_audio == step_audio diff --git a/tests/engine/duplex/test_duplex_deploy_config.py b/tests/engine/duplex/test_duplex_deploy_config.py index 800070efb..e9168729c 100644 --- a/tests/engine/duplex/test_duplex_deploy_config.py +++ b/tests/engine/duplex/test_duplex_deploy_config.py @@ -33,6 +33,7 @@ def test_duplex_session_runtime_defaults_are_typed_and_immutable(tmp_path) -> No assert deploy.duplex_session.max_pending_turns_per_session == 4 assert deploy.duplex_session.max_sessions == 1 assert deploy.duplex_session.completed_append_cache_size == 256 + assert deploy.duplex_session.server_vad_model_path is None with pytest.raises(FrozenInstanceError): deploy.duplex_session.idle_ttl_s = 1.0 # type: ignore[misc] @@ -46,6 +47,29 @@ def test_duplex_session_runtime_accepts_disabled_idle_expiry(tmp_path) -> None: assert deploy.duplex_session.idle_ttl_s is None +def test_duplex_session_runtime_accepts_local_server_vad_model(tmp_path) -> None: + deploy_path = tmp_path / "duplex.yaml" + deploy_path.write_text( + "duplex_session:\n server_vad_model_path: /models/silero_vad.onnx\nstages: []\n", + encoding="utf-8", + ) + + deploy = load_deploy_config(deploy_path) + + assert deploy.duplex_session.server_vad_model_path == "/models/silero_vad.onnx" + + +def test_duplex_session_runtime_rejects_empty_server_vad_model_path(tmp_path) -> None: + deploy_path = tmp_path / "duplex.yaml" + deploy_path.write_text( + "duplex_session:\n server_vad_model_path: ''\nstages: []\n", + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="server_vad_model_path"): + load_deploy_config(deploy_path) + + @pytest.mark.parametrize( ("name", "value"), [ diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index fb8231a75..0e88a00ef 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Tests for Omni config utils. For stability, these tests should largely be invariant to the specific attributes of vLLM config except in cases where we @@ -63,7 +66,7 @@ def test_full_payload_capability_reaches_omni_model_config(monkeypatch): class ConnectorWorker: model_runner_cls = ConnectorRunner - worker_module.ConnectorWorker = ConnectorWorker + setattr(worker_module, "ConnectorWorker", ConnectorWorker) monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda _self: None) @@ -95,7 +98,7 @@ def test_stage_without_connector_configuration_accepts_plain_runner(monkeypatch) class WorkerWithoutConnector: model_runner_cls = object - worker_module.WorkerWithoutConnector = WorkerWithoutConnector + setattr(worker_module, "WorkerWithoutConnector", WorkerWithoutConnector) monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) args = OmniEngineArgs( @@ -113,7 +116,7 @@ def test_full_payload_capability_requires_selected_worker_connector(monkeypatch) class WorkerWithoutConnector: model_runner_cls = object - worker_module.WorkerWithoutConnector = WorkerWithoutConnector + setattr(worker_module, "WorkerWithoutConnector", WorkerWithoutConnector) monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) with pytest.raises(ValueError, match="does not provide an Omni connector model runner"): @@ -130,7 +133,7 @@ def test_full_payload_capability_validates_platform_selected_worker(monkeypatch) class WorkerWithoutConnector: model_runner_cls = object - worker_module.WorkerWithoutConnector = WorkerWithoutConnector + setattr(worker_module, "WorkerWithoutConnector", WorkerWithoutConnector) monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) monkeypatch.setattr( current_omni_platform, @@ -293,7 +296,7 @@ def test_remote_tokenizer_subfolder_download_does_not_report_failure(tmp_path, m baseline_config = Mock() warning = mocker.patch("vllm_omni.engine.arg_utils.logger.warning") - monkeypatch.setattr("huggingface_hub.snapshot_download", lambda *args, **kwargs: str(tmp_path)) + monkeypatch.setattr("huggingface_hub.HfApi.snapshot_download", lambda *args, **kwargs: str(tmp_path)) monkeypatch.setattr(OmniEngineArgs, "_patch_empty_hf_config", lambda *args, **kwargs: None) monkeypatch.setattr(EngineArgs, "create_model_config", lambda _self: baseline_config) monkeypatch.setattr( @@ -339,6 +342,7 @@ def test_patch_missing_local_hf_config(tmp_path): def test_non_missing_local_hf_config_error_reaches_parent_loader(tmp_path, monkeypatch, config_entry): """Non-missing config errors must reach vLLM's normal loader.""" config_path = tmp_path / "config.json" + loader_error: Exception if config_entry == "malformed": config_path.write_text("{not valid json", encoding="utf-8") loader_error = json.JSONDecodeError("invalid config", "{not valid json", 1) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 528e46f3a..b1771da28 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -11,7 +11,7 @@ import types import pytest -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, normalize_omni_diffusion_kwargs +from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.engine import async_omni_engine as async_omni_engine_module from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_init_utils import ( @@ -1732,133 +1732,6 @@ def test_inject_kv_stage_info_infers_receiver_tp_topology(): assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2} -def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - stage_cfg = types.SimpleNamespace( - stage_type="diffusion", - engine_args=types.SimpleNamespace( - diffusion_attention_config=None, - lora_path=None, - lora_scale=None, - enable_sleep_mode=None, - quantization_config=None, - ), - ) - - monkeypatch.setattr( - engine_mod, - "load_and_resolve_stage_configs", - 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 - assert isinstance(diffusion_attention_config, AttentionConfig) - assert diffusion_attention_config.default is not None - assert diffusion_attention_config.default.backend == "FLASH_ATTN" - - -def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - existing_attention = AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")) - stage_cfg = types.SimpleNamespace( - stage_type="diffusion", - engine_args=types.SimpleNamespace( - diffusion_attention_config=existing_attention, - lora_path=None, - lora_scale=None, - enable_sleep_mode=None, - quantization_config=None, - ), - ) - - monkeypatch.setattr( - engine_mod, - "load_and_resolve_stage_configs", - 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 - - -def test_resolve_stage_configs_does_not_inject_over_stage_diffusion_attention_backend(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - stage_cfg = types.SimpleNamespace( - stage_type="diffusion", - engine_args=types.SimpleNamespace( - diffusion_attention_backend="TORCH_SDPA", - diffusion_attention_config=None, - lora_path=None, - lora_scale=None, - enable_sleep_mode=None, - quantization_config=None, - ), - ) - - monkeypatch.setattr( - engine_mod, - "load_and_resolve_stage_configs", - 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, - ) - - engine_args = stage_configs[0].engine_args - assert engine_args.diffusion_attention_backend == "TORCH_SDPA" - assert engine_args.diffusion_attention_config is None - normalized = normalize_omni_diffusion_kwargs(vars(engine_args)) - assert normalized["diffusion_attention_config"].default.backend == "TORCH_SDPA" - - -def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stage(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - stage_cfg = types.SimpleNamespace( - stage_type="llm", - engine_args=types.SimpleNamespace( - attention_config={"backend": "FLASH_ATTN"}, - enable_sleep_mode=None, - ), - ) - - monkeypatch.setattr( - engine_mod, - "load_and_resolve_stage_configs", - 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"} - assert not hasattr(stage_configs[0].engine_args, "diffusion_attention_config") - - def test_extract_legacy_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch): """ROCm default attention logic only applies to LLM stages, not diffusion.""" from vllm_omni.engine.stage_init_utils import extract_legacy_stage_metadata diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 99b81aeb9..aa4f25ded 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -2665,7 +2665,10 @@ async def test_duplex_reaper_loop_waits_between_ticks(): @pytest.mark.asyncio -async def test_duplex_reaper_loop_survives_one_cleanup_failure(): +@pytest.mark.parametrize("first_cleanup_delay", [0.0, 0.05], ids=["immediate", "delayed"]) +async def test_duplex_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: float) -> None: + recovered = asyncio.Event() + class _Plane: def __init__(self) -> None: self.calls = 0 @@ -2673,7 +2676,9 @@ async def test_duplex_reaper_loop_survives_one_cleanup_failure(): async def reap_expired(self) -> int: self.calls += 1 if self.calls == 1: + await asyncio.sleep(first_cleanup_delay) raise RuntimeError("transient cleanup failure") + recovered.set() return 0 orchestrator = object.__new__(Orchestrator) @@ -2682,11 +2687,13 @@ async def test_duplex_reaper_loop_survives_one_cleanup_failure(): orchestrator._shutdown_event = asyncio.Event() task = asyncio.create_task(orchestrator._duplex_reaper_loop()) - await asyncio.sleep(0.035) - orchestrator._shutdown_event.set() - await task - - assert orchestrator.duplex_control_plane.calls >= 2 + try: + # Wait for recovery itself, including when cleanup exceeds the old 35 ms window. + await asyncio.wait_for(recovered.wait(), timeout=5.0) + assert orchestrator.duplex_control_plane.calls >= 2 + finally: + orchestrator._shutdown_event.set() + await asyncio.wait_for(task, timeout=5.0) @pytest.mark.asyncio diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 90db55edb..eaa246072 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -16,7 +16,7 @@ from types import SimpleNamespace import janus import pytest -from vllm.sampling_params import SamplingParams +from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.serial_utils import MsgpackEncoder @@ -33,11 +33,16 @@ from vllm_omni.engine.orchestrator import ( cleanup_request_artifact_dirs, ) from vllm_omni.engine.stage_pool import StageUnavailableError +from vllm_omni.errors import OmniClientError from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.stage_input_processors import ( + joyai_vl_interaction as joyai_bridge, +) from vllm_omni.outputs import OmniRequestOutput from .test_orchestrator import ( FakeOutputProcessor, + FakeRunningCounter, FakeStageClient, OrchestratorFixture, _build_harness, @@ -45,6 +50,7 @@ from .test_orchestrator import ( _build_stage_pools, _engine_core_outputs, _enqueue_add_request, + _get_output_message, _wait_for, ) @@ -935,6 +941,139 @@ async def test_diffusion_client_error_output_propagates_status_code(orchestrator orchestrator_fixture.thread.join(timeout=5) +@pytest.mark.asyncio +@pytest.mark.parametrize("event_driven", [False, True], ids=["legacy", "event-driven"]) +@pytest.mark.parametrize("first_speaker", ["Ryan", "not-a-real-speaker"]) +async def test_joyai_tts_request_validation_keeps_orchestrator_running( + orchestrator_factory, + monkeypatch: pytest.MonkeyPatch, + first_speaker: str, + event_driven: bool, +) -> None: + """Check that the same orchestrator can handle the next valid request.""" + monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", "1" if event_driven else "0") + + # Use a fixed token count so the test does not need tokenizer files. + class FakeTokenizer: + def encode(self, _text: str, *, add_special_tokens: bool) -> list[int]: + return list(range(12)) + + talker_model_config = SimpleNamespace( + max_model_len=64, + hf_config=SimpleNamespace( + talker_config=SimpleNamespace(spk_id={"vivian": 0, "ryan": 1}, codec_language_id={}, spk_is_dialect={}) + ), + ) + monkeypatch.setattr(joyai_bridge, "cached_tokenizer_from_config", lambda _: FakeTokenizer()) + + # Use the real bridge so speaker validation is part of the test. + class JoyAITalkerStage(FakeStageClient): + def process_engine_inputs(self, source_outputs, prompt=None, streaming_context=None): + self.bridge_inputs = joyai_bridge.joyai_action_to_tts( + source_outputs, prompt, self.requires_multimodal_data, target_model_config=talker_model_config + ) + return self.bridge_inputs + + stage0 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = JoyAITalkerStage(stage_type="llm", final_output=True) + stage0_processor, stage1_processor = FakeOutputProcessor(), FakeOutputProcessor() + fixture = orchestrator_factory( + [stage0, stage1], + output_processors=[stage0_processor, stage1_processor], + stage_vllm_configs=[ + SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + SimpleNamespace(model_config=talker_model_config), + ], + ) + counter = FakeRunningCounter() + fixture.orchestrator._running_counter = counter + # With a valid speaker, task_type should be changed from Base to CustomVoice. + # An unknown speaker should cause a request error. + first_info = {"tts_task_type": "Base", "speaker": "Vivian", "tts_speaker": first_speaker} + talker_submissions = 0 + # Send both requests to the same orchestrator. The second uses the default Vivian voice. + for index, (request_id, info) in enumerate((("first", first_info), ("good", {})), start=1): + await _enqueue_add_request( + fixture, + request_id=request_id, + prompt=SimpleNamespace(request_id=request_id, prompt_token_ids=[1, 2]), + original_prompt={"additional_information": info}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + await _wait_for(lambda: len(stage0.add_request_calls) == index and counter.value == 1) + stage0_processor.request_outputs = [_build_request_output(request_id, text=" Hello.")] + stage0.push_engine_core_outputs(_engine_core_outputs(f"{request_id}-raw", float(index))) + + if request_id == "first" and first_speaker == "not-a-real-speaker": + # Reject this request before Talker runs, while keeping the engine alive. + error = await _wait_for_error_message(fixture, request_id=request_id) + assert error.status_code == 400 and error.error_type == "BadRequestError" + assert error.fatal is False and error.stage_id == 1 + assert "not-a-real-speaker" in error.error + assert stage1.add_request_calls == [] + else: + talker_submissions += 1 + await _wait_for(lambda: len(stage1.add_request_calls) == talker_submissions or fixture.result_future.done()) + assert not fixture.result_future.done() + assert stage1.add_request_calls[-1][0].request_id == request_id + metadata = stage1.bridge_inputs[0]["additional_information"] + assert metadata["task_type"] == ["CustomVoice"] + assert metadata["speaker"] == [first_speaker if request_id == "first" else "Vivian"] + # Finish the valid request and check its final output. + stage1_processor.request_outputs = [_build_request_output(request_id, text="Completed.")] + stage1.push_engine_core_outputs(_engine_core_outputs(f"{request_id}-talker", float(index))) + output = await _get_output_message(fixture) + assert output.request_id == request_id and output.stage_id == 1 + assert output.finished is True + assert output.engine_outputs.outputs[0].text == "Completed." + + # Wait for full cleanup before sending the next request. + await _wait_for(lambda: request_id not in fixture.orchestrator.request_states and counter.value == 0) + assert all(pool.get_bound_replica_id(request_id) is None for pool in fixture.orchestrator.stage_pools) + assert fixture.thread.is_alive() and not fixture.result_future.done() + + +@pytest.mark.asyncio +async def test_streaming_input_processor_client_error_does_not_forward_terminal_update() -> None: + """Do not send a second update after the first one fails and clears the request.""" + processor_calls = [] + + class FailingStage(FakeStageClient): + def process_engine_inputs(self, *args, **kwargs): + processor_calls.append((args, kwargs)) + raise OmniClientError("Input rejected") + + orchestrator, queues = _build_bare_orchestrator( + _build_stage_pools([[FakeStageClient()], [FailingStage(final_output=True)]]) + ) + state = OrchestratorRequestState( + request_id="bad-stream", + prompt={"prompt_token_ids": [1]}, + sampling_params_list=[SamplingParams(output_kind=RequestOutputKind.CUMULATIVE), _sampling_params()], + final_stage_id=1, + ) + state.streaming.enabled = True + assert state.duplex_identity is None + orchestrator.request_states[state.request_id] = state + try: + await orchestrator._route_output(0, 0, _build_request_output(state.request_id), state, stage_metrics=None) + + assert len(processor_calls) == 1 + error = queues[1].sync_q.get_nowait() + assert isinstance(error, ErrorMessage) + assert error.request_id == state.request_id + assert error.stage_id == 1 + assert error.status_code == 400 + assert error.fatal is False + assert queues[1].sync_q.empty() + assert state.request_id not in orchestrator.request_states + assert all(pool.get_bound_replica_id(state.request_id) is None for pool in orchestrator.stage_pools) + finally: + for q in queues: + q.close() + + @pytest.mark.asyncio async def test_duplex_input_processor_failure_is_request_scoped(orchestrator_factory, monkeypatch) -> None: class FailingStage(FakeStageClient): diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index cacda1f60..675a0650b 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -376,14 +376,26 @@ class TestSingleStageModeDetection: mocker: MockerFixture, *, stage_cfgs: list[Any] | None = None, + resolved_config_path: str = "/fake/path", + patch_deploy_config: bool = True, **kwargs: Any, ) -> AsyncOmniEngine: mock_stage_configs = stage_cfgs or [_make_stage_cfg(0)] + if patch_deploy_config: + mocker.patch.object( + StageConfigFactory, + "get_pipeline_config", + return_value=None, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_deploy_config", + return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()), + ) mocker.patch.object( AsyncOmniEngine, "_resolve_stage_configs", - return_value=("/fake/path", mock_stage_configs), + return_value=(resolved_config_path, mock_stage_configs), ) mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator") mock_thread_cls = mocker.patch("threading.Thread") @@ -445,6 +457,8 @@ class TestSingleStageModeDetection: engine = self._make_engine_no_thread( mocker, deploy_config="/fake/duplex.yaml", + resolved_config_path="/fake/duplex.yaml", + patch_deploy_config=False, ) get_pipeline_config.assert_called_once_with( @@ -455,6 +469,30 @@ class TestSingleStageModeDetection: load_deploy_config.assert_called_once_with("/fake/duplex.yaml") assert engine.duplex_session_config is duplex_session + def test_auto_discovered_deploy_loads_duplex_runtime_config( + self, + mocker: MockerFixture, + ) -> None: + deploy_path = "/resolved/qwen3_omni_moe.yaml" + duplex_session = DuplexSessionRuntimeConfig(server_vad_model_path="/models/silero_vad.onnx") + mocker.patch( + "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config", + return_value=None, + ) + load_deploy_config = mocker.patch( + "vllm_omni.engine.async_omni_engine.load_deploy_config", + return_value=SimpleNamespace(duplex_session=duplex_session), + ) + + engine = self._make_engine_no_thread( + mocker, + resolved_config_path=deploy_path, + patch_deploy_config=False, + ) + + load_deploy_config.assert_called_once_with(deploy_path) + assert engine.duplex_session_config is duplex_session + def test_single_stage_mode_without_stage_id_has_no_filter(self, mocker: MockerFixture): engine = self._make_engine_no_thread( mocker, @@ -518,6 +556,10 @@ class TestEndpointRestrictionsTrustRemoteCode: StageConfigFactory.try_infer_model_type.cache_clear() def _make_engine_no_thread(self, mocker: MockerFixture, **kwargs: Any) -> AsyncOmniEngine: + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_deploy_config", + return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()), + ) mocker.patch.object( AsyncOmniEngine, "_resolve_stage_configs", diff --git a/tests/entrypoints/duplex/test_runtime_adapter_boundary.py b/tests/entrypoints/duplex/test_runtime_adapter_boundary.py index dc7fffccd..cb90e4bd5 100644 --- a/tests/entrypoints/duplex/test_runtime_adapter_boundary.py +++ b/tests/entrypoints/duplex/test_runtime_adapter_boundary.py @@ -61,20 +61,19 @@ if loaded: assert result.returncode == 0, result.stderr -def test_generic_handler_requires_explicit_serving_runtime_adapter() -> None: +def test_generic_handler_supports_adapterless_chat_fallback() -> None: repo_root = Path(__file__).resolve().parents[3] script = """ from types import SimpleNamespace from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler -chat_service = SimpleNamespace(engine_client=SimpleNamespace()) -try: - OmniDuplexSessionHandler(chat_service=chat_service) -except ValueError as exc: - if "serving runtime adapter" not in str(exc).lower(): - raise -else: +chat_service = SimpleNamespace( + engine_client=SimpleNamespace(), + model_config=SimpleNamespace(model="test-model"), +) +handler = OmniDuplexSessionHandler(chat_service=chat_service) +if handler._serving_runtime_adapter is not None: raise SystemExit("generic handler silently selected a model serving adapter") """ diff --git a/tests/entrypoints/openai/test_audio8_tts_adapter.py b/tests/entrypoints/openai/test_audio8_tts_adapter.py new file mode 100644 index 000000000..ffd22d9f3 --- /dev/null +++ b/tests/entrypoints/openai/test_audio8_tts_adapter.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Serving-side registration invariants for the Audio8 TTS adapter. + +The engine never names ``audio8_tts`` directly: ``/v1/audio/speech`` resolves the +adapter by matching the deployed stage's ``model_stage`` key against the +``stage_keys`` each adapter declares. That makes the pipeline definition and the +adapter two halves of one contract, with nothing in the type system holding them +together -- rename either side and the endpoint silently stops recognising the +model (falling through to "unknown TTS model") instead of failing at startup. + +These tests pin the contract from both ends. +""" + +import pytest + +from vllm_omni.entrypoints.openai.tts_adapters import ( + all_tts_model_types, + all_tts_stage_keys, + detect_tts_model_type, + resolve_adapter, +) +from vllm_omni.entrypoints.openai.tts_adapters.audio8_tts import Audio8TTSAdapter +from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SLOW_AR_STAGE_KEY = "audio8_tts_slow_ar" + + +def _pipeline_stage_keys() -> list[str]: + return [stage.model_stage for stage in AUDIO8_TTS_PIPELINE.stages] + + +def test_pipeline_stage_key_is_the_one_the_adapter_claims(): + """The shipped pipeline must name the stage the adapter matches on.""" + assert SLOW_AR_STAGE_KEY in _pipeline_stage_keys() + assert SLOW_AR_STAGE_KEY in Audio8TTSAdapter.stage_keys + + +def test_detection_resolves_slow_ar_stage_to_the_audio8_adapter(): + assert detect_tts_model_type(SLOW_AR_STAGE_KEY, "Audio8TTSSlowARForConditionalGeneration") == "audio8_tts" + assert resolve_adapter("audio8_tts") is Audio8TTSAdapter + + +def test_adapter_is_registered_in_the_global_tables(): + """Importing the package must register the adapter, not just define it.""" + assert "audio8_tts" in all_tts_model_types() + assert SLOW_AR_STAGE_KEY in all_tts_stage_keys() + + +def test_codec_decoder_stage_does_not_serve_speech(): + """Only stage 0 answers /v1/audio/speech; the codec stage must not match.""" + codec_stage_key = _pipeline_stage_keys()[1] + assert codec_stage_key != SLOW_AR_STAGE_KEY + assert detect_tts_model_type(codec_stage_key, "Audio8TTSCodecDecoder") != "audio8_tts" + + +def test_apply_sampling_overrides_defaults_tts_local_seed_for_fast_ar(): + """The deploy seed must reach the Fast AR through ``tts_local_seed``. + + Only the Slow AR is seeded via ``sampling_metadata.generators``; the residual + codebooks read ``tts_local_seed`` and otherwise fall back to global RNG, which + makes an unseeded request's audio depend on batch composition. The adapter + defaults it from the stage-0 seed without clobbering an explicit value. + """ + from types import SimpleNamespace + + from vllm.sampling_params import SamplingParams + + adapter = Audio8TTSAdapter.__new__(Audio8TTSAdapter) + request = SimpleNamespace(max_new_tokens=None) + + # Deploy default seed propagates to tts_local_seed. + out = adapter.apply_sampling_overrides([SamplingParams(seed=42)], request) + assert out[0].extra_args["tts_local_seed"] == 42 + + # An explicit tts_local_seed is left intact (setdefault, not overwrite). + pinned = [SamplingParams(seed=42, extra_args={"tts_local_seed": 7})] + assert adapter.apply_sampling_overrides(pinned, request)[0].extra_args["tts_local_seed"] == 7 + + # No seed at all -> nothing injected (non-deterministic by request contract). + out_unseeded = adapter.apply_sampling_overrides([SamplingParams(seed=None)], request) + assert (out_unseeded[0].extra_args or {}).get("tts_local_seed") is None diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index 707b4b961..71f93f6dc 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -568,6 +568,48 @@ def test_websocket_routes_emit_stable_unavailable_frames_and_close(path: str, pa websocket.receive_text() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("duplex_query", "expected_handler"), + [ + (None, "duplex"), + ("1", "duplex"), + ("true", "duplex"), + ("on", "duplex"), + ("0", "legacy"), + ("false", "legacy"), + ], +) +async def test_realtime_route_defaults_to_configured_duplex_handler( + monkeypatch, duplex_query: str | None, expected_handler: str +) -> None: + calls: list[str] = [] + + class _DuplexHandler: + async def handle_realtime_session(self, _websocket) -> None: + calls.append("duplex") + + class _LegacyConnection: + async def handle_connection(self) -> None: + calls.append("legacy") + + monkeypatch.setattr(api_server, "RealtimeConnection", lambda _websocket, _serving: _LegacyConnection()) + query_params = {} if duplex_query is None else {"duplex": duplex_query} + websocket = SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + openai_serving_duplex=_DuplexHandler(), + openai_serving_realtime=object(), + ) + ), + query_params=query_params, + ) + + await api_server.realtime_websocket(websocket) + + assert calls == [expected_handler] + + def test_health_without_engine_returns_stable_unhealthy_response() -> None: """Lock ``/health`` with no engine initialized. diff --git a/tests/entrypoints/openai_api/test_audex_serving_guards.py b/tests/entrypoints/openai_api/test_audex_serving_guards.py index 27ed11356..87524825e 100644 --- a/tests/entrypoints/openai_api/test_audex_serving_guards.py +++ b/tests/entrypoints/openai_api/test_audex_serving_guards.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Serving-layer regressions for the Audex zero-codec-token contract. A zero-codec-token (or phase-invalid TTA) request reaches the serving layer @@ -171,7 +171,7 @@ class TestAudexCfgValidation: class TestAudexCfgInjection: def _serving_with_fake_tokenizer(self): serving = _serving() - serving._get_audex_tokenizer = lambda: _MarkerTokenizer() + serving._adapter._get_tokenizer = lambda: _MarkerTokenizer() return serving def _cond_prompt(self, text: str = "Hello world.") -> str: @@ -333,8 +333,8 @@ class TestAdapterRouting: class TestAudexTTAInjection: def _serving(self): serving = _serving(model_type="audex_tta") - serving._get_audex_tokenizer = lambda: _TTAVocabTokenizer() - serving._audex_tta_rvq = None + serving._adapter._get_tokenizer = lambda: _TTAVocabTokenizer() + serving._adapter._tta_rvq = None return serving def _prompt(self): diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py index e47885d9b..a9bf76319 100644 --- a/tests/entrypoints/openai_api/test_audio_format.py +++ b/tests/entrypoints/openai_api/test_audio_format.py @@ -163,6 +163,8 @@ class TestStreamingAudioResampler: chunked = StreamingAudioResampler(24000, 8000) pieces = [chunked.process(chunk) for chunk in np.split(waveform, [137, 2048, 9001, 17003])] + # Bounded FIR history must own its data, not retain the last large chunk. + assert chunked._history.base is None pieces.append(chunked.process(np.empty(0), final=True)) actual = np.concatenate(pieces) @@ -193,9 +195,15 @@ class TestStreamingAudioResampler: output_rms = np.sqrt(np.mean(output[100:-100] ** 2)) assert output_rms < input_rms * 0.01 - def test_rejects_non_integer_downsampling_ratio(self): - with pytest.raises(ValueError, match="integer downsampling ratio"): - StreamingAudioResampler(24000, 16000) + def test_supports_rational_downsampling_ratio(self): + waveform = np.sin(np.linspace(0, 200 * np.pi, 24000, endpoint=False)).astype(np.float32) + resampler = StreamingAudioResampler(24000, 16000) + + output = resampler.process(waveform, final=True) + reference = torchaudio.functional.resample(torch.from_numpy(waveform), 24000, 16000).numpy() + + assert output.shape == (16000,) + np.testing.assert_allclose(output[100:-100], reference[100:-100], atol=2e-3, rtol=1e-3) def test_rejects_multichannel_audio(self): resampler = StreamingAudioResampler(24000, 8000) diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index c69bc5c47..d3bbc5d01 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -5,8 +5,10 @@ from __future__ import annotations import asyncio import base64 +import io import json import struct +import wave from types import SimpleNamespace from typing import Any @@ -30,8 +32,8 @@ from vllm_omni.engine.duplex.control_client import DuplexControlRequestError from vllm_omni.engine.duplex.lease import DuplexLeaseActivity from vllm_omni.engine.duplex.messages import DuplexFence, DuplexSessionLifecycleMessage from vllm_omni.engine.duplex.runtime import duplex_resource_request_id -from vllm_omni.entrypoints.duplex import vad as realtime_vad from vllm_omni.entrypoints.duplex.protocol import ( + DuplexCapabilities, DuplexOverlapPolicy, DuplexPlaybackCommitPolicy, DuplexSession, @@ -69,6 +71,11 @@ from vllm_omni.outputs import OmniRequestOutput from vllm_omni.outputs.duplex import attach_duplex_output_decision pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +INVALID_VALIDATED_SAMPLE_RATES = ( + pytest.param(16_000.5, id="fractional"), + pytest.param(float("inf"), id="infinite"), + pytest.param(float("nan"), id="nan"), +) def test_native_input_append_supports_explicit_session_opt_out(): @@ -327,6 +334,46 @@ class FakeChatService: return SimpleNamespace(audio_data=f"wav-{int(audio_obj.audio_tensor.shape[0])}") +class TurnBasedFakeChatService(FakeChatService): + """Chat-fallback service without a model-native duplex adapter.""" + + duplex_serving_adapter_path = None + + def __init__(self, engine_client: FakeEngineClient) -> None: + super().__init__(engine_client) + self.requests: list[Any] = [] + + async def create_chat_completion(self, request, raw_request=None): + del raw_request + self.requests.append(request) + self.seen_request_ids.append(request.request_id) + + async def _gen(): + yield "data: [DONE]\n\n" + + return _gen() + + +class GateChatService(FakeChatService): + def __init__(self, engine_client: FakeEngineClient) -> None: + super().__init__(engine_client) + self.first_response_gate = asyncio.Event() + self.response_count = 0 + + async def create_chat_completion(self, request, raw_request=None): + del raw_request + self.seen_request_ids.append(request.request_id) + self.response_count += 1 + response_number = self.response_count + + async def _gen(): + if response_number == 1: + await self.first_response_gate.wait() + yield "data: [DONE]\n\n" + + return _gen() + + def _test_data_plane() -> MiniCPMO45DataPlaneSession: def encode_audio( audio: object, @@ -405,6 +452,105 @@ class TimedWebSocket: return [m.get("type", "") for m in self.sent] +class FakeServerVADBackend: + frame_samples = 160 + + def __init__(self, probabilities: list[float]) -> None: + self.probabilities = probabilities + self.calls = 0 + + def new_state(self) -> int: + return 0 + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: + del frame + self.calls += 1 + index = int(state) + probability = self.probabilities[index] if index < len(self.probabilities) else 0.0 + return probability, index + 1 + + +class FakeServerVADProvider: + def __init__(self, backend: FakeServerVADBackend) -> None: + self.backend = backend + + def get(self) -> FakeServerVADBackend: + return self.backend + + +class FailingServerVADProvider: + def __init__(self) -> None: + self.calls = 0 + + def get(self) -> FakeServerVADBackend: + self.calls += 1 + raise RuntimeError("server VAD backend unavailable") + + +class FailOnceServerVADBackend(FakeServerVADBackend): + def __init__(self, probabilities: list[float]) -> None: + super().__init__(probabilities) + self._failed = False + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: + if not self._failed and int(state) == 2: + self._failed = True + raise RuntimeError("server VAD inference failed") + return super().infer(frame, state) + + +def _server_vad_turn_detection(*, create_response: bool = True) -> dict[str, Any]: + return { + "type": "server_vad", + "prefix_padding_ms": 10, + "silence_duration_ms": 20, + "create_response": create_response, + "interrupt_response": False, + } + + +def _server_vad_session_update( + turn_detection: dict[str, Any] | None, + *, + event_id: str | None = None, + sample_rate_hz: int = 16_000, +) -> dict[str, Any]: + event: dict[str, Any] = { + "type": "session.update", + "session": { + "model": "test-model", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": sample_rate_hz}, + "turn_detection": turn_detection, + } + }, + }, + } + if event_id is not None: + event["event_id"] = event_id + return event + + +def _server_vad_audio_append( + frame_count: int = 5, + *, + event_id: str | None = None, + is_speech: bool | None = None, + sample_rate_hz: int = 16_000, +) -> dict[str, Any]: + sample_count = frame_count * 160 * sample_rate_hz // 16_000 + event: dict[str, Any] = { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(np.zeros(sample_count, dtype=" None: + if data.get("type") == "response.done": + ws.put({"type": "session.close"}) + + ws = TimedWebSocket(on_send=on_send) + if not at_creation: + ws.put(_server_vad_session_update(_server_vad_turn_detection())) + update = _server_vad_session_update(_server_vad_turn_detection(), event_id="bad-vad-audio-format") + update["session"]["audio"]["input"].update(input_update) + ws.put(update) + if not at_creation: + ws.put(_server_vad_audio_append()) + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0])), + ) + + await handler.handle_realtime_session(ws) # type: ignore[arg-type] + + (error,) = [event["error"] for event in ws.sent if event.get("type") == "error"] + assert error["code"] == "unsupported_turn_detection" + assert error["event_id"] == "bad-vad-audio-format" + assert "mono PCM16" in error["message"] + assert ("session.created" in ws.sent_types()) is not at_creation + assert ("response.done" in ws.sent_types()) is not at_creation + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sample_rate_hz", [16_000, 24_000]) +@pytest.mark.parametrize("create_response", [False, True], ids=["commit-only", "minimal-config"]) +@pytest.mark.parametrize("at_creation", [True, False], ids=["create", "update"]) +async def test_realtime_server_vad_commits_without_client_commit( + sample_rate_hz: int, create_response: bool, at_creation: bool +): + def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: + terminal_event = "response.done" if create_response else "input_audio_buffer.committed" + if data.get("type") == terminal_event: + ws.put({"type": "session.close"}) + + ws = TimedWebSocket(on_send=on_send) + if not at_creation: + ws.put(_server_vad_session_update(None, sample_rate_hz=sample_rate_hz)) + ws.put( + _server_vad_session_update( + {"type": "server_vad"} if create_response else _server_vad_turn_detection(create_response=False), + sample_rate_hz=sample_rate_hz, + ) + ) + # Cover the default 500 ms silence timeout, plus streaming FIR right context. + ws.put(_server_vad_audio_append(54, sample_rate_hz=sample_rate_hz)) + engine = FakeEngineClient() + chat_service = TurnBasedFakeChatService(engine) + handler = OmniDuplexSessionHandler( + chat_service=chat_service, + server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0])), + ) + + await handler.handle_realtime_session(ws) # type: ignore[arg-type] + + event_types = ws.sent_types() + assert "error" not in event_types + effective_session = [ + event["session"] for event in ws.sent if event["type"] in {"session.created", "session.updated"} + ][-1] + assert effective_session["turn_detection"]["interrupt_response"] is False + assert effective_session["audio"]["input"]["turn_detection"] == effective_session["turn_detection"] + assert effective_session["overlap_policy"] == "listen_only" + assert event_types.index("input_audio_buffer.speech_started") < event_types.index( + "input_audio_buffer.speech_stopped" + ) + assert event_types.index("input_audio_buffer.speech_stopped") < event_types.index("input_audio_buffer.committed") + assert event_types.index("input_audio_buffer.committed") < event_types.index("conversation.item.added") + assert event_types.index("conversation.item.added") < event_types.index("conversation.item.done") + assert "conversation.item.created" not in event_types + committed = next(event for event in ws.sent if event.get("type") == "input_audio_buffer.committed") + added = next(event for event in ws.sent if event.get("type") == "conversation.item.added") + done = next(event for event in ws.sent if event.get("type") == "conversation.item.done") + assert committed["item_id"] == added["item"]["id"] == done["item"]["id"] + assert event_types.count("response.created") == int(create_response) + assert len(chat_service.seen_request_ids) == int(create_response) + if create_response: + assert event_types.index("input_audio_buffer.committed") < event_types.index("response.created") + assert event_types.index("response.created") < event_types.index("response.done") + assert engine.opened == [] + assert engine.signals == [] + assert engine.closed == [] + + +@pytest.mark.asyncio +async def test_realtime_server_vad_create_response_false_accepts_response_create(): + def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: + if data.get("type") == "input_audio_buffer.committed": + ws.put({"type": "response.create"}) + if data.get("type") == "response.done": + ws.put({"type": "session.close"}) + + ws = TimedWebSocket(on_send=on_send) + ws.put(_server_vad_session_update(_server_vad_turn_detection(create_response=False))) + ws.put(_server_vad_audio_append()) + engine = FakeEngineClient() + chat_service = TurnBasedFakeChatService(engine) + handler = OmniDuplexSessionHandler( + chat_service=chat_service, + server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0])), + ) + + await handler.handle_realtime_session(ws) # type: ignore[arg-type] + + event_types = ws.sent_types() + assert event_types.index("input_audio_buffer.committed") < event_types.index("response.created") + assert event_types.index("response.created") < event_types.index("response.done") + assert event_types.count("response.created") == 1 + assert len(chat_service.seen_request_ids) == 1 + assert engine.opened == engine.signals == engine.closed == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sample_rate_hz", [16_000, 24_000]) +async def test_realtime_server_vad_complete_audio_item_bypasses_vad(sample_rate_hz: int): + def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: + item = data.get("item") + if ( + data.get("type") == "conversation.item.done" + and isinstance(item, dict) + and item.get("id") == "item-complete-audio" + ): + ws.put({"type": "response.create"}) + if data.get("type") == "response.done": + ws.put({"type": "session.close"}) + + ws = TimedWebSocket(on_send=on_send) + ws.put(_server_vad_session_update(_server_vad_turn_detection(), sample_rate_hz=sample_rate_hz)) + samples = np.arange(sample_rate_hz // 10, dtype=np.float32) + pcm16 = np.rint(0.5 * np.sin(2 * np.pi * 1_000 * samples / sample_rate_hz) * 32_767).astype(" None: + if data.get("type") == "response.done": + ws.put({"type": "session.close"}) + + ws = TimedWebSocket(on_send=on_send) + ws.put(_server_vad_session_update(_server_vad_turn_detection())) + first_append = _server_vad_audio_append(frame_count=2) + first_append["audio"] = base64.b64encode(np.full(320, 12_000, dtype=" None: + nonlocal committed_count, response_count + if data.get("type") == "input_audio_buffer.committed": + committed_count += 1 + if data.get("type") == "error" and data.get("error", {}).get("code") == "input_backpressure": + chat_service.first_response_gate.set() + if data.get("type") == "response.created": + response_count += 1 + if response_count == 2: + ws.put({"type": "session.close"}) + + ws = TimedWebSocket(on_send=on_send) + ws.put(_server_vad_session_update(_server_vad_turn_detection())) + for _ in range(3): + ws.put(_server_vad_audio_append()) + handler = OmniDuplexSessionHandler( + chat_service=chat_service, + server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0] * 3)), + ) + + await handler.handle_realtime_session(ws) # type: ignore[arg-type] + + assert committed_count == 2 + assert response_count == 2 + event_types = ws.sent_types() + committed_indices = [index for index, event in enumerate(event_types) if event == "input_audio_buffer.committed"] + response_created_indices = [index for index, event in enumerate(event_types) if event == "response.created"] + assert committed_indices[1] < event_types.index("response.done") < response_created_indices[1] + backpressure = [ + event + for event in ws.sent + if event.get("type") == "error" and event.get("error", {}).get("code") == "input_backpressure" + ] + assert len(backpressure) == 1 + + @pytest.mark.asyncio async def test_native_session_update_rejects_client_runtime_config(): engine = FakeEngineClient() @@ -1110,6 +1967,7 @@ async def test_minicpmo_native_session_update_requires_ref_audio_before_enabling def test_personaplex_candidate_update_does_not_leak_minicpmo_ref_audio_check(): handler = OmniDuplexSessionHandler( chat_service=SimpleNamespace(engine_client=SimpleNamespace()), + served_model_name="test-model", serving_runtime_adapter=PersonaPlexServingRuntimeAdapter(lambda *_: None), ) session = DuplexSession( @@ -1124,6 +1982,51 @@ def test_personaplex_candidate_update_does_not_leak_minicpmo_ref_audio_check(): assert error is None +@pytest.mark.parametrize( + ("native", "turn_detection", "expected_code"), + [ + (False, {"type": "server_vad", "interrupt_response": False}, None), + (False, {"type": "server_vad"}, None), + (False, {"type": "server_vad", "interrupt_response": True}, "server_vad_requires_native_duplex"), + (True, {"type": "server_vad", "interrupt_response": False}, "unsupported_turn_detection"), + (True, {"type": "server_vad"}, None), + ], + ids=[ + "turn-based-endpointing", + "turn-based-default-endpointing", + "turn-based-explicit-interruption", + "native-non-interrupting", + "native-default-interruption", + ], +) +def test_server_vad_runtime_capability_matrix( + native: bool, + turn_detection: dict[str, object], + expected_code: str | None, +): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + ) + capabilities = minicpmo45_native_capabilities() if native else DuplexCapabilities() + session = DuplexSession( + "sid-server-vad-capability", + DuplexSessionConfig(), + capabilities=capabilities, + ) + candidate = DuplexSessionConfig.from_event({"session": {"turn_detection": turn_detection}}) + handler._resolve_server_vad_defaults(session, candidate, {"turn_detection": turn_detection}) + + error = handler._runtime_session_candidate_update_error(session, candidate) + + assert (error or {}).get("code") == expected_code + if expected_code is None: + assert candidate.server_vad is not None + assert candidate.server_vad.interrupt_response is native + assert candidate.overlap_policy == ("barge_in_on_speech" if native else "listen_only") + elif expected_code == "server_vad_requires_native_duplex": + assert "set interrupt_response=false" in error["error"] + + @pytest.mark.asyncio async def test_minicpmo_native_session_update_rejects_instructions_after_native_context_locked(): engine = FakeEngineClient() @@ -1550,6 +2453,23 @@ def test_duplex_session_playback_commit_uses_multi_delta_audio_text_marks(): assert session.history == (committed,) +def test_duplex_session_preserves_response_history_order_when_user_item_arrives_during_generation(): + session = DuplexSession(session_id="sid-history-order", config=DuplexSessionConfig()) + first_user = {"role": "user", "content": "first turn"} + next_user = { + "role": "user", + "content": [{"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,AAAA"}}], + } + session.append_history_message(first_user) + session.begin_response() + session.append_assistant_text("first answer") + session.append_history_message(next_user) + + assistant = session.end_response(commit_text=True) + + assert session.history == (first_user, assistant, next_user) + + @pytest.mark.asyncio async def test_duplex_session_actor_preserves_wire_order_before_control(): ws = TimedWebSocket() @@ -1616,7 +2536,7 @@ def _native_realtime_session_update( } -def _server_vad_update(event_id: str, **session: object) -> dict[str, object]: +def _native_server_vad_update(event_id: str, **session: object) -> dict[str, object]: return { "type": "session.update", "event_id": event_id, @@ -2356,7 +3276,7 @@ def test_auto_response_force_barge_in_has_no_waiting_turn_variant(): ) -def test_duplex_endpoint_requires_explicit_session_mode_duplex(): +def test_duplex_endpoint_requires_explicit_realtime_configuration(): assert should_enable_duplex_endpoint(None) is False assert should_enable_duplex_endpoint([]) is False assert should_enable_duplex_endpoint([SimpleNamespace(session_mode="turn")]) is False @@ -2388,6 +3308,22 @@ stage_args: assert should_enable_duplex_endpoint([], config_path=str(config_path)) is True +def test_duplex_endpoint_supports_inherited_duplex_session(tmp_path): + base_path = tmp_path / "base.yaml" + base_path.write_text( + """ +session_mode: turn +duplex_session: {} +stages: [] +""", + encoding="utf-8", + ) + config_path = tmp_path / "deploy.yaml" + config_path.write_text("base_config: base.yaml\n", encoding="utf-8") + + assert should_enable_duplex_endpoint([], config_path=str(config_path)) is True + + def test_duplex_handler_splits_data_plane_audio_list_into_deltas(): import torch @@ -4254,7 +5190,7 @@ async def test_duplex_handler_runtime_open_failure_is_reported_to_client(): chat_service = FakeChatService(engine) handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) ws = TimedWebSocket() - ws.put(_session_create("sid-open-fail")) + ws.put(_native_session_create("sid-open-fail")) ws.put({"type": "session.close"}) await handler.handle_session(ws) @@ -4289,7 +5225,7 @@ async def test_duplex_handler_preserves_typed_admission_error(): idle_timeout_s=1, ) ws = TimedWebSocket() - ws.put(_session_create("sid-admission-rejected")) + ws.put(_native_session_create("sid-admission-rejected")) await handler.handle_session(ws) @@ -4903,13 +5839,37 @@ async def test_cancel_active_native_data_plane_request_aborts_stage_request_id() assert "audio.cancelled" in ws.sent_types() +@pytest.mark.asyncio +async def test_cancel_chat_fallback_response_aborts_request_and_task(): + engine = FakeEngineClient() + handler = OmniDuplexSessionHandler(chat_service=TurnBasedFakeChatService(engine)) + session = DuplexSession("sid-chat-cancel", DuplexSessionConfig()) + response_id = session.begin_response() + session.bind_request("chatcmpl-cancel") + ws = TimedWebSocket() + task = asyncio.create_task(asyncio.sleep(60)) + try: + cancelled = await handler._cancel_active_response(session, task, ws.send_json, reason="response.cancel") + + assert cancelled is True + assert task.cancelled() + assert engine.internal_abort_batches == [["chatcmpl-cancel"]] + assert "error" not in ws.sent_types() + (event,) = [event for event in ws.sent if event["type"] == "audio.cancelled"] + assert event["response_id"] == response_id + assert event["reason"] == "response.cancel" + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.asyncio async def test_duplex_handler_explicit_close_closes_runtime_once_with_client_reason(): engine = FakeEngineClient() chat_service = FakeChatService(engine) handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) ws = TimedWebSocket() - ws.put(_session_create("sid-close")) + ws.put(_native_session_create("sid-close")) ws.put({"type": "session.close"}) await handler.handle_session(ws) @@ -5600,7 +6560,7 @@ async def test_duplex_handler_idle_timeout_close_does_not_emit_runtime_control() chat_service = FakeChatService(engine) handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=0.1) ws = TimedWebSocket() - ws.put(_session_create("sid-disconnect")) + ws.put(_native_session_create("sid-disconnect")) await handler.handle_session(ws) @@ -5633,7 +6593,7 @@ async def test_duplex_handler_runtime_close_failure_is_reported_without_closed_a chat_service = FakeChatService(engine) handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) ws = TimedWebSocket() - ws.put(_session_create("sid-close-fail")) + ws.put(_native_session_create("sid-close-fail")) ws.put({"type": "session.close"}) await handler.handle_session(ws) @@ -5662,7 +6622,7 @@ async def test_duplex_handler_control_close_failure_is_reported_without_closed_a chat_service = FakeChatService(engine) handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) ws = TimedWebSocket() - ws.put(_session_create("sid-control-close-fail")) + ws.put(_native_session_create("sid-control-close-fail")) ws.put({"type": "session.close"}) await handler.handle_session(ws) @@ -5701,7 +6661,7 @@ async def test_turn_signal_input_cancel_uses_epoch_fence_transition(): idle_timeout_s=1, ) ws = TimedWebSocket() - ws.put(_session_create("sid-turn-signal-cancel")) + ws.put(_native_session_create("sid-turn-signal-cancel")) ws.put({"type": "input.text.append", "text": "discard me"}) ws.put({"type": "turn.signal", "event": "input.cancel"}) ws.put({"type": "session.close"}) @@ -5731,59 +6691,6 @@ async def test_duplex_handler_local_turn_signal_does_not_round_trip_runtime(): assert "runtime_signal_failed" not in {m.get("code") for m in ws.sent} -@pytest.mark.asyncio -async def test_duplex_barge_in_aborts_active_response_when_runtime_signal_fails(): - engine = FakeEngineClient(fail_signal_events={"barge_in"}) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "response.created": - ws.put({"type": "input.cancel", "reason": "test_barge_in"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_session_create("sid-barge-signal-fail")) - ws.put({"type": "input.text.append", "text": "hello"}) - ws.put({"type": "input.commit"}) - - await handler.handle_session(ws) - - assert "audio.cancelled" in ws.sent_types() - assert engine.aborted == ["chatcmpl-duplex-sid-barge-signal-fail-0-1"] - assert engine.internal_abort_batches == [["chatcmpl-duplex-sid-barge-signal-fail-0-1"]] - error = next(m for m in ws.sent if m.get("type") == "error") - assert error["code"] == "runtime_signal_failed" - - -@pytest.mark.asyncio -async def test_duplex_handler_surfaces_stage_unsupported_result_to_client(): - control_result = { - "operation": "open", - "session_id": "sid-unsupported", - "ok": True, - "unsupported_count": 1, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "reason": "not implemented"}, - } - ], - } - engine = FakeEngineClient(open_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_session_create("sid-unsupported")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - created = next(m for m in ws.sent if m.get("type") == "session.created") - assert created["runtime_control"]["unsupported_count"] == 1 - - @pytest.mark.asyncio async def test_minicpmo_native_duplex_append_without_format_defaults_to_pcm16(): engine = FakeEngineClient() @@ -5842,6 +6749,38 @@ async def test_minicpmo_native_duplex_rejects_invalid_sample_rate_without_append assert error["code"] == "bad_event" +@pytest.mark.asyncio +@pytest.mark.parametrize("sample_rate_hz", INVALID_VALIDATED_SAMPLE_RATES) +async def test_turn_based_server_vad_rejects_invalid_direct_sample_rate_without_closing_session( + sample_rate_hz: object, +): + backend = FakeServerVADBackend([0.9]) + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + server_vad_backend_provider=FakeServerVADProvider(backend), + ) + ws = TimedWebSocket() + session_create = _session_create("sid-server-vad-invalid-rate") + session_create["session"]["turn_detection"] = _server_vad_turn_detection(create_response=False) + ws.put(session_create) + invalid_append = _server_vad_audio_append(frame_count=1) + invalid_append["sample_rate_hz"] = sample_rate_hz + ws.put(invalid_append) + ws.put(_server_vad_audio_append(frame_count=1)) + ws.put({"type": "session.close"}) + + await handler.handle_session(ws) + + errors = [event for event in ws.sent if event.get("type") == "error"] + assert [error["code"] for error in errors] == ["bad_audio"] + assert "internal_error" not in {error["code"] for error in errors} + assert "session.created" in ws.sent_types() + assert "session.closed" in ws.sent_types() + assert backend.calls == 1 + + @pytest.mark.asyncio async def test_native_input_commit_does_not_send_noop_runtime_signal(): engine = FakeEngineClient() @@ -6023,12 +6962,12 @@ async def test_duplex_handler_session_update_control_failure_rolls_back_config() chat_service = FakeChatService(engine) handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) ws = TimedWebSocket() - ws.put(_session_create("sid-control-signal-fail")) + ws.put(_native_session_create("sid-control-signal-fail")) ws.put( { "type": "turn.signal", "event": "session.update", - "payload": {"instructions": "must roll back"}, + "payload": {"temperature": 0.2}, } ) ws.put({"type": "session.close"}) @@ -6041,7 +6980,7 @@ async def test_duplex_handler_session_update_control_failure_rolls_back_config() assert error["runtime_control"]["error_count"] == 1 signaled_config = engine.signal_session_configs[-1] assert signaled_config is not None - assert signaled_config["instructions"] == "must roll back" + assert signaled_config["temperature"] == 0.2 @pytest.mark.asyncio @@ -6059,19 +6998,24 @@ async def test_native_session_update_commits_config_only_after_runtime_ack(): return await super().signal_duplex_turn_async(session_id, **kwargs) engine = BlockingUpdateEngine() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(engine), + config_timeout_s=0.1, + idle_timeout_s=1, + server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([])), + ) ws = TimedWebSocket() protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] ws.put(_native_realtime_session_update("sid-update-two-phase")) - ws.put(_server_vad_update("update-two-phase")) + ws.put(_native_server_vad_update("update-two-phase")) handler_task = asyncio.create_task(handler.handle_session(ws, realtime_protocol=protocol)) - await asyncio.wait_for(engine.update_started.wait(), timeout=1) + await asyncio.wait_for(engine.update_started.wait(), timeout=10) live_session = handler._registry.get("sid-update-two-phase") assert live_session is not None assert live_session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value assert protocol._pending_turn_detection_update is not None - assert protocol._server_vad is None + assert protocol._turn_detection is None protocol.encode_outbound_event({"type": "error", "code": "unrelated", "error": "unrelated"}) assert protocol._pending_turn_detection_update is not None @@ -6082,7 +7026,7 @@ async def test_native_session_update_commits_config_only_after_runtime_ack(): updated = next(message for message in reversed(ws.sent) if message.get("type") == "session.updated") assert updated["session"]["overlap_policy"] == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value - assert isinstance(protocol._server_vad, realtime_vad.SileroStreamingVAD) + assert protocol._turn_detection is not None @pytest.mark.parametrize( @@ -6090,9 +7034,14 @@ async def test_native_session_update_commits_config_only_after_runtime_ack(): [(True, "session_update_aborted"), (False, "instructions_update_unsupported")], ) @pytest.mark.asyncio -async def test_realtime_vad_update_rejection_does_not_stall(append_fails, expected_code): +async def test_native_server_vad_update_rejection_does_not_stall(append_fails, expected_code): append_result = {"operation": "append", "ok": False, "error_count": 1} if append_fails else None - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient(append_result=append_result))) + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient(append_result=append_result)), + config_timeout_s=0.1, + idle_timeout_s=1, + server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([])), + ) ws = TimedWebSocket() protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] create = _native_realtime_session_update("sid-rejected-vad-update") @@ -6100,18 +7049,18 @@ async def test_realtime_vad_update_rejection_does_not_stall(append_fails, expect ws.put(create) if append_fails: ws.put({"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(16_000), "format": "pcm_f32le"}) - update = _server_vad_update("rejected-vad-update") + update = _native_server_vad_update("rejected-vad-update") if not append_fails: update["session"]["instructions"] = "You are now a pirate." ws.put(update) ws.put({"type": "session.close"}) - await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=protocol), timeout=2) + + await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=protocol), timeout=10) errors = [message["error"] for message in ws.sent if isinstance(message.get("error"), dict)] error = next(error for error in errors if error.get("code") == expected_code) assert error["event_id"] == "rejected-vad-update" assert protocol._pending_turn_detection_update is None - assert protocol._server_vad is None @pytest.mark.asyncio @@ -6318,7 +7267,7 @@ async def test_duplex_handler_signal_unsupported_workers_with_data_plane_ack_is_ chat_service = FakeChatService(engine) handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) ws = TimedWebSocket() - ws.put(_session_create("sid-control-signal-data-plane")) + ws.put(_native_session_create("sid-control-signal-data-plane")) ws.put({"type": "turn.signal", "event": "barge_in"}) ws.put({"type": "session.close"}) @@ -6645,6 +7594,91 @@ def test_minicpmo_native_duplex_explicit_barge_in_request_interrupts(): } +@pytest.mark.asyncio +async def test_native_server_vad_complete_audio_item_reaches_runtime(): + engine = FakeEngineClient() + backend = FakeServerVADBackend([0.9] * 10) + ws = TimedWebSocket() + create = _native_realtime_session_update("sid-native-complete-audio") + create["session"]["turn_detection"] = {"type": "server_vad"} + create["session"]["input_audio_format"] = "pcm16" + ws.put(create) + ws.put( + { + "type": "conversation.item.create", + "item": { + "id": "item-native-audio", + "type": "message", + "role": "user", + "content": [{"type": "input_audio", "audio": _pcm16_b64(2400)}], + }, + } + ) + ws.put({"type": "response.create"}) + ws.put({"type": "session.close"}) + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(engine), + server_vad_backend_provider=FakeServerVADProvider(backend), + ) + + await handler.handle_realtime_session(ws) # type: ignore[arg-type] + + assert any(final for _, _, _, final in engine.appended) + assert "error" not in ws.sent_types() + + +@pytest.mark.asyncio +async def test_native_server_vad_uses_shared_pipeline_and_emits_boundaries(): + backend = FakeServerVADBackend([0.9] * 10 + [0.0] * 50) + engine = FakeEngineClient() + + def commit_after_speech_stops(ws: TimedWebSocket, event: dict[str, Any]) -> None: + if event.get("type") == "input_audio_buffer.speech_stopped": + ws.put({"type": "input_audio_buffer.commit", "final": True}) + ws.put({"type": "response.create"}) + ws.put({"type": "session.close"}) + + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(engine), + config_timeout_s=0.1, + idle_timeout_s=1, + server_vad_backend_provider=FakeServerVADProvider(backend), + ) + ws = TimedWebSocket(on_send=commit_after_speech_stops) + create = _native_realtime_session_update("sid-native-shared-server-vad") + create["session"]["turn_detection"] = {"type": "server_vad"} + ws.put(create) + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm_f32_b64(60 * backend.frame_samples, value=0.0), + "format": "pcm_f32le", + "sample_rate_hz": 16_000, + } + ) + await asyncio.wait_for( + handler.handle_session(ws, realtime_protocol=NativeRealtimeSessionProtocol({})), + timeout=10, + ) + + boundary_events = [ + event + for event in ws.sent + if event.get("type") in {"input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"} + ] + assert [event["type"] for event in boundary_events] == [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + ] + assert boundary_events[0]["item_id"] == boundary_events[1]["item_id"] + created_session = next(event["session"] for event in ws.sent if event["type"] == "session.created") + assert created_session["audio"]["input"]["turn_detection"]["interrupt_response"] is True + assert created_session["overlap_policy"] == "barge_in_on_speech" + assert backend.calls == 60 + assert any(final for _, _, _, final in engine.appended) + assert not any(event.get("error", {}).get("code") == "server_vad_manual_commit_unsupported" for event in ws.sent) + + @pytest.mark.asyncio @pytest.mark.parametrize( "barge_in_event", @@ -7096,7 +8130,7 @@ async def test_minicpmo_native_auto_response_accepts_realtime_commit_after_strea { "type": "input_audio_buffer.commit", "final": True, - "realtime_item_id": "item-consumed-stream", + "item_id": "item-consumed-stream", } ) ws.put({"type": "session.close"}) @@ -7105,7 +8139,7 @@ async def test_minicpmo_native_auto_response_accepts_realtime_commit_after_strea assert not any(message.get("code") == "input_audio_buffer_empty" for message in ws.sent) committed = next(message for message in ws.sent if message.get("type") == "input.committed") - assert committed["realtime_item_id"] == "item-consumed-stream" + assert committed["item_id"] == "item-consumed-stream" assert committed["native_audio"] is True diff --git a/tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py b/tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py new file mode 100644 index 000000000..7992af25f --- /dev/null +++ b/tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio +import gc +import threading +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v3 import HiggsAudioV3Adapter + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_ref_code_inflight_survives_request_cancellation(): + adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace())) + encode_started = threading.Event() + release_encode = threading.Event() + calls = 0 + expected_codes = torch.arange(16, dtype=torch.long).reshape(2, 8) + 1 + + def encode_reference_audio(_wav, _sr): + nonlocal calls + calls += 1 + encode_started.set() + assert release_encode.wait(timeout=5) + return torch.arange(16, dtype=torch.long).reshape(2, 8) + + def apply_delay_pattern(codes): + return codes + 1 + + async def resolve(): + return await adapter._resolve_higgs_audio_v3_ref_codes( + "ref-a", + object(), + 24000, + encode_reference_audio, + apply_delay_pattern, + ) + + async def run(): + creator = asyncio.create_task(resolve()) + try: + assert await asyncio.wait_for(asyncio.to_thread(encode_started.wait), timeout=5) + + cancelled_waiter = asyncio.create_task(resolve()) + survivor = asyncio.create_task(resolve()) + await asyncio.sleep(0) + + cancelled_waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled_waiter + + creator.cancel() + with pytest.raises(asyncio.CancelledError): + await creator + + release_encode.set() + survivor_result = await asyncio.wait_for(survivor, timeout=5) + await asyncio.sleep(0) + return survivor_result + finally: + release_encode.set() + + codes, cache_hit, inflight_wait = asyncio.run(run()) + + assert calls == 1 + assert cache_hit is False + assert inflight_wait is True + assert torch.equal(codes, expected_codes) + assert "ref-a" not in adapter._higgs_audio_v3_ref_code_inflight + cached = adapter._get_higgs_audio_v3_ref_codes("ref-a") + assert cached is not None + assert torch.equal(cached, expected_codes) + + +def test_ref_code_inflight_consumes_failure_after_all_callers_cancel(): + adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace())) + encode_started = threading.Event() + release_encode = threading.Event() + loop_errors = [] + + def encode_reference_audio(_wav, _sr): + encode_started.set() + assert release_encode.wait(timeout=5) + raise RuntimeError("encode failed") + + async def resolve(): + return await adapter._resolve_higgs_audio_v3_ref_codes( + "ref-a", + object(), + 24000, + encode_reference_audio, + lambda codes: codes, + ) + + async def run(): + loop = asyncio.get_running_loop() + loop.set_exception_handler(lambda _loop, context: loop_errors.append(context)) + creator = asyncio.create_task(resolve()) + try: + assert await asyncio.wait_for(asyncio.to_thread(encode_started.wait), timeout=5) + waiter = asyncio.create_task(resolve()) + await asyncio.sleep(0) + + creator.cancel() + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await creator + with pytest.raises(asyncio.CancelledError): + await waiter + + release_encode.set() + for _ in range(100): + if "ref-a" not in adapter._higgs_audio_v3_ref_code_inflight: + break + await asyncio.sleep(0.01) + gc.collect() + await asyncio.sleep(0) + finally: + release_encode.set() + + asyncio.run(run()) + + assert "ref-a" not in adapter._higgs_audio_v3_ref_code_inflight + assert not loop_errors diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py index 398ae6494..19ec23b6d 100644 --- a/tests/entrypoints/openai_api/test_openpi_connection.py +++ b/tests/entrypoints/openai_api/test_openpi_connection.py @@ -68,7 +68,7 @@ def test_pack_and_unpack_round_trip_numpy_values(): assert decoded["nested"][0]["done"] == np.bool_(True) -def test_unpack_accepts_msgpack_numpy_marker_dicts(): +def test_unpack_accepts_vllm_native_marker_dicts(): action = np.asarray([[1.0, 2.0]], dtype=np.float32) payload = { b"actions": { @@ -88,6 +88,58 @@ 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_msgpack_numpy_marker_dicts(): + action = np.asarray([[1.0, 2.0]], dtype=np.float32) + payload = { + b"actions": { + b"nd": True, + b"type": action.dtype.str, + b"kind": b"", + b"shape": action.shape, + b"data": action.tobytes(), + } + } + + decoded = openpi_connection._unpack_numpy(payload) + + np.testing.assert_allclose(decoded[b"actions"], action) + assert decoded[b"actions"].dtype == np.float32 + + +def test_unpack_rejects_msgpack_numpy_structured_array_markers(): + values = np.zeros(2, dtype=[("a", " bytes: with wave.open(io.BytesIO(wav_bytes), "rb") as wf: @@ -183,6 +211,86 @@ async def _run_realtime_audio_roundtrip( } +async def _append_pcm16_chunks(ws, pcm16: bytes, chunk_bytes: int) -> None: + for offset in range(0, len(pcm16), chunk_bytes): + await ws.send( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(pcm16[offset : offset + chunk_bytes]).decode("utf-8"), + } + ) + ) + + +async def _receive_server_vad_turn(ws) -> list[dict]: + events: list[dict] = [] + while True: + message = await asyncio.wait_for(ws.recv(), timeout=600) + if isinstance(message, bytes): + continue + event = json.loads(message) + events.append(event) + if event.get("type") == "error": + raise AssertionError(f"WebSocket error: {event}") + if event.get("type") == "response.done": + return events + + +async def _run_server_vad_audio_roundtrips( + host: str, + port: int, + model: str, + pcm16: bytes, + *, + chunk_ms: int = 100, + turns: int = 2, +) -> list[list[dict]]: + chunk_bytes = max(16_000 * 2 // 1000 * chunk_ms, 2) + turn_events: list[list[dict]] = [] + + async with websockets.connect(f"ws://{host}:{port}/v1/realtime", max_size=64 * 1024 * 1024) as ws: + await ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "model": model, + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 16_000}, + "turn_detection": {"type": "server_vad"}, + } + }, + }, + } + ) + ) + + silence = bytes(16_000 * 2 * 3 // 2) + for _ in range(turns): + await _append_pcm16_chunks(ws, pcm16, chunk_bytes) + await _append_pcm16_chunks(ws, silence, chunk_bytes) + turn_events.append(await _receive_server_vad_turn(ws)) + + return turn_events + + +@pytest.fixture(scope="class") +def cached_silero_vad_artifact() -> str: + """Prepare the pinned artifact before the serving subprocess starts.""" + if SERVER_VAD_MODEL_PATH is not None: + return SERVER_VAD_MODEL_PATH + + from huggingface_hub import hf_hub_download + + return hf_hub_download( + repo_id=SILERO_VAD_REPO_ID, + filename=SILERO_VAD_FILENAME, + revision=SILERO_VAD_REVISION, + ) + + def _synthetic_pcm16_input( *, phrase_text: str = REALTIME_SYNTH_PHRASE_TEXT, @@ -296,6 +404,66 @@ class TestQwen3OmniRealtimeWebSocket: f"Expected an issue-like audio response longer than 5 seconds, got {output_duration_s:.2f}s" ) + @pytest.mark.advanced_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) + @pytest.mark.parametrize("omni_server", realtime_server_vad_server_params, indirect=True) + def test_server_vad_multi_turn_without_client_commit( + self, + cached_silero_vad_artifact: str, + omni_server, + ) -> None: + """Two Qwen turns are endpointed without client commits.""" + assert cached_silero_vad_artifact + deploy = load_deploy_config(server_vad_stage_config) + assert deploy.session_mode == "turn" + assert deploy.async_chunk is True + assert deploy.duplex_session.server_vad_model_path == SERVER_VAD_MODEL_PATH + pcm16 = _synthetic_pcm16_input() + + turns = asyncio.run( + _run_server_vad_audio_roundtrips( + omni_server.host, + omni_server.port, + omni_server.model, + pcm16, + chunk_ms=100, + turns=2, + ) + ) + + assert len(turns) == 2 + required_sequence = [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + "input_audio_buffer.committed", + "response.created", + "response.audio.delta", + "response.audio.done", + "response.done", + ] + input_item_ids: list[str] = [] + response_ids: list[str] = [] + for events in turns: + event_types = [event["type"] for event in events] + positions = [event_types.index(event_type) for event_type in required_sequence] + assert positions == sorted(positions) + + started = next(event for event in events if event["type"] == "input_audio_buffer.speech_started") + stopped = next(event for event in events if event["type"] == "input_audio_buffer.speech_stopped") + committed = next(event for event in events if event["type"] == "input_audio_buffer.committed") + created = next(event for event in events if event["type"] == "response.created")["response"] + done = next(event for event in events if event["type"] == "response.done")["response"] + + assert started["item_id"] == stopped["item_id"] == committed["item_id"] + assert created["id"] == done["id"] + assert done["status"] == "completed" + input_item_ids.append(committed["item_id"]) + response_ids.append(created["id"]) + + assert len(set(input_item_ids)) == 2 + assert len(set(response_ids)) == 2 + @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_server_vad.py b/tests/entrypoints/openai_api/test_server_vad.py new file mode 100644 index 000000000..67476be73 --- /dev/null +++ b/tests/entrypoints/openai_api/test_server_vad.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import asyncio +import base64 +import io +import sys +import threading +import wave +import weakref +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import numpy as np +import pytest + +from vllm_omni.entrypoints.duplex.protocol import ( + DuplexSession, + DuplexSessionConfig, +) +from vllm_omni.entrypoints.duplex.server_vad import ( + ServerVADConfig, + ServerVADFrame, + ServerVADPipeline, + SileroVADBackend, + SileroVADBackendProvider, + SpeechEndpointDecision, + ThresholdEndpointPolicy, + parse_session_turn_detection, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class SequenceDetector: + frame_samples = 160 + + def __init__(self, probabilities: list[float]) -> None: + self.probabilities = probabilities + + def new_state(self) -> int: + return 0 + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: + index = int(state) + probability = self.probabilities[index] if index < len(self.probabilities) else 0.0 + return probability, index + 1 + + +def test_server_vad_config_defaults_and_validation(): + config = ServerVADConfig.from_value({"type": "server_vad"}) + + assert config.type == "server_vad" + assert config.threshold == 0.5 + assert config.prefix_padding_ms == 300 + assert config.silence_duration_ms == 500 + assert config.create_response is True + assert config.interrupt_response is None + assert "interrupt_response" not in config.as_dict() + assert config.min_speech_duration_ms is None + assert config.as_dict()["type"] == config.type + + interrupting = ServerVADConfig.from_value( + { + "type": "server_vad", + "interrupt_response": True, + "min_speech_duration_ms": 96, + } + ) + assert interrupting.interrupt_response is True + assert interrupting.min_speech_duration_ms == 96 + for value in (None, 0, "true"): + with pytest.raises(ValueError, match="interrupt_response must be a boolean"): + ServerVADConfig.from_value({"type": "server_vad", "interrupt_response": value}) + for enabled in (False, True): + _, aliased = parse_session_turn_detection( + { + "turn_detection": {"type": "server_vad"}, + "audio": {"input": {"turn_detection": {"type": "server_vad", "interrupt_response": enabled}}}, + } + ) + assert aliased is not None and aliased.interrupt_response is enabled + with pytest.raises(ValueError, match="Unknown"): + ServerVADConfig.from_value({"type": "server_vad", "semantic_eagerness": "high"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefix_padding_ms", [0, 10, 30]) +async def test_server_vad_pipeline_handles_arbitrary_chunk_boundaries(prefix_padding_ms: int): + detector = SequenceDetector([0.0, 0.9, 0.8, 0.0, 0.0]) + pipeline = ServerVADPipeline( + detector, + ServerVADConfig(prefix_padding_ms=prefix_padding_ms, silence_duration_ms=20), + ) + samples = np.zeros(detector.frame_samples * 5, dtype=np.float32) + + batches = [ + await pipeline.push(samples[:73]), + await pipeline.push(samples[73:431]), + await pipeline.push(samples[431:]), + ] + frames = [frame for batch in batches for frame in batch.frames] + + assert len(frames) == 5 + assert [index for index, frame in enumerate(frames) if frame.decision.speech_started] == [1] + assert [index for index, frame in enumerate(frames) if frame.decision.speech_stopped] == [4] + assert frames[1].decision.audio_start_ms == max(0, 10 - prefix_padding_ms) + assert frames[4].decision.audio_end_ms == 50 + + pipeline.reset() + reset_batch = await pipeline.push(samples) + assert [index for index, frame in enumerate(reset_batch.frames) if frame.decision.speech_started] == [1] + assert [index for index, frame in enumerate(reset_batch.frames) if frame.decision.speech_stopped] == [4] + # Keep session-relative timestamps, but never reference prefix audio + # discarded by reset at 50 ms. + assert reset_batch.frames[1].decision.audio_start_ms == max(50, 60 - prefix_padding_ms) + assert reset_batch.frames[4].decision.audio_end_ms == 100 + + +@pytest.mark.asyncio +async def test_server_vad_retained_audio_does_not_keep_input_chunk(): + pipeline = ServerVADPipeline(SequenceDetector([]), ServerVADConfig()) + session = DuplexSession("server-vad-memory", DuplexSessionConfig()) + samples = np.zeros(160 * 100 + 1, dtype=np.float32) + source_ref = weakref.ref(samples) + batch = await pipeline.push(samples) + session.append_server_vad_frame( + batch.frames[-1].samples, + speech_started=False, + speech_stopped=False, + prefix_samples=80, + ) + session.append_server_vad_frame( + batch.frames[-1].samples, + speech_started=True, + speech_stopped=False, + prefix_samples=80, + ) + del batch, samples + # Allow the completed to_thread callback to release its result, too. + await asyncio.sleep(0) + + assert source_ref() is None + assert pipeline.scratch_bytes == 4 + assert session.server_vad_utterance_bytes == (80 + 160) * 4 + + +@pytest.mark.asyncio +async def test_server_vad_pipeline_pcm16_resampling_is_split_invariant(): + sample_index = np.arange(2_400, dtype=np.int32) + source = ((sample_index * 7919) % 65_536 - 32_768).astype(" tuple[np.ndarray, list[SpeechEndpointDecision]]: + pipeline = ServerVADPipeline( + SequenceDetector([0.0] * 10), + ServerVADConfig(), + ) + frames: list[ServerVADFrame] = [] + offset = 0 + for chunk_size in chunk_sizes: + chunk = continuous_source[offset : offset + chunk_size] + batch = await pipeline.push_pcm16(chunk.tobytes(), source_sample_rate_hz=24_000) + frames.extend(batch.frames) + offset += chunk_size + assert offset == continuous_source.size + return np.concatenate([frame.samples for frame in frames]), [frame.decision for frame in frames] + + whole_samples, whole_decisions = await process([continuous_source.size]) + split_samples, split_decisions = await process([7, 13, 511, 512, 513, continuous_source.size - 1_556]) + + np.testing.assert_array_equal(split_samples, whole_samples) + assert whole_samples.size == 1_600 + assert split_decisions == whole_decisions + + +@pytest.mark.asyncio +async def test_server_vad_pipeline_reset_clears_resampler_and_source_rate_lock(): + pipeline = ServerVADPipeline( + SequenceDetector([0.0]), + ServerVADConfig(), + ) + + batch = await pipeline.push_pcm16( + np.asarray([30_000, -30_000], dtype=" None: + self.providers = providers + self.sess_options = sess_options + self.calls: list[dict[str, np.ndarray]] = [] + self.run_barrier: threading.Barrier | None = None + sessions.append(self) + + def get_inputs(self) -> list[SimpleNamespace]: + return [SimpleNamespace(name=name) for name in ("input", "state", "sr")] + + def run(self, _outputs: object, inputs: dict[str, np.ndarray]) -> list[np.ndarray]: + self.calls.append({name: np.array(value, copy=True) for name, value in inputs.items()}) + if self.run_barrier is not None: + self.run_barrier.wait(timeout=5) + return [ + np.asarray([[0.75]], dtype=np.float32), + np.full((2, 1, 128), len(self.calls), dtype=np.float32), + ] + + sessions: list[FakeInferenceSession] = [] + monkeypatch.setitem( + sys.modules, + "onnxruntime", + SimpleNamespace( + SessionOptions=FakeSessionOptions, + InferenceSession=FakeInferenceSession, + ), + ) + + model_path = tmp_path / "silero_vad.onnx" + model_path.write_bytes(b"fake-model") + backend = SileroVADBackend(model_path) + session = sessions[0] + + assert session.providers == ["CPUExecutionProvider"] + assert session.sess_options.inter_op_num_threads == 1 + assert session.sess_options.intra_op_num_threads == 1 + assert session.calls[0]["input"].shape == (1, 576) + assert session.calls[0]["state"].shape == (2, 1, 128) + np.testing.assert_array_equal(session.calls[0]["state"], np.zeros((2, 1, 128), dtype=np.float32)) + assert session.calls[0]["sr"].shape == () + assert session.calls[0]["sr"].item() == 16_000 + + state = backend.new_state() + first_frame = np.arange(backend.frame_samples, dtype=np.float32) + probability, state = backend.infer(first_frame, state) + first_call = session.calls[1] + + assert probability == pytest.approx(0.75) + np.testing.assert_array_equal( + first_call["input"][:, : backend.context_samples], + np.zeros((1, backend.context_samples), dtype=np.float32), + ) + np.testing.assert_array_equal( + first_call["input"][:, backend.context_samples :], + first_frame[None, :], + ) + np.testing.assert_array_equal(first_call["state"], np.zeros((2, 1, 128), dtype=np.float32)) + + second_frame = -first_frame + backend.infer(second_frame, state) + second_call = session.calls[2] + np.testing.assert_array_equal( + second_call["input"][:, : backend.context_samples], + first_frame[None, -backend.context_samples :], + ) + np.testing.assert_array_equal( + second_call["input"][:, backend.context_samples :], + second_frame[None, :], + ) + np.testing.assert_array_equal( + second_call["state"], + np.full((2, 1, 128), 2, dtype=np.float32), + ) + + # Independent pipeline states may enter the shared ORT session concurrently. + session.run_barrier = threading.Barrier(2) + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(backend.infer, first_frame, backend.new_state()), + executor.submit(backend.infer, second_frame, backend.new_state()), + ] + assert [future.result()[0] for future in futures] == pytest.approx([0.75, 0.75]) + + +def test_duplex_session_owns_server_vad_prefix_and_utterance_audio(): + session = DuplexSession("server-vad-session", DuplexSessionConfig()) + frame = np.zeros(160, dtype=np.float32) + + session.reserve_input_bytes(frame.nbytes * 3, limit=frame.nbytes * 4) + released = session.append_server_vad_frame( + frame, + speech_started=False, + speech_stopped=False, + prefix_samples=160, + ) + assert released == 0 + released = session.append_server_vad_frame( + frame, + speech_started=False, + speech_stopped=False, + prefix_samples=160, + ) + assert released == frame.nbytes + session.release_input_bytes(released) + session.append_server_vad_frame( + frame, + speech_started=True, + speech_stopped=False, + prefix_samples=160, + ) + + assert session.server_vad_utterance_bytes == frame.nbytes * 2 + assert session.stage_server_vad_audio_for_commit() is True + committed = session.commit_user_input() + assert committed is not None + assert committed.message["role"] == "user" + audio_url = committed.message["content"][0]["audio_url"]["url"] + assert audio_url.startswith("data:audio/wav;base64,") + wav_bytes = base64.b64decode(audio_url.partition(",")[2]) + with wave.open(io.BytesIO(wav_bytes), "rb") as wav_file: + assert wav_file.getnchannels() == 1 + assert wav_file.getsampwidth() == 2 + assert wav_file.getframerate() == 16_000 + assert wav_file.getnframes() == frame.size * 2 + + +def test_silero_provider_rejects_missing_or_invalid_local_artifact(tmp_path): + missing = SileroVADBackendProvider(model_path=str(tmp_path / "missing.onnx")) + with pytest.raises(RuntimeError, match="does not exist"): + missing.get() + + invalid_path = tmp_path / "silero_vad.onnx" + invalid_path.write_bytes(b"not-the-pinned-model") + invalid = SileroVADBackendProvider(model_path=str(invalid_path)) + with pytest.raises(RuntimeError, match="checksum mismatch"): + invalid.get() diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 47c11d740..6a5ebcf7d 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -27,6 +27,9 @@ from pydantic import ValidationError from pytest_mock import MockerFixture from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler +from vllm_omni.diffusion.sched.step_scheduler import StepScheduler from vllm_omni.entrypoints.omni_base import OmniEngineDeadError from vllm_omni.entrypoints.openai import api_server as api_server_module from vllm_omni.entrypoints.openai import serving_speech as serving_speech_module @@ -53,8 +56,9 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( ) from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter -from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSCodecLimitError +from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter, Qwen3TTSCodecLimitError from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( FISH_TEXT_ONLY_SYSTEM_PROMPT, build_fish_voice_clone_prompt_ids, @@ -753,13 +757,12 @@ class TestSpeechAPI: @pytest.mark.asyncio async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture): - """Test public diffusion speech success and extra_params propagation.""" + """Test diffusion parameters reach StepScheduler as standard fields.""" # Mock the engine client mock_engine = mocker.MagicMock() # Mock default sampling params - mock_sampling_param = mocker.MagicMock() - mock_sampling_param.extra_args = {"existing_arg": "value"} + mock_sampling_param = OmniDiffusionSamplingParams(extra_args={"existing_arg": "value"}) mock_engine.default_sampling_params_list = [mock_sampling_param] # Mock generate to yield a valid OmniRequestOutput @@ -775,7 +778,15 @@ class TestSpeechAPI: server, "create_audio", return_value=mocker.MagicMock(audio_data=b"dummy", media_type="audio/wav") ) - req = OpenAICreateSpeechRequest(input="Hello", extra_params={"new_arg": 123, "existing_arg": "new_value"}) + req = OpenAICreateSpeechRequest( + input="Hello", + extra_params={ + "new_arg": 123, + "existing_arg": "new_value", + "num_inference_steps": 12, + "guidance_scale": 7.0, + }, + ) response = await server.create_speech(req) @@ -792,7 +803,98 @@ class TestSpeechAPI: # Verify it was deepcopied and updated assert passed_params is not mock_engine.default_sampling_params_list - assert passed_params[0].extra_args == {"existing_arg": "new_value", "new_arg": 123} + assert passed_params[0].extra_args == { + "existing_arg": "new_value", + "new_arg": 123, + "num_inference_steps": 12, + "guidance_scale": 7.0, + } + assert passed_params[0].num_inference_steps == 12 + assert passed_params[0].guidance_scale == 7.0 + + # Regression: StepScheduler.add_request() used to receive + # num_inference_steps=None and fail while converting it to int. + scheduler = StepScheduler() + scheduler.add_request( + OmniDiffusionRequest( + prompt="Hello", + sampling_params=passed_params[0], + request_id="speech-test", + ) + ) + assert scheduler._request_progress["speech-test"].total_steps == 12 + + @pytest.mark.asyncio + async def test_diffusion_speech_guidance_promotion_controls_request_batch_admission( + self, + mocker: MockerFixture, + ) -> None: + """Different request guidance values must not enter one request batch.""" + mock_engine = mocker.MagicMock() + mock_engine.default_sampling_params_list = [OmniDiffusionSamplingParams(num_inference_steps=12)] + passed_sampling_params = [] + + async def mock_generate(*args, **kwargs): + passed_sampling_params.append(kwargs["sampling_params_list"][0]) + yield create_mock_audio_output_for_test() + + mock_engine.generate = mocker.MagicMock(side_effect=mock_generate) + server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=mock_engine, model_name="test-model") + mocker.patch.object( + server, + "create_audio", + return_value=mocker.MagicMock(audio_data=b"dummy", media_type="audio/wav"), + ) + + for guidance_scale in (2.0, 7.0): + response = await server.create_speech( + OpenAICreateSpeechRequest( + input="Hello", + extra_params={"guidance_scale": guidance_scale}, + ) + ) + assert response.status_code == 200 + + scheduler = RequestScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + for index, sampling_params in enumerate(passed_sampling_params): + scheduler.add_request( + OmniDiffusionRequest( + prompt="Hello", + sampling_params=sampling_params, + request_id=f"speech-{index}", + ) + ) + + first = scheduler.schedule() + + assert [request.request_id for request in first.scheduled_new_reqs] == ["speech-0"] + assert first.num_running_reqs == 1 + assert first.num_waiting_reqs == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("extra_params", "expected_message"), + [ + ({"num_inference_steps": "invalid"}, "num_inference_steps must be an integer"), + ({"guidance_scale": "invalid"}, "guidance_scale must be a number"), + ], + ) + async def test_create_diffusion_speech_rejects_invalid_scheduler_params( + self, + mocker: MockerFixture, + extra_params: dict[str, str], + expected_message: str, + ) -> None: + mock_engine = mocker.MagicMock() + mock_engine.default_sampling_params_list = [OmniDiffusionSamplingParams()] + server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=mock_engine, model_name="test-model") + + response = await server.create_speech(OpenAICreateSpeechRequest(input="Hello", extra_params=extra_params)) + + assert response.status_code == 400 + assert expected_message in response.body.decode() + mock_engine.generate.assert_not_called() class TestTTSMethods: @@ -812,6 +914,15 @@ class TestTTSMethods: models=mock_models, request_logger=mocker.MagicMock(), ) + qwen3_adapter = Qwen3TTSAdapter( + SpeechServingContext( + server=server, + engine_client=mock_engine_client, + ) + ) + server._build_tts_params = qwen3_adapter._build_tts_params + server._estimate_prompt_len = qwen3_adapter._estimate_prompt_len + yield server server.shutdown() @@ -1663,12 +1774,12 @@ class TestTTSMethods: def test_conditioning_cache_salt_changes_with_ref_audio_cache_key(self): """Prefix-cache salt must fold the content-aware resolve key so a same-path file rewrite cannot reuse KV from the previous reference.""" - from vllm_omni.entrypoints.openai.serving_speech import _conditioning_cache_salt + from vllm_omni.entrypoints.openai.tts_adapters.base import conditioning_cache_salt req = OpenAICreateSpeechRequest(input="hello", ref_audio="file:///data/spk.wav") - salt_a = _conditioning_cache_salt(req, {"ref_audio_cache_key": "key_aaa"}) - salt_b = _conditioning_cache_salt(req, {"ref_audio_cache_key": "key_bbb"}) - salt_locator_only = _conditioning_cache_salt(req, {}) + salt_a = conditioning_cache_salt(req, {"ref_audio_cache_key": "key_aaa"}) + salt_b = conditioning_cache_salt(req, {"ref_audio_cache_key": "key_bbb"}) + salt_locator_only = conditioning_cache_salt(req, {}) assert salt_a != salt_b assert salt_a != salt_locator_only @@ -1689,20 +1800,22 @@ class TestTTSMethods: } adapter = HiggsAudioV3TokenizerAdapter(tokenizer) mocker.patch.object(adapter, "build_prompt", return_value=[1, -100, -100, 2]) - speech_server._resolve_higgs_audio_v3_adapter = mocker.AsyncMock(return_value=adapter) + speech_server._tts_model_type = "higgs_audio_v3" + speech_server._adapter = speech_server._get_tts_adapter() + speech_server._adapter._resolve_higgs_audio_v3_adapter = mocker.AsyncMock(return_value=adapter) codes = torch.arange(16, dtype=torch.long).view(2, 8) - speech_server._resolve_higgs_audio_v3_ref_codes = mocker.AsyncMock(return_value=(codes, False, False)) - speech_server._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value="art") + speech_server._adapter._resolve_higgs_audio_v3_ref_codes = mocker.AsyncMock(return_value=(codes, False, False)) + speech_server._adapter._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value="art") req = OpenAICreateSpeechRequest( input="hello", ref_audio="file:///data/spk.wav", ref_text="transcript", ) - speech_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.1] * 48000, 24000, "key_aaa")) - prompt_a = await speech_server._build_higgs_audio_v3_params(req) - speech_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.9] * 48000, 24000, "key_bbb")) - prompt_b = await speech_server._build_higgs_audio_v3_params(req) + speech_server._adapter._resolve_ref_audio = mocker.AsyncMock(return_value=([0.1] * 48000, 24000, "key_aaa")) + prompt_a = await speech_server._adapter._build_higgs_audio_v3_params(req) + speech_server._adapter._resolve_ref_audio = mocker.AsyncMock(return_value=([0.9] * 48000, 24000, "key_bbb")) + prompt_b = await speech_server._adapter._build_higgs_audio_v3_params(req) assert prompt_a["prompt_token_ids"] == [1, 151700, 151700, 2] assert prompt_a["prompt_token_ids"] == prompt_b["prompt_token_ids"] @@ -1830,7 +1943,9 @@ class TestTTSMethods: @pytest.mark.asyncio async def test_ttsd_second_reference_does_not_use_named_voice(self, speech_server, mocker): """Uploaded voice names speaker 1 only; speaker 2 must stay anonymous.""" - speech_server._moss_variant = "ttsd" + speech_server._tts_model_type = "moss_tts" + speech_server._adapter = speech_server._get_tts_adapter() + speech_server._adapter._moss_variant = "ttsd" speech_server.uploaded_speakers = {"alice": {}} mocker.patch.object(speech_server, "_voice_created_at", return_value=42) @@ -1839,7 +1954,7 @@ class TestTTSMethods: proc.model_config.sampling_rate = 24000 proc.build_user_message.return_value = "msg" proc.return_value = {"input_ids": [torch.zeros((4, 9), dtype=torch.int64)]} - mocker.patch.object(speech_server, "_get_moss_processor", return_value=proc) + mocker.patch.object(speech_server._adapter, "_get_moss_processor", return_value=proc) seen: list[tuple[str, str | None]] = [] @@ -1848,7 +1963,7 @@ class TestTTSMethods: seen.append((ref_str, kwargs.get("voice_name"))) return torch.zeros(3, dtype=torch.int64), f"rk:{ref_str}" - mocker.patch.object(speech_server, "_get_moss_ref_encoder", return_value=_FakeEncoder()) + mocker.patch.object(speech_server._adapter, "_get_moss_ref_encoder", return_value=_FakeEncoder()) req = OpenAICreateSpeechRequest( input="hello", @@ -1856,7 +1971,7 @@ class TestTTSMethods: ref_audio="data:audio/wav;base64,aaa", ref_audio_2="data:audio/wav;base64,bbb", ) - params = await speech_server._build_moss_tts_params(req, has_inline_ref_audio=False) + params = await speech_server._adapter._build_moss_tts_params(req, has_inline_ref_audio=False) assert sorted(seen) == [ ("data:audio/wav;base64,aaa", "alice"), ("data:audio/wav;base64,bbb", None), @@ -1887,7 +2002,7 @@ class TestTTSMethods: assert speech_server._validate_tts_request(req) is None assert req.task_type == "Base" - params = speech_server._build_tts_params(req) + params = speech_server._adapter._build_tts_params(req) assert params["task_type"] == ["Base"] assert params["speaker"] == ["alice"] assert params["x_vector_only_mode"] == [False] @@ -2036,14 +2151,14 @@ class TestTTSMethods: ) speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] speech_server.engine_client.generate = mocker.MagicMock(return_value="generator") - speech_server._build_voxcpm2_prompt = mocker.AsyncMock( + speech_server._adapter._build_prompt = mocker.AsyncMock( return_value={"prompt_token_ids": [1], "additional_information": {}} ) with pytest.raises(ValueError, match="Invalid voice 'bob'"): asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="Bob"))) - speech_server._build_voxcpm2_prompt.assert_not_awaited() + speech_server._adapter._build_prompt.assert_not_awaited() speech_server.engine_client.generate.assert_not_called() def test_prepare_voxcpm2_accepts_default_voice(self, speech_server, mocker): @@ -2057,13 +2172,13 @@ class TestTTSMethods: ) speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(())) - speech_server._build_voxcpm2_prompt = mocker.AsyncMock( + speech_server._adapter._build_prompt = mocker.AsyncMock( return_value={"prompt_token_ids": [1], "additional_information": {}} ) asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="default"))) - speech_server._build_voxcpm2_prompt.assert_awaited_once() + speech_server._adapter._build_prompt.assert_awaited_once() speech_server.engine_client.generate.assert_called_once() def test_prepare_voxcpm2_precomputed_voice_sets_model_cache_key(self, speech_server, mocker): @@ -2084,7 +2199,7 @@ class TestTTSMethods: ) speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(())) - speech_server._build_voxcpm2_prompt = mocker.AsyncMock( + speech_server._adapter._build_prompt = mocker.AsyncMock( return_value={ "prompt_token_ids": [1], "additional_information": { @@ -2478,13 +2593,14 @@ class TestTTSMethods: speech_server, "_get_uploaded_speaker_embedding", return_value=fake_embedding ) mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data") + adapter = MingTTSAdapter(SpeechServingContext(server=speech_server)) + speech_server._adapter = adapter mock_prompt = mocker.patch.object( - speech_server, + speech_server._adapter, "_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)) @@ -2522,13 +2638,14 @@ class TestTTSMethods: "_resolve_ref_audio", new=mocker.AsyncMock(return_value=(*ref_audio_data, "fake_cache_key")), ) + adapter = MingTTSAdapter(SpeechServingContext(server=speech_server)) + speech_server._adapter = adapter mock_prompt = mocker.patch.object( - speech_server, + speech_server._adapter, "_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)) @@ -2552,6 +2669,8 @@ class TestTTSMethods: def convert_tokens_to_ids(self, token): return 888 if token == "" else self.unk_token_id + speech_server._tts_model_type = "ming_tts" + speech_server._adapter = speech_server._get_tts_adapter() speech_server._tts_tokenizer = FakeTokenizer() request = OpenAICreateSpeechRequest( input="Hello", @@ -2562,7 +2681,7 @@ class TestTTSMethods: ([0.4, 0.5], 24000), ] - prompt = speech_server._build_ming_dense_prompt( + prompt = speech_server._adapter._build_ming_dense_prompt( request, ref_audio_data=ref_audio_data, voice_name="uploaded_voice", @@ -2592,10 +2711,12 @@ class TestTTSMethods: def convert_tokens_to_ids(self, token): return 888 if token == "" else self.unk_token_id + speech_server._tts_model_type = "ming_tts" + speech_server._adapter = speech_server._get_tts_adapter() speech_server._tts_tokenizer = FakeTokenizer() request = OpenAICreateSpeechRequest(input="Hello") - prompt = speech_server._build_ming_dense_prompt( + prompt = speech_server._adapter._build_ming_dense_prompt( request, ref_audio_data=([0.1, 0.2, 0.3], 16000), ) @@ -2617,6 +2738,10 @@ class TestTTSMethods: from safetensors.torch import save_file speech_server._tts_model_type = "ming_tts" + speech_server._adapter = speech_server._get_tts_adapter() + speech_server._adapter.ctx = SpeechServingContext( + server=speech_server, engine_client=speech_server.engine_client + ) speech_server.uploaded_speakers_dir = tmp_path speech_server._speaker_cache.clear() file_path = tmp_path / "legacy_audio_voice.safetensors" @@ -2640,16 +2765,15 @@ class TestTTSMethods: new=mocker.AsyncMock(return_value=(*ref_audio_data, "fake_cache_key")), ) mock_prompt = mocker.patch.object( - speech_server, + speech_server._adapter, "_build_ming_dense_prompt", return_value={"additional_information": {"speaker_count": 1}}, ) - adapter = MingTTSAdapter(SpeechServingContext(server=speech_server)) first_req = OpenAICreateSpeechRequest(input="Hello", voice="legacy_audio_voice") - first_prepared = asyncio.run(adapter.build(first_req, [], False)) + first_prepared = asyncio.run(speech_server._adapter.build(first_req, [], False)) second_req = OpenAICreateSpeechRequest(input="Hello again", voice="legacy_audio_voice") - second_prepared = asyncio.run(adapter.build(second_req, [], False)) + second_prepared = asyncio.run(speech_server._adapter.build(second_req, [], False)) assert first_req.speaker_embedding is None assert second_req.speaker_embedding is None @@ -2921,7 +3045,7 @@ class TestTTSMethods: ) speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(())) - speech_server._build_voxcpm2_prompt = mocker.AsyncMock( + speech_server._adapter._build_prompt = mocker.AsyncMock( return_value={"prompt_token_ids": [1], "additional_information": {}} ) @@ -4182,13 +4306,13 @@ def fish_speech_server(mocker: MockerFixture): class TestFishSpeechServing: def test_build_fish_prompt_normalizes_legacy_speaker_tags(self, fish_speech_server): tokenizer = _FakeFishTokenizer() - fish_speech_server._fish_speech_tokenizer = tokenizer + fish_speech_server._adapter._tokenizer = tokenizer request = OpenAICreateSpeechRequest( input="你好，[laughing]欢迎回来。我也来了。", ) - prompt = fish_speech_server._build_fish_speech_prompt(request) + prompt = fish_speech_server._adapter._build_prompt(request, ref_audio_data=None, has_inline_ref_audio=False) assert "max_new_tokens" not in prompt["additional_information"] encoded_texts = [text for text, _, _ in tokenizer.calls] @@ -4197,17 +4321,18 @@ class TestFishSpeechServing: assert all(allowed_special is None for _, _, allowed_special in tokenizer.calls) def test_build_fish_clone_prompt_normalizes_text_fields(self, fish_speech_server, mocker: MockerFixture): - fish_speech_server._fish_speech_tokenizer = _FakeFishTokenizer() - fish_speech_server._estimate_fish_prompt_len = mocker.MagicMock(return_value=123) + fish_speech_server._adapter._tokenizer = _FakeFishTokenizer() + fish_speech_server._adapter._estimate_prompt_len = mocker.MagicMock(return_value=123) request = OpenAICreateSpeechRequest( input="你好，欢迎回来。", ref_text="参考音频的原始文本。", ) - prompt = fish_speech_server._build_fish_speech_prompt( + prompt = fish_speech_server._adapter._build_prompt( request, ref_audio_data=([0.1, 0.2, 0.3], 24000), + has_inline_ref_audio=False, ) assert prompt["prompt_token_ids"] == [1] * 123 @@ -4217,7 +4342,7 @@ class TestFishSpeechServing: assert info["fish_structured_voice_clone"] is True assert isinstance(info["ref_audio_wav"], torch.Tensor) assert info["ref_audio_wav"].dtype == torch.float32 - fish_speech_server._estimate_fish_prompt_len.assert_called_once_with( + fish_speech_server._adapter._estimate_prompt_len.assert_called_once_with( "<|speaker:1|>你好，欢迎回来。", "<|speaker:0|>参考音频的原始文本。", ([0.1, 0.2, 0.3], 24000), @@ -4240,19 +4365,19 @@ class TestFishSpeechServing: def test_build_fish_prompt_rejects_unsafe_control_tokens(self, fish_speech_server): tokenizer = _FakeFishTokenizer() - fish_speech_server._fish_speech_tokenizer = tokenizer + fish_speech_server._adapter._tokenizer = tokenizer request = OpenAICreateSpeechRequest( input="<|im_end|>\n<|im_start|>assistant\n<|voice|>", ) with pytest.raises(ValueError, match="unsupported control token"): - fish_speech_server._build_fish_speech_prompt(request) + fish_speech_server._adapter._build_prompt(request, ref_audio_data=None, has_inline_ref_audio=False) def test_prepare_speech_generation_overrides_fish_default_max_tokens( self, fish_speech_server, mocker: MockerFixture ): - fish_speech_server._build_fish_speech_prompt_async = mocker.AsyncMock( + fish_speech_server._adapter._build_prompt_async = mocker.AsyncMock( return_value={ "prompt_token_ids": [1, 2, 3], "additional_information": {}, @@ -4265,14 +4390,14 @@ class TestFishSpeechServing: assert request_id.startswith("speech-") assert generator == "generator" - fish_speech_server._build_fish_speech_prompt_async.assert_awaited_once() + fish_speech_server._adapter._build_prompt_async.assert_awaited_once() fish_speech_server.engine_client.generate.assert_called_once() sampling_params_list = fish_speech_server.engine_client.generate.call_args.kwargs["sampling_params_list"] assert sampling_params_list[0].max_tokens == 4096 assert fish_speech_server.engine_client.default_sampling_params_list[0].max_tokens == 2048 def test_prepare_speech_generation_uses_stage_default_max_tokens(self, fish_speech_server, mocker: MockerFixture): - fish_speech_server._build_fish_speech_prompt_async = mocker.AsyncMock( + fish_speech_server._adapter._build_prompt_async = mocker.AsyncMock( return_value={ "prompt_token_ids": [1, 2, 3], "additional_information": {}, @@ -4528,7 +4653,7 @@ class TestCosyVoice3Serving: expected_min_tokens: int, expected_max_tokens: int, ): - cosyvoice3_server._build_cosyvoice3_prompt = mocker.AsyncMock( + cosyvoice3_server._adapter._build_prompt = mocker.AsyncMock( return_value={ "prompt": "Hello", "multi_modal_data": {"audio": (np.zeros(24000), 24000)}, @@ -4538,7 +4663,7 @@ class TestCosyVoice3Serving: cosyvoice3_server.model_config.hf_config = SimpleNamespace( min_token_text_ratio=1, max_token_text_ratio=200, allowed_special=True ) - cosyvoice3_server._cosyvoice3_tokenizer = mocker.MagicMock() + cosyvoice3_server._adapter._tokenizer = mocker.MagicMock() mocker.patch("vllm_omni.model_executor.models.cosyvoice3.utils.extract_text_token", return_value=[None, 10]) request = OpenAICreateSpeechRequest( @@ -4555,7 +4680,7 @@ class TestCosyVoice3Serving: assert request_id.startswith("speech-") assert generator == "generator" assert tts_params == {} - cosyvoice3_server._build_cosyvoice3_prompt.assert_awaited_once() + cosyvoice3_server._adapter._build_prompt.assert_awaited_once() # ---- GLM-TTS Serving Tests ---- @@ -4635,13 +4760,13 @@ class TestGLMTTSServing: return_value="normalized target", ) - text_token_len = glm_tts_server._estimate_glm_tts_text_token_len("abcdef") + text_token_len = glm_tts_server._adapter._estimate_text_token_len("abcdef") assert text_token_len == 3 load_tokenizer.assert_called_once() def test_prepare_speech_generation_glm_tts(self, glm_tts_server, mocker: MockerFixture): - glm_tts_server._build_glm_tts_prompt = mocker.AsyncMock( + glm_tts_server._adapter._build_prompt = mocker.AsyncMock( return_value={ "prompt": "Hello", "multi_modal_data": {"audio": (np.zeros(24000), 24000)}, @@ -4654,7 +4779,7 @@ class TestGLMTTSServing: ref_text="Reference text", ) glm_tts_server.model_config.hf_config = SimpleNamespace(min_token_text_ratio=1, max_token_text_ratio=200) - glm_tts_server._estimate_glm_tts_text_token_len = mocker.MagicMock(return_value=10) + glm_tts_server._adapter._estimate_text_token_len = mocker.MagicMock(return_value=10) request_id, generator, tts_params = asyncio.run(glm_tts_server._prepare_speech_generation(request)) sampling_params = glm_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0] @@ -4663,7 +4788,7 @@ class TestGLMTTSServing: assert request_id.startswith("speech-") assert generator == "generator" assert tts_params == {} - glm_tts_server._build_glm_tts_prompt.assert_awaited_once() + glm_tts_server._adapter._build_prompt.assert_awaited_once() @pytest.fixture @@ -4782,16 +4907,14 @@ class TestMingFlashOmniTTSServing: assert error is not None assert "ref_audio" in error - def test_ming_flash_omni_tts_adapter_builds_prompt(self, ming_flash_omni_tts_server, mocker: MockerFixture): - ming_flash_omni_tts_server._build_ming_flash_omni_prompt = mocker.MagicMock( - return_value={ - "prompt_token_ids": [1, 2, 3], - "additional_information": {"voice": ["test"]}, - } - ) + def test_ming_flash_omni_tts_adapter_builds_prompt(self, ming_flash_omni_tts_server): request = OpenAICreateSpeechRequest(input="Hello", voice="test") asyncio.run(ming_flash_omni_tts_server._prepare_speech_generation(request)) - ming_flash_omni_tts_server._build_ming_flash_omni_prompt.assert_called_once() + + prompt = ming_flash_omni_tts_server.engine_client.generate.call_args.kwargs["prompt"] + assert prompt["prompt_token_ids"] == [0] + assert prompt["additional_information"]["text"] == "Hello" + assert prompt["additional_information"]["voice_name"] == "test" @pytest.fixture @@ -4880,10 +5003,6 @@ class TestDotsTTSServing: class TestTTSAsyncOffloading: """Tests for event-loop-safe offloading of blocking TTS operations.""" - def test_build_voxtral_prompt_is_sync(self): - """_build_voxtral_prompt should be a regular function, not a coroutine.""" - assert not asyncio.iscoroutinefunction(OmniOpenAIServingSpeech._build_voxtral_prompt) - @pytest.fixture def voxtral_server(self, mocker: MockerFixture): mocker.patch( @@ -4977,7 +5096,7 @@ class TestTTSAsyncOffloading: def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server, mocker: MockerFixture): """Voxtral path in _prepare_speech_generation should call the async wrapper.""" - voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock( + voxtral_server._adapter._build_prompt_async = mocker.AsyncMock( return_value={ "prompt_token_ids": [1, 2, 3], "additional_information": {"voice": ["test"]}, @@ -4985,19 +5104,19 @@ class TestTTSAsyncOffloading: ) request = OpenAICreateSpeechRequest(input="hello", voice="test") asyncio.run(voxtral_server._prepare_speech_generation(request)) - voxtral_server._build_voxtral_prompt_async.assert_awaited_once() + voxtral_server._adapter._build_prompt_async.assert_awaited_once() 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._adapter.validate = mocker.MagicMock(return_value=None) - qwen3_tts_server._build_tts_params = mocker.MagicMock( + qwen3_tts_server._adapter._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) + qwen3_tts_server._adapter._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) request = OpenAICreateSpeechRequest(input="hello") asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) - qwen3_tts_server._build_tts_params.assert_called_once() - qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once() + qwen3_tts_server._adapter._build_tts_params.assert_called_once() + qwen3_tts_server._adapter._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 @@ -5007,10 +5126,10 @@ class TestTTSAsyncOffloading: 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( + qwen3_tts_server._adapter._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) + qwen3_tts_server._adapter._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) request = OpenAICreateSpeechRequest(input="hello") asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) @@ -5058,10 +5177,10 @@ class TestTTSAsyncOffloading: 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( + qwen3_tts_server._adapter._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) + qwen3_tts_server._adapter._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, @@ -5075,10 +5194,10 @@ class TestTTSAsyncOffloading: 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( + qwen3_tts_server._adapter._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) + qwen3_tts_server._adapter._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, @@ -5092,10 +5211,10 @@ class TestTTSAsyncOffloading: 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( + qwen3_tts_server._adapter._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) + qwen3_tts_server._adapter._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, @@ -5112,10 +5231,10 @@ class TestTTSAsyncOffloading: """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( + qwen3_tts_server._adapter._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) + qwen3_tts_server._adapter._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, @@ -5132,7 +5251,7 @@ class TestTTSAsyncOffloading: """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( + voxtral_server._adapter._build_prompt_async = mocker.AsyncMock( return_value={ "prompt_token_ids": [1, 2, 3], "additional_information": {"voice": ["test"]}, @@ -5153,7 +5272,7 @@ class TestTTSAsyncOffloading: ): """VoiceDesign explicit false should reach the model prompt additional_information.""" qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) - qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + qwen3_tts_server._adapter._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) request = OpenAICreateSpeechRequest( input="hello", @@ -5176,7 +5295,7 @@ class TestTTSAsyncOffloading: qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000, "fake_cache_key")) qwen3_tts_server._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value=None) - qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + qwen3_tts_server._adapter._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) request = OpenAICreateSpeechRequest( input="hello", @@ -5491,7 +5610,7 @@ class TestTTSAsyncOffloading: 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 + assert qwen3_tts_server._adapter._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") @@ -5499,7 +5618,7 @@ class TestTTSAsyncOffloading: 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 + assert qwen3_tts_server._adapter._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") @@ -5507,7 +5626,7 @@ class TestTTSAsyncOffloading: 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 + assert qwen3_tts_server._adapter._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_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 7c553af8f..6414ab4ff 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -197,7 +197,7 @@ class TestStreamingSpeechWebSocket: error = ws.receive_json() assert error["type"] == "error" - assert "while input is buffered" in error["message"] + assert "while an utterance is in progress" in error["message"] # The buffered text survives the rejected reconfiguration. ws.send_json({"type": "input.done"}) @@ -742,6 +742,8 @@ class TestStreamingSpeechWebSocket: config.speaker_embedding = None config.stream_audio = True config.word_timestamps = False + config.seed = None + config.non_streaming_mode = None with pytest.raises(WebSocketDisconnect): asyncio.run( @@ -758,6 +760,163 @@ class TestStreamingSpeechWebSocket: assert websocket.send_json.await_count == 2 +class TestWebSocketSentenceSplitting: + def test_sentence_granularity_emits_one_request_per_sentence(self, mocker: MockerFixture): + app, speech_service = _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", + "split_granularity": "sentence", + } + ) + ws.send_json({"type": "input.text", "text": "Hello world. How are you? "}) + ws.send_json({"type": "input.done"}) + + first = ws.receive_json() + assert first["sentence_index"] == 0 + assert first["sentence_text"] == "Hello world." + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + second = ws.receive_json() + assert second["sentence_index"] == 1 + assert second["sentence_text"] == "How are you?" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": 0, + "total_sentences": 2, + } + + assert speech_service._generate_audio_bytes.await_count == 2 + assert [call.args[0].input for call in speech_service._generate_audio_bytes.await_args_list] == [ + "Hello world.", + "How are you?", + ] + + def test_sentence_granularity_emits_before_input_done(self, mocker: MockerFixture): + app, speech_service = _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", + "split_granularity": "sentence", + } + ) + ws.send_json({"type": "input.text", "text": "First sentence. "}) + + start = ws.receive_json() + assert start["sentence_text"] == "First sentence." + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert speech_service._generate_audio_bytes.await_count == 1 + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": 0, + "total_sentences": 1, + } + + def test_indic_danda_splits_without_latin_period(self, mocker: MockerFixture): + app, speech_service = _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", + "language": "Auto", + "split_granularity": "sentence", + } + ) + ws.send_json({"type": "input.text", "text": "नमस्ते। कैसे हो?"}) + ws.send_json({"type": "input.done"}) + + first = ws.receive_json() + assert first["sentence_text"] == "नमस्ते।" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + second = ws.receive_json() + assert second["sentence_text"] == "कैसे हो?" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json()["total_sentences"] == 2 + + def test_session_config_rejected_after_a_split_unit_was_emitted(self, mocker: MockerFixture): + """The splitter buffer is empty here, but the utterance is still open.""" + app, speech_service = _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", + "split_granularity": "sentence", + } + ) + ws.send_json({"type": "input.text", "text": "First sentence. "}) + ws.receive_json() + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + ws.send_json({"type": "session.config", "voice": "Serena"}) + error = ws.receive_json() + assert error["type"] == "error" + assert "while an utterance is in progress" in error["message"] + + ws.send_json({"type": "input.text", "text": "Second sentence. "}) + assert ws.receive_json()["sentence_index"] == 1 + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": 0, + "total_sentences": 2, + } + + # Reconfiguration is allowed again once the utterance closed. + ws.send_json({"type": "session.config", "voice": "Serena"}) + ws.send_json({"type": "input.text", "text": "Third."}) + ws.send_json({"type": "input.done"}) + ws.receive_json() + ws.receive_bytes() + ws.receive_json() + ws.receive_json() + + voices = [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] + assert voices == ["Vivian", "Vivian", "Serena"] + + def test_seed_is_forwarded_to_speech_request(self, mocker: MockerFixture): + app, speech_service = _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", "seed": 42}) + ws.send_json({"type": "input.text", "text": "Hello."}) + ws.send_json({"type": "input.done"}) + ws.receive_json() + ws.receive_bytes() + ws.receive_json() + ws.receive_json() + + assert speech_service._generate_audio_bytes.await_args_list[0].args[0].seed == 42 + + class TestGeneratePcmChunksContract: """Guard: _generate_pcm_chunks must exist on OmniOpenAIServingSpeech. diff --git a/tests/entrypoints/openai_api/test_speech_text_splitter.py b/tests/entrypoints/openai_api/test_speech_text_splitter.py new file mode 100644 index 000000000..ca2eb76d8 --- /dev/null +++ b/tests/entrypoints/openai_api/test_speech_text_splitter.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from vllm_omni.entrypoints.openai import speech_text_splitter +from vllm_omni.entrypoints.openai.speech_text_splitter import SpeechTextSplitter, extract_complete_units + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestExtractCompleteUnits: + def test_ascii_sentence_needs_whitespace_or_flush(self): + units, remainder, _ = extract_complete_units("Hello world.", frozenset(".!?。"), flush=False) + assert units == [] + assert remainder == "Hello world." + + units, remainder, _ = extract_complete_units("Hello world.", frozenset(".!?。"), flush=True) + assert units == ["Hello world."] + assert remainder == "" + + def test_ascii_sentence_with_trailing_space(self): + units, remainder, _ = extract_complete_units("Hello world. More", frozenset(".!?"), flush=False) + assert units == ["Hello world."] + assert remainder == "More" + + def test_decimal_is_not_a_sentence_boundary(self): + units, remainder, _ = extract_complete_units("It costs 3.14 dollars. ", frozenset(".!?"), flush=False) + assert units == ["It costs 3.14 dollars."] + assert remainder == "" + + def test_thousands_separator_is_not_a_clause_boundary(self): + units, _, _ = extract_complete_units("It is 1,000 rupees, sir. ", frozenset(".!?,;"), flush=False) + assert units == ["It is 1,000 rupees,", "sir."] + + def test_punctuation_run_stays_one_unit(self): + units, _, _ = extract_complete_units("Wait... Really?! Yes.", frozenset(".!?"), flush=True) + assert units == ["Wait...", "Really?!", "Yes."] + + def test_closing_quote_stays_with_its_sentence(self): + units, _, _ = extract_complete_units('He said "Hello." Bye.', frozenset(".!?"), flush=True) + assert units == ['He said "Hello."', "Bye."] + + def test_abbreviations_and_initials_do_not_split(self): + units, _, _ = extract_complete_units("Dr. Smith, e.g. J. R. Tolkien. Next.", frozenset(".!?"), flush=True) + assert units == ["Dr. Smith, e.g. J. R. Tolkien.", "Next."] + + def test_long_token_before_period_is_not_an_abbreviation(self): + units, _, _ = extract_complete_units("We landed in Washington. Next.", frozenset(".!?"), flush=True) + assert units == ["We landed in Washington.", "Next."] + + def test_abbreviation_lookback_is_bounded(self, monkeypatch): + # One WebSocket frame can reach ~128 KiB; an unbounded lookback walked + # the whole prefix at every `.`, quadratic in the frame size. + checks = 0 + is_token_char = speech_text_splitter._is_token_char + + def counting_is_token_char(ch): + nonlocal checks + checks += 1 + return is_token_char(ch) + + monkeypatch.setattr(speech_text_splitter, "_is_token_char", counting_is_token_char) + text = "a." * 20000 + units, _, _ = extract_complete_units(text, frozenset(".!?"), flush=True) + assert units == [text] + assert checks <= len(text) * (speech_text_splitter._MAX_ABBREVIATION_LEN + 1) + + def test_cjk_period_completes_once_the_next_char_is_known(self): + # The trailing 。 waits: a closing delimiter could still follow it. + units, remainder, _ = extract_complete_units("你好。世界。", frozenset(".!?。！？"), flush=False) + assert units == ["你好。"] + assert remainder == "世界。" + + units, remainder, _ = extract_complete_units("你好。世界。", frozenset(".!?。！？"), flush=True) + assert units == ["你好。", "世界。"] + assert remainder == "" + + def test_indic_danda_and_question_mark(self): + units, remainder, _ = extract_complete_units("नमस्ते। कैसे हो?", frozenset(".!?।॥؟"), flush=True) + assert units == ["नमस्ते।", "कैसे हो?"] + assert remainder == "" + + def test_scan_cursor_skips_already_examined_text(self): + units, remainder, scan = extract_complete_units("no terminator yet", frozenset(".!?"), flush=False) + assert units == [] + assert remainder == "no terminator yet" + assert scan == len(remainder) + + units, remainder, scan = extract_complete_units( + "no terminator yet. Next", frozenset(".!?"), flush=False, scan_from=scan + ) + assert units == ["no terminator yet."] + assert remainder == "Next" + + +class TestSpeechTextSplitter: + def test_none_buffers_until_flush(self): + splitter = SpeechTextSplitter("none") + assert splitter.feed("Hello world. ") == [] + assert splitter.flush() == ["Hello world."] + assert splitter.flush() == [] + + def test_sentence_mode_incremental_words(self): + splitter = SpeechTextSplitter("sentence") + assert splitter.feed("Hello ") == [] + assert splitter.feed("world. ") == ["Hello world."] + assert splitter.feed("How are you?") == [] + assert splitter.flush() == ["How are you?"] + + def test_sentence_mode_resumes_scan_across_feeds(self): + splitter = SpeechTextSplitter("sentence") + assert splitter.feed("It costs 3") == [] + assert splitter.feed(".") == [] + assert splitter.feed("14 dollars. ") == ["It costs 3.14 dollars."] + + def test_clause_mode_splits_on_comma(self): + splitter = SpeechTextSplitter("clause") + assert splitter.feed("Hello, world. ") == ["Hello,", "world."] + + def test_clause_mode_splits_on_arabic_semicolon(self): + splitter = SpeechTextSplitter("clause") + assert splitter.feed("مرحبا؛ كيف حالك") == ["مرحبا؛"] + + def test_empty_flush(self): + splitter = SpeechTextSplitter("sentence") + assert splitter.flush() == [] diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index ed71167a3..a17a6ccbd 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -22,12 +22,14 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( detect_tts_model_type, resolve_adapter, ) +from vllm_omni.entrypoints.openai.tts_adapters.covo_audio import CovoAudioAdapter from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import ( IndexTTS2Adapter, IndexTTS25Adapter, indextts2_conditioning_cache_salt, ) +from vllm_omni.entrypoints.openai.tts_adapters.ming_flash_omni_tts import MingFlashOmniTTSAdapter from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import ( MossTTSAdapter, MossTTSNanoAdapter, @@ -37,6 +39,8 @@ from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import ( Qwen3TTSAdapter, Qwen3TTSCodecLimitError, ) +from vllm_omni.entrypoints.openai.tts_adapters.step_audio2 import StepAudio2Adapter +from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter from vllm_omni.model_executor.models.indextts2 import prompt_utils from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import ( INDEXTTS25_TOKENIZER_FILE, @@ -125,8 +129,8 @@ def test_moss_tts_applies_request_max_new_tokens(adapter_cls): def _build_moss_tts_request(adapter_cls, mocker, *, request_seed): server = mocker.Mock() - server._build_moss_tts_params = mocker.AsyncMock(return_value={}) adapter = adapter_cls(SpeechServingContext(server=server)) + adapter._build_moss_tts_params = mocker.AsyncMock(return_value={}) request = OpenAICreateSpeechRequest(input="hello", seed=request_seed) return asyncio.run( @@ -162,46 +166,45 @@ def test_moss_tts_seed_falls_back_to_stage_default(adapter_cls, mocker): assert prepared.tts_params["seed"] == [42] -@pytest.mark.parametrize("variant,ref_text", [("local", None), ("local", " "), ("tts", "Reference.")]) -def test_moss_tts_generation_fallback(variant, ref_text, mocker): - server = mocker.Mock(_moss_variant=variant) - server._build_moss_tts_params = mocker.AsyncMock(return_value={"prompt_token_ids": [7]}) - request = OpenAICreateSpeechRequest(input="Target.", ref_text=ref_text) - adapter = MossTTSAdapter(SpeechServingContext(server=server)) - - prepared = asyncio.run(adapter.build(request, [], has_inline_ref_audio=True)) - - server._build_moss_tts_params.assert_awaited_once_with(request, has_inline_ref_audio=True) - server._get_moss_processor.assert_not_called() - assert prepared.prompt["prompt_token_ids"] == [7] - - -def test_moss_local_reference_transcript_continuation(mocker): +@pytest.mark.parametrize( + "variant,ref_text,mode", + [ + ("local", " Reference. ", "continuation"), + ("local", None, "generation"), + ("local", " ", "generation"), + ("tts", "Reference.", "generation"), + ], +) +def test_moss_reference_transcript_mode(variant, ref_text, mode, mocker): import torch reference = [torch.ones((3, 12), dtype=torch.int64)] unified = torch.arange(52, dtype=torch.int64).reshape(1, 4, 13) processor = mocker.Mock(return_value={"input_ids": unified}) server = mocker.Mock(_moss_variant="local", uploaded_speakers={"speaker": {}}) - server._build_moss_tts_params = mocker.AsyncMock(return_value={"prompt_token_ids": [7]}) - server._get_moss_processor.return_value = processor server._voice_created_at.return_value = 123 - server._encode_moss_references = mocker.AsyncMock(return_value=(reference, {0: "reference-key"})) request = OpenAICreateSpeechRequest( - input="Target.", ref_text=" Reference. ", language="English", voice="Speaker", seed=0, max_new_tokens=2048 + input="Target.", ref_text=ref_text, language="English", voice="Speaker", seed=0, max_new_tokens=2048 ) adapter = MossTTSAdapter(SpeechServingContext(server=server)) + adapter._moss_variant = variant + mocker.patch.object(adapter, "_get_moss_processor", return_value=processor) + mocker.patch.object( + adapter, "_encode_moss_references", new=mocker.AsyncMock(return_value=(reference, {0: "reference-key"})) + ) + prepared = asyncio.run(adapter.build(request, [], has_inline_ref_audio=False)) - processor.build_user_message.assert_called_once_with(text="Reference. Target.", language="English") - processor.build_assistant_message.assert_called_once_with(audio_codes_list=reference) - processor.assert_called_once_with( - conversations=[[processor.build_user_message.return_value, processor.build_assistant_message.return_value]], - mode="continuation", - ) - server._encode_moss_references.assert_awaited_once_with(request, has_inline_ref_audio=False, two_speaker=False) - server._build_moss_tts_params.assert_not_called() + if mode == "continuation": + processor.build_user_message.assert_called_once_with(text="Reference. Target.", language="English") + processor.build_assistant_message.assert_called_once_with(audio_codes_list=reference) + conversation = [processor.build_user_message.return_value, processor.build_assistant_message.return_value] + else: + processor.build_user_message.assert_called_once_with(text="Target.", language="English", reference=reference) + processor.build_assistant_message.assert_not_called() + conversation = [processor.build_user_message.return_value] + processor.assert_called_once_with(conversations=[conversation], mode=mode) assert prepared.prompt["prompt_token_ids"] == unified[0, :, 0].tolist() assert torch.equal(prepared.tts_params["codes"]["ref"], unified[0, :, 1:]) assert prepared.tts_params["max_new_frames"] == [2048] @@ -218,6 +221,81 @@ def test_qwen3_tts_metadata(): assert issubclass(Qwen3TTSAdapter, ARTTSAdapter) +def test_qwen3_tts_build_constructs_prepared_request(): + server = SimpleNamespace(_tts_executor=None, _tts_tokenizer=None, uploaded_speakers={}) + engine_client = SimpleNamespace(model_config=SimpleNamespace()) + adapter = Qwen3TTSAdapter(SimpleNamespace(server=server, engine_client=engine_client)) + + async def estimate_prompt_len(_tts_params): + return 3 + + adapter._estimate_prompt_len_async = estimate_prompt_len + request = OpenAICreateSpeechRequest(input="hello") + + prepared = asyncio.run(adapter.build(request, [], False)) + + assert prepared.prompt["prompt_token_ids"] == [1, 1, 1] + assert prepared.prompt["additional_information"] is prepared.tts_params + assert prepared.tts_params["text"] == ["hello"] + assert prepared.tts_params["speaker"] == ["Vivian"] + assert prepared.model_type == "CustomVoice" + + +def test_covo_audio_build_constructs_tokenized_prompt(): + class FakeTokenizer: + def apply_chat_template(self, messages, *, tokenize, add_generation_prompt): + assert messages[-1] == {"role": "user", "content": "hello"} + assert tokenize is True + assert add_generation_prompt is True + return [11, 12, 13] + + adapter = CovoAudioAdapter(SimpleNamespace(server=SimpleNamespace(), engine_client=SimpleNamespace())) + adapter._tokenizer = FakeTokenizer() + + prepared = asyncio.run(adapter.build(OpenAICreateSpeechRequest(input="hello"), [], False)) + + assert prepared.prompt == {"prompt_token_ids": [11, 12, 13]} + assert prepared.tts_params == {} + assert prepared.model_type == "covo_audio" + + +def test_ming_flash_omni_build_constructs_prepared_request(): + server = SimpleNamespace(uploaded_speakers={}) + adapter = MingFlashOmniTTSAdapter(SimpleNamespace(server=server, engine_client=SimpleNamespace())) + request = OpenAICreateSpeechRequest(input="hello", voice="test") + + prepared = asyncio.run(adapter.build(request, [], False)) + + assert prepared.prompt["prompt_token_ids"] == [0] + info = prepared.prompt["additional_information"] + assert info["ming_task"] == "instruct" + assert info["text"] == "hello" + assert info["voice_name"] == "test" + assert prepared.model_type == "ming_flash_omni_tts" + + +def test_step_audio2_build_constructs_open_assistant_turn(): + adapter = StepAudio2Adapter(SimpleNamespace(server=SimpleNamespace(), engine_client=SimpleNamespace())) + request = OpenAICreateSpeechRequest(input="hello", instructions="Read warmly.") + + prepared = asyncio.run(adapter.build(request, [], False)) + + assert "<|im_start|>system\nRead warmly.<|im_end|>" in prepared.prompt["prompt"] + assert prepared.prompt["prompt"].endswith("<|im_start|>assistant\n") + assert prepared.tts_params == {} + assert prepared.model_type == "step_audio2" + + +def test_voxtral_prompt_builder_is_sync(): + ctx = SimpleNamespace( + server=SimpleNamespace(_tts_executor=None), + engine_client=SimpleNamespace(), + ) + adapter = VoxtralTTSAdapter(ctx) + + assert not asyncio.iscoroutinefunction(adapter._build_prompt) + + @pytest.mark.parametrize( ("task_type", "text_tokens", "request_cap", "expected_cap"), [ diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index 3b8d17997..d54ce36d5 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -132,6 +132,8 @@ _PIPELINE_STAGES = [ "audex_thinker", "audex_tta_thinker", "audex_xcodec", + "audio8_tts_codec_decoder", + "audio8_tts_slow_ar", "audio_generation", "audio_tokenizer", "audio_vae", diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 9f643d8ca..bf395b813 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -23,9 +23,9 @@ from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _png_bytes() -> bytes: +def _png_bytes(size: tuple[int, int] = (2, 1)) -> bytes: buffer = BytesIO() - Image.new("RGB", (2, 1), color=(12, 34, 56)).save(buffer, format="PNG") + Image.new("RGB", size, color=(12, 34, 56)).save(buffer, format="PNG") return buffer.getvalue() @@ -122,6 +122,52 @@ async def test_decode_image_url_keeps_data_urls_local(monkeypatch): assert image.mode == "RGB" +def test_decode_image_bytes_rejects_image_over_pixel_limit_before_conversion(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100) + + def _unexpected_convert(*args, **kwargs): + pytest.fail("over-limit images must be rejected before conversion") + + monkeypatch.setattr(Image.Image, "convert", _unexpected_convert) + + with pytest.raises( + InvalidInputReferenceError, + match=r"Image dimensions 20x20 \(400 pixels\) exceed the maximum of 100 pixels", + ): + video_api_utils._decode_image_bytes(_png_bytes((20, 20)), source="image reference") + + +def test_decode_image_bytes_maps_pillow_pixel_limit_error(monkeypatch): + monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 100) + + with pytest.raises(InvalidInputReferenceError, match="decoder pixel limit") as exc_info: + video_api_utils._decode_image_bytes(_png_bytes((20, 20)), source="image reference") + + assert isinstance(exc_info.value.__cause__, Image.DecompressionBombError) + + +@pytest.mark.parametrize("max_pixels", [0, 400]) +def test_decode_image_bytes_allows_disabled_or_inclusive_pixel_limit(monkeypatch, max_pixels): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", max_pixels) + + image = video_api_utils._decode_image_bytes(_png_bytes((20, 20)), source="image reference") + + assert image.size == (20, 20) + assert image.mode == "RGB" + + +@pytest.mark.asyncio +async def test_decode_input_reference_preserves_image_pixel_limit_error(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100) + + with pytest.raises(InvalidInputReferenceError, match="VLLM_MAX_IMAGE_PIXELS"): + await video_api_utils.decode_input_reference( + image_reference=None, + video_reference=None, + input_reference_bytes=_png_bytes((20, 20)), + ) + + def _install_fake_video_mux(monkeypatch, mux_calls): def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None): mux_calls.append( diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 969691ebc..1b18a9ae5 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -75,6 +75,7 @@ class FakeAsyncOmni: self.model_class_name = "WanPipeline" self.captured_prompt = None self.captured_reference_video_bytes = None + self.captured_control_reference_bytes = {} self.captured_sampling_params_list = None def get_diffusion_od_config(self): @@ -83,6 +84,10 @@ class FakeAsyncOmni: async def generate(self, prompt, request_id, sampling_params_list): self.captured_prompt = prompt self.captured_sampling_params_list = sampling_params_list + for control_type in ("edge", "blur", "depth", "seg", "wsm"): + control_params = sampling_params_list[0].extra_args.get(control_type) + if isinstance(control_params, dict) and isinstance(control_params.get("control_path"), str): + self.captured_control_reference_bytes[control_type] = Path(control_params["control_path"]).read_bytes() reference_videos = prompt.get("multi_modal_data", {}).get("video") if ( isinstance(reference_videos, list) @@ -338,6 +343,8 @@ def _cosmos3_stage_configs(): return [ SimpleNamespace( stage_type="diffusion", + final_output=True, + final_output_type="video", engine_args=SimpleNamespace(model_class_name="Cosmos3OmniDiffusersPipeline"), ) ] @@ -919,6 +926,17 @@ def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_fal assert handler.supports_mixed_reference_inputs +@pytest.mark.parametrize("model_class_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"]) +def test_control_upload_capability_is_declared_only_by_cosmos3(test_client, model_class_name): + handler = test_client.app.state.openai_serving_video + handler._engine_client.model_class_name = model_class_name + + assert handler.supported_control_upload_types == frozenset({"edge", "blur", "depth", "seg", "wsm"}) + + handler._engine_client.model_class_name = "WanPipeline" + assert handler.supported_control_upload_types == frozenset() + + def test_decode_video_bytes_can_keep_first_frames(): from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes @@ -1612,7 +1630,8 @@ def test_generic_video_model_rejects_mixed_image_and_video_references(test_clien assert "does not support mixed image and video" in response.json()["detail"].lower() -def test_h3_multipart_rejects_bmp_image_reference(test_client): +def test_h3_multipart_rejects_bmp_image_reference(test_client, monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100) image = Image.new("RGB", (64, 64), color="blue") image_buffer = io.BytesIO() image.save(image_buffer, format="BMP") @@ -1628,6 +1647,37 @@ def test_h3_multipart_rejects_bmp_image_reference(test_client): assert "must use jpg" in response.json()["detail"].lower() +@pytest.mark.parametrize("field", ["input_reference", "input_references"]) +def test_h3_multipart_rejects_image_over_pixel_limit(field, test_client, monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100) + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "reject oversized image", "extra_params": '{"task":"ref2va"}'}, + files=[(field, ("reference.png", _make_test_image_bytes((20, 20)), "image/png"))], + ) + + assert response.status_code == 400 + assert "VLLM_MAX_IMAGE_PIXELS" in response.json()["detail"] + + +@pytest.mark.parametrize("field", ["input_reference", "input_references"]) +def test_h3_multipart_maps_pillow_pixel_limit_error(field, test_client, monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 0) + monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 100) + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "reject decoder bomb", "extra_params": '{"task":"ref2va"}'}, + files=[(field, ("reference.png", _make_test_image_bytes((20, 20)), "image/png"))], + ) + + assert response.status_code == 400 + assert "decoder pixel limit" in response.json()["detail"] + + @pytest.mark.asyncio async def test_h3_upload_limit_checks_declared_size_before_read(): class OversizedUpload: @@ -1800,6 +1850,24 @@ def test_video_request_validation(): VideoGenerationRequest(prompt="test", quality="medium") +def test_async_create_accepts_fractional_fps(test_client, mocker: MockerFixture): + """Queued VideoResponse must accept fractional fps from the request path.""" + _mock_encode_video_bytes(mocker) + response = test_client.post( + "/v1/videos", + data={"prompt": "fractional fps", "fps": "12.5", "num_frames": "5"}, + ) + assert response.status_code == 200, response.text + body = response.json() + assert body["fps"] == 12.5 + assert body["num_frames"] == 5 + video_id = body["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + engine = test_client.app.state.openai_serving_video._engine_client + assert engine.captured_sampling_params_list[0].fps == 12.5 + assert engine.captured_sampling_params_list[0].frame_rate == 12.5 + + def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture): mocker.patch( "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", @@ -2218,6 +2286,171 @@ def test_sync_v2v_returns_video_bytes(test_client, mocker: MockerFixture): assert input_video[0].size == (32, 24) +@pytest.mark.parametrize("endpoint", ["/v1/videos", "/v1/videos/sync"]) +@pytest.mark.parametrize("control_type", ["wsm", "depth"]) +def test_cosmos3_accepts_optional_uploaded_control( + endpoint, + control_type, + test_client, + mocker: MockerFixture, +): + control_bytes = f"{control_type}-control".encode() + _mock_encode_video_bytes(mocker, b"controlled-video") + engine = test_client.app.state.openai_serving_video._engine_client + engine.model_class_name = "Cosmos3OmniDiffusersPipeline" + + response = test_client.post( + endpoint, + data={ + "prompt": "Follow the uploaded control.", + "control_type": control_type, + "extra_params": json.dumps({control_type: {"control_weight": 0.75}}), + }, + files=[ + ("input_reference", ("input.mp4", _make_test_video_bytes(), "video/mp4")), + ("control_reference", (f"{control_type}.mp4", control_bytes, "video/mp4")), + ], + ) + + assert response.status_code == 200 + if endpoint.endswith("/sync"): + assert response.content == b"controlled-video" + else: + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + captured = engine.captured_sampling_params_list[0].extra_args[control_type] + assert captured["control_weight"] == 0.75 + assert engine.captured_control_reference_bytes[control_type] == control_bytes + assert len(engine.captured_prompt["multi_modal_data"]["video"]) == 3 + assert not Path(captured["control_path"]).exists() + + +@pytest.mark.parametrize("endpoint", ["/v1/videos", "/v1/videos/sync"]) +def test_cosmos3_uploaded_control_selects_transfer_reference_video_decode_policy( + endpoint, + test_client, + mocker: MockerFixture, +): + _mock_encode_video_bytes(mocker, b"controlled-video") + test_client.app.state.stage_configs = _cosmos3_stage_configs() + + response = test_client.post( + endpoint, + data={ + "prompt": "Preserve all conditioning motion.", + "control_type": "wsm", + "num_frames": "9", + "extra_params": json.dumps({"num_first_chunk_conditional_frames": 9}), + }, + files=[ + ( + "input_reference", + ("input.mp4", _make_test_video_bytes(num_frames=6), "video/mp4"), + ), + ("control_reference", ("wsm.mp4", b"control", "video/mp4")), + ], + ) + + assert response.status_code == 200 + if not endpoint.endswith("/sync"): + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + engine = test_client.app.state.openai_serving_video._engine_client + assert len(engine.captured_prompt["multi_modal_data"]["video"]) == 6 + + +def test_cosmos3_control_upload_is_optional(test_client, mocker: MockerFixture): + _mock_encode_video_bytes(mocker) + engine = test_client.app.state.openai_serving_video._engine_client + engine.model_class_name = "Cosmos3OmniDiffusersPipeline" + + response = test_client.post("/v1/videos/sync", data={"prompt": "No control for this request."}) + + assert response.status_code == 200 + assert "wsm" not in engine.captured_sampling_params_list[0].extra_args + + +def test_control_upload_does_not_affect_models_without_capability(test_client): + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "Unsupported control.", "control_type": "wsm"}, + files={"control_reference": ("wsm.mp4", b"control", "video/mp4")}, + ) + + assert response.status_code == 400 + assert "not supported by this model" in response.json()["detail"] + assert test_client.app.state.openai_serving_video._engine_client.captured_prompt is None + + +@pytest.mark.parametrize( + ("data", "files", "message"), + [ + ( + {"prompt": "Missing type."}, + {"control_reference": ("wsm.mp4", b"control", "video/mp4")}, + "requires control_type", + ), + ( + {"prompt": "Missing file.", "control_type": "wsm"}, + None, + "requires a control_reference", + ), + ( + {"prompt": "Unknown type.", "control_type": "unknown"}, + {"control_reference": ("control.mp4", b"control", "video/mp4")}, + "not supported by this model", + ), + ], +) +def test_cosmos3_control_upload_validates_contract(data, files, message, test_client): + test_client.app.state.openai_serving_video._engine_client.model_class_name = "Cosmos3OmniDiffusersPipeline" + + response = test_client.post("/v1/videos/sync", data=data, files=files) + + assert response.status_code == 400 + assert message in response.json()["detail"] + + +def test_cosmos3_control_upload_rejects_existing_control_source(test_client): + test_client.app.state.openai_serving_video._engine_client.model_class_name = "Cosmos3OmniDiffusersPipeline" + response = test_client.post( + "/v1/videos/sync", + data={ + "prompt": "Ambiguous control.", + "control_type": "wsm", + "extra_params": json.dumps({"wsm": {"control_path": "/already/present.mp4"}}), + }, + files={"control_reference": ("wsm.mp4", b"control", "video/mp4")}, + ) + + assert response.status_code == 400 + assert "not both" in response.json()["detail"] + + +@pytest.mark.parametrize( + ("control_bytes", "message"), + [ + (b"", "must not be empty"), + (b"control", "size limit"), + ], +) +def test_cosmos3_control_upload_rejects_invalid_size(control_bytes, message, test_client, monkeypatch): + test_client.app.state.openai_serving_video._engine_client.model_class_name = "Cosmos3OmniDiffusersPipeline" + monkeypatch.setattr(api_server, "CONTROL_REFERENCE_MAX_BYTES", 3) + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "Invalid control.", "control_type": "wsm"}, + files={"control_reference": ("wsm.mp4", control_bytes, "video/mp4")}, + ) + + assert response.status_code == 400 + assert message in response.json()["detail"] + assert test_client.app.state.openai_serving_video._engine_client.captured_prompt is None + + def test_sync_missing_handler_returns_503(): app = FastAPI() app.include_router(router) diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 616d5486e..dda9ab38c 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -4,8 +4,12 @@ from types import SimpleNamespace import pytest +import torch +from pydantic import ValidationError -from vllm_omni.diffusion.data import AttentionConfig +from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.resolver import OmniConfigResolution +from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig from vllm_omni.engine import stage_init_utils from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.entrypoints.cli.serve import OmniServeCommand @@ -14,9 +18,13 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def _terminal_config(stage_cfg: dict) -> OmniDiffusionConfig: + return OmniDiffusionConfig.from_kwargs(**stage_cfg["engine_args"]) + + def test_default_stage_config_includes_cache_backend(): """Ensure cache knobs survive the default diffusion-stage builder.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "cache_backend": "cache_dit", "cache_config": '{"Fn_compute_blocks": 2}', @@ -30,12 +38,12 @@ def test_default_stage_config_includes_cache_backend(): assert engine_args["cache_backend"] == "cache_dit" assert engine_args["cache_config"]["Fn_compute_blocks"] == 2 assert engine_args["vae_use_slicing"] is True - assert engine_args["parallel_config"].ulysses_degree == 2 + assert engine_args["parallel_config"]["ulysses_degree"] == 2 assert engine_args["model_stage"] == "diffusion" def test_default_stage_config_preserves_ulysses_a2a_permute() -> None: - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "ulysses_degree": 4, "ulysses_a2a_permute": True, @@ -43,18 +51,18 @@ def test_default_stage_config_preserves_ulysses_a2a_permute() -> None: )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] - assert parallel_config.ulysses_degree == 4 - assert parallel_config.ulysses_a2a_permute is True + assert parallel_config["ulysses_degree"] == 4 + assert parallel_config["ulysses_a2a_permute"] is True def test_default_stage_config_preserves_model_extras(): - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"extras": {"ltx2_use_conv_vae": True}})[0] + stage_cfg = StageConfigFactory.create_default_diffusion({"extras": {"ltx2_use_conv_vae": True}})[0] assert stage_cfg["engine_args"]["extras"]["ltx2_use_conv_vae"] is True def test_default_stage_config_preserves_and_overrides_promoted_extras(): - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "extras": { "auxiliary_text_encoder": "/models/extras-llama", @@ -68,7 +76,7 @@ def test_default_stage_config_preserves_and_overrides_promoted_extras(): assert extras["auxiliary_text_encoder"] == "/models/extras-llama" assert extras["default_llama_model_id"] == "extras/default-llama" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "extras": { "auxiliary_text_encoder": "/models/extras-llama", @@ -84,43 +92,57 @@ def test_default_stage_config_preserves_and_overrides_promoted_extras(): assert extras["default_llama_model_id"] == "top-level/default-llama" -def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker): +@pytest.mark.parametrize( + "stage_overrides", + [ + {"0": {"extras": {"ltx2_use_conv_vae": True}}}, + '{"0":{"extras":{"ltx2_use_conv_vae":true}}}', + ], +) +def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker, stage_overrides): """Local/unregistered Diffusers checkpoints still honor stage-0 extras.""" - - def resolve_with_default(*_args, default_stage_cfg_factory, **_kwargs): - return None, default_stage_cfg_factory(), None - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", - side_effect=resolve_with_default, + "vllm_omni.config.resolver.StageConfigFactory.create_from_model", + return_value=None, + ) + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "LTX2Pipeline"), ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) _, stage_configs = engine._resolve_stage_configs( "/models/LTX-2.5-Diffusers", - {"stage_overrides": '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'}, + {"stage_overrides": stage_overrides}, trust_remote_code=False, ) assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True +def test_default_stage_rejects_unknown_nested_parallel_config_key(): + unknown_key = "unknown_parallel_field" + with pytest.raises(ValidationError, match=unknown_key): + StageConfigFactory.create_default_diffusion( + {"parallel_config": {unknown_key: 2}}, + ) + + def test_default_cache_config_used_when_missing(): """Ensure default cache_config is synthesized when only backend is given.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "cache_backend": "cache_dit", } )[0] - cache_config = stage_cfg["engine_args"]["cache_config"] - assert cache_config is not None - assert cache_config["Fn_compute_blocks"] == 1 + cache_config = _terminal_config(stage_cfg).cache_config + assert cache_config.Fn_compute_blocks == 1 def test_default_stage_devices_from_sequence_parallel(): """Ensure runtime devices reflect computed diffusion world size.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "ulysses_degree": 2, "ring_degree": 2, @@ -132,7 +154,7 @@ def test_default_stage_devices_from_sequence_parallel(): def test_default_stage_devices_and_dp_from_num_gpus(): """Resolve omitted DP before deriving devices from an explicit GPU count.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "num_gpus": 8, "tensor_parallel_size": 2, @@ -141,15 +163,14 @@ def test_default_stage_devices_and_dp_from_num_gpus(): )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] - assert parallel_config.data_parallel_size == 2 - assert parallel_config.world_size == 8 + assert parallel_config["data_parallel_size"] == 2 assert stage_cfg["engine_args"]["num_gpus"] == 8 assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7" def test_default_stage_config_uses_parallel_size_kwargs(): """Ensure default diffusion parallel_config uses CLI/API parallel sizes.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "pipeline_parallel_size": 2, "data_parallel_size": 3, @@ -159,15 +180,15 @@ def test_default_stage_config_uses_parallel_size_kwargs(): )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] - assert parallel_config.pipeline_parallel_size == 2 - assert parallel_config.data_parallel_size == 3 - assert parallel_config.tensor_parallel_size == 4 - assert parallel_config.enable_expert_parallel is True + assert parallel_config["pipeline_parallel_size"] == 2 + assert parallel_config["data_parallel_size"] == 3 + assert parallel_config["tensor_parallel_size"] == 4 + assert parallel_config["enable_expert_parallel"] is True def test_default_stage_config_preserves_omitted_dp_for_runtime_inference(): """Keep omitted DP unresolved until the runtime WORLD size is known.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "pipeline_parallel_size": None, "data_parallel_size": None, @@ -180,18 +201,19 @@ def test_default_stage_config_preserves_omitted_dp_for_runtime_inference(): )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] - assert parallel_config.pipeline_parallel_size == 1 - assert parallel_config.data_parallel_size is None - assert parallel_config.tensor_parallel_size == 1 - assert parallel_config.enable_expert_parallel is False - assert stage_cfg["engine_args"]["enforce_eager"] is False - assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "regional" - assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is True + assert parallel_config["pipeline_parallel_size"] == 1 + assert parallel_config["data_parallel_size"] is None + assert parallel_config["tensor_parallel_size"] == 1 + assert parallel_config["enable_expert_parallel"] is False + terminal_config = _terminal_config(stage_cfg) + assert terminal_config.enforce_eager is False + assert terminal_config.diffusion_compile_granularity == "regional" + assert terminal_config.diffusion_compile_dynamic is True def test_default_stage_config_propagates_ulysses_mode(): """Ensure UAA mode survives default diffusion-stage creation.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "ulysses_degree": 4, "ulysses_mode": "advanced_uaa", @@ -199,13 +221,13 @@ def test_default_stage_config_propagates_ulysses_mode(): )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] - assert parallel_config.ulysses_degree == 4 - assert parallel_config.ulysses_mode == "advanced_uaa" + assert parallel_config["ulysses_degree"] == 4 + assert parallel_config["ulysses_mode"] == "advanced_uaa" def test_default_stage_config_includes_default_sampling_params(): """Ensure default sampling params survive the default diffusion-stage builder.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "default_sampling_params": '{"0": {"generator_device":"cpu", "guidance_scale":7.5}}', } @@ -219,13 +241,13 @@ def test_default_stage_config_includes_default_sampling_params(): def test_default_stage_config_includes_diffusion_attention_backend(): """Ensure diffusion attention shorthand lands in engine_args.diffusion_attention_config.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "diffusion_attention_backend": "FLASH_ATTN", } )[0] - diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"] + diffusion_attention_config = _terminal_config(stage_cfg).diffusion_attention_config assert isinstance(diffusion_attention_config, AttentionConfig) assert diffusion_attention_config.default is not None assert diffusion_attention_config.default.backend == "FLASH_ATTN" @@ -233,7 +255,7 @@ def test_default_stage_config_includes_diffusion_attention_backend(): def test_default_stage_config_includes_diffusion_attention_config(): """Ensure structured diffusion attention config survives default stage creation.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "diffusion_attention_config": { "default": {"backend": "FLASH_ATTN"}, @@ -242,7 +264,7 @@ def test_default_stage_config_includes_diffusion_attention_config(): } )[0] - diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"] + diffusion_attention_config = _terminal_config(stage_cfg).diffusion_attention_config assert isinstance(diffusion_attention_config, AttentionConfig) assert diffusion_attention_config.default is not None assert diffusion_attention_config.default.backend == "FLASH_ATTN" @@ -252,7 +274,7 @@ def test_default_stage_config_includes_diffusion_attention_config(): def test_default_stage_config_rejects_conflicting_diffusion_attention_inputs(): """Ensure shorthand and default.backend stay mutually exclusive.""" with pytest.raises(ValueError, match="mutually exclusive"): - AsyncOmniEngine._create_default_diffusion_stage_cfg( + StageConfigFactory.create_default_diffusion( { "diffusion_attention_backend": "FLASH_ATTN", "diffusion_attention_config": { @@ -264,7 +286,7 @@ def test_default_stage_config_rejects_conflicting_diffusion_attention_inputs(): def test_default_stage_config_engine_args(): """Ensure default diffusion-stage builder sets and propagates engine_args.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { "distributed_executor_backend": "ray", "boundary_ratio": 0.875, @@ -283,7 +305,7 @@ def test_default_stage_config_engine_args(): def test_default_stage_config_whitelist_none_fallback(): """DeployConfig / StageDeployConfig whitelist fields with value None fall back to OmniDiffusionConfig dataclass defaults.""" - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + stage_cfg = StageConfigFactory.create_default_diffusion( { # DeployConfig pipeline-wide "trust_remote_code": None, @@ -294,12 +316,11 @@ def test_default_stage_config_whitelist_none_fallback(): } )[0] - engine_args = stage_cfg["engine_args"] - - assert engine_args["trust_remote_code"] is False - assert engine_args["distributed_executor_backend"] is None - assert engine_args["dtype"] == "auto" - assert engine_args["enforce_eager"] is False + terminal = _terminal_config(stage_cfg) + assert terminal.trust_remote_code is False + assert terminal.distributed_executor_backend is None + assert terminal.dtype == torch.bfloat16 + assert terminal.enforce_eager is False def test_serve_cli_accepts_ulysses_mode(): @@ -321,12 +342,12 @@ def test_serve_cli_accepts_ulysses_mode(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] assert args.ulysses_mode == "advanced_uaa" - assert parallel_config.ulysses_degree == 4 - assert parallel_config.ulysses_mode == "advanced_uaa" + assert parallel_config["ulysses_degree"] == 4 + assert parallel_config["ulysses_mode"] == "advanced_uaa" def test_serve_cli_accepts_text_encoder_tp_size(): @@ -345,11 +366,11 @@ def test_serve_cli_accepts_text_encoder_tp_size(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] assert args.text_encoder_tp_size == 4 - assert parallel_config.text_encoder_tp_size == 4 + assert parallel_config["text_encoder_tp_size"] == 4 def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage(): @@ -368,7 +389,7 @@ def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] assert args.task_type == "fl2va" assert stage_cfg["engine_args"]["task_type"] == "fl2va" @@ -390,7 +411,7 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] assert args.enable_diffusion_pipeline_profiler is True assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True @@ -416,7 +437,7 @@ def test_serve_cli_forwards_distilled_lora_to_diffusion_stage(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] engine_args = stage_cfg["engine_args"] assert explicit_kwargs["lora_backend"] == "distill" @@ -447,7 +468,7 @@ def test_serve_cli_forwards_compact_diffusion_offload_config(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] engine_args = stage_cfg["engine_args"] expected = { @@ -513,7 +534,7 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] engine_args = stage_cfg["engine_args"] assert explicit_kwargs["host_weight_runtime_mode"] == "preferred" @@ -542,7 +563,7 @@ def test_serve_cli_accepts_diffusion_compile_controls(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] assert args.diffusion_compile_granularity == "full" assert args.diffusion_compile_dynamic is False @@ -569,7 +590,7 @@ def test_serve_cli_accepts_diffusion_attention_backend(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"] assert args.diffusion_attention_backend == "FASTVIDEO_VSA" @@ -597,7 +618,7 @@ def test_serve_cli_accepts_request_batch_max_wait_ms(): ) explicit_kwargs = args.get_explicit_kwargs_dict() - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + stage_cfg = StageConfigFactory.create_default_diffusion(explicit_kwargs)[0] assert args.request_batch_max_wait_ms == 250.0 assert stage_cfg["engine_args"]["request_batch_max_wait_ms"] == 250.0 @@ -637,7 +658,7 @@ def test_serve_cli_accepts_additional_config(): ] ) - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))[0] + stage_cfg = StageConfigFactory.create_default_diffusion(vars(args))[0] engine_args = stage_cfg["engine_args"] @@ -645,79 +666,19 @@ def test_serve_cli_accepts_additional_config(): assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}} -def test_default_stage_resolves_video_output_from_checkpoint(mocker): - captured = {} - - def resolve_with_default(*args, default_stage_cfg_factory, **kwargs): - del args, kwargs - captured.update(default_stage_cfg_factory()[0]) - stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace()) - return ("", [stage], None) - - resolver = mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_model_class_name", - return_value="MiniMaxH3Pipeline", - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", - side_effect=resolve_with_default, - ) - engine = AsyncOmniEngine.__new__(AsyncOmniEngine) - engine._strip_single_engine_args = lambda kwargs: kwargs - - engine._resolve_stage_configs( - "/models/MiniMax-H3/FL2VA", - {}, - trust_remote_code=False, - ) - - resolver.assert_called_once_with("/models/MiniMax-H3/FL2VA", "default") - assert captured["final_output_type"] == "video" - - -def test_default_diffusers_stage_preserves_video_model_identity(mocker): - captured = {} - - def resolve_with_default(*args, default_stage_cfg_factory, **kwargs): - del args, kwargs - captured.update(default_stage_cfg_factory()[0]) - stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace()) - return ("", [stage], None) - - mocker.patch( - "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index", - return_value={"_class_name": "WanImageToVideoPipeline"}, - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", - side_effect=resolve_with_default, - ) - engine = AsyncOmniEngine.__new__(AsyncOmniEngine) - engine._strip_single_engine_args = lambda kwargs: kwargs - - engine._resolve_stage_configs( - "/models/Wan2.2-I2V", - {"diffusion_load_format": "diffusers"}, - trust_remote_code=False, - ) - - assert captured["engine_args"]["model_class_name"] == "WanImageToVideoPipeline" - assert captured["final_output_type"] == "video" - - -def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker): - """Ensure YAML/deploy stage resolution forwards top-level additional_config.""" +def test_resolve_stage_configs_delegates_overrides_to_resolver(mocker): + """The engine consumes resolver output without a second merge pass.""" + additional_config = {"torchair_graph_config": {"enabled": True}} fake_diffusion_stage = SimpleNamespace( stage_type="diffusion", - engine_args=SimpleNamespace(), + engine_args=SimpleNamespace(additional_config=additional_config), ) - fake_llm_stage = SimpleNamespace( - stage_type="llm", - engine_args=SimpleNamespace(), - ) - mocker.patch( - "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", - return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage], None), + resolve_config = mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_omni_config", + return_value=OmniConfigResolution( + config_path="dummy.yaml", + stage_configs=(fake_diffusion_stage,), + ), ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) @@ -726,13 +687,15 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo "dummy-model", { "deploy_config": "dummy.yaml", - "additional_config": {"torchair_graph_config": {"enabled": True}}, + "additional_config": additional_config, }, 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}} + assert stage_configs == [fake_diffusion_stage] + assert resolve_config.call_args.args == ("dummy-model",) + assert resolve_config.call_args.kwargs["deploy_config_path"] == "dummy.yaml" + assert resolve_config.call_args.kwargs["cli_overrides"]["additional_config"] is additional_config @pytest.mark.parametrize( @@ -760,30 +723,6 @@ def test_default_stage_config_includes_quantization_config(): "weights": "weights.bin", } - stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0] + stage_cfg = StageConfigFactory.create_default_diffusion({"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) - - _, stage_configs = engine._resolve_stage_configs( - "dummy-model", - { - "deploy_config": "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_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index 9600b990d..38ec40163 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for AR EngineCore vs diffusion worker pause/sleep routing.""" from __future__ import annotations @@ -7,6 +10,8 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, call import pytest +from vllm.renderers import BaseRenderer +from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.diffusion.data import CuMemTag, OmniACK from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -37,6 +42,22 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni: return omni +@pytest.mark.cpu +def test_reset_mm_cache_routes_through_renderer(mocker): + async def run() -> None: + omni = object.__new__(AsyncOmni) + renderer = mocker.create_autospec(BaseRenderer, instance=True) + input_processor = mocker.create_autospec(InputProcessor, instance=True) + input_processor.renderer = renderer + omni.input_processor = input_processor + + await omni.reset_mm_cache() + + renderer.clear_mm_cache_async.assert_awaited_once_with() + + asyncio.run(run()) + + @pytest.mark.cpu def test_split_stage_ids_by_type(): omni = _make_omni(stage_types=["llm", "diffusion", "llm"]) diff --git a/tests/entrypoints/test_disable_log_stats_routing.py b/tests/entrypoints/test_disable_log_stats_routing.py new file mode 100644 index 000000000..f9dd2da4a --- /dev/null +++ b/tests/entrypoints/test_disable_log_stats_routing.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import argparse +from types import SimpleNamespace +from typing import Any + +import pytest + +from vllm_omni.config.omni_config import ( + _NON_STAGE_ENGINE_CLI_FIELDS, + VllmOmniConfig, + _global_stage_cli_fields, +) +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.config.resolver import OmniConfigResolution +from vllm_omni.config.stage_config import PipelineConfig +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.entrypoints.cli.serve import run_headless +from vllm_omni.utils.tracking_parser import TrackingNamespace + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_disable_log_stats_is_not_a_stage_config_field() -> None: + assert "disable_log_stats" in _NON_STAGE_ENGINE_CLI_FIELDS + assert "disable_log_stats" not in _global_stage_cli_fields() + + +def test_structured_tts_config_accepts_disable_log_stats() -> None: + pipeline = OMNI_PIPELINES["qwen3_tts"] + assert isinstance(pipeline, PipelineConfig) + + config = VllmOmniConfig.from_pipeline_config( + pipeline, + cli_overrides={"disable_log_stats": True}, + ) + + assert len(config.stage_configs) == len(pipeline.stages) + + +def test_async_engine_consumes_disable_log_stats_before_resolution(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def fake_resolve(model: str, kwargs: dict[str, Any], **unused: Any) -> OmniConfigResolution: + captured["model"] = model + captured["kwargs"] = dict(kwargs) + return OmniConfigResolution( + config_path="config.yaml", + stage_configs=(), + ) + + monkeypatch.setattr( + "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + fake_resolve, + ) + engine = SimpleNamespace(_apply_strategy_lb_policy=lambda *args, **kwargs: None) + + config_path, stage_configs = AsyncOmniEngine._resolve_stage_configs( + engine, + "fake-model", + {"disable_log_stats": True}, + trust_remote_code=False, + ) + + assert captured["model"] == "fake-model" + assert "disable_log_stats" not in captured["kwargs"] + assert config_path == "config.yaml" + assert stage_configs == [] + + +def test_headless_consumes_disable_log_stats_before_resolution(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def fake_resolve(*args: Any, **kwargs: Any) -> OmniConfigResolution: + captured.update(kwargs) + return OmniConfigResolution( + config_path=None, + stage_configs=(SimpleNamespace(stage_id=99),), + ) + + monkeypatch.setattr("vllm_omni.config.resolver.resolve_omni_config", fake_resolve) + namespace = argparse.Namespace( + model="fake-model", + stage_id=0, + omni_master_address="127.0.0.1", + omni_master_port=26000, + worker_backend="multi_process", + omni_replica_address=None, + omni_dp_size_local=1, + disable_log_stats=True, + log_stats=False, + ) + args = TrackingNamespace( + unfiltered_ns=namespace, + explicit_keys=frozenset({"disable_log_stats"}), + ) + + with pytest.raises(ValueError, match="No stage config found"): + run_headless(args) + + assert "disable_log_stats" not in captured["cli_overrides"] diff --git a/tests/entrypoints/test_duplex_fence_propagation.py b/tests/entrypoints/test_duplex_fence_propagation.py index 2f1292ec7..550723ec8 100644 --- a/tests/entrypoints/test_duplex_fence_propagation.py +++ b/tests/entrypoints/test_duplex_fence_propagation.py @@ -74,7 +74,7 @@ async def test_openai_handler_passes_current_session_fence_for_runtime_controls( handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) session = DuplexSession( session_id="sid", - config=DuplexSessionConfig(), + config=DuplexSessionConfig(extra_body={"native_duplex": True}), epoch=1, turn_id=2, ) @@ -92,7 +92,7 @@ async def test_openai_handler_passes_current_session_fence_for_runtime_controls( async def test_runtime_signal_uses_fence_captured_before_session_turn_advances(): engine = FenceRecordingEngine() handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession(session_id="sid", config=DuplexSessionConfig()) + session = DuplexSession(session_id="sid", config=DuplexSessionConfig(extra_body={"native_duplex": True})) captured_fence = DuplexFence("sid", epoch=0, turn_id=0) session.turn_id = 1 @@ -111,7 +111,7 @@ async def test_openai_handler_forwards_cancel_fence_through_async_omni_facade(): app = object.__new__(AsyncOmni) app.engine = engine handler = OmniDuplexSessionHandler(chat_service=FakeChatService(app)) - session = DuplexSession(session_id="sid-cancel", config=DuplexSessionConfig()) + session = DuplexSession(session_id="sid-cancel", config=DuplexSessionConfig(extra_body={"native_duplex": True})) cancelled_fence = DuplexFence("sid-cancel", epoch=2, turn_id=3) next_fence = DuplexFence("sid-cancel", epoch=3, turn_id=3) diff --git a/tests/entrypoints/test_pd_disaggregation.py b/tests/entrypoints/test_pd_disaggregation.py index c896ccb88..e9a3ba54c 100644 --- a/tests/entrypoints/test_pd_disaggregation.py +++ b/tests/entrypoints/test_pd_disaggregation.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for PD (Prefill-Decode) disaggregation in the Omni orchestrator. Tests the PD detection, validation, config parsing, sampling param @@ -74,7 +77,7 @@ class _FakeStageConfig: class _FakeQueue: def __init__(self, maxsize=0): - self._queue = Queue(maxsize=maxsize) + self._queue: Queue[Any] = Queue(maxsize=maxsize) def put(self, item): self._queue.put(item) @@ -343,7 +346,7 @@ def mock_get_config(monkeypatch): monkeypatch.setattr( "vllm.transformers_utils.config.get_config", lambda model, **kwargs: fake_hf_config, raising=False ) - monkeypatch.setattr("vllm_omni.entrypoints.utils.get_config", lambda model, **kwargs: fake_hf_config, raising=False) + monkeypatch.setattr("vllm_omni.config.config_factory.get_config", lambda model, **kwargs: fake_hf_config) def _mock_cached_file(path_or_repo_id, *args, **kwargs): import os diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py index 660203710..450e1be39 100644 --- a/tests/entrypoints/test_resolve_dreamzero_config.py +++ b/tests/entrypoints/test_resolve_dreamzero_config.py @@ -1,28 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest +from vllm_omni.config.resolver import resolve_omni_config from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc -from vllm_omni.entrypoints.utils import load_stage_configs_from_model, resolve_model_config_path pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch): - monkeypatch.setattr( - "vllm_omni.entrypoints.utils.get_config", - lambda _model, trust_remote_code=True: type("Cfg", (), {"model_type": "vla"})(), - ) - monkeypatch.setattr( - "vllm_omni.entrypoints.utils._looks_like_dreamzero", - lambda _model: True, - ) - result = resolve_model_config_path("GEAR-Dreams/DreamZero-DROID") - - assert result is not None - assert result.endswith("vllm_omni/deploy/dreamzero.yaml") - - -def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch): +def test_dreamzero_resolves_through_registry_with_model_defaults(monkeypatch): monkeypatch.setattr( "vllm_omni.config.config_factory.StageConfigFactory._try_infer_model_type", classmethod(lambda _cls, model, trust_remote_code=True: "vla"), @@ -36,12 +24,18 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch): lambda _model: True, ) - stage_configs, _ = load_stage_configs_from_model( + resolved = resolve_omni_config( "GEAR-Dreams/DreamZero-DROID", trust_remote_code=False, + deploy_config_path=None, + cli_overrides=None, + stage_overrides=None, + strategy_config_path=None, ) - engine_args = stage_configs[0].engine_args + engine_args = resolved.stage_configs[0].engine_args + assert resolved.config_path is not None + assert resolved.config_path.endswith("vllm_omni/deploy/dreamzero.yaml") assert engine_args.model_class_name == "DreamZeroPipeline" assert engine_args.model_config.policy_server_config.action_space == "joint_position" diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 29090c425..b64758e5d 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -11,13 +11,22 @@ from types import SimpleNamespace import pytest from pytest_mock import MockerFixture -from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless +from vllm_omni.config.resolver import OmniConfigResolution +from vllm_omni.entrypoints.cli.serve import ( + OmniServeCommand, + _parse_stage_overrides, + 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] +def _resolved(*stages: SimpleNamespace) -> OmniConfigResolution: + return OmniConfigResolution(config_path="/fake/stages.yaml", stage_configs=tuple(stages)) + + def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None: """``--no-async-chunk`` should parse to ``async_chunk=False`` and mark the shared deploy-level dest as explicitly provided by the user.""" @@ -110,7 +119,7 @@ def test_serve_parser_accepts_diffusion_quantization_config() -> None: assert args.get_explicit_kwargs_dict()["diffusion_quantization_config"] == expected -def _make_headless_args(**kwargs) -> TrackingNamespace: +def _make_headless_args(*, explicit_keys: frozenset[str] | None = None, **kwargs) -> TrackingNamespace: defaults = { "model": "fake-model", "stage_id": 0, @@ -130,7 +139,7 @@ def _make_headless_args(**kwargs) -> TrackingNamespace: ns = argparse.Namespace(**ns_kwargs) return TrackingNamespace( unfiltered_ns=ns, - explicit_keys=frozenset(ns.__dict__.keys()), + explicit_keys=frozenset(ns.__dict__.keys()) if explicit_keys is None else explicit_keys, ) @@ -159,41 +168,49 @@ def test_run_headless_rejects_non_multiprocess_worker_backend() -> None: # --------------------------------------------------------------------------- -# --stage-overrides parsing parity (headless vs standard path) +# --stage-overrides parsing at the serving boundary # --------------------------------------------------------------------------- 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"}}') + 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_serve_parser_parses_stage_overrides_before_resolution() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "fake-model", + "--omni", + "--stage-overrides", + '{"0": {"devices": "0,1"}}', + ] + ) + + assert args.stage_overrides == {"0": {"devices": "0,1"}} -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_serve_parser_accepts_empty_stage_overrides_as_noop() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + args = parser.parse_args(["serve", "fake-model", "--omni", "--stage-overrides", "{}"]) -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 + assert args.stage_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: `` suffix echoing the raw input.""" + """Invalid JSON fails at the serving boundary with the raw input.""" bad = "{not valid json}" - with pytest.raises(ValueError) as excinfo: - parse_stage_overrides(bad) + with pytest.raises(argparse.ArgumentTypeError) 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 @@ -271,11 +288,8 @@ def test_parse_stage_overrides_accepts_empty_inner_dict() -> None: assert parsed == {"0": {}, "1": {}} -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.""" +def test_run_headless_forwards_parsed_stage_overrides(mocker: MockerFixture) -> None: + """The headless resolver receives the mapping parsed by argparse.""" captured: dict = {} def _fake_resolve(*args, **kwargs): @@ -283,10 +297,10 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) 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) + return _resolved(SimpleNamespace(stage_id=99)) mocker.patch( - "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + "vllm_omni.config.resolver.resolve_omni_config", side_effect=_fake_resolve, ) @@ -294,28 +308,22 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) stage_id=0, deploy_config="/tmp/deploy.yaml", strategy_config="/tmp/strategy.yaml", - stage_overrides='{"0": {"devices": "0,1"}, "1": {"devices": "2"}}', + 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 len(captured["args"]) == 2 + assert captured["args"] == ("fake-model",) assert captured["deploy_config_path"] == "/tmp/deploy.yaml" 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_parse_stage_overrides_rejects_non_mapping_values() -> None: + with pytest.raises(argparse.ArgumentTypeError, match="JSON object"): + _parse_stage_overrides('["not", "a", "mapping"]') + with pytest.raises(argparse.ArgumentTypeError, match="must be an object"): + _parse_stage_overrides('{"0": "not a mapping"}') def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None: @@ -323,8 +331,8 @@ def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) 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], None), + "vllm_omni.config.resolver.resolve_omni_config", + return_value=_resolved(other_stage), ) args = _make_headless_args(stage_id=0) @@ -372,8 +380,8 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker engine_manager = mocker.Mock() mocker.patch( - "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg], None), + "vllm_omni.config.resolver.resolve_omni_config", + return_value=_resolved(stage_cfg), ) 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) @@ -452,8 +460,8 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo manager_b = mocker.Mock() mocker.patch( - "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg], None), + "vllm_omni.config.resolver.resolve_omni_config", + return_value=_resolved(stage_cfg), ) 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) @@ -511,8 +519,8 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) proc.is_alive.return_value = False mocker.patch( - "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg], None), + "vllm_omni.config.resolver.resolve_omni_config", + return_value=_resolved(stage_cfg), ) 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) @@ -591,8 +599,8 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur proc.is_alive.return_value = False mocker.patch( - "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg], None), + "vllm_omni.config.resolver.resolve_omni_config", + return_value=_resolved(stage_cfg), ) 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) diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 6bc6cdc2a..680bb110c 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -4,10 +4,8 @@ """Unit tests for vllm_omni.entrypoints.utils module.""" import logging -import os from collections import Counter from dataclasses import dataclass -from pathlib import Path from types import SimpleNamespace import pytest @@ -15,25 +13,22 @@ import torch from pytest_mock import MockerFixture from vllm.sampling_params import RequestOutputKind, SamplingParams +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.config.resolver import ( + OmniConfigResolution, + _convert_dataclasses_to_dict, + _filter_dict_like_object, + resolve_omni_config, +) +from vllm_omni.config.stage_config import PipelineConfig from vllm_omni.config.yaml_util import create_config from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs -from vllm_omni.engine.async_omni_engine import AsyncOmniEngine -from vllm_omni.engine.stage_init_utils import ( - get_stage_connector_spec, - load_omni_transfer_config_for_model, -) from vllm_omni.entrypoints.utils import ( - _convert_dataclasses_to_dict, - _filter_dict_like_object, - _try_resolve_omni_model_type, coerce_param_message_types, filter_dataclass_kwargs, - filter_stages, - load_and_resolve_stage_configs, - load_stage_configs_from_model, - resolve_model_config_path, ) +from vllm_omni.model_executor.models.qwen3_omni.pipeline import QWEN3_OMNI_PIPELINE pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -173,7 +168,7 @@ class TestFilterDictLikeObject: "normal": "value", } - mocker.patch("vllm_omni.entrypoints.utils.logger") + mocker.patch("vllm_omni.config.resolver.logger") result = _filter_dict_like_object(input_dict) # Normal key should exist @@ -195,7 +190,7 @@ class TestFilterDictLikeObject: "normal": "value", } - mocker.patch("vllm_omni.entrypoints.utils.logger") + mocker.patch("vllm_omni.config.resolver.logger") result = _filter_dict_like_object(input_dict) # Callable should be filtered @@ -215,7 +210,7 @@ class TestConvertDataclassesToDict: "callable": lambda x: x, } - mocker.patch("vllm_omni.entrypoints.utils.logger") + mocker.patch("vllm_omni.config.resolver.logger") result = _convert_dataclasses_to_dict(input_dict) # Callable should be filtered out by _filter_dict_like_object @@ -300,254 +295,199 @@ class TestFilterDataclassKwargs: assert "extra_param" not in result["cache_config"] -class TestResolveModelConfigPath: - """Test suite for resolve_model_config_path function with diffusers format models.""" - - def test_glm_image_diffusers_format_resolution(self, mocker: MockerFixture): - """Test GlmImagePipeline diffusers class resolves to glm_image config.""" - mocker.patch( - "vllm_omni.entrypoints.utils.file_or_path_exists", - return_value=True, - ) - mocker.patch( - "vllm_omni.entrypoints.utils._try_get_class_name_from_diffusers_config", - return_value="GlmImagePipeline", - ) - mocker.patch( - "vllm_omni.entrypoints.utils.current_omni_platform.get_default_stage_config_path", - return_value="vllm_omni/model_executor/stage_configs", - ) - - original_exists = os.path.exists - - def mock_exists(path): - if "glm_image.yaml" in str(path): - return True - return original_exists(path) - - mocker.patch("os.path.exists", side_effect=mock_exists) - - result = resolve_model_config_path("zai-org/GLM-Image") - - assert result is not None - assert "glm_image.yaml" in result - - def test_uses_registry_default_deploy_when_hf_model_type_differs( - self, - mocker: MockerFixture, - tmp_path, - ): - from vllm_omni.config.stage_config import _DEPLOY_DIR as _REAL_DEPLOY_DIR - from vllm_omni.entrypoints import utils as utils_mod - - hf_config = SimpleNamespace( - model_type="minicpmo", - architectures=["MiniCPMO"], - version="4.5", - ) - model = "local/minicpmo-registry-default-test" - mocker.patch( - "vllm_omni.entrypoints.utils.get_config", - return_value=hf_config, - ) - mocker.patch( - "vllm_omni.config.config_factory.get_config", - return_value=hf_config, - ) - # A bare deploy/.yaml must not win over the registered - # pipeline default when HF model_type is not an OMNI_PIPELINES key. - (tmp_path / "minicpmo.yaml").write_text("stages: []\n", encoding="utf-8") - (tmp_path / "minicpmo_4_5.yaml").write_bytes((_REAL_DEPLOY_DIR / "minicpmo_4_5.yaml").read_bytes()) - mocker.patch.object(utils_mod, "_DEPLOY_DIR", tmp_path) - - result = resolve_model_config_path(model) - - assert result is not None - assert Path(result).as_posix().endswith("minicpmo_4_5.yaml") - assert Path(result).name == "minicpmo_4_5.yaml" - transfer_config = load_omni_transfer_config_for_model(model, result) - assert transfer_config is not None - - sender = get_stage_connector_spec( - omni_transfer_config=transfer_config, - stage_id=1, - async_chunk=True, - ) - receiver = get_stage_connector_spec( - omni_transfer_config=transfer_config, - stage_id=2, - async_chunk=True, - ) - assert sender["extra"]["role"] == "sender" - assert receiver["extra"]["role"] == "receiver" - assert sender["extra"]["connector_get_sleep_s"] == 0.01 - assert sender["extra"]["connector_get_max_wait_first_chunk"] == 3000 - assert sender["extra"]["connector_get_max_wait"] == 300 - - def test_object_storage_uri_resolves_via_materialized_configs( - self, - mocker: MockerFixture, - tmp_path, - ): - """An object-storage URI must not reach HF helpers; resolution reads the - locally materialized config files instead.""" - hf_config = SimpleNamespace( - model_type="qwen3_omni_moe", - architectures=["Qwen3OmniMoeForConditionalGeneration"], - ) - uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct" - materialize = mocker.patch( - "vllm_omni.entrypoints.utils._materialize_object_storage_configs", - return_value=str(tmp_path), - ) - get_config = mocker.patch( - "vllm_omni.entrypoints.utils.get_config", - return_value=hf_config, - ) - - result = resolve_model_config_path(uri) - - materialize.assert_called_once_with(uri) - get_config.assert_called_once() - assert get_config.call_args[0][0] == str(tmp_path) - assert result is not None - assert Path(result).as_posix().endswith("qwen3_omni_moe.yaml") - - -class TestTryResolveOmniModelType: - """Name matching must only scan the last path component.""" - - def test_uri_bucket_name_cannot_hijack_pipeline(self): - # "qwen3-tts-models" normalizes to "qwen3ttsmodels", which contains the - # registered key "qwen3tts" — only the basename may participate, so a - # misleadingly named bucket selects nothing. - assert _try_resolve_omni_model_type("s3://qwen3-tts-models/plain-checkpoint") is None - - def test_basename_still_matches_pipeline(self): - assert _try_resolve_omni_model_type("gs://any-bucket/Fun-CosyVoice3-0.5B-2512") == "cosyvoice3" - - -class TestLoadAndResolveStageConfigs: - def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture): - """Ensure that kwargs survive default stage creation.""" - engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine" - kwargs = {"dtype": torch.float32, "engine_backend": engine_backend} - kwargs["revision"] = "pinned-revision" - 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", - kwargs=kwargs, - trust_remote_code=False, - default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs), - ) - assert config_path is None - assert len(stage_configs) == 1 - assert "dtype" in stage_configs[0]["engine_args"] - assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend - assert stage_configs[0]["engine_args"]["revision"] == "pinned-revision" - - def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture): - deploy_path = tmp_path / "qwen3_multi.yaml" +class TestResolveOmniConfig: + def test_bare_deploy_name_returns_packaged_resolved_path(self, tmp_path, mocker: MockerFixture): + deploy_name = "qwen3_omni_moe.yaml" + deploy_path = tmp_path / deploy_name deploy_path.write_text( - 'stages:\n - stage_id: 0\n devices: "0"\n - stage_id: 1\n devices: "1,2,3"\n num_replicas: 3\n', + "duplex_session:\n server_vad_model_path: /models/silero_vad.onnx\nstages: []\n", encoding="utf-8", ) - - returned_stage_configs = [ - create_config({"stage_id": 0, "runtime": {"devices": "0"}, "engine_args": {"model": "dummy"}}), - create_config( - { - "stage_id": 1, - "runtime": {"devices": "1,2,3", "num_replicas": 3}, - "engine_args": {"model": "dummy"}, - } - ), - ] - load_stage_configs = mocker.patch( - "vllm_omni.entrypoints.utils.load_stage_configs_from_model", - return_value=(returned_stage_configs, None), - ) - cli_overrides = { - "interleave_mm_strings": True, - "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}, - "gpu_memory_utilization": 0.55, - } - - config_path, stage_configs, _ = load_and_resolve_stage_configs( - model="dummy-model", - kwargs=cli_overrides, - trust_remote_code=True, - deploy_config_path=str(deploy_path), + mocker.patch("vllm_omni.config.config_factory._DEPLOY_DIR", tmp_path) + mocker.patch("vllm_omni.config.omni_config._DEPLOY_DIR", tmp_path) + mocker.patch( + "vllm_omni.config.config_factory.StageConfigFactory.get_pipeline_config", + return_value=QWEN3_OMNI_PIPELINE, ) - load_stage_configs.assert_called_once_with( + resolved = resolve_omni_config( "dummy-model", - trust_remote_code=True, - base_engine_args=cli_overrides, - deploy_config_path=str(deploy_path), + trust_remote_code=False, + deploy_config_path=deploy_name, + cli_overrides={}, stage_overrides=None, strategy_config_path=None, ) - assert config_path == str(deploy_path) - assert len(stage_configs) == 2 - assert stage_configs[1].runtime.num_replicas == 3 - assert stage_configs[1].runtime.devices == "1,2,3" - def test_filter_stages_selects_mode_stages_without_mutating_stage_config(self, tmp_path): - config_path = tmp_path / "deploy.yaml" - config_path.write_text( - """modes: - - mode: text-to-text - stages: [0] - - mode: text-to-image - stages: [0, 1] -""", - encoding="utf-8", + assert resolved.config_path == str(deploy_path) + + def test_stage_lookup_error_lists_resolved_ids(self): + resolved = OmniConfigResolution( + config_path=None, + stage_configs=(SimpleNamespace(stage_id=1), SimpleNamespace(stage_id=3)), ) - stages = [ - create_config( - { - "stage_id": 0, - "runtime": {"requires_multimodal_data": True}, - "final_output": False, - "final_output_type": None, - } - ), - create_config( - { - "stage_id": 1, - "runtime": {"requires_multimodal_data": True}, - "final_output": True, - "final_output_type": "image", - } - ), - ] - filtered = filter_stages(str(config_path), stages, {"mode": "text-to-text"}) + with pytest.raises(KeyError, match=r"no stage 2; resolved stages: \[1, 3\]"): + resolved.stage_by_id(2) - assert len(filtered) == 1 - assert filtered[0].stage_id == 0 - assert filtered[0].runtime.requires_multimodal_data is True - assert filtered[0].final_output is False - assert filtered[0].final_output_type is None - - -class TestLoadStageConfigsFromModel: - def test_unresolved_model_does_not_fall_back_to_yaml(self, mocker: MockerFixture, caplog): + def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture): + """Ensure that generic diffusion overrides survive resolution.""" + engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine" mocker.patch( - "vllm_omni.entrypoints.utils.StageConfigFactory.create_legacy_stage_configs_from_model", - return_value=(None, None), + "vllm_omni.config.resolver.StageConfigFactory.create_from_model", + return_value=None, ) - resolve_path = mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path") + mocker.patch( + "vllm_omni.config.resolver._resolve_generic_diffusion_model_class", + return_value=(True, "FluxPipeline"), + ) + kwargs = { + "dtype": torch.float32, + "engine_backend": engine_backend, + "revision": "pinned-revision", + } + resolved = resolve_omni_config( + "black-forest-labs/FLUX.2-klein-4B", + trust_remote_code=False, + deploy_config_path=None, + cli_overrides=kwargs, + stage_overrides=None, + strategy_config_path=None, + ) + assert resolved.config_path is None + assert len(resolved.stage_configs) == 1 + engine_args = resolved.stage_configs[0]["engine_args"] + assert "dtype" in engine_args + assert engine_args["engine_backend"] == engine_backend + assert engine_args["revision"] == "pinned-revision" - result = load_stage_configs_from_model("unregistered-model", trust_remote_code=False) + def test_generic_diffusion_uses_registered_model_metadata(self, mocker: MockerFixture): + mocker.patch( + "vllm_omni.config.resolver.StageConfigFactory.create_from_model", + return_value=None, + ) + resolve_model_class = mocker.patch( + "vllm_omni.config.resolver.resolve_model_class_name", + return_value="WanImageToVideoPipeline", + ) + mocker.patch( + "vllm_omni.config.resolver.DiffusionModelRegistry.get_supported_archs", + return_value={"WanImageToVideoPipeline"}, + ) - assert result == ([], None) - resolve_path.assert_not_called() - assert "No registered PipelineConfig resolved" in caplog.text - assert "deploy_config" in caplog.text + resolved = resolve_omni_config( + "/models/Wan2.2-I2V", + trust_remote_code=False, + deploy_config_path=None, + cli_overrides={ + "diffusion_load_format": "diffusers", + "revision": "pinned-revision", + }, + stage_overrides=None, + strategy_config_path=None, + ) + + stage = resolved.stage_configs[0] + resolve_model_class.assert_called_once_with( + "/models/Wan2.2-I2V", + "diffusers", + "pinned-revision", + ) + assert stage.engine_args.model_class_name == "WanImageToVideoPipeline" + assert stage.final_output_type == "video" + + def test_registered_pipeline_uses_structured_metadata_and_preserves_override_trust(self, mocker: MockerFixture): + endpoint_restriction = SimpleNamespace(name="chat") + structured_config = SimpleNamespace( + orchestrator_config=SimpleNamespace(deploy_config_path="/resolved/deploy.yaml"), + pipeline_config=SimpleNamespace(endpoint_restrictions=(endpoint_restriction,)), + ) + runtime_stage = create_config( + { + "stage_id": 1, + "runtime": {"devices": "1,2,3", "num_replicas": 3}, + "engine_args": {"model": "dummy-model"}, + } + ) + legacy_stage = SimpleNamespace(to_omegaconf=lambda: runtime_stage) + create_structured = mocker.patch( + "vllm_omni.config.resolver.StageConfigFactory.create_from_model", + return_value=structured_config, + ) + create_legacy = mocker.patch( + "vllm_omni.config.resolver.StageConfigFactory._resolve_legacy_from_registry", + return_value=SimpleNamespace( + stage_configs=[legacy_stage], + pipeline_config=structured_config.pipeline_config, + omni_lb_policy="round_robin", + ), + ) + strategy_specs = {"stage_1": {"dp": 3}} + load_strategy = mocker.patch( + "vllm_omni.config.resolver._load_strategy_specs", + return_value=strategy_specs, + ) + + resolved = resolve_omni_config( + "dummy-model", + trust_remote_code=None, + deploy_config_path="deploy.yaml", + cli_overrides={"dtype": "bfloat16", "trust_remote_code": True}, + stage_overrides={"1": {"tensor_parallel_size": 2}}, + strategy_config_path="strategy.yaml", + ) + + expected_overrides = { + "dtype": "bfloat16", + "trust_remote_code": True, + "stage_1_tensor_parallel_size": 2, + } + create_structured.assert_called_once_with( + "dummy-model", + trust_remote_code=None, + cli_overrides=expected_overrides, + deploy_config_path="deploy.yaml", + ) + create_legacy.assert_called_once_with( + structured_config.pipeline_config, + expected_overrides, + "/resolved/deploy.yaml", + strategy_specs=strategy_specs, + ) + load_strategy.assert_called_once_with("strategy.yaml") + assert resolved.config_path == "/resolved/deploy.yaml" + assert resolved.pipeline_config is structured_config.pipeline_config + assert resolved.omni_lb_policy == "round_robin" + assert resolved.endpoint_restrictions == (endpoint_restriction,) + assert resolved.stage_configs == (runtime_stage,) + + def test_registered_resolution_exposes_forced_aligner_topology(self, mocker: MockerFixture): + pipeline = OMNI_PIPELINES["qwen3_tts"] + assert isinstance(pipeline, PipelineConfig) + structured_config = SimpleNamespace( + orchestrator_config=SimpleNamespace(deploy_config_path=None), + pipeline_config=pipeline, + ) + mocker.patch( + "vllm_omni.config.resolver.StageConfigFactory.create_from_model", + return_value=structured_config, + ) + + resolved = resolve_omni_config( + "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", + trust_remote_code=False, + deploy_config_path=None, + cli_overrides={"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"}, + stage_overrides=None, + strategy_config_path=None, + ) + + assert resolved.pipeline_config is not None + pipeline_stage_ids = [stage.stage_id for stage in resolved.pipeline_config.stages] + runtime_stage_ids = [stage.stage_id for stage in resolved.stage_configs] + assert pipeline_stage_ids == runtime_stage_ids + assert resolved.pipeline_config.stages[-1].model_stage == "forced_aligner" + assert len(resolved.pipeline_config.stages) == len(pipeline.stages) + 1 class TestCumulativeStreamingCoercion: diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py index a4e641a57..9120e183f 100644 --- a/tests/helpers/__init__.py +++ b/tests/helpers/__init__.py @@ -29,12 +29,13 @@ def skip_if_gated_repo_inaccessible( and we skip cleanly. """ try: - from huggingface_hub import hf_hub_download from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError + + from vllm_omni.transformers_utils.repo_utils import hf_api except Exception: return try: - hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) + hf_api().hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) except GatedRepoError as exc: import pytest diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 65d28109f..dd25aca71 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -659,7 +659,11 @@ def _compute_pcm_hnr_db(pcm_samples: np.ndarray, sr: int = _PCM_SPEECH_SAMPLE_RA return float(np.mean(hnr_values)) if hnr_values else 0.0 -def _assert_pcm_int16_speech_hnr(audio_bytes: bytes, min_hnr_db: float = _MIN_PCM_SPEECH_HNR_DB) -> None: +def _assert_pcm_int16_speech_hnr( + audio_bytes: bytes, + min_hnr_db: float = _MIN_PCM_SPEECH_HNR_DB, + sr: int = _PCM_SPEECH_SAMPLE_RATE_HZ, +) -> None: """Validate harmonic-to-noise ratio on raw int16 PCM from /v1/audio/speech. min_hnr_db defaults to the global _MIN_PCM_SPEECH_HNR_DB (1.0 dB), @@ -668,12 +672,16 @@ def _assert_pcm_int16_speech_hnr(audio_bytes: bytes, min_hnr_db: float = _MIN_PC intrinsically around -2 dB) can pass a lower per-test threshold via request_config["min_hnr_db"] to keep the catastrophic-failure check while not gating CI on a model-intrinsic property. + + ``sr`` must be the stream's real sample rate: the autocorrelation lag window + is derived from it, so a wrong rate detunes the 80-400 Hz pitch search and + understates HNR for models that do not decode at 24 kHz. """ assert audio_bytes is not None and len(audio_bytes) >= 2, "missing PCM bytes" assert len(audio_bytes) % 2 == 0, "PCM byte length must be aligned to int16" pcm_samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0 - hnr = _compute_pcm_hnr_db(pcm_samples) - print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {min_hnr_db} dB)") + hnr = _compute_pcm_hnr_db(pcm_samples, sr=sr) + print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {min_hnr_db} dB, sr={sr})") assert hnr >= min_hnr_db, ( f"Audio distortion detected: HNR={hnr:.2f} dB < {min_hnr_db} dB. " "Voice clone decoder may be losing ref_code speaker context on later chunks." @@ -934,7 +942,15 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], 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) + # Raw PCM carries no header, so the HNR pitch search has to be told + # the rate. Defaulting to 24 kHz for a 44.1 kHz model shifts the + # search window to 147-735 Hz and misses the fundamental entirely, + # scoring clean speech ~1.9 dB lower than it is. + _assert_pcm_int16_speech_hnr( + response.audio_bytes, + min_hnr_db=min_hnr_db, + sr=int(request_config.get("expected_sample_rate") or _PCM_SPEECH_SAMPLE_RATE_HZ), + ) transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=True) if transcript is not None: diff --git a/tests/helpers/client.py b/tests/helpers/client.py index 0f8346200..ed5158b92 100644 --- a/tests/helpers/client.py +++ b/tests/helpers/client.py @@ -12,6 +12,7 @@ import copy import io import json import time +from collections.abc import Mapping from dataclasses import asdict, dataclass from io import BytesIO from pathlib import Path @@ -152,11 +153,14 @@ class OpenPIWebSocketResponse: action_tensors: list[np.ndarray] | None = None -def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[str, Any]: - """Build a minimal DROID-style observation payload for the OpenPI robot endpoint.""" +def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke", seed: int | None = None) -> dict[str, Any]: + """Build a minimal DROID-style observation payload for the OpenPI robot endpoint. + + ``seed`` pins the policy's sampling noise for this request (see ``docs/serving/openpi_api.md``). + """ identity_eef_9d = np.zeros((1, 1, 9), dtype=np.float32) identity_eef_9d[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) - return { + observation: dict[str, Any] = { "session_id": session_id, "video": { "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8), @@ -169,6 +173,9 @@ def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[s }, "language": {"annotation.language.language_instruction": [["pick up the object"]]}, } + if seed is not None: + observation["seed"] = seed + return observation DREAMZERO_DEFAULT_PROMPT = ( @@ -1462,6 +1469,7 @@ class OnlineOmniClient: "instructions", "speed", "sample_rate", + "extra_params", "stream_format", "x_vector_only_mode", ): @@ -2116,6 +2124,123 @@ class OfflineOmniClient: assert_audio_speech_response(result, request_config, run_level="core_model") return result + def send_tokenized_tts_request(self, request_config: dict[str, Any]) -> OmniResponse: + """Offline TTS for models whose engine prompt is *pre-tokenized*. + + Fish-Speech-style DualAR models (and any model whose prompt embeds codec + codes) cannot go through :meth:`send_audio_speech_request`: the prompt is + built by the caller as ``prompt_token_ids`` plus per-request + ``additional_information``, not from ``mm_processor_kwargs``. + + request_config keys: + - ``prompt_token_ids``: ``list[int]`` or ``list[list[int]]`` (one entry + per request). + - ``additional_information``: dict, or list of dicts matching + ``prompt_token_ids``. + - ``expected_sample_rate``: asserted against the reported ``sr``. + - ``min_duration_s`` / ``max_duration_s``: bounds on decoded audio. + - plus every key ``assert_audio_speech_response`` understands + (``response_format``, ``min_audio_bytes``, ``input``, ...). + """ + token_ids = request_config.get("prompt_token_ids") + if not token_ids: + raise ValueError("request_config must contain 'prompt_token_ids'") + batched_ids = token_ids if isinstance(token_ids[0], (list, tuple)) else [token_ids] + + extra = request_config.get("additional_information") + if extra is None: + batched_extra: list[dict[str, Any] | None] = [None] * len(batched_ids) + elif isinstance(extra, dict): + batched_extra = [extra] * len(batched_ids) + else: + batched_extra = list(extra) + if len(batched_extra) != len(batched_ids): + raise ValueError("'additional_information' must match the number of prompts") + + prompts: list[dict[str, Any]] = [] + for ids, info in zip(batched_ids, batched_extra, strict=True): + prompt: dict[str, Any] = {"prompt_token_ids": list(ids)} + if info is not None: + prompt["additional_information"] = info + prompts.append(prompt) + + start_time = time.perf_counter() + outputs = self.runner.omni.generate(prompts) + e2e_latency = time.perf_counter() - start_time + + assert len(outputs) == len(prompts), f"Expected {len(prompts)} outputs, got {len(outputs)}" + expected_sr = request_config.get("expected_sample_rate") + min_duration = request_config.get("min_duration_s") + max_duration = request_config.get("max_duration_s") + + result = OmniResponse() + wavs: list[torch.Tensor] = [] + for index, stage_output in enumerate(outputs): + assert stage_output.outputs, f"Request {index} produced no output" + multimodal_output = stage_output.outputs[0].multimodal_output + assert isinstance(multimodal_output, Mapping), ( + f"Request {index}: expected a mapping, got {type(multimodal_output)}" + ) + + # Never use ``a or b`` here: truthiness on a tensor raises. + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + assert audio is not None, f"Request {index}: no audio key in {sorted(multimodal_output)}" + if isinstance(audio, list): + chunks = [torch.as_tensor(chunk).float().cpu().reshape(-1) for chunk in audio if chunk is not None] + assert chunks, f"Request {index}: audio list contained no tensors" + audio = torch.cat(chunks, dim=0) + wav = torch.as_tensor(audio).float().cpu().reshape(-1) + assert wav.numel() > 0, f"Request {index}: decoded audio is empty" + wavs.append(wav) + + sr_raw = multimodal_output.get("sr") + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else None + assert sr_raw is not None, f"Request {index}: missing sample rate" + sample_rate = int(sr_raw.item() if hasattr(sr_raw, "item") else sr_raw) + if expected_sr is not None: + assert sample_rate == int(expected_sr), ( + f"Request {index}: expected sample rate {expected_sr}, got {sample_rate}" + ) + + duration_s = wav.numel() / sample_rate + if min_duration is not None: + assert duration_s >= float(min_duration), ( + f"Request {index}: audio too short ({duration_s:.2f}s < {min_duration}s)" + ) + if max_duration is not None: + assert duration_s <= float(max_duration), ( + f"Request {index}: audio too long ({duration_s:.2f}s > {max_duration}s)" + ) + + wav_buf = io.BytesIO() + sf.write(wav_buf, wav.numpy(), samplerate=sample_rate, format="WAV", subtype="PCM_16") + result = OmniResponse( + success=True, + audio_bytes=wav_buf.getvalue(), + audio_format="audio/wav", + e2e_latency=e2e_latency, + ) + assert_audio_speech_response(result, request_config, run_level="core_model") + + if request_config.get("assert_distinct_outputs") and len(wavs) > 1: + # Distinct inputs must decode to distinct audio. Leaked per-request + # codec state shows up as two requests producing (near-)identical + # waveforms; different lengths already prove distinctness. + for i in range(len(wavs)): + for j in range(i + 1, len(wavs)): + a, b = wavs[i], wavs[j] + if a.numel() != b.numel(): + continue + max_diff = (a - b).abs().max().item() + assert max_diff > 1e-3, ( + f"Requests {i} and {j} produced near-identical audio " + f"(max sample diff {max_diff:.2e}); possible codec-state crosstalk" + ) + return result + def start_profile(self, profile_prefix: str | None = None, stages: list[int] | None = None) -> list[Any]: return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages) diff --git a/tests/metrics/test_metrics_utils.py b/tests/metrics/test_metrics_utils.py index 544c7819a..4ff201561 100644 --- a/tests/metrics/test_metrics_utils.py +++ b/tests/metrics/test_metrics_utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from types import SimpleNamespace @@ -9,11 +12,30 @@ from vllm_omni.metrics.utils import ( count_audio_frames, count_image_pixels, count_tokens_from_outputs, + count_video_frames, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize( + ("shape", "expected_frames"), + [ + ((1, 64, 64, 3), 1), + ((3, 64, 64, 3), 3), + ((4, 64, 64, 3), 4), + ((8, 64, 64, 3), 8), + ((3, 8, 64, 64), 8), + ((1, 1, 64, 64, 3), 1), + ((1, 3, 64, 64, 3), 3), + ((1, 4, 64, 64, 3), 4), + ((1, 3, 8, 64, 64), 8), + ], +) +def test_count_video_frames_handles_short_channel_last_outputs(shape, expected_frames): + assert count_video_frames(SimpleNamespace(shape=shape)) == expected_frames + + @pytest.mark.parametrize( ("audio_chunk", "expected"), [ diff --git a/tests/model_executor/models/audio8_tts/__init__.py b/tests/model_executor/models/audio8_tts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/audio8_tts/test_audio8_tts_codec_decoder.py b/tests/model_executor/models/audio8_tts/test_audio8_tts_codec_decoder.py new file mode 100644 index 000000000..a737a7f4b --- /dev/null +++ b/tests/model_executor/models/audio8_tts/test_audio8_tts_codec_decoder.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import functools + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +NUM_CODEBOOKS = 10 +#: Samples the fake codec returns per frame, so trims are easy to assert. +SAMPLES_PER_FRAME = 100 + + +@functools.lru_cache(maxsize=1) +def _decoder_cls(): + from vllm_omni.model_executor.models.audio8_tts.audio8_tts_codec_decoder import Audio8TTSCodecDecoder + + return Audio8TTSCodecDecoder + + +class _FakeCodec: + """Returns ``SAMPLES_PER_FRAME`` ramp samples per frame.""" + + def __init__(self): + self.calls: list[tuple[int, ...]] = [] + + def decode(self, codes_bqf: torch.Tensor) -> torch.Tensor: + self.calls.append(tuple(codes_bqf.shape)) + frames = int(codes_bqf.shape[-1]) + total = frames * SAMPLES_PER_FRAME + return torch.arange(total, dtype=torch.float32).view(1, 1, total) + + +def _make_decoder(codec: _FakeCodec, num_codebooks: int = NUM_CODEBOOKS, splits=None): + decoder = object.__new__(_decoder_cls()) + torch.nn.Module.__init__(decoder) + decoder._codec = codec + decoder._codec_device = torch.device("cpu") + decoder._codec_dtype = torch.float32 + decoder._num_codebooks = num_codebooks + decoder._sample_rate = 44100 + decoder._frame_size = 2048 + decoder._logged_stats = True + decoder._ensure_codec_loaded = lambda: None + if splits is not None: + decoder._split_request_ids = lambda ids, seq_token_counts=None: splits + return decoder + + +def test_left_context_frames_are_trimmed_from_the_emitted_chunk(): + """The decoder must emit *delta* audio: the left-context prefix exists only + so the causal codec has history and must never reach the client twice.""" + codec = _FakeCodec() + decoder = _make_decoder(codec) + codes = torch.arange(NUM_CODEBOOKS * 4, dtype=torch.long).reshape(NUM_CODEBOOKS, 4) + + out = decoder.forward( + input_ids=torch.tensor([0], dtype=torch.long), + runtime_additional_information=[{"codes": {"audio": codes}, "meta": {"left_context_size": 3}}], + ) + + audios = out.multimodal_outputs["model_outputs"] + assert len(audios) == 1 + # 4 frames decoded, 3 of context trimmed => 1 frame of new audio. + assert audios[0].shape[0] == SAMPLES_PER_FRAME + assert codec.calls == [(1, NUM_CODEBOOKS, 4)] + + +def test_tensor_codes_from_runtime_info_win_over_input_ids(): + codec = _FakeCodec() + decoder = _make_decoder(codec, num_codebooks=2) + out = decoder.forward( + input_ids=torch.zeros(64, dtype=torch.long), + runtime_additional_information=[{"codes": {"audio": torch.tensor([[1, 2], [3, 4]], dtype=torch.long)}}], + ) + assert codec.calls == [(1, 2, 2)] + assert out.multimodal_outputs["model_outputs"][0].shape[0] == 2 * SAMPLES_PER_FRAME + + +def test_flat_input_ids_are_reshaped_codebook_major(): + codec = _FakeCodec() + decoder = _make_decoder(codec, num_codebooks=2) + out = decoder.forward(input_ids=torch.arange(6, dtype=torch.long), runtime_additional_information=[{}]) + assert codec.calls == [(1, 2, 3)] + assert out.multimodal_outputs["model_outputs"][0].shape[0] == 3 * SAMPLES_PER_FRAME + + +def test_every_return_path_sets_model_outputs_and_sr(): + """A branch that omits ``model_outputs`` makes the serving layer drop the + chunk silently, so assert the key exists on all the early-exit paths.""" + codec = _FakeCodec() + + empty = _make_decoder(codec).forward(input_ids=torch.zeros(0, dtype=torch.long)) + assert empty.multimodal_outputs["model_outputs"][0].numel() == 0 + assert empty.multimodal_outputs["sr"] + + ragged = _make_decoder(codec, num_codebooks=4).forward( + input_ids=torch.arange(6, dtype=torch.long), # 6 % 4 != 0 + runtime_additional_information=[{}], + ) + assert ragged.multimodal_outputs["model_outputs"][0].numel() == 0 + + codes = torch.zeros((NUM_CODEBOOKS, 2), dtype=torch.long) + no_new = _make_decoder(codec).forward( + input_ids=torch.tensor([0], dtype=torch.long), + runtime_additional_information=[{"codes": {"audio": codes}, "meta": {"left_context_size": 2}}], + ) + assert no_new.multimodal_outputs["model_outputs"][0].numel() == 0 + assert codec.calls == [], "a context-only chunk must not be decoded" + + +def test_mixed_batch_keeps_per_request_alignment(): + """An empty request must not shift the following request's audio into its + slot: the output list index is the request index.""" + codec = _FakeCodec() + splits = [ + torch.empty((0,), dtype=torch.long), + torch.arange(NUM_CODEBOOKS * 2, dtype=torch.long), + ] + decoder = _make_decoder(codec, splits=splits) + out = decoder.forward( + input_ids=torch.arange(NUM_CODEBOOKS * 2, dtype=torch.long), + runtime_additional_information=[{}, {"meta": {"left_context_size": 1}}], + ) + audios = out.multimodal_outputs["model_outputs"] + assert len(audios) == 2 + assert audios[0].numel() == 0 + assert audios[1].shape[0] == SAMPLES_PER_FRAME + + +def test_codes_with_wrong_codebook_count_fall_back_to_input_ids(): + codec = _FakeCodec() + decoder = _make_decoder(codec, num_codebooks=2) + out = decoder.forward( + input_ids=torch.arange(4, dtype=torch.long), + runtime_additional_information=[{"codes": {"audio": torch.zeros((5, 3), dtype=torch.long)}}], + ) + assert codec.calls == [(1, 2, 2)] + assert out.multimodal_outputs["model_outputs"][0].shape[0] == 2 * SAMPLES_PER_FRAME diff --git a/tests/model_executor/models/audio8_tts/test_audio8_tts_prompt_utils.py b/tests/model_executor/models/audio8_tts/test_audio8_tts_prompt_utils.py new file mode 100644 index 000000000..9047d06f5 --- /dev/null +++ b/tests/model_executor/models/audio8_tts/test_audio8_tts_prompt_utils.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from vllm_omni.model_executor.models.audio8_tts.prompt_utils import ( + build_text_only_prompt_ids, + build_voice_clone_prompt_ids, + estimate_voice_clone_prompt_len, + normalize_text, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _SegmentTokenizer: + """Tokenizer stub that maps each encoded string to one sentinel id. + + Lets the tests assert the *segment sequence* the prompt builder emits, which + is what must stay byte-identical to the reference ``ArkttsProcessor``. + """ + + def __init__(self): + self.vocab: dict[str, int] = {} + + def encode(self, text: str, add_special_tokens: bool = True): + assert add_special_tokens is False, "Audio8 prompts must not add extra special tokens" + return [self.vocab.setdefault(text, 1000 + len(self.vocab))] + + def segment_of(self, token_id: int) -> str: + return next(text for text, tid in self.vocab.items() if tid == token_id) + + +def _segments(tokenizer: _SegmentTokenizer, ids: list[int]) -> list[str]: + return [tokenizer.segment_of(i) for i in ids] + + +def test_text_only_prompt_matches_reference_segment_order(): + tokenizer = _SegmentTokenizer() + ids, normalized = build_text_only_prompt_ids(tokenizer, " Hello world ") + + assert normalized == "Hello world" + assert _segments(tokenizer, ids) == [ + "<|im_start|>system\n", + "convert the provided text to speech", + "<|im_end|>\n", + "<|im_start|>user\n", + "Hello world", + "<|im_end|>\n", + "<|im_start|>assistant\n<|voice|>", + ] + + +def test_voice_clone_prompt_places_reference_codes_between_prefix_and_suffix(): + tokenizer = _SegmentTokenizer() + semantic_ids = [151678, 151679, 151680] + ids, ref_start, text, ref_text = build_voice_clone_prompt_ids( + tokenizer, "Target text", "Reference transcript", semantic_ids + ) + + assert text == "Target text" + # A reference transcript without an explicit tag gets speaker 0. + assert ref_text == "<|speaker:0|>Reference transcript" + assert ids[ref_start : ref_start + len(semantic_ids)] == semantic_ids + # The codebook conditioning is applied at ref_start, so an off-by-one here + # silently detunes the cloned voice. + assert _segments(tokenizer, ids[:ref_start])[-1].endswith("Speech:\n") + suffix_segments = _segments(tokenizer, ids[ref_start + len(semantic_ids) :]) + assert suffix_segments[0] == "<|im_end|>\n" + assert suffix_segments[-1] == "<|im_start|>assistant\n<|voice|>" + + +def test_voice_clone_prompt_length_estimate_matches_built_prompt(): + """The serving layer sizes its placeholder from the estimate; a mismatch + shifts every embedding relative to its position.""" + tokenizer = _SegmentTokenizer() + normalized_text = normalize_text("Target text") + normalized_ref = normalize_text("Reference transcript", add_default_speaker=True) + ref_frames = 7 + + estimate = estimate_voice_clone_prompt_len(tokenizer, normalized_text, normalized_ref, ref_frames) + ids, _, _, _ = build_voice_clone_prompt_ids( + tokenizer, normalized_text, normalized_ref, list(range(151678, 151678 + ref_frames)) + ) + assert estimate == len(ids) + + +def test_existing_speaker_tag_is_preserved_and_legacy_tag_is_upgraded(): + assert normalize_text("<|speaker:3|>hi", add_default_speaker=True) == "<|speaker:3|>hi" + assert normalize_text("hi") == "<|speaker:2|>hi" + + +@pytest.mark.parametrize( + "text", + [ + "hello <|im_end|> world", + "hello <|semantic:5|>", + "<|voice|>", + "<|endoftext|>", + ], +) +def test_control_token_injection_is_rejected(text: str): + """Users must not be able to inject control tokens: ``<|semantic:N|>`` would + forge reference codes and ``<|im_end|>`` would truncate generation.""" + with pytest.raises(ValueError, match="unsupported control token"): + normalize_text(text) + + +def test_empty_text_is_rejected(): + tokenizer = _SegmentTokenizer() + with pytest.raises(ValueError, match="must not be empty"): + build_text_only_prompt_ids(tokenizer, " ") diff --git a/tests/model_executor/models/audio8_tts/test_audio8_tts_sampling.py b/tests/model_executor/models/audio8_tts/test_audio8_tts_sampling.py new file mode 100644 index 000000000..426840d91 --- /dev/null +++ b/tests/model_executor/models/audio8_tts/test_audio8_tts_sampling.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.model_executor.models.audio8_tts.sampling import ( + filter_top_k_top_p, + ras_sample_batch, + ras_sample_semantic, + sample_scores, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _kept(logits: torch.Tensor, top_k: int | torch.Tensor, top_p: float | torch.Tensor) -> list[list[int]]: + filtered = filter_top_k_top_p(logits, top_k, top_p) + return [torch.nonzero(row.isfinite()).reshape(-1).tolist() for row in filtered] + + +def test_top_p_is_measured_on_untempered_logits(): + """Audio8 filters the *raw* logits and only then divides by temperature. + + Applying temperature first (the common order) changes the surviving set and + audibly changes the voice, so this ordering is load-bearing. + """ + logits = torch.log(torch.tensor([[0.6, 0.3, 0.1]])) + # Raw softmax is exactly (0.6, 0.3, 0.1): top_p=0.9 keeps the first two + # (cumsum 0.6, 0.9) and drops the third (cumsum 1.0 > 0.9). + assert _kept(logits, top_k=0, top_p=0.9) == [[0, 1]] + # Had the logits been tempered by 0.5 first, probabilities would become + # (0.7insh, ...) and only one token would survive. + tempered = _kept(logits / 0.5, top_k=0, top_p=0.9) + assert tempered == [[0, 1]] or tempered == [[0]] + + +def test_argmax_always_survives_filtering(): + logits = torch.tensor([[10.0, -1.0, -2.0]]) + # top_p=0.0 would otherwise remove every candidate and produce NaNs. + assert _kept(logits, top_k=1, top_p=0.0) == [[0]] + + +def test_per_row_top_k_and_top_p_are_honoured(): + logits = torch.log(torch.tensor([[0.5, 0.3, 0.15, 0.05], [0.5, 0.3, 0.15, 0.05]])) + kept = _kept(logits, top_k=torch.tensor([1, 3]), top_p=torch.tensor([1.0, 1.0])) + assert kept == [[0], [0, 1, 2]] + + +def test_zero_temperature_row_falls_back_to_argmax(): + logits = torch.tensor([[1.0, 5.0, 2.0], [1.0, 5.0, 2.0]]) + out = sample_scores( + logits, + temperature=torch.tensor([0.0, 1.0]), + top_k=0, + top_p=1.0, + generator=torch.Generator().manual_seed(0), + ) + assert int(out[0]) == 1 + + +def test_ras_substitutes_only_repeated_semantic_ids(): + """RAS re-rolls a repeat from a flatter distribution; the EOS slot and + non-repeats must pass through untouched, or utterances get truncated.""" + num_semantic = 2 + # Row 0 will draw id 0 (a repeat) -> substituted. + # Row 1 will draw id 2 (the EOS slot) -> never substituted even if recent. + logits = torch.tensor( + [ + [10.0, -20.0, -20.0], + [-20.0, -20.0, 10.0], + ] + ) + recent = torch.tensor([[0, -1], [2, -1]]) + out = ras_sample_semantic( + logits, + recent, + temperature=1e-6, # deterministic: argmax within each draw + top_k=0, + top_p=1.0, + ras_temperature=1e-6, + ras_top_p=1.0, + num_semantic_ids=num_semantic, + ) + # Both draws are argmax, so substitution is observable only through which + # branch was taken; assert the EOS row is preserved verbatim. + assert int(out[1]) == 2 + assert int(out[0]) == 0 + + +def test_ras_without_history_returns_the_plain_draw(): + logits = torch.tensor([[10.0, 0.0, -5.0]]) + out = ras_sample_semantic( + logits, + None, + temperature=1e-6, + top_k=0, + top_p=1.0, + ras_temperature=1.0, + ras_top_p=0.9, + num_semantic_ids=2, + ) + assert int(out[0]) == 0 + + +def test_sampling_is_reproducible_for_a_seeded_generator(): + logits = torch.randn(4, 32, generator=torch.Generator().manual_seed(7)) + first = sample_scores(logits, temperature=0.7, top_k=8, top_p=0.9, generator=torch.Generator().manual_seed(1234)) + second = sample_scores(logits, temperature=0.7, top_k=8, top_p=0.9, generator=torch.Generator().manual_seed(1234)) + assert torch.equal(first, second) + + +def test_batch_routes_each_row_to_its_own_generator(): + """Regression for the shared slot-0 generator bug. + + A shared generator over ``[num_reqs, vocab]`` made every row draw from slot + 0's RNG stream, so a request's audio depended on its batch neighbours and + per-request ``seed`` was ignored for slots > 0. Flat logits make the draw + purely noise-driven so the wrong generator changes the token. + """ + # Flat logits: the sampled id is determined only by the generator's noise, + # so routing the generator to the wrong slot changes the result (coincidence ~1/256). + logits = torch.zeros(2, 256) + kw = dict(temperature=1.0, top_k=0, top_p=1.0, ras_temperature=1.0, ras_top_p=0.9, num_semantic_ids=200) + + # (1) Each batched row equals sampling that row alone with its own generator. + batched = ras_sample_batch( + logits, + None, + generators={0: torch.Generator().manual_seed(11), 1: torch.Generator().manual_seed(22)}, + num_reqs=2, + **kw, + ) + row0 = ras_sample_semantic(logits[0:1], None, generator=torch.Generator().manual_seed(11), **kw) + row1 = ras_sample_semantic(logits[1:2], None, generator=torch.Generator().manual_seed(22), **kw) + assert int(batched[0]) == int(row0) + assert int(batched[1]) == int(row1) + + # (2) Row 1 must not change when only slot 0's seed changes. + other = ras_sample_batch( + logits, + None, + generators={0: torch.Generator().manual_seed(9999), 1: torch.Generator().manual_seed(22)}, + num_reqs=2, + **kw, + ) + assert int(other[1]) == int(batched[1]) + + +def test_single_request_batch_keeps_the_fast_path(): + """``num_reqs == 1`` must match a direct single-row draw (no per-row loop).""" + logits = torch.randn(1, 48, generator=torch.Generator().manual_seed(5)) + kw = dict(temperature=0.8, top_k=16, top_p=0.9, ras_temperature=1.0, ras_top_p=0.9, num_semantic_ids=40) + batched = ras_sample_batch(logits, None, generators={0: torch.Generator().manual_seed(99)}, num_reqs=1, **kw) + direct = ras_sample_semantic(logits, None, generator=torch.Generator().manual_seed(99), **kw) + assert torch.equal(batched, direct) diff --git a/tests/model_executor/models/fish_speech/test_dac_modules.py b/tests/model_executor/models/fish_speech/test_dac_modules.py new file mode 100644 index 000000000..0896f5533 --- /dev/null +++ b/tests/model_executor/models/fish_speech/test_dac_modules.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU tests for the vendored DAC codec modules (dac_modules). + +Guards the checkpoint contract: the module tree built by build_dac_codec() +must keep the exact parameter/buffer names and shapes that codec.pth was +saved with (verified against the upstream fish-speech 0.1.0 package), and +the encode/decode surfaces used by dac_encoder.py / fish_speech_dac_decoder.py +must keep their shapes. +""" + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Spot checks covering every submodule family and both weight-norm styles. +# Derived from the upstream fish-speech 0.1.0 state dict (541 entries total); +# legacy weight_norm (weight_g/weight_v) comes from the descript quantizer +# projections, parametrizations.weight.original0/1 from the causal convs. +EXPECTED_KEY_SAMPLES = [ + # encoder: first conv + deepest encoder-block transformer + "encoder.block.0.conv.parametrizations.weight.original0", + "encoder.block.4.block.5.layers.3.attention.wqkv.weight", + "encoder.block.4.block.5.freqs_cis", + "encoder.block.1.block.0.block.0.alpha", + # quantizer: semantic + residual VQ (legacy weight_norm keys) + "quantizer.semantic_quantizer.quantizers.0.codebook.weight", + "quantizer.quantizer.quantizers.8.in_proj.weight_g", + "quantizer.quantizer.quantizers.8.in_proj.weight_v", + # quantizer pre/post transformers + down/upsample convnext + "quantizer.pre_module.layers.7.feed_forward.w1.weight", + "quantizer.post_module.norm.weight", + "quantizer.downsample.1.1.dwconv.conv.weight", + "quantizer.upsample.0.0.conv.weight", + # decoder: transposed conv + final conv + "decoder.model.1.block.1.conv.parametrizations.weight.original1", + "decoder.model.6.conv.parametrizations.weight.original0", +] + +NUM_STATE_DICT_ENTRIES = 541 +DAC_NUM_CODEBOOKS = 10 # 1 semantic + 9 residual + + +@pytest.fixture(scope="module") +def codec(): + from vllm_omni.model_executor.models.fish_speech.dac_utils import build_dac_codec + + model = build_dac_codec() + model.eval() + return model + + +def test_state_dict_contract(codec): + sd = codec.state_dict() + assert len(sd) == NUM_STATE_DICT_ENTRIES + for key in EXPECTED_KEY_SAMPLES: + assert key in sd, f"checkpoint-contract key missing: {key}" + + +def test_no_fish_speech_import_required(codec): + import sys + + assert "fish_speech" not in sys.modules, "vendored codec must not import the fish-speech package" + + +@torch.no_grad() +def test_encode_shapes(codec): + frame_length = codec.frame_length # 2048 samples per code frame + wav = torch.randn(1, 1, frame_length * 3 + 100) + lengths = torch.tensor([wav.shape[-1]], dtype=torch.long) + codes, code_lengths = codec.encode(wav, lengths) + assert codes.shape[0] == 1 + assert codes.shape[1] == DAC_NUM_CODEBOOKS + assert codes.shape[2] == 4 # ceil(3.x frames) + assert code_lengths.tolist() == [4] + assert codes.dtype in (torch.int64, torch.int32) + + +@torch.no_grad() +def test_decode_shapes(codec): + frames = 8 + codes = torch.randint(0, 1024, (1, DAC_NUM_CODEBOOKS, frames), dtype=torch.long) + feature_lengths = torch.tensor([frames], dtype=torch.long) + wav, audio_lengths = codec.decode(codes, feature_lengths) + assert wav.shape[0] == 1 and wav.shape[1] == 1 + assert audio_lengths.tolist() == [frames * codec.frame_length] + assert wav.shape[-1] >= audio_lengths[0] - codec.frame_length + + +@torch.no_grad() +def test_decoder_runner_mutations(codec): + """The stage-1 runner prunes encode-only parts and bakes weight norm; + decode must still work afterwards (mirrors fish_speech_dac_decoder).""" + from torch.nn.utils.parametrize import remove_parametrizations + + codec.encoder = None + codec.quantizer.pre_module = None + codec.quantizer.downsample = None + for module in codec.modules(): + parametrizations = getattr(module, "parametrizations", None) + if not parametrizations: + continue + for name in list(parametrizations.keys()): + remove_parametrizations(module, name, leave_parametrized=True) + + codes = torch.randint(0, 1024, (2, DAC_NUM_CODEBOOKS, 5), dtype=torch.long) + wav, audio_lengths = codec.decode(codes, torch.tensor([5, 3], dtype=torch.long)) + assert wav.shape[0] == 2 + assert audio_lengths.tolist() == [5 * codec.frame_length, 3 * codec.frame_length] diff --git a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py index 3fba046dc..3d43078ce 100755 --- a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py @@ -11,13 +11,15 @@ import asyncio import json import os import time -from collections import OrderedDict, defaultdict +from collections import defaultdict from types import SimpleNamespace import pytest import torch from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v3 import HiggsAudioV3Adapter pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -1683,31 +1685,21 @@ class TestVoiceCloneReferenceCache: 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 - - serving = object.__new__(OmniOpenAIServingSpeech) - serving._higgs_audio_v3_ref_code_cache = OrderedDict() - serving._higgs_audio_v3_ref_code_cache_bytes = 0 - serving._higgs_audio_v3_ref_code_inflight = {} + adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace())) codes = torch.arange(16, dtype=torch.long).reshape(2, 8) - serving._put_higgs_audio_v3_ref_codes("ref-a", codes) + adapter._put_higgs_audio_v3_ref_codes("ref-a", codes) - cached = serving._get_higgs_audio_v3_ref_codes("ref-a") + cached = adapter._get_higgs_audio_v3_ref_codes("ref-a") assert cached is not None cached.fill_(0) - cached_again = serving._get_higgs_audio_v3_ref_codes("ref-a") + cached_again = adapter._get_higgs_audio_v3_ref_codes("ref-a") assert cached_again is not None assert torch.equal(cached_again, codes) def test_higgs_v3_ref_code_inflight_deduplicates_concurrent_encode(self): - from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech - - serving = object.__new__(OmniOpenAIServingSpeech) - serving._higgs_audio_v3_ref_code_cache = OrderedDict() - serving._higgs_audio_v3_ref_code_cache_bytes = 0 - serving._higgs_audio_v3_ref_code_inflight = {} + adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace())) calls = 0 def encode_reference_audio(_wav, _sr): @@ -1722,7 +1714,7 @@ class TestVoiceCloneReferenceCache: async def run(): return await asyncio.gather( *[ - serving._resolve_higgs_audio_v3_ref_codes( + adapter._resolve_higgs_audio_v3_ref_codes( "ref-a", object(), 24000, @@ -1741,7 +1733,7 @@ class TestVoiceCloneReferenceCache: assert torch.equal(codes, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1) cached, cache_hit, inflight_wait = asyncio.run( - serving._resolve_higgs_audio_v3_ref_codes( + adapter._resolve_higgs_audio_v3_ref_codes( "ref-a", object(), 24000, diff --git a/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py b/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py index 194a024db..c20ee7f48 100644 --- a/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py +++ b/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py @@ -91,7 +91,11 @@ def _minimal_model(mocker: MockerFixture): decoder=decoder, config=SimpleNamespace(avg_pooler=2, stride_size=2, hop_length=240), ) - model._tokenizer_service = SimpleNamespace(audio_tokenizer=audio_tok, cuda_graph_wrapper=None) + model._tokenizer_service = SimpleNamespace( + audio_tokenizer=audio_tok, + cuda_graph_wrapper=None, + device="cpu", + ) return model, audio_tok diff --git a/tests/model_executor/models/mimo_audio/test_mimo_audio_tokenizer_worker_cache.py b/tests/model_executor/models/mimo_audio/test_mimo_audio_tokenizer_worker_cache.py new file mode 100644 index 000000000..06b16f853 --- /dev/null +++ b/tests/model_executor/models/mimo_audio/test_mimo_audio_tokenizer_worker_cache.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for the MiMo-Audio tokenizer-worker cache key. + +The worker cache key used to keep only the device *type*, so an explicit +``MIMO_AUDIO_TOKENIZER_DEVICE=cuda:1`` silently loaded the tokenizer on the +process's current CUDA device and shared one cache entry with ``cuda:0``. +The index must survive both the cache key and the worker constructor. + +CPU-only and weight-free: the worker class is replaced by a stub so no +tokenizer weights are loaded. +""" + +from __future__ import annotations + +import functools + +import pytest + +torch = pytest.importorskip("torch") + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_TOK_PATH = "/models/mimo-audio-tokenizer" + + +@functools.lru_cache(maxsize=1) +def _code2wav_module(): + """Defer the import (pulls vLLM model_executor) until first use.""" + from vllm_omni.model_executor.models.mimo_audio import mimo_audio_code2wav + + return mimo_audio_code2wav + + +class _StubWorker: + def __init__(self, device_str: str, config_path: str | None, audio_tokenizer_path: str): + self.device = device_str + self.config_path = config_path + self.audio_tokenizer_path = audio_tokenizer_path + + +@pytest.fixture +def isolated_worker_cache(monkeypatch): + """Empty worker cache with the real worker replaced by a stub.""" + mod = _code2wav_module() + monkeypatch.setattr(mod, "MiMoAudioTokenizerWorker", _StubWorker) + monkeypatch.setattr(mod, "_TOKENIZER_WORKER_CACHE", {}) + return mod + + +@pytest.mark.parametrize("device", [torch.device("cuda:1"), "cuda:1"]) +def test_cache_key_keeps_cuda_index(device): + """Both torch.device and string spellings retain the index.""" + mod = _code2wav_module() + key = mod._normalize_tokenizer_worker_cache_key(device, None, _TOK_PATH) + assert key[0] == "cuda:1" + + +def test_cache_key_cpu_unchanged(): + mod = _code2wav_module() + key = mod._normalize_tokenizer_worker_cache_key(torch.device("cpu"), None, _TOK_PATH) + assert key[0] == "cpu" + + +def test_cache_key_bare_cuda_resolves_to_current_device(monkeypatch): + """A bare "cuda" shares the entry with its equivalent indexed spelling.""" + mod = _code2wav_module() + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 3) + + bare = mod._normalize_tokenizer_worker_cache_key(torch.device("cuda"), None, _TOK_PATH) + indexed = mod._normalize_tokenizer_worker_cache_key(torch.device("cuda:3"), None, _TOK_PATH) + assert bare[0] == "cuda:3" + assert bare == indexed + + +def test_worker_constructed_on_requested_index(isolated_worker_cache): + mod = isolated_worker_cache + worker = mod.get_tokenizer_worker( + device=torch.device("cuda:1"), + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + assert worker.device == "cuda:1" + + +def test_distinct_cuda_indices_get_distinct_workers(isolated_worker_cache): + mod = isolated_worker_cache + first = mod.get_tokenizer_worker( + device=torch.device("cuda:0"), + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + second = mod.get_tokenizer_worker( + device=torch.device("cuda:1"), + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + assert first is not second + assert (first.device, second.device) == ("cuda:0", "cuda:1") + + +def test_same_device_reuses_cached_worker(isolated_worker_cache): + mod = isolated_worker_cache + first = mod.get_tokenizer_worker( + device=torch.device("cuda:1"), + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + second = mod.get_tokenizer_worker( + device="cuda:1", + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + assert first is second 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 21c8dfd14..3ec6662cc 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 @@ -1,12 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace import pytest import torch -from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import ( + MingTTSAdapter, + parse_ming_instruction_fields, +) from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import CFMGraphExecutor, CFMSampler @@ -174,19 +178,19 @@ def test_ming_dense_validation_rejects_semantic_audio_vae_config(): def test_ming_instruction_parser_preserves_dense_and_flash_defaults(): """Ming dense and Ming flash keep distinct instruction defaults.""" - serving = object.__new__(OmniOpenAIServingSpeech) - serving.uploaded_speakers = {"uploaded": {}} + serving = SimpleNamespace(uploaded_speakers={"uploaded": {}}) + adapter = MingTTSAdapter(SpeechServingContext(server=serving)) - dense_plain = serving._parse_ming_instruction(SimpleNamespace(instructions="calm", language=None, voice=None)) + dense_plain = adapter._parse_ming_instruction(SimpleNamespace(instructions="calm", language=None, voice=None)) assert dense_plain == "calm" - dense_with_fields = serving._parse_ming_instruction( + dense_with_fields = adapter._parse_ming_instruction( SimpleNamespace(instructions="calm", language="Auto", voice="灵小甄") ) assert dense_with_fields == {"IP": "灵小甄", "风格": "calm"} - flash_fields = serving._parse_ming_instruction_fields( - SimpleNamespace(instructions="calm", language="粤语", voice="灵小甄") + flash_fields = parse_ming_instruction_fields( + serving, SimpleNamespace(instructions="calm", language="粤语", voice="灵小甄") ) assert flash_fields == {"风格": "calm"} diff --git a/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py b/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py new file mode 100644 index 000000000..d25b58e48 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for the MiniCPM-o 4.5 ``audio_embeds`` input path. + +Two defects made the OpenAI ``audio_embeds`` content part unusable with this model: + +1. ``MiniCPMOAudioEmbeddingItems`` registered itself under ``modality="image"`` + (copy-pasted from ``MiniCPMVImageEmbeddingItems``), so every request failed with + ``Modality 'image' not found. Available modalities: {'audio'}``. +2. ``_get_prompt_updates`` sized the placeholder from + ``sum(map(len, single_audio_embeds))``. ``MiniCPMOAudioEmbeddingInputs`` declares + ``TensorShape("bn", "s", "h")``, so one item is ``(s, h)`` and that expression + iterates rows and sums the hidden size -- ``s * h`` instead of ``s``. A one-second + slice of shape ``(10, 4096)`` asked for 40960 placeholder tokens instead of 10, + pushing every real request past ``max_model_len``. +""" + +from __future__ import annotations + +import math +from types import SimpleNamespace + +import pytest +import torch +from vllm.multimodal.parse import MultiModalDataItems + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMMultiModalProcessor, + MiniCPMO45OmniLLMProcessingInfo, + MiniCPMOAudioEmbeddingItems, + _minicpmo_field_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +POOL_STEP = 5 +HIDDEN_SIZE = 4096 +#: Audio embedding vectors produced per second: 100 mel -> 50 cnn -> (50-5)//5+1. +EMBEDS_PER_SECOND = 10 + + +class _FakeProcessor: + """Minimal processor surface used by ``get_audio_placeholder``.""" + + def __init__(self) -> None: + self.pool_step = POOL_STEP + self.audio_processor = SimpleNamespace(hop_length=160) + self.image_processor = SimpleNamespace(mean=[0.0], std=[1.0]) + + def get_audio_placeholder( + self, + audio_lens: int, + chunk_input: bool = True, + chunk_length: int = 1, + ) -> str: + del chunk_input, chunk_length + feature_lens = math.ceil(audio_lens / self.audio_processor.hop_length) + cnn_feature_lens = (feature_lens - 1) // 2 + 1 + output_lens = (cnn_feature_lens - self.pool_step) // self.pool_step + 1 + return "" * output_lens + + +class _IdentityTokenizer: + """``_get_prompt_updates`` round-trips the image/video patterns through the + tokenizer to detect encode/decode drift; an identity pair keeps that a no-op.""" + + def encode(self, text: str, add_special_tokens: bool = True) -> str: + del add_special_tokens + return text + + def decode(self, tokens: str) -> str: + return tokens + + +def _processing_info() -> MiniCPMO45OmniLLMProcessingInfo: + processor = _FakeProcessor() + tokenizer = _IdentityTokenizer() + info = object.__new__(MiniCPMO45OmniLLMProcessingInfo) + info.ctx = SimpleNamespace( + tokenizer=tokenizer, + get_tokenizer=lambda: tokenizer, + get_hf_config=lambda: SimpleNamespace(audio_pool_step=POOL_STEP), + get_hf_processor=lambda **kwargs: processor, + ) + return info + + +def _audio_placeholder_for(single_audio_embeds: torch.Tensor) -> str: + """Run the real prompt-update path and return the audio placeholder text. + + Goes through ``_get_prompt_updates`` rather than recomputing the length, so the + ``sum(map(len, ...))`` regression is actually exercised. + """ + info = _processing_info() + proc = object.__new__(MiniCPMO45OmniLLMMultiModalProcessor) + proc.info = info + + items = MiniCPMOAudioEmbeddingItems( + {"audio_embeds": single_audio_embeds.unsqueeze(0)}, + fields_factory=_minicpmo_field_config, + ) + mm_items = MultiModalDataItems({"audio": items}) + + updates = proc._get_prompt_updates(mm_items, {}, {}) + audio_update = next(u for u in updates if u.modality == "audio") + return audio_update.content(0).full + + +def test_audio_embedding_items_use_the_audio_modality() -> None: + """Registering under ``image`` makes every ``audio_embeds`` request fail.""" + items = MiniCPMOAudioEmbeddingItems( + {"audio_embeds": torch.zeros(1, EMBEDS_PER_SECOND, HIDDEN_SIZE)}, + fields_factory=_minicpmo_field_config, + ) + + assert items.modality == "audio" + + +@pytest.mark.parametrize("num_seconds", [1, 3, 30]) +def test_placeholder_length_matches_embedding_count(num_seconds: int) -> None: + """An ``(s, h)`` item must yield exactly ``s`` placeholder tokens.""" + num_embeds = num_seconds * EMBEDS_PER_SECOND + + placeholder = _audio_placeholder_for(torch.zeros(num_embeds, HIDDEN_SIZE)) + + assert placeholder.count("") == num_embeds + + +def test_placeholder_does_not_scale_with_hidden_size() -> None: + """The count must depend on ``s`` alone, never on ``h``. + + ``sum(map(len, ...))`` returned ``s * h``, so the placeholder grew with the model's + hidden size. Two items with the same ``s`` but different ``h`` must agree. + """ + narrow = _audio_placeholder_for(torch.zeros(EMBEDS_PER_SECOND, 8)) + wide = _audio_placeholder_for(torch.zeros(EMBEDS_PER_SECOND, HIDDEN_SIZE)) + + assert narrow.count("") == wide.count("") == EMBEDS_PER_SECOND diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py index 5aa78262a..8d5aeade6 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for MiniCPM-o 4.5 code2wav model-dir resolution (#5442). In hub/CI deployments ``model_config.model`` is a repo id rather than a local @@ -27,7 +27,7 @@ def _no_hub(monkeypatch): def _fail(*args, **kwargs): # pragma: no cover - must not be reached raise AssertionError("snapshot_download must not be called here") - monkeypatch.setattr(huggingface_hub, "snapshot_download", _fail) + monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", _fail) def test_local_directory_is_returned_unchanged(tmp_path, monkeypatch): @@ -38,13 +38,13 @@ def test_local_directory_is_returned_unchanged(tmp_path, monkeypatch): def test_repo_id_resolves_via_snapshot_download(tmp_path, monkeypatch): calls = {} - def _fake_snapshot_download(model_ref, revision=None, allow_patterns=None): + def _fake_snapshot_download(self, model_ref, revision=None, allow_patterns=None): calls["model_ref"] = model_ref calls["revision"] = revision calls["allow_patterns"] = allow_patterns return str(tmp_path / "snapshot") - monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot_download) + monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", _fake_snapshot_download) resolved = _resolve_model_dir("openbmb/MiniCPM-o-4_5", revision="abc123") assert resolved == str(tmp_path / "snapshot") assert calls["model_ref"] == "openbmb/MiniCPM-o-4_5" @@ -56,7 +56,7 @@ def test_snapshot_download_failure_propagates(monkeypatch): def _raise(*args, **kwargs): raise FileNotFoundError("offline and not cached") - monkeypatch.setattr(huggingface_hub, "snapshot_download", _raise) + monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", _raise) with pytest.raises(FileNotFoundError): _resolve_model_dir("openbmb/MiniCPM-o-4_5") diff --git a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py b/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py index ba4015b98..3e8ee2a63 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py +++ b/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py @@ -10,6 +10,10 @@ Covers: - ``process_audios``: >30s audio unpads per chunk without ``TypeError`` - multiple >30s audios unpad in chunk order - <=30s audio unpadding byte-identical to the pre-fix behavior + - ``MiniCPMOAudioFeatureInputs``: a batch mixing a >30s and a <=30s audio + (different slice counts) validates instead of raising, whether from a + single request with multiple audios of different lengths or from + multiple requests sharing an encoder batch """ from __future__ import annotations @@ -26,6 +30,7 @@ from vllm.multimodal.inputs import ( from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( MiniCPMO45OmniLLMMultiModalProcessor, + MiniCPMOAudioFeatureInputs, MiniCPMOMultiModalDataParser, _minicpmo_field_config, ) @@ -152,3 +157,43 @@ class TestProcessAudiosUnpadding: assert len(features) == 1 assert features[0].shape == (_FEAT_DIM, 980) assert torch.equal(features[0], fake_outputs["audio_features"][0][:, :980]) + + +class TestMixedChunkCountBatchValidation: + def test_batch_mixes_different_slice_counts(self) -> None: + # A >30s audio (3 slices) and a <=30s audio (2 slices) scheduled in + # the same encoder batch: audio_feature_lens is a list of per-audio + # tensors with different lengths (3 vs 2). This is the shape produced + # both by a single request with two audios of different lengths and + # by two separate requests batched together. Before the fix this + # raised ValueError: "audio_feature_lens contains inconsistent shapes". + audio_feature_lens = [ + torch.tensor([_CHUNK_FRAMES, _CHUNK_FRAMES, 1]), + torch.tensor([_CHUNK_FRAMES, 1]), + ] + audio_features = [torch.zeros(_FEAT_DIM, _CHUNK_FRAMES) for _ in range(5)] + + inputs = MiniCPMOAudioFeatureInputs( + type="audio_features", + audio_features=audio_features, + audio_feature_lens=audio_feature_lens, + ) + + assert inputs.audio_feature_lens is audio_feature_lens + + def test_single_slice_count_batch_still_validated(self) -> None: + # Same slice count per audio (the common case) must still validate, + # unaffected by marking "s" dynamic. + audio_feature_lens = [ + torch.tensor([_CHUNK_FRAMES, 1]), + torch.tensor([_CHUNK_FRAMES, 1]), + ] + audio_features = [torch.zeros(_FEAT_DIM, _CHUNK_FRAMES) for _ in range(4)] + + inputs = MiniCPMOAudioFeatureInputs( + type="audio_features", + audio_features=audio_features, + audio_feature_lens=audio_feature_lens, + ) + + assert inputs.audio_feature_lens is audio_feature_lens diff --git a/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py b/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py index b55f8279c..ef05bf738 100644 --- a/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py +++ b/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py @@ -6,19 +6,27 @@ Tests for OmniVoice CUDA Graph generator wrapper numerical equivalence. Verifies that _OmniVoiceCUDAGraphForward produces results equivalent to eager mode across three scenarios: - Exact-size inputs (no padding) → bit-identical - - Padded inputs (padded to nearest bucket) → correct slicing, exact match - for position-independent models - - Oversized inputs (fallback to eager) → bit-identical + - Padded inputs (padded to nearest bucket) → correct slicing and exact match + - Oversized inputs (128-aligned lazy capture) → bit-identical -Uses a lightweight synthetic generator to keep warmup fast and avoid -loading actual model weights. +Uses a small randomly initialized OmniVoiceGenerator so every test exercises +the production embedding, transformer, varlen-attention, and logits paths +without loading checkpoint weights. """ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch -import torch.nn as nn +from vllm.utils.math_utils import round_up + +from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( + OmniVoiceGenerator, + _OmniVoiceCUDAGraphForward, +) +from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig pytestmark = [ pytest.mark.core_model, @@ -29,53 +37,7 @@ pytestmark = [ DEVICE = torch.device("cuda:0") NUM_CB = 8 VOCAB = 1025 -HIDDEN = 64 CAPTURE_SIZES = [32, 64, 128] -HEAD_DIM = 64 - - -# --------------------------------------------------------------------------- -# Synthetic generator: matches _OmniVoiceCUDAGraphForward's interface -# --------------------------------------------------------------------------- - - -class _SyntheticGenerator: - """Position-independent Embedding+Linear replacing OmniVoiceGenerator's 28-layer transformer. - - Each token position is processed independently (no attention), so padded - runs produce bit-identical results to non-padded runs on the same positions. - This makes the padding/slicing logic easy to verify exactly. - """ - - class _Cfg: - num_audio_codebook = NUM_CB - - def __init__(self, device: torch.device): - self.config = self._Cfg() - self.text_embedding = nn.Embedding(1000, HIDDEN).to(device).eval() - self._linear = nn.Linear(HIDDEN, NUM_CB * VOCAB, bias=False).to(device).eval() - self._rope_table: torch.Tensor | None = None - - @property - def model_dtype(self) -> torch.dtype: - """Part of the interface _OmniVoiceCUDAGraphForward expects of a generator.""" - return self.text_embedding.weight.dtype - - def _rope_table_for(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: - """Part of the interface _OmniVoiceCUDAGraphForward expects of a generator.""" - return torch.zeros(seq_len, HEAD_DIM, device=device, dtype=dtype) - - def _step_forward( - self, - input_ids: torch.Tensor, - audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, - rope_table: torch.Tensor, - ) -> torch.Tensor: - two_b, _, S = input_ids.shape - x = self.text_embedding(input_ids[:, 0, :].clamp(0, 999)) # [two_b, S, H] - logits = self._linear(x) # [two_b, S, 8*1025] - return logits.view(two_b, S, NUM_CB, VOCAB).permute(0, 2, 1, 3) # [two_b, 8, S, 1025] # --------------------------------------------------------------------------- @@ -83,18 +45,37 @@ class _SyntheticGenerator: # --------------------------------------------------------------------------- -def _eager(gen: _SyntheticGenerator, ids: torch.Tensor, mask: torch.Tensor, attn: torch.Tensor) -> torch.Tensor: - """Run _step_forward with a zero RoPE table (synthetic gen ignores it).""" - two_b, _, S = ids.shape - return gen._step_forward(ids, mask, attn, gen._rope_table_for(S, ids.device, torch.float32)) +def _eager(gen: OmniVoiceGenerator, ids: torch.Tensor, mask: torch.Tensor, cu_seqs) -> torch.Tensor: + """Run the production varlen step outside CUDA Graph.""" + seq_len = ids.shape[0] + rope_table = gen._rope_table_for(seq_len, ids.device, gen.model_dtype) + return gen._step_forward(ids, mask, cu_seqs, rope_table) + + +def _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, batch_size) -> torch.Tensor: + """Run eager with the exact padded shape and metadata used by Graph.""" + seq_len = ids.shape[0] + bucket = wrapper._find_bucket(batch_size, seq_len) + if bucket is None: + bucket = round_up(seq_len, wrapper._LAZY_CAPTURE_ALIGNMENT) + padded_ids, padded_mask = wrapper._pad_inputs(ids, mask, bucket) + padded_cu_seqs = cu_seqs.clone() + padded_cu_seqs[-1] = bucket + return _eager(gen, padded_ids, padded_mask, padded_cu_seqs)[:, :seq_len, :] def _make_inputs(seq_len: int, device: torch.device = DEVICE): - two_b = 2 - ids = torch.randint(0, 100, (two_b, NUM_CB, seq_len), dtype=torch.long, device=device) - mask = torch.ones(two_b, seq_len, dtype=torch.bool, device=device) - attn = torch.ones(two_b, 1, seq_len, seq_len, dtype=torch.bool, device=device) - return ids, mask, attn + ids = torch.randint(0, 100, (seq_len, NUM_CB), dtype=torch.long, device=device) + mask = torch.ones(seq_len, dtype=torch.bool, device=device) + cond_len = (seq_len + 1) // 2 + uncond_len = seq_len - cond_len + + cu_seqs = torch.tensor( + [0, cond_len, cond_len + uncond_len, cond_len + uncond_len], + dtype=torch.int32, + device=device, + ) + return ids, mask, cu_seqs # --------------------------------------------------------------------------- @@ -105,15 +86,26 @@ def _make_inputs(seq_len: int, device: torch.device = DEVICE): @pytest.fixture(scope="module") def gen(): torch.manual_seed(42) - return _SyntheticGenerator(DEVICE) + config = OmniVoiceConfig( + audio_vocab_size=VOCAB, + num_audio_codebook=NUM_CB, + llm_config={ + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "intermediate_size": 128, + "vocab_size": 128, + "max_position_embeddings": 4096, + "head_dim": 16, + }, + enable_cuda_graph=False, + ) + return OmniVoiceGenerator(config, SimpleNamespace(max_num_seqs=2)).to(DEVICE).eval() @pytest.fixture(scope="module") def wrapper(gen): - from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( - _OmniVoiceCUDAGraphForward, - ) - w = _OmniVoiceCUDAGraphForward(gen, capture_sizes=CAPTURE_SIZES) w.warmup(DEVICE) return w @@ -127,10 +119,10 @@ def wrapper(gen): @pytest.mark.parametrize("seq_len", CAPTURE_SIZES) def test_exact_size_bit_identical(gen, wrapper, seq_len): """When input exactly matches a captured bucket, output must be bit-identical to eager.""" - ids, mask, attn = _make_inputs(seq_len) + ids, mask, cu_seqs = _make_inputs(seq_len) with torch.no_grad(): - eager_out = _eager(gen, ids, mask, attn) - graph_out = wrapper(ids, mask, attn) + eager_out = _eager(gen, ids, mask, cu_seqs) + graph_out = wrapper(ids, mask, cu_seqs, 1) torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) @@ -140,40 +132,36 @@ def test_exact_size_bit_identical(gen, wrapper, seq_len): @pytest.mark.parametrize("seq_len", [1, 15, 33, 60, 100]) -def test_padded_output_shape(gen, wrapper, seq_len): +def test_padded_output_shape(wrapper, seq_len): """Graph output must be sliced back to actual seq_len, not the bucket size.""" - ids, mask, attn = _make_inputs(seq_len) + ids, mask, cu_seqs = _make_inputs(seq_len) with torch.no_grad(): - graph_out = wrapper(ids, mask, attn) - assert graph_out.shape == (2, NUM_CB, seq_len, VOCAB) + graph_out = wrapper(ids, mask, cu_seqs, 1) + assert graph_out.shape == (NUM_CB, seq_len, VOCAB) @pytest.mark.parametrize("seq_len", [15, 33, 60, 100]) def test_padded_output_matches_eager(gen, wrapper, seq_len): - """Padded graph output must equal eager output at actual positions. - - The synthetic model has no attention across positions, so zero-padding - does not affect non-padded positions — exact match is expected. - """ - ids, mask, attn = _make_inputs(seq_len) + """Padded graph output must equal eager output at actual positions.""" + ids, mask, cu_seqs = _make_inputs(seq_len) with torch.no_grad(): - eager_out = _eager(gen, ids, mask, attn) - graph_out = wrapper(ids, mask, attn) + eager_out = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, 1) + graph_out = wrapper(ids, mask, cu_seqs, 1) torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) # --------------------------------------------------------------------------- -# 3. Oversized inputs → fallback to eager (lazy capture), bit-identical +# 3. Oversized inputs → aligned lazy capture, bit-identical # --------------------------------------------------------------------------- @pytest.mark.parametrize("seq_len", [129, 200, 256]) -def test_fallback_eager_bit_identical(gen, wrapper, seq_len): - """Sequences exceeding the largest bucket fall back to lazy capture → bit-identical.""" - ids, mask, attn = _make_inputs(seq_len) +def test_aligned_lazy_capture_bit_identical(gen, wrapper, seq_len): + """Sequences beyond the static plan use aligned lazy capture and remain bit-identical.""" + ids, mask, cu_seqs = _make_inputs(seq_len) with torch.no_grad(): - eager_out = _eager(gen, ids, mask, attn) - graph_out = wrapper(ids, mask, attn) + eager_out = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, 1) + graph_out = wrapper(ids, mask, cu_seqs, 1) torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) @@ -184,13 +172,73 @@ def test_fallback_eager_bit_identical(gen, wrapper, seq_len): def test_deterministic_across_calls(wrapper): """Same input must produce identical output on repeated CUDA graph replays.""" - ids, mask, attn = _make_inputs(32) + ids, mask, cu_seqs = _make_inputs(32) with torch.no_grad(): - out1 = wrapper(ids, mask, attn) - out2 = wrapper(ids, mask, attn) + out1 = wrapper(ids, mask, cu_seqs, 1).clone() + out2 = wrapper(ids, mask, cu_seqs, 1).clone() torch.testing.assert_close(out1, out2, atol=0, rtol=0) +def test_replay_uses_updated_cu_seqs(gen, wrapper): + """The same graph key must honor new sequence boundaries on replay.""" + ids, mask, _ = _make_inputs(48) + cu_seqs_a = torch.tensor([0, 32, 48, 48], dtype=torch.int32, device=DEVICE) + cu_seqs_b = torch.tensor([0, 17, 48, 48], dtype=torch.int32, device=DEVICE) + + with torch.no_grad(): + eager_a = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs_a, 1) + graph_a = wrapper(ids, mask, cu_seqs_a, 1).clone() + eager_b = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs_b, 1) + graph_b = wrapper(ids, mask, cu_seqs_b, 1).clone() + + torch.testing.assert_close(graph_a, eager_a, atol=0, rtol=0) + torch.testing.assert_close(graph_b, eager_b, atol=0, rtol=0) + assert not torch.equal(graph_a, graph_b) + + +def test_batch_two_cu_seqs_matches_eager(gen, wrapper): + """B=2 uses four real sequence segments plus the fixed tail slot.""" + ids, mask, _ = _make_inputs(48) + cu_seqs = torch.tensor([0, 10, 18, 32, 48, 48], dtype=torch.int32, device=DEVICE) + + assert cu_seqs.numel() == 2 * 2 + 2 + with torch.no_grad(): + eager_out = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, 2) + graph_out = wrapper(ids, mask, cu_seqs, 2).clone() + + torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) + + +def test_lazy_graph_cache_uses_lru_eviction(wrapper): + """A lazy-cache hit must protect that graph from the next eviction.""" + original_limit = wrapper._MAX_LAZY_GRAPHS + wrapper._lazy_graphs.clear() + wrapper._MAX_LAZY_GRAPHS = 2 + try: + # Static B=1 coverage ends at 128. These lengths round to three + # distinct 128-aligned lazy keys: 256, 384, and 512. + for seq_len in (129, 257): + ids, mask, cu_seqs = _make_inputs(seq_len) + with torch.no_grad(): + wrapper(ids, mask, cu_seqs, 1) + assert list(wrapper._lazy_graphs) == [(1, 256), (1, 384)] + + # Refresh key 256, making key 384 the least recently used entry. + ids, mask, cu_seqs = _make_inputs(129) + with torch.no_grad(): + wrapper(ids, mask, cu_seqs, 1) + assert list(wrapper._lazy_graphs) == [(1, 384), (1, 256)] + + # Inserting key 512 must evict key 384, not the recently hit key 256. + ids, mask, cu_seqs = _make_inputs(385) + with torch.no_grad(): + wrapper(ids, mask, cu_seqs, 1) + assert list(wrapper._lazy_graphs) == [(1, 256), (1, 512)] + finally: + wrapper._MAX_LAZY_GRAPHS = original_limit + wrapper._lazy_graphs.clear() + + # --------------------------------------------------------------------------- # 5. _find_bucket logic (CPU, no CUDA graph) # --------------------------------------------------------------------------- @@ -203,58 +251,17 @@ def test_find_bucket_returns_nearest_bucket(): ) w = _OmniVoiceCUDAGraphForward.__new__(_OmniVoiceCUDAGraphForward) - w._capture_sizes = [32, 64, 128] + w.capture_bucket_sizes_by_batch = { + 1: [32, 64, 128], + 2: [64, 128, 256], + } w._graphs = {} - assert w._find_bucket(1) == 32 - assert w._find_bucket(32) == 32 - assert w._find_bucket(33) == 64 - assert w._find_bucket(64) == 64 - assert w._find_bucket(128) == 128 - assert w._find_bucket(129) is None - assert w._find_bucket(1000) is None - - -# --------------------------------------------------------------------------- -# 6. enable_cuda_graph=False produces same tokens as enable_cuda_graph=True -# --------------------------------------------------------------------------- - - -def test_cuda_graph_disabled_matches_eager_generator(): - """OmniVoiceGenerator with enable_cuda_graph=False must produce the same - _step_forward output as one with enable_cuda_graph=True (after warmup). - - Uses a 2-layer config so warmup completes quickly without loading weights. - """ - from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator - from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig - - cfg = OmniVoiceConfig() - cfg.llm_num_hidden_layers = 2 - cfg.num_hidden_layers = 2 - - torch.manual_seed(0) - gen_eager = OmniVoiceGenerator(cfg).to(DEVICE).eval() - gen_eager.config.enable_cuda_graph = False - gen_eager._cuda_graph_fwd = None - - torch.manual_seed(0) - gen_graph = OmniVoiceGenerator(cfg).to(DEVICE).eval() - gen_graph.config.cuda_graph_capture_sizes = [64] - from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import _OmniVoiceCUDAGraphForward - - gen_graph._cuda_graph_fwd = _OmniVoiceCUDAGraphForward(gen_graph, capture_sizes=[64]) - gen_graph._cuda_graph_fwd.warmup(DEVICE) - - seq_len = 64 - ids = torch.zeros(2, cfg.num_audio_codebook, seq_len, dtype=torch.long, device=DEVICE) - mask = torch.ones(2, seq_len, dtype=torch.bool, device=DEVICE) - attn = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=DEVICE) - - rope_table = gen_eager._rope_table_for(seq_len, DEVICE, gen_eager.model_dtype) - - with torch.no_grad(): - eager_logits = gen_eager._step_forward(ids, mask, attn, rope_table) - graph_logits = gen_graph._cuda_graph_fwd(ids, mask, attn) - - torch.testing.assert_close(graph_logits, eager_logits, atol=0, rtol=0) + assert w._find_bucket(1, 1) == 32 + assert w._find_bucket(1, 32) == 32 + assert w._find_bucket(1, 33) == 64 + assert w._find_bucket(1, 128) == 128 + assert w._find_bucket(1, 129) is None + assert w._find_bucket(2, 33) == 64 + assert w._find_bucket(2, 129) == 256 + assert w._find_bucket(3, 64) is None diff --git a/tests/model_executor/models/omnivoice/test_fused_projection_load.py b/tests/model_executor/models/omnivoice/test_fused_projection_load.py index fc5db01c7..c80f0e65b 100644 --- a/tests/model_executor/models/omnivoice/test_fused_projection_load.py +++ b/tests/model_executor/models/omnivoice/test_fused_projection_load.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """The OmniVoice generator packs q/k/v and gate/up into fused projections. The HF checkpoint stores those five tensors separately, so packing them is a @@ -14,6 +14,8 @@ that would leave one partly written raises instead. from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -49,6 +51,13 @@ def _config() -> OmniVoiceConfig: ) +def _generator() -> OmniVoiceGenerator: + return OmniVoiceGenerator( + _config(), + SimpleNamespace(max_num_seqs=1), + ) + + def _checkpoint_shards() -> dict[str, torch.Tensor]: """The five per-layer tensors an OmniVoice checkpoint actually stores.""" torch.manual_seed(0) @@ -64,7 +73,7 @@ def _checkpoint_shards() -> dict[str, torch.Tensor]: def test_every_fused_parameter_is_written() -> None: - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() loaded = generator._load_fused_projections(shards) @@ -77,7 +86,7 @@ def test_every_fused_parameter_is_written() -> None: def test_packed_weight_matches_the_shards_it_came_from() -> None: """Packing order is q,k,v and gate,up -- the split in forward() assumes it.""" - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() generator._load_fused_projections(shards) @@ -101,7 +110,7 @@ def test_packed_weight_matches_the_shards_it_came_from() -> None: def test_no_fused_parameter_keeps_its_random_init() -> None: """The corruption this guards against: a fused param never written at all.""" - generator = OmniVoiceGenerator(_config()) + generator = _generator() before = { name: param.detach().clone() for name, param in generator.named_parameters() @@ -121,7 +130,7 @@ def test_no_fused_parameter_keeps_its_random_init() -> None: ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj", "mlp.gate_proj", "mlp.up_proj"], ) def test_a_missing_shard_raises_instead_of_loading_partially(dropped: str) -> None: - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() del shards[f"llm.layers.0.{dropped}.weight"] @@ -130,7 +139,7 @@ def test_a_missing_shard_raises_instead_of_loading_partially(dropped: str) -> No def test_a_wrong_shaped_shard_raises() -> None: - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() shards["llm.layers.0.self_attn.q_proj.weight"] = torch.randn(NUM_HEADS * HEAD_DIM + 1, HIDDEN) @@ -140,7 +149,7 @@ def test_a_wrong_shaped_shard_raises() -> None: def test_a_checkpoint_missing_a_whole_layer_raises() -> None: """Half-loaded is the dangerous state: it neither errors nor works.""" - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = {k: v for k, v in _checkpoint_shards().items() if not k.startswith("llm.layers.1.")} with pytest.raises(ValueError, match="fused projections"): @@ -149,7 +158,7 @@ def test_a_checkpoint_missing_a_whole_layer_raises() -> None: def test_a_checkpoint_with_no_fused_shards_is_left_alone() -> None: """An unrelated state_dict must not trip the completeness check.""" - generator = OmniVoiceGenerator(_config()) + generator = _generator() assert generator._load_fused_projections({"audio_heads.weight": torch.randn(4, 4)}) == set() @@ -162,7 +171,7 @@ def test_load_weights_wires_the_packing_in(tmp_path, monkeypatch: pytest.MonkeyP """ import safetensors.torch - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() before = generator.layers[0].self_attn.qkv_proj.weight.detach().clone() diff --git a/tests/model_executor/models/omnivoice/test_mask_dtype.py b/tests/model_executor/models/omnivoice/test_mask_dtype.py index ff899e2a2..ff8dda49d 100644 --- a/tests/model_executor/models/omnivoice/test_mask_dtype.py +++ b/tests/model_executor/models/omnivoice/test_mask_dtype.py @@ -1,183 +1,79 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""OmniVoice attention masks must follow the model dtype, not a hardcoded float32. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OmniVoice SDPA fallback masks must not depend on the model dtype. -SDPA requires the additive attention mask to carry the same dtype as the query. -The generator used to materialize that mask as float32 unconditionally, in three -places (the shared per-forward mask, the CUDA-graph capture buffer, and the -replay-time normalization), so every half-precision deployment died on the first -attention call with:: +The previous padded attention path used an additive floating-point mask. SDPA +requires such a mask to have the same dtype as the query, and constructing it +as float32 caused half-precision OmniVoice inference to fail on CUDA. - RuntimeError: invalid dtype for bias - should match query's dtype - -(the exact wording depends on which SDPA backend is selected for the shape; the -math backend words the same rejection as ``attn_mask.dtype``.) - -That left OmniVoice servable only in float32, even though the upstream k2-fsa -implementation runs it in float16. - -Note the split between the CPU and CUDA tests below: the CPU SDPA backend -silently accepts the mismatched mask, so only the CUDA tests can pin the actual -crash. The CPU tests cover the mask contract itself, which is what the fix -changes and what a future regression would break first. +The packed-varlen path no longer needs an additive mask. When SDPA fallback is +required, ``_attention_metadata_from_cu_seqs`` builds a boolean block mask. +Boolean SDPA masks have no dtype coupling with the query, so no mask coercion +is required for float16 or bfloat16 inference. """ from __future__ import annotations import pytest import torch +import torch.nn.functional as F from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( - OmniVoiceAttention, - OmniVoiceGenerator, - _additive_float_mask, + _attention_metadata_from_cu_seqs, ) -from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig HALF_DTYPES = [torch.float16, torch.bfloat16] -ALL_DTYPES = HALF_DTYPES + [torch.float32] cpu_test = pytest.mark.core_model cuda_test = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def _tiny_config() -> OmniVoiceConfig: - """A real OmniVoiceConfig sized so the module fits in a unit test.""" - return OmniVoiceConfig( - llm_config={ - "hidden_size": 32, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 8, - "intermediate_size": 64, - "vocab_size": 64, - "max_position_embeddings": 128, - }, - enable_cuda_graph=False, +@cpu_test +@pytest.mark.cpu +def test_sdpa_fallback_mask_is_dtype_independent() -> None: + """The fallback mask is boolean rather than an additive model-dtype mask.""" + cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32) + + metadata = _attention_metadata_from_cu_seqs( + cu_seqs, + 6, + needs_sdpa_mask=True, ) - -def _inputs(dtype: torch.dtype, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - hidden = torch.randn(2, 6, 32, device=device, dtype=dtype) - bool_mask = torch.ones(2, 1, 6, 6, dtype=torch.bool, device=device) - bool_mask[:, :, :, 4:] = False - rope_table = torch.zeros(6, 8, device=device, dtype=dtype) - rope_table[:, :4] = 1.0 # cos = 1, sin = 0: no rotation, so the test is about the mask - return hidden, bool_mask, rope_table - - -# -------------------------------------------------------------------------- -# The mask contract (CPU) -# -------------------------------------------------------------------------- - - -@cpu_test -@pytest.mark.cpu -@pytest.mark.parametrize("dtype", ALL_DTYPES) -def test_additive_float_mask_uses_requested_dtype(dtype: torch.dtype) -> None: - out = _additive_float_mask(torch.tensor([[True, False]]), dtype) - assert out.dtype == dtype - - -@cpu_test -@pytest.mark.cpu -def test_additive_float_mask_keeps_mask_semantics() -> None: - """Attend positions stay at 0.0 and masked ones at -inf, in half precision too.""" - out = _additive_float_mask(torch.tensor([[True, False]]), torch.float16) - assert out[0, 0].item() == 0.0 - assert out[0, 1].item() == float("-inf") - - -@cpu_test -@pytest.mark.cpu -def test_additive_float_mask_requires_an_explicit_dtype() -> None: - """The float32 default is the bug; keeping the parameter required is the fix.""" - with pytest.raises(TypeError): - _additive_float_mask(torch.tensor([[True]])) # type: ignore[call-arg] - - -@cpu_test -@pytest.mark.cpu -@pytest.mark.parametrize("dtype", ALL_DTYPES) -def test_generator_model_dtype_tracks_its_weights(dtype: torch.dtype) -> None: - """The single source of truth the three former float32 literals now defer to.""" - assert OmniVoiceGenerator(_tiny_config()).to(dtype).model_dtype == dtype - - -@cpu_test -@pytest.mark.cpu -@pytest.mark.parametrize("dtype", HALF_DTYPES) -def test_attention_accepts_a_model_dtype_mask(dtype: torch.dtype) -> None: - attn = OmniVoiceAttention(_tiny_config()).to(dtype).eval() - hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cpu")) - - with torch.inference_mode(): - out = attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, dtype)) - - assert out.shape == hidden.shape - assert out.dtype == dtype - assert torch.isfinite(out).all() - - -# -------------------------------------------------------------------------- -# The crash itself (CUDA only — the CPU backend does not reject the mismatch) -# -------------------------------------------------------------------------- + assert metadata.attn_mask is not None + assert metadata.attn_mask.dtype == torch.bool @cuda_test @pytest.mark.parametrize("dtype", HALF_DTYPES) -def test_float32_mask_is_what_broke_half_precision(dtype: torch.dtype) -> None: - """Pin the original failure so a float32 default cannot come back unnoticed.""" - attn = OmniVoiceAttention(_tiny_config()).to(device="cuda:0", dtype=dtype).eval() - hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cuda:0")) - - with pytest.raises(RuntimeError, match=r"(invalid dtype for bias|attn_mask)"), torch.inference_mode(): - attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, torch.float32)) - - -@cuda_test -@pytest.mark.parametrize("dtype", HALF_DTYPES) -def test_real_path_forward_runs_in_half_precision(dtype: torch.dtype) -> None: - """The regression, on the production path with the Triton norm kernels live.""" - attn = OmniVoiceAttention(_tiny_config()).to(device="cuda:0", dtype=dtype).eval() - hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cuda:0")) - - with torch.inference_mode(): - out = attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, dtype)) - - assert out.dtype == dtype - assert torch.isfinite(out).all() - - -@cuda_test -@pytest.mark.parametrize("dtype", HALF_DTYPES) -def test_generator_forward_runs_in_half_precision(dtype: torch.dtype) -> None: - """End-to-end over the real iterative loop, which is where the float32 mask was built. - - This is the test that fails on an unfixed tree: ``forward`` materialized the - shared SDPA mask as float32 regardless of the weights' dtype. - """ +def test_sdpa_fallback_mask_needs_no_half_precision_coercion( + dtype: torch.dtype, +) -> None: + """Half-precision SDPA accepts the production bool mask without casting.""" device = torch.device("cuda:0") - config = _tiny_config() - generator = OmniVoiceGenerator(config).to(device=device, dtype=dtype).eval() + seq_len = 6 + cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32, device=device) - seq_len, target_len = 12, 4 - input_ids = torch.zeros(2, config.num_audio_codebook, seq_len, dtype=torch.long, device=device) - input_ids[:, 1:, :] = config.audio_mask_id - audio_mask = torch.zeros(2, seq_len, dtype=torch.bool, device=device) - audio_mask[:, seq_len - target_len :] = True - attention_mask = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=device) + metadata = _attention_metadata_from_cu_seqs( + cu_seqs, + seq_len, + needs_sdpa_mask=True, + ) + + assert metadata.attn_mask is not None + assert metadata.attn_mask.dtype == torch.bool + + query = torch.randn(1, 2, seq_len, 8, device=device, dtype=dtype) + key = torch.randn_like(query) + value = torch.randn_like(query) with torch.inference_mode(): - tokens = generator( - input_ids, - audio_mask, - attention_mask, - target_lens=[target_len], - seed=0, - num_step=2, + output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=metadata.attn_mask, ) - assert tokens.shape == (1, config.num_audio_codebook, target_len) - assert tokens.dtype == torch.long + assert output.dtype == dtype + assert torch.isfinite(output).all() diff --git a/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py b/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py index fbf5ac0b0..95524fdce 100644 --- a/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py +++ b/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py @@ -1,12 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import pytest -from huggingface_hub import hf_hub_download from tokenizers import Tokenizer as HFTokenizer from vllm_omni.diffusion.models.omnivoice.pipeline_omnivoice import ( _tokenize_with_nonverbal_tags, ) +from vllm_omni.transformers_utils.repo_utils import hf_api pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -18,7 +21,7 @@ def _ids(tok: HFTokenizer, s: str) -> list[int]: class TestNonVerbalTags: @classmethod def setup_class(cls): - tokenizer_path = hf_hub_download(repo_id="k2-fsa/OmniVoice", filename="tokenizer.json") + tokenizer_path = hf_api().hf_hub_download(repo_id="k2-fsa/OmniVoice", filename="tokenizer.json") cls.tokenizer = HFTokenizer.from_file(tokenizer_path) def test_plain_text_fallback(self): diff --git a/tests/model_executor/models/omnivoice/test_pipeline_batching.py b/tests/model_executor/models/omnivoice/test_pipeline_batching.py new file mode 100644 index 000000000..6e611f2ff --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_pipeline_batching.py @@ -0,0 +1,228 @@ +# 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_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.models.omnivoice.pipeline_omnivoice import ( + OmniVoicePipeline, + _PreparedOmniVoiceRequest, +) +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.models.omnivoice.omnivoice_generator import ( + _attention_metadata_from_cu_seqs, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_sdpa_fallback_mask_preserves_packed_sequence_boundaries() -> None: + cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32) + + metadata = _attention_metadata_from_cu_seqs(cu_seqs, 6, needs_sdpa_mask=True) + + assert metadata.attn_mask is not None + expected = torch.tensor( + [ + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [0, 0, 1, 0, 0, 0], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + ], + dtype=torch.bool, + ) + torch.testing.assert_close(metadata.attn_mask[0, 0], expected) + + +def test_eager_attention_metadata_honors_exact_max_seqlen() -> None: + cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32) + + metadata = _attention_metadata_from_cu_seqs( + cu_seqs, + 6, + needs_sdpa_mask=False, + max_seqlen=3, + ) + + assert metadata.extra["max_seqlen_q"] == 3 + assert metadata.extra["max_seqlen_k"] == 3 + + +def test_request_batch_prepare_error_does_not_skip_valid_request() -> None: + """A malformed request must not prevent another request from completing.""" + prepared = _PreparedOmniVoiceRequest( + input_ids=torch.zeros(5, 8, dtype=torch.long), + audio_mask=torch.ones(5, dtype=torch.bool), + cond_len=3, + target_len=2, + seed=42, + ) + + def prepare_request(prompt, extra): + del extra + if prompt == "bad": + return DiffusionOutput(error="invalid prompt") + return prepared + + def collate_requests(requests): + assert requests == [prepared] + return prepared.input_ids, prepared.audio_mask, [prepared.cond_len] + + generator_calls = [] + + def generate_tokens(**kwargs): + generator_calls.append(kwargs) + assert kwargs["target_lens"] == [prepared.target_len] + return torch.ones(1, 8, prepared.target_len, dtype=torch.long) + + pipeline = SimpleNamespace( + _prepare_request_input=prepare_request, + _collate_request_inputs=collate_requests, + generator=generate_tokens, + decoder=lambda tokens: tokens.float(), + num_step=32, + guidance_scale=7.0, + t_shift=1.0, + layer_penalty_factor=0.0, + position_temperature=0.0, + class_temperature=0.0, + ) + sampling = OmniDiffusionSamplingParams(num_inference_steps=32, guidance_scale=7.0) + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest(prompt="bad", sampling_params=sampling, request_id="bad"), + OmniDiffusionRequest( + prompt="good", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=32, guidance_scale=7.0), + request_id="good", + ), + ] + ) + + outputs = OmniVoicePipeline.forward(pipeline, batch) + + assert len(outputs) == 2 + assert outputs[0].error == "invalid prompt" + assert outputs[1].error is None + torch.testing.assert_close(outputs[1].output, torch.ones(1, 8, 2)) + assert generator_calls[0]["guidance_scale"] == 7.0 + + +def test_request_batch_honors_explicit_sampling_overrides() -> None: + """Request steps and guidance must override the OmniVoice defaults.""" + prepared = _PreparedOmniVoiceRequest( + input_ids=torch.zeros(5, 8, dtype=torch.long), + audio_mask=torch.ones(5, dtype=torch.bool), + cond_len=3, + target_len=2, + seed=42, + ) + captured_sampling = [] + + def generate_tokens(**kwargs): + captured_sampling.append((kwargs["num_step"], kwargs["guidance_scale"])) + return torch.ones(1, 8, prepared.target_len, dtype=torch.long) + + pipeline = SimpleNamespace( + _prepare_request_input=lambda prompt, extra: prepared, + _collate_request_inputs=lambda requests: (prepared.input_ids, prepared.audio_mask, [prepared.cond_len]), + generator=generate_tokens, + decoder=lambda tokens: tokens.float(), + num_step=32, + guidance_scale=2.0, + t_shift=1.0, + layer_penalty_factor=0.0, + position_temperature=0.0, + class_temperature=0.0, + ) + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt="hello", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=12, guidance_scale=6.5), + request_id="request", + ) + ] + ) + + outputs = OmniVoicePipeline.forward(pipeline, batch) + + assert outputs[0].error is None + assert captured_sampling == [(12, 6.5)] + + +def test_request_batch_prepare_error_preserves_surrounding_output_indices() -> None: + """A middle prepare error must not shift valid outputs into the wrong slots.""" + prepared_before = _PreparedOmniVoiceRequest( + input_ids=torch.zeros(3, 8, dtype=torch.long), + audio_mask=torch.ones(3, dtype=torch.bool), + cond_len=2, + target_len=1, + seed=1, + ) + prepared_after = _PreparedOmniVoiceRequest( + input_ids=torch.zeros(5, 8, dtype=torch.long), + audio_mask=torch.ones(5, dtype=torch.bool), + cond_len=3, + target_len=2, + seed=2, + ) + + def prepare_request(prompt, extra): + del extra + if prompt == "before": + return prepared_before + if prompt == "after": + return prepared_after + return DiffusionOutput(error="invalid middle request") + + def collate_requests(requests): + assert requests == [prepared_before, prepared_after] + return ( + torch.cat([prepared_before.input_ids, prepared_after.input_ids]), + torch.cat([prepared_before.audio_mask, prepared_after.audio_mask]), + [prepared_before.cond_len, prepared_after.cond_len], + ) + + def generate_tokens(**kwargs): + assert kwargs["target_lens"] == [1, 2] + before = torch.full((1, 8, 1), 11, dtype=torch.long) + after = torch.full((1, 8, 2), 22, dtype=torch.long) + return torch.cat([before, after], dim=-1) + + pipeline = SimpleNamespace( + _prepare_request_input=prepare_request, + _collate_request_inputs=collate_requests, + generator=generate_tokens, + decoder=lambda tokens: tokens.float(), + num_step=32, + guidance_scale=2.0, + t_shift=1.0, + layer_penalty_factor=0.0, + position_temperature=0.0, + class_temperature=0.0, + ) + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt=prompt, + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=32), + request_id=prompt, + ) + for prompt in ("before", "bad", "after") + ] + ) + + outputs = OmniVoicePipeline.forward(pipeline, batch) + + assert len(outputs) == 3 + torch.testing.assert_close(outputs[0].output, torch.full((1, 8, 1), 11.0)) + assert outputs[1].error == "invalid middle request" + torch.testing.assert_close(outputs[2].output, torch.full((1, 8, 2), 22.0)) 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 619f2169a..438a5ec1d 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 @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests for code predictor dtype alignment (fix for #2385). @@ -20,6 +20,9 @@ import pytest import torch import torch.nn.functional as F from pytest_mock import MockerFixture +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeTalkerCodePredictorConfig, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -41,6 +44,7 @@ _COMMON = os.path.join(_MODELS, "common") def _load_module(name: str, filename: str): path = os.path.abspath(os.path.join(_BASE, filename)) spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None mod = importlib.util.module_from_spec(spec) sys.modules[name] = mod # register before exec (needed for dataclasses etc.) spec.loader.exec_module(mod) @@ -83,7 +87,7 @@ 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 + setattr(custom_op_mock, "CustomOp", NativeCustomOp) return { "vllm_omni": mocker.MagicMock(), @@ -119,6 +123,7 @@ def _load_target_classes(mocker: MockerFixture): common_spec = importlib.util.spec_from_file_location( "vllm_omni.model_executor.models.common.qwen3_code_predictor", common_cp_path ) + assert common_spec is not None and common_spec.loader is not None common_cp_mod = importlib.util.module_from_spec(common_spec) sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] = common_cp_mod common_spec.loader.exec_module(common_cp_mod) @@ -209,6 +214,80 @@ def test_npu_custom_ops_use_fused_norm_and_cached_rope(mocker: MockerFixture, lo torch.testing.assert_close(sin, rotary.sin_cached[position_ids].to(torch.float16)) +@pytest.mark.parametrize( + ("rope_kwargs", "expected_theta"), + [ + pytest.param( + {"rope_theta": 1_000_000.0}, + 1_000_000.0, + id="serialized_checkpoint", + ), + pytest.param( + { + "rope_theta": 10_000.0, + "rope_parameters": {"rope_type": "default", "rope_theta": 1_000_000.0}, + }, + 1_000_000.0, + id="nested_precedence", + ), + pytest.param({}, 10_000.0, id="default"), + ], +) +def test_code_predictor_rotary_uses_qwen3_omni_rope_parameters( + loaded_target_classes, + rope_kwargs, + expected_theta, +) -> None: + """Follow Transformers 5 deserialization and precedence for Qwen3-Omni.""" + _ = loaded_target_classes + config = Qwen3OmniMoeTalkerCodePredictorConfig( + hidden_size=16, + num_attention_heads=1, + head_dim=16, + num_code_groups=16, + **rope_kwargs, + ) + assert not hasattr(config, "rope_theta") + assert config.rope_parameters["rope_theta"] == expected_theta + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + + rotary = common_mod._RotaryEmbedding(config) + + expected = 1.0 / (expected_theta ** (torch.arange(0, 16, 2, dtype=torch.float32) / 16)) + torch.testing.assert_close(rotary.inv_freq, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("theta", [10_000.0, 1_000_000.0]) +@pytest.mark.parametrize( + "legacy_top_level_only", + [False, True], + ids=["normalized", "legacy_top_level"], +) +def test_code_predictor_rotary_preserves_qwen3_tts_rope_theta( + loaded_target_classes, + theta, + legacy_top_level_only, +) -> None: + """Keep the shared predictor compatible with current and legacy Qwen3-TTS.""" + config_class = loaded_target_classes[0] + config = config_class( + hidden_size=16, + num_attention_heads=1, + head_dim=16, + num_code_groups=16, + rope_theta=theta, + ) + assert config.rope_parameters["rope_theta"] == theta + if legacy_top_level_only: + delattr(config, "rope_parameters") + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + + rotary = common_mod._RotaryEmbedding(config) + + expected = 1.0 / (theta ** (torch.arange(0, 16, 2, dtype=torch.float32) / 16)) + torch.testing.assert_close(rotary.inv_freq, expected, rtol=0, atol=0) + + class TestCodePredictorDtypeAlignment: """Test that code predictor buffers match model parameter dtype.""" diff --git a/tests/model_executor/stage_input_processors/test_audio8_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_audio8_tts_async_chunk.py new file mode 100644 index 000000000..9190a33f6 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_audio8_tts_async_chunk.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.audio8_tts import ( + extract_last_frame, + slow_ar_to_codec_decoder_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +NUM_CODEBOOKS = 3 + + +def _make_transfer_manager(*, chunk_frames=2, left_context=1, initial_chunk_frames=0): + return SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + put_req_chunk=defaultdict(int), + connector=SimpleNamespace( + config={ + "extra": { + "codec_chunk_frames": chunk_frames, + "codec_left_context_frames": left_context, + "initial_codec_chunk_frames": initial_chunk_frames, + } + } + ), + ) + + +def _make_request(req_id="req", *, finished=False): + return SimpleNamespace( + external_req_id=req_id, + additional_information=None, + is_finished=lambda: finished, + ) + + +def _frame(value: int) -> torch.Tensor: + return torch.full((NUM_CODEBOOKS,), value, dtype=torch.long) + + +def _step(transfer_manager, request, frame_value: int | None, *, finished: bool = False): + multimodal_output = {} if frame_value is None else {"audio_codes": _frame(frame_value).unsqueeze(0)} + return slow_ar_to_codec_decoder_async_chunk(transfer_manager, multimodal_output, request, is_finished=finished) + + +def test_codes_accumulate_across_steps_and_chunk_is_withheld_until_full(): + """Codes must accumulate; emitting per-step would starve the codec of the + context it needs and resetting would truncate the utterance.""" + transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1) + request = _make_request() + + assert _step(transfer_manager, request, 1) is None + payload = _step(transfer_manager, request, 2) + + assert payload is not None + assert len(transfer_manager.code_prompt_token_ids["req"]) == 2 + # [num_codebooks, frames], codebook-major. + assert tuple(payload.codes.audio.shape) == (NUM_CODEBOOKS, 2) + assert payload.codes.audio[0].tolist() == [1, 2] + # Nothing precedes the first chunk, so there is no left context yet. + assert payload.meta.left_context_size == 0 + + +def test_second_chunk_carries_left_context_and_reports_its_size(): + transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1) + request = _make_request() + for value in (1, 2, 3): + _step(transfer_manager, request, value) + payload = _step(transfer_manager, request, 4) + + assert payload is not None + # 1 context frame + 2 new frames. + assert tuple(payload.codes.audio.shape) == (NUM_CODEBOOKS, 3) + assert payload.codes.audio[0].tolist() == [2, 3, 4] + assert payload.meta.left_context_size == 1 + + +def test_partial_tail_is_flushed_when_the_request_finishes(): + transfer_manager = _make_transfer_manager(chunk_frames=4, left_context=2) + request = _make_request() + _step(transfer_manager, request, 1) + _step(transfer_manager, request, 2) + payload = _step(transfer_manager, request, 3, finished=True) + + assert payload is not None + assert payload.codes.audio[0].tolist() == [1, 2, 3] + assert bool(payload.meta.finished) is True + + +def test_finished_request_without_any_frame_emits_a_terminal_payload(): + transfer_manager = _make_transfer_manager() + payload = slow_ar_to_codec_decoder_async_chunk( + transfer_manager, None, _make_request(finished=True), is_finished=True + ) + assert payload is not None + assert payload.codes.audio.numel() == 0 + assert bool(payload.meta.finished) is True + + +def test_terminal_call_on_a_boundary_does_not_re_emit_the_last_chunk(): + """A finished call carrying no new frame, arriving when the stream sits + exactly on a chunk boundary, must close the stream rather than re-ship the + last window. Before the guard, the final chunk_frames frames shipped twice.""" + transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1) + request = _make_request() + for value in (1, 2, 3, 4): # length ends exactly on a chunk boundary + _step(transfer_manager, request, value) + # Separate terminal call, no new frame: the boundary chunk is already sent. + payload = _step(transfer_manager, request, None, finished=True) + assert payload is not None + assert payload.codes.audio.numel() == 0 + assert bool(payload.meta.finished) is True + + +def test_final_frame_on_a_boundary_still_emits_when_it_arrives_finished(): + """If the last frame both lands on a boundary and carries finished, its chunk + has not been shipped yet, so it must still be emitted (the guard must key on + 'no new frame this call', not merely on being finished at a boundary).""" + transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1) + request = _make_request() + _step(transfer_manager, request, 1) + payload = _step(transfer_manager, request, 2, finished=True) + assert payload is not None + assert payload.codes.audio[0].tolist() == [1, 2] + assert bool(payload.meta.finished) is True + + +def test_initial_chunk_shortens_time_to_first_audio(): + transfer_manager = _make_transfer_manager(chunk_frames=4, left_context=2, initial_chunk_frames=1) + request = _make_request() + payload = _step(transfer_manager, request, 1) + + assert payload is not None + assert payload.codes.audio[0].tolist() == [1] + assert payload.meta.left_context_size == 0 + + +def test_prefill_placeholder_frames_are_not_forwarded(): + """Prefill emits an all-zero placeholder frame per position; forwarding it + would prepend silence and desynchronise the chunk boundaries.""" + assert extract_last_frame({"audio_codes": torch.zeros((3, NUM_CODEBOOKS), dtype=torch.long)}) is None + assert extract_last_frame({}) is None + assert extract_last_frame({"audio_codes": torch.zeros((0, NUM_CODEBOOKS), dtype=torch.long)}) is None + + +def test_explicit_validity_flag_overrides_the_all_zero_heuristic(): + frame = torch.zeros((1, NUM_CODEBOOKS), dtype=torch.long) + out = extract_last_frame({"audio_codes": frame, "audio_code_valid": torch.tensor([True])}) + assert out is not None and out.tolist() == [0] * NUM_CODEBOOKS + + +def test_invalid_chunk_config_fails_loudly(): + transfer_manager = _make_transfer_manager(chunk_frames=0) + request = _make_request() + with pytest.raises(ValueError, match="codec chunk config"): + _step(transfer_manager, request, 1) diff --git a/tests/model_executor/stage_input_processors/test_joyai_vl_interaction.py b/tests/model_executor/stage_input_processors/test_joyai_vl_interaction.py new file mode 100644 index 000000000..0822cfcd2 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_joyai_vl_interaction.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Tests for JoyAI action-to-TTS conversion.""" + +from types import SimpleNamespace + +import pytest + +from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.model_executor.stage_input_processors import ( + joyai_vl_interaction as joyai_bridge, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_joyai_action_to_tts_bridge(monkeypatch: pytest.MonkeyPatch) -> None: + """Convert each JoyAI action through the real stage-client entry point.""" + target_model_config = SimpleNamespace( + hf_config=SimpleNamespace(talker_config=SimpleNamespace(spk_id={"vivian": 1, "ryan": 2})) + ) + prompt_length_calls: list[tuple[dict[str, list[str]], object]] = [] + prompt_lengths = { + "Colored squares are visible.": 12, + "I am checking this request now.": 18, + } + + def compute_prompt_length(metadata, model_config): + prompt_length_calls.append((metadata, model_config)) + return prompt_lengths[metadata["text"][0]] + + monkeypatch.setattr( + joyai_bridge, + "_compute_talker_prompt_length", + compute_prompt_length, + ) + + client = object.__new__(StageEngineCoreClient) + client.custom_process_input_func = joyai_bridge.joyai_action_to_tts + client.requires_multimodal_data = False + client.vllm_config = SimpleNamespace(model_config=target_model_config) + + cases = [ + ( + " Colored squares are visible.", + None, + { + "tts_task_type": "Base", + "language": ["Chinese"], + "speaker": ["Ryan"], + }, + ("Colored squares are visible.", "Chinese", "Ryan"), + ), + ("", None, {}, None), + ( + "", + ( + " I am checking this request now. " + " Why is the purple elephant dancing beside a volcano?" + ), + { + "language": "Chinese", + "speaker": "Ryan", + "tts_language": "English", + "tts_speaker": "Vivian", + }, + ("I am checking this request now.", "English", "Vivian"), + ), + ] + + for text, cumulative_text, additional_information, expected in cases: + completion = SimpleNamespace(text=text) + if cumulative_text is not None: + completion.cumulative_text = cumulative_text + + call_count = len(prompt_length_calls) + talker_inputs = client.process_engine_inputs( + [SimpleNamespace(outputs=[completion])], + prompt={"additional_information": additional_information}, + ) + + if expected is None: + assert talker_inputs == [] + assert len(prompt_length_calls) == call_count + continue + + spoken_text, language, speaker = expected + [talker_input] = talker_inputs + metadata = talker_input["additional_information"] + assert talker_input["prompt_token_ids"] == [0] * prompt_lengths[spoken_text] + assert metadata == { + "task_type": ["CustomVoice"], + "language": [language], + "speaker": [speaker], + "instruct": [""], + "text": [spoken_text], + } + assert prompt_length_calls[-1] == (metadata, target_model_config) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 509f113b7..183ef8659 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -47,6 +47,7 @@ EXCLUDED_MODELS = [ "LongCatVideoAvatarPipeline", "BagelPipeline", "BooguImagePipeline", + "BooguImageTurboPipeline", "LancePipeline", "MingImagePipeline", "InternVLAA1Pipeline", diff --git a/tests/utils/test_qwen3_force_align_processor.py b/tests/utils/test_qwen3_force_align_processor.py index 4a928d0e5..f846ac7a8 100644 --- a/tests/utils/test_qwen3_force_align_processor.py +++ b/tests/utils/test_qwen3_force_align_processor.py @@ -10,7 +10,21 @@ def test_build_prompt_has_boundary_timestamp_markers(): assert prompt.count("") == 4 assert "helloworld" in prompt - assert prompt.endswith("<|im_start|>assistant\n") + # Official aligner format: audio placeholder + words, no chat template + # (a leading "<|im_start|>user\n" shifts the predicted markers one bin late). + assert prompt == f"{processor.AUDIO_PLACEHOLDER}helloworld" + assert "<|im_start|>" not in prompt + assert processor.build_prompt([]) == f"{processor.AUDIO_PLACEHOLDER}" + + +def test_build_prompt_matches_official_encode_timestamp(): + pytest.importorskip("qwen_asr") + from qwen_asr.inference.qwen3_forced_aligner import Qwen3ForceAlignProcessor + + text = "It's 3 o'clock, 你好 world." + word_list, official_prompt = Qwen3ForceAlignProcessor().encode_timestamp(text, "english") + + assert processor.build_prompt(list(word_list)) == official_prompt @pytest.mark.parametrize( diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 5e3cfe865..be282da81 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -1829,6 +1829,7 @@ class TestPreferModelSamplerNoneFallback: if _declares_prefer_model_sampler(p.read_text(encoding="utf-8", errors="ignore")) } expected = { + "audio8_tts", "cosyvoice3", "glm_tts", "higgs_audio_v2", diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 22ba693d4..c5a96148e 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -17,7 +17,7 @@ import regex as re _HF_NAMES = ( r"HfApi|HfFileSystem|hf_hub_download|snapshot_download" r"|list_repo_files|file_exists|try_to_load_from_cache" - r"|list_repo_refs|repo_exists" + r"|list_repo_refs|repo_exists|get_safetensors_metadata" ) # Non-library trees: examples, tests, and tooling may keep stdlib imports. @@ -169,56 +169,20 @@ CHECK_IMPORTS = { r"^\s*from\s+huggingface_hub\s+import\b[^\n]*\b(?:" + _HF_NAMES + r")\b" ), tip=( - "Use the shared, vLLM-tagged helpers from " - "vllm.transformers_utils.repo_utils (e.g. hf_api(), hf_fs(), " - "list_repo_files, file_exists) instead of calling " - "huggingface_hub directly." + "Use 'hf_api()' from 'vllm_omni.transformers_utils.repo_utils' (or " + "add a tagged helper there) instead, so requests are tagged with " + "vLLM-Omni's library info." ), allowed_dirs=_NON_LIBRARY_DIRS, allowed_files={ - "vllm_omni/benchmarks/data_modules/daily_omni_dataset.py", - "vllm_omni/benchmarks/data_modules/seed_tts_dataset.py", - "vllm_omni/benchmarks/data_modules/seed_tts_eval.py", - "vllm_omni/diffusion/lora/loader.py", - "vllm_omni/diffusion/model_loader/hub_prefetch.py", - "vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py", - "vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py", - "vllm_omni/diffusion/models/helios/pipeline_helios.py", - "vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py", - "vllm_omni/diffusion/models/ltx2/ltx2_components.py", - "vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py", - "vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py", - "vllm_omni/diffusion/models/pi0/pipeline_pi0.py", - "vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py", - "vllm_omni/diffusion/models/soulx_singer/utils.py", - "vllm_omni/diffusion/models/utils.py", - "vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py", - "vllm_omni/engine/arg_utils.py", - "vllm_omni/entrypoints/openai/api_server.py", - "vllm_omni/entrypoints/openai/serving_speech.py", - "vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py", - "vllm_omni/model_executor/model_loader/weight_utils.py", - "vllm_omni/model_executor/models/audex/checkpoint.py", - "vllm_omni/model_executor/models/bagel/bagel.py", - "vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py", - "vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py", - "vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py", - "vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py", - "vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py", - "vllm_omni/model_executor/models/glm_tts/glm_tts.py", + # ``try_to_load_from_cache`` is a local cache lookup that issues no + # Hub request, so there is nothing for ``hf_api()`` to tag. "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py", "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py", "vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py", - "vllm_omni/model_executor/models/indextts2/preprocess_utils.py", - "vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py", - "vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py", - "vllm_omni/model_executor/models/ming_tts/speaker_extractor.py", - "vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py", - "vllm_omni/model_executor/models/omnivoice/omnivoice.py", - "vllm_omni/model_executor/models/personaplex/personaplex_mimi.py", - "vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py", - "vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py", "vllm_omni/transformers_utils/configs/higgs_audio_v3.py", + # Defines the tagged helper itself. + "vllm_omni/transformers_utils/repo_utils.py", }, ), } @@ -330,6 +294,7 @@ def test_regex(): ("from huggingface_hub import HfFileSystem", True), ("from huggingface_hub import list_repo_files", True), ("from huggingface_hub import try_to_load_from_cache", True), + ("from huggingface_hub import get_safetensors_metadata", True), (" from huggingface_hub import snapshot_download", True), ("from huggingface_hub import PyTorchModelHubMixin, hf_hub_download", True), ("from huggingface_hub import (snapshot_download)", True), @@ -346,6 +311,7 @@ def test_regex(): ("from huggingface_hub.constants import HF_HUB_CACHE", False), ("from huggingface_hub.utils import EntryNotFoundError", False), ("from vllm.transformers_utils.repo_utils import hf_api", False), + ("from vllm_omni.transformers_utils.repo_utils import hf_api", False), ("from huggingface_hub import (\n PyTorchModelHubMixin,\n)", False), ("# from huggingface_hub import snapshot_download", False), ] diff --git a/tools/pre_commit/check_tts_adapter.py b/tools/pre_commit/check_tts_adapter.py index d46928662..199c46a30 100644 --- a/tools/pre_commit/check_tts_adapter.py +++ b/tools/pre_commit/check_tts_adapter.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Ratchet on the TTS adapter migration (RFC #4327, #4855). Per-model behaviour is being moved out of ``serving_speech.py`` and into @@ -28,7 +29,7 @@ REPO_ROOT = Path(__file__).resolve().parents[2] SERVING_SPEECH = Path("vllm_omni/entrypoints/openai/serving_speech.py") # Budgets. These may only go DOWN. See the module docstring. -MAX_MODEL_TYPE_BRANCHES = 9 +MAX_MODEL_TYPE_BRANCHES = 6 _GUIDANCE = """ Move the per-model behaviour into its adapter under diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py index abb02b1d5..3bf852e8c 100644 --- a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py +++ b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py @@ -170,7 +170,7 @@ def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: raise ValueError("repo_id is required to download Daily-Omni Videos.tar") try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api except ImportError as e: raise ImportError( "Daily-Omni Hub media download requires huggingface_hub. " @@ -180,7 +180,7 @@ def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: tar_path: Path | None = None for fname in ("Videos.tar", "videos.tar"): try: - tar_path = Path(hf_hub_download(repo_id=rid, filename=fname, repo_type="dataset")) + tar_path = Path(hf_api().hf_hub_download(repo_id=rid, filename=fname, repo_type="dataset")) break except Exception: continue @@ -237,14 +237,14 @@ def ensure_daily_omni_hub_root(repo_id: str) -> Path: raise ValueError("repo_id is required to download Daily-Omni from Hugging Face") try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except ImportError as e: raise ImportError( "Install huggingface_hub to load Daily-Omni from the Hub, or pass a local " "--daily-omni-qa-json / --daily-omni-video-dir mirror." ) from e - cache = snapshot_download( + cache = hf_api().snapshot_download( repo_id=rid, repo_type="dataset", allow_patterns=_DAILY_OMNI_QA_ALLOW_PATTERNS, diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py index 23e8b8ade..996f4f248 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Seed-TTS zero-shot evaluation-style prompts for ``vllm bench serve``. Loads rows from the `meta.lst` format used in `BytedanceSpeech/seed-tts-eval`_ (or any @@ -121,7 +124,7 @@ def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None repo_id = dataset_path.strip() try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except ImportError as e: raise ImportError( "Install huggingface_hub to download Seed-TTS from the Hub, or clone the dataset " @@ -130,7 +133,7 @@ def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None # Download only the requested locale subtree instead of the whole dataset # repo. This avoids large, flaky nightly downloads when we only need e.g. # ``en/meta.lst`` + ``en/prompt-wavs/**``. - cache = snapshot_download( + cache = hf_api().snapshot_download( repo_id=repo_id, repo_type="dataset", allow_patterns=[f"{locale}/**"], diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index 18a72a6c1..100933870 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -332,7 +332,8 @@ def _ensure_utmos_jit_model() -> Any | None: return _utmos_jit_model try: import torch - from huggingface_hub import hf_hub_download + + from vllm_omni.transformers_utils.repo_utils import hf_api repo = os.environ.get("SEED_TTS_UTMOS_HF_REPO", "balacoon/utmos").strip() or "balacoon/utmos" fname = os.environ.get("SEED_TTS_UTMOS_JIT_FILE", "utmos.jit").strip() or "utmos.jit" @@ -341,7 +342,7 @@ def _ensure_utmos_jit_model() -> Any | None: repo, fname, ) - path = hf_hub_download(repo_id=repo, filename=fname, repo_type="model") + path = hf_api().hf_hub_download(repo_id=repo, filename=fname, repo_type="model") # TODO The model weights in UTMOS must be loaded in cuda:0; otherwise, the model execution will fail. want = "cuda:0" diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index 3d7c983f3..52eec86b9 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/vllm_omni/benchmarks/metrics/metrics.py @@ -52,6 +52,21 @@ _MULTIMODAL_BENCHMARK_FIELDS = [ (defs.MEDIAN_IMAGE_GENERATION_MS, float, field(default=0.0)), (defs.STD_IMAGE_GENERATION_MS, float, field(default=0.0)), (defs.PERCENTILES_IMAGE_GENERATION_MS, _PERCENTILE_ROWS_TYPE, field(default=None)), + (defs.TOTAL_VIDEO_DURATION_S, float, field(default=0.0)), + (defs.TOTAL_VIDEO_FRAMES, int, field(default=0)), + (defs.VIDEO_THROUGHPUT, float, field(default=0.0)), + (defs.MEAN_VIDEO_RTF, float, field(default=0.0)), + (defs.MEDIAN_VIDEO_RTF, float, field(default=0.0)), + (defs.STD_VIDEO_RTF, float, field(default=0.0)), + (defs.PERCENTILES_VIDEO_RTF, _PERCENTILE_ROWS_TYPE, field(default=None)), + (defs.MEAN_VIDEO_GENERATION_MS, float, field(default=0.0)), + (defs.MEDIAN_VIDEO_GENERATION_MS, float, field(default=0.0)), + (defs.STD_VIDEO_GENERATION_MS, float, field(default=0.0)), + (defs.PERCENTILES_VIDEO_GENERATION_MS, _PERCENTILE_ROWS_TYPE, field(default=None)), + (defs.MEAN_PEAK_MEMORY_MB, float, field(default=0.0)), + (defs.MEDIAN_PEAK_MEMORY_MB, float, field(default=0.0)), + (defs.STD_PEAK_MEMORY_MB, float, field(default=0.0)), + (defs.PERCENTILES_PEAK_MEMORY_MB, _PERCENTILE_ROWS_TYPE, field(default=None)), ] MultiModalsBenchmarkMetrics = make_dataclass( @@ -171,8 +186,10 @@ def _stage_modality_flags( ) -> tuple[bool, bool, bool, bool, bool]: is_text_stage = final_output_type == "text" or output_unit_type == "text" is_audio_stage = final_output_type == "audio" or output_unit_type == "audio" - is_image_stage = final_output_type in {"image", "images"} or output_unit_type == "image" is_video_stage = final_output_type in {"video", "videos"} or output_unit_type == "video" + # Video diffusion may still report output_unit_type="image" when frames are + # stored in ``images``; prefer video when final_output_type says so. + is_image_stage = (not is_video_stage) and (final_output_type in {"image", "images"} or output_unit_type == "image") is_internal_stream_stage = ( output_unit_type in _STREAMING_OUTPUT_UNIT_TYPES and not is_text_stage and not is_audio_stage ) @@ -222,6 +239,7 @@ def print_metrics( print("{:<40} {:<10.2f}".format("Request goodput (req/s):", metrics.request_goodput)) if isinstance(metrics, MultiModalsBenchmarkMetrics): print("{:<40} {:<10.2f}".format("Peak concurrent requests:", metrics.max_concurrent_requests)) + print_peak_memory_metrics(metrics) if task_type != TaskType.GENERATION or "e2el" in selected_percentile_metrics: process_one_metric("e2el", metrics) print_text_metrics(task_type, selected_percentile_metrics, metrics) @@ -230,9 +248,13 @@ def print_metrics( print_audio_metrics(selected_percentile_metrics, metrics) if _has_image_output(metrics): print_image_metrics(selected_percentiles or [], metrics) + if _has_video_output(metrics): + print_video_metrics(selected_percentiles or [], metrics) if print_stage and outputs and selected_percentiles is not None: - print("\n{s:{c}^{n}}".format(s=" Stage Benchmark Result ", n=50, c="=")) - for sm in _build_stage_metrics_from_outputs(outputs): + stage_metrics = _build_stage_metrics_from_outputs(outputs) + if stage_metrics: + print("\n{s:{c}^{n}}".format(s=" Stage Benchmark Result ", n=50, c="=")) + for sm in stage_metrics: print_stage_metrics( task_type, selected_percentile_metrics, @@ -243,6 +265,11 @@ def print_metrics( def print_text_metrics(task_type, selected_percentile_metrics, metrics: MultiModalsBenchmarkMetrics): + # Pure image/video runs have no user-facing text tokens; skip the whole + # Text Result section (token throughput / peak would be misleading). + if metrics.total_output <= 0 and (_has_image_output(metrics) or _has_video_output(metrics)): + return + print("{s:{c}^{n}}".format(s=" Text Result ", n=50, c="=")) print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) if isinstance(metrics, MultiModalsBenchmarkMetrics): @@ -275,6 +302,14 @@ def _has_image_output(metrics: MultiModalsBenchmarkMetrics) -> bool: return int(getattr(metrics, defs.TOTAL_IMAGES, 0) or 0) > 0 +def _has_video_output(metrics: MultiModalsBenchmarkMetrics) -> bool: + return ( + float(getattr(metrics, defs.TOTAL_VIDEO_DURATION_S, 0.0) or 0.0) > 0.0 + or int(getattr(metrics, defs.TOTAL_VIDEO_FRAMES, 0) or 0) > 0 + or float(getattr(metrics, defs.MEAN_VIDEO_GENERATION_MS, 0.0) or 0.0) > 0.0 + ) + + def print_audio_metrics(selected_percentile_metrics, metrics: MultiModalsBenchmarkMetrics): print("{s:{c}^{n}}".format(s=" Audio Result ", n=50, c="=")) print( @@ -293,6 +328,16 @@ def print_audio_metrics(selected_percentile_metrics, metrics: MultiModalsBenchma process_one_metric(metric, metrics) +def print_peak_memory_metrics(metrics: MultiModalsBenchmarkMetrics): + if getattr(metrics, defs.MEAN_PEAK_MEMORY_MB) <= 0: + return + print("{s:{c}^{n}}".format(s="Peak Memory", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean PEAK_MEMORY_MB (MB):", getattr(metrics, defs.MEAN_PEAK_MEMORY_MB))) + print("{:<40} {:<10.2f}".format("Median PEAK_MEMORY_MB (MB):", getattr(metrics, defs.MEDIAN_PEAK_MEMORY_MB))) + for p, value in getattr(metrics, defs.PERCENTILES_PEAK_MEMORY_MB) or []: + print("{:<40} {:<10.2f}".format(f"P{_p_label(p)} PEAK_MEMORY_MB (MB):", value)) + + def print_image_metrics(selected_percentiles: list[float], metrics: MultiModalsBenchmarkMetrics): print("{s:{c}^{n}}".format(s=" Image Result ", n=50, c="=")) print("{:<40} {:<10}".format("Total images generated:", getattr(metrics, defs.TOTAL_IMAGES))) @@ -311,7 +356,7 @@ def print_image_metrics(selected_percentiles: list[float], metrics: MultiModalsB ) ) if getattr(metrics, defs.MEAN_IMAGE_GENERATION_MS) > 0: - print("-----------------Image Generation-----------------") + print("{s:{c}^{n}}".format(s=" Image Generation ", n=50, c="-")) print("{:<40} {:<10.2f}".format("Mean IMAGE_GENERATION (ms):", getattr(metrics, defs.MEAN_IMAGE_GENERATION_MS))) print( "{:<40} {:<10.2f}".format( @@ -324,6 +369,42 @@ def print_image_metrics(selected_percentiles: list[float], metrics: MultiModalsB print("{:<40} {:<10.2f}".format(f"P{_p_label(p)} IMAGE_GENERATION (ms):", value)) +def print_video_metrics(selected_percentiles: list[float], metrics: MultiModalsBenchmarkMetrics): + print("{s:{c}^{n}}".format(s=" Video Result ", n=50, c="=")) + print( + "{:<40} {:<10.2f}".format( + "Total video duration generated(s):", + getattr(metrics, defs.TOTAL_VIDEO_DURATION_S), + ) + ) + print("{:<40} {:<10}".format("Total video frames generated:", getattr(metrics, defs.TOTAL_VIDEO_FRAMES))) + print( + "{:<40} {:<10.2f}".format( + "Video throughput(video duration/s):", + getattr(metrics, defs.VIDEO_THROUGHPUT), + ) + ) + if getattr(metrics, defs.MEAN_VIDEO_RTF) > 0: + print("{s:{c}^{n}}".format(s=" Video RTF ", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean VIDEO_RTF:", getattr(metrics, defs.MEAN_VIDEO_RTF))) + print("{:<40} {:<10.2f}".format("Median VIDEO_RTF:", getattr(metrics, defs.MEDIAN_VIDEO_RTF))) + for p, value in getattr(metrics, defs.PERCENTILES_VIDEO_RTF) or []: + if p in selected_percentiles: + print("{:<40} {:<10.2f}".format(f"P{_p_label(p)} VIDEO_RTF:", value)) + if getattr(metrics, defs.MEAN_VIDEO_GENERATION_MS) > 0: + print("{s:{c}^{n}}".format(s=" Video Generation ", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean VIDEO_GENERATION (ms):", getattr(metrics, defs.MEAN_VIDEO_GENERATION_MS))) + print( + "{:<40} {:<10.2f}".format( + "Median VIDEO_GENERATION (ms):", + getattr(metrics, defs.MEDIAN_VIDEO_GENERATION_MS), + ) + ) + for p, value in getattr(metrics, defs.PERCENTILES_VIDEO_GENERATION_MS) or []: + if p in selected_percentiles: + print("{:<40} {:<10.2f}".format(f"P{_p_label(p)} VIDEO_GENERATION (ms):", value)) + + def process_one_metric( metric_attribute_name: str, metrics: MultiModalsBenchmarkMetrics, @@ -495,6 +576,21 @@ def _print_image_stage_metrics( ) +def _print_video_stage_metrics( + selected_percentiles: list[float], + sm: StageBenchmarkMetrics, +) -> None: + print("{s:{c}^{n}}".format(s=" Video Result ", n=50, c="=")) + _print_percentile_metric( + "Video Generation", + "VIDEO_GENERATION", + getattr(sm, defs.STAGE_GEN_TIMES_MS), + selected_percentiles, + values_are_ms=True, + to_ms=True, + ) + + def _print_internal_stream_stage_metrics( selected_percentile_metrics: list[str], selected_percentiles: list[float], @@ -549,11 +645,12 @@ def print_stage_metrics( ) = _stage_modality_flags(getattr(sm, "final_output_type"), getattr(sm, "output_unit_type")) print("{s:{c}^{n}}".format(s=title, n=50, c="=")) - if is_image_stage: - _print_image_stage_metrics(selected_percentiles, sm) + if is_video_stage: + _print_video_stage_metrics(selected_percentiles, sm) return - if is_video_stage: + if is_image_stage: + _print_image_stage_metrics(selected_percentiles, sm) return _print_stage_timing(sm, selected_percentiles) @@ -754,6 +851,11 @@ def calculate_metrics( denoise_step_latencies_ms: list[float] = [] total_images = 0 total_image_pixels = 0 + video_durations: list[float] = [] + video_rtfs: list[float] = [] + total_video_frames = 0 + video_generation_times_ms: list[float] = [] + peak_memories_mb: list[float] = [] audio_underruns: list[float] = [] audio_continuity_ok: list[bool] = [] input_audio_duration = 0.0 @@ -762,7 +864,18 @@ def calculate_metrics( output_len = outputs[i].output_tokens if not output_len: - if tokenizer is None: + # Pure image/video (no generated text) must stay at 0 tokens. + # Do not invent output_len=1 via the tokenizer-is-None placeholder. + generated_text = getattr(outputs[i], "generated_text", None) or "" + has_image = int(getattr(outputs[i], defs.IMAGE_COUNT, 0) or 0) > 0 + has_video = ( + float(getattr(outputs[i], defs.VIDEO_DURATION, 0.0) or 0.0) > 0.0 + or int(getattr(outputs[i], defs.VIDEO_FRAMES, 0) or 0) > 0 + or float(getattr(outputs[i], defs.VIDEO_GENERATION_TIME_MS, 0.0) or 0.0) > 0.0 + ) + if not generated_text and (has_image or has_video): + output_len = 0 + elif tokenizer is None: output_len = 1 else: # We use the tokenizer to count the number of output tokens @@ -827,6 +940,19 @@ def calculate_metrics( denoise_step_latency_ms = float(getattr(outputs[i], defs.DENOISE_STEP_LATENCY_MS, 0.0) or 0.0) if denoise_step_latency_ms > 0: denoise_step_latencies_ms.append(denoise_step_latency_ms) + video_duration = float(getattr(outputs[i], defs.VIDEO_DURATION, 0.0) or 0.0) + if video_duration > 0: + video_durations.append(video_duration) + total_video_frames += int(getattr(outputs[i], defs.VIDEO_FRAMES, 0) or 0) + video_rtf = float(getattr(outputs[i], defs.VIDEO_RTF, 0.0) or 0.0) + if video_rtf > 0: + video_rtfs.append(video_rtf) + video_generation_time_ms = float(getattr(outputs[i], defs.VIDEO_GENERATION_TIME_MS, 0.0) or 0.0) + if video_generation_time_ms > 0: + video_generation_times_ms.append(video_generation_time_ms) + peak_memory_mb = float(getattr(outputs[i], defs.PEAK_MEMORY_MB, 0.0) or 0.0) + if peak_memory_mb > 0: + peak_memories_mb.append(peak_memory_mb) audio_underruns.append(getattr(outputs[i], f"{defs.AUDIO_UNDERRUN}_s", 0.0)) audio_continuity_ok.append(bool(getattr(outputs[i], defs.AUDIO_CONTINUITY_OK, True))) e2els.append(outputs[i].latency) @@ -892,18 +1018,24 @@ def calculate_metrics( token_timeline_available = False else: # Calculate token generation timestamp using - # start_time, ttft, and itl - token_times = [output.start_time + output.ttft] - current_time = token_times[0] - for itl_value in output.itl: - current_time += itl_value - token_times.append(current_time) + # start_time, ttft, and itl. Skip requests with no text tokens + # (empty itl, zero output_tokens, and no generated_text) so pure + # image/video runs do not seed a fake peak of 1 tok/s from + # start_time+ttft alone. + output_token_count = int(getattr(output, "output_tokens", 0) or 0) + has_generated_text = bool(getattr(output, "generated_text", None) or "") + if output_token_count > 0 or output.itl or has_generated_text: + token_times = [output.start_time + output.ttft] + current_time = token_times[0] + for itl_value in output.itl: + current_time += itl_value + token_times.append(current_time) - # Add tokens to second buckets - for token_time in token_times: - second_bucket = int(token_time - min_start_time) - if 0 <= second_bucket < duration_seconds: - tokens_per_second[second_bucket] += 1 + # Add tokens to second buckets + for token_time in token_times: + second_bucket = int(token_time - min_start_time) + if 0 <= second_bucket < duration_seconds: + tokens_per_second[second_bucket] += 1 # Track concurrent requests for each second this request was active request_start_second = int(output.start_time - min_start_time) @@ -1009,6 +1141,25 @@ def calculate_metrics( defs.PERCENTILES_IMAGE_GENERATION_MS: [ (p, np.percentile(image_generation_times_ms or 0, p)) for p in selected_percentiles ], + defs.TOTAL_VIDEO_DURATION_S: sum(video_durations), + defs.TOTAL_VIDEO_FRAMES: total_video_frames, + defs.VIDEO_THROUGHPUT: sum(video_durations) / dur_s, + defs.MEAN_VIDEO_RTF: np.mean(video_rtfs or 0), + defs.STD_VIDEO_RTF: np.std(video_rtfs or 0), + defs.MEDIAN_VIDEO_RTF: np.median(video_rtfs or 0), + defs.PERCENTILES_VIDEO_RTF: [(p, np.percentile(video_rtfs or 0, p)) for p in selected_percentiles], + defs.MEAN_VIDEO_GENERATION_MS: np.mean(video_generation_times_ms or 0), + defs.STD_VIDEO_GENERATION_MS: np.std(video_generation_times_ms or 0), + defs.MEDIAN_VIDEO_GENERATION_MS: np.median(video_generation_times_ms or 0), + defs.PERCENTILES_VIDEO_GENERATION_MS: [ + (p, np.percentile(video_generation_times_ms or 0, p)) for p in selected_percentiles + ], + defs.MEAN_PEAK_MEMORY_MB: np.mean(peak_memories_mb or 0), + defs.STD_PEAK_MEMORY_MB: np.std(peak_memories_mb or 0), + defs.MEDIAN_PEAK_MEMORY_MB: np.median(peak_memories_mb or 0), + defs.PERCENTILES_PEAK_MEMORY_MB: [ + (p, np.percentile(peak_memories_mb or 0, p)) for p in selected_percentiles + ], defs.MEAN_AUDIO_UNDERRUN_S: np.mean(audio_underruns or 0), defs.STD_AUDIO_UNDERRUN_S: np.std(audio_underruns or 0), defs.MEDIAN_AUDIO_UNDERRUN_S: np.median(audio_underruns or 0), diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index b51dd2ee4..10f0853e0 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -14,7 +14,7 @@ import time import traceback import uuid import wave -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from dataclasses import dataclass, replace from datetime import datetime from pathlib import Path @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any, Literal import aiohttp import numpy as np import pybase64 as base64 +from PIL import Image from tqdm.asyncio import tqdm from vllm.benchmarks import datasets from vllm.benchmarks.datasets import SampleRequest @@ -85,7 +86,7 @@ from vllm_omni.benchmarks.omniinteract import ( write_batch_artifacts as write_omniinteract_batch_artifacts, ) from vllm_omni.metrics import definitions as defs -from vllm_omni.metrics.utils import coerce_positive_int_scalar +from vllm_omni.metrics.utils import coerce_bool, coerce_positive_float_scalar, coerce_positive_int_scalar if TYPE_CHECKING: from vllm_omni.clients.duplex import DuplexClient @@ -100,7 +101,13 @@ 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"}) +_IMAGE_STAGE_METRICS_BACKENDS = frozenset( + { + "/v1/images/generations", + "/v1/images/edits", + "openai-image-edits-omni", + } +) _PRINT_STAGE = False @@ -204,6 +211,61 @@ def _seed_tts_capture_pcm_for_wer() -> bool: ) +_DEFAULT_REQUEST_TIMEOUT_S = 900.0 +_LEGACY_REQUEST_TIMEOUT_S = 6 * 60 * 60.0 + +# Set from the ``--omni-request-timeout-s`` CLI flag by ``vllm bench serve`` +# before the benchmark session is built (``None`` = use the default above). +_REQUEST_TIMEOUT_OVERRIDE_S: float | None = None + + +def set_request_timeout_s(value: float) -> None: + """Record the explicitly requested per-request timeout (from the CLI).""" + global _REQUEST_TIMEOUT_OVERRIDE_S + _REQUEST_TIMEOUT_OVERRIDE_S = float(value) + + +def _omni_request_timeout_s() -> float: + """Per-request total timeout for the shared benchmark ``aiohttp`` session. + + An explicit ``--omni-request-timeout-s`` value wins over the 900 s default; + ``<= 0`` restores the legacy 6 h cap. A bounded per-request timeout makes a + hung server surface as ``failed`` requests once the deadline fires instead + of pinning the benchmark slot indefinitely. + """ + value = _REQUEST_TIMEOUT_OVERRIDE_S + if value is None: + return _DEFAULT_REQUEST_TIMEOUT_S + if value <= 0: + return _LEGACY_REQUEST_TIMEOUT_S + return value + + +def _build_benchmark_session( + max_concurrency: int | None, + ssl_setting: ssl.SSLContext | bool, +) -> aiohttp.ClientSession: + """Build the session shared by every benchmark request. + + Connections are reused across requests to reduce TLS handshake overhead; + the per-request total timeout comes from ``_omni_request_timeout_s()``. + """ + connector = aiohttp.TCPConnector( + limit=max_concurrency or 0, + limit_per_host=max_concurrency or 0, + ttl_dns_cache=300, + use_dns_cache=True, + enable_cleanup_closed=True, + force_close=True, + ssl=ssl_setting, + ) + return aiohttp.ClientSession( + connector=connector, + trust_env=True, + timeout=aiohttp.ClientTimeout(total=_omni_request_timeout_s()), + ) + + def _merge_extra_body_mm_kwargs(base: dict | None, overlay: dict | None) -> dict | None: """Shallow-merge ``extra_body`` dicts; deep-merge ``mm_processor_kwargs`` if both set.""" if not base and not overlay: @@ -662,6 +724,11 @@ class MixRequestFuncOutput(RequestFuncOutput): image_generation_time_ms: float = 0.0 image_pixels: int = 0 denoise_step_latency_ms: float = 0.0 + video_duration: float = 0.0 + video_rtf: float = 0.0 + video_frames: int = 0 + video_generation_time_ms: float = 0.0 + peak_memory_mb: float = 0.0 text_latency: float = 0.0 tpot_measured: bool = True #: Worst-case streaming-audio underrun (wall-clock seconds the player @@ -849,7 +916,7 @@ def _record_text_token_stream_intervals( def _update_output_stage_metrics_from_payload( output: MixRequestFuncOutput, - data: dict[str, Any], + data: Mapping[str, object], *, update_output_tokens: bool = True, ) -> None: @@ -897,7 +964,55 @@ def _apply_chat_stage0_token_timings(output: MixRequestFuncOutput) -> bool: ) -def _image_metrics_from_stage_metrics(metrics: dict[str, Any] | None) -> tuple[int, float, int, float]: +def _peak_memory_mb_from_payload(data: Mapping[str, object]) -> float: + peak_memory_mb = coerce_positive_float_scalar(data.get(defs.PEAK_MEMORY_MB)) + if peak_memory_mb is not None: + return peak_memory_mb + + for key in ("metrics", "usage"): + nested = data.get(key) + if isinstance(nested, dict): + peak_memory_mb = coerce_positive_float_scalar(nested.get(defs.PEAK_MEMORY_MB)) + if peak_memory_mb is not None: + return peak_memory_mb + + response_data = data.get("data") + if isinstance(response_data, list): + for item in response_data: + if isinstance(item, dict): + peak_memory_mb = coerce_positive_float_scalar(item.get(defs.PEAK_MEMORY_MB)) + if peak_memory_mb is not None: + return peak_memory_mb + + choices = data.get("choices") + if isinstance(choices, list): + for choice in choices: + if not isinstance(choice, dict): + continue + contents = [] + message = choice.get("message") + if isinstance(message, dict): + contents.append(message.get("content")) + delta = choice.get("delta") + if isinstance(delta, dict): + contents.append(delta.get("content")) + for content in contents: + if isinstance(content, list): + for item in content: + if isinstance(item, dict): + peak_memory_mb = coerce_positive_float_scalar(item.get(defs.PEAK_MEMORY_MB)) + if peak_memory_mb is not None: + return peak_memory_mb + return 0.0 + + +def _update_output_peak_memory_from_payload(output: MixRequestFuncOutput, data: Mapping[str, object]) -> None: + peak_memory_mb = _peak_memory_mb_from_payload(data) + if peak_memory_mb > output.peak_memory_mb: + output.peak_memory_mb = peak_memory_mb + + +def _image_metrics_from_stage_metrics(metrics: object) -> tuple[int, float, int, float]: if not isinstance(metrics, dict): return 0, 0.0, 0, 0.0 stage_snapshot = metrics.get("stage_metrics") @@ -924,7 +1039,7 @@ def _image_metrics_from_stage_metrics(metrics: dict[str, Any] | None) -> tuple[i return image_count, image_generation_ms, image_pixels, denoise_step_latency_ms -def _image_generation_ms_from_content(content: Any) -> float: +def _image_generation_ms_from_content(content: object) -> float: if not isinstance(content, list): return 0.0 for item in content: @@ -943,6 +1058,263 @@ def _image_generation_ms_from_content(content: Any) -> float: return 0.0 +def _image_info_from_response_data(content: object) -> tuple[int, int]: + if not isinstance(content, list): + return 0, 0 + image_count = 0 + total_pixels = 0 + for item in content: + if not isinstance(item, dict): + continue + b64_json = item.get("b64_json") + if not isinstance(b64_json, str) or not b64_json: + continue + try: + with Image.open(io.BytesIO(base64.b64decode(b64_json, validate=True))) as img: + width, height = img.size + img.verify() + image_count += 1 + total_pixels += int(width) * int(height) + except Exception: + logger.debug("Failed to decode generated image payload", exc_info=True) + return image_count, total_pixels + + +def _apply_image_metrics_from_payload(output: MixRequestFuncOutput, data: Mapping[str, object]) -> int: + """Populate image benchmark fields from an OpenAI-compatible image payload.""" + _update_output_stage_metrics_from_payload(output, data, update_output_tokens=False) + _update_output_peak_memory_from_payload(output, data) + + payload_image_count = 0 + response_data = data.get("data") + if isinstance(response_data, list): + payload_image_count, content_image_pixels = _image_info_from_response_data(response_data) + output.image_count = max(output.image_count, payload_image_count) + content_image_ms = _image_generation_ms_from_content(response_data) + if content_image_ms > 0: + output.image_generation_time_ms = max(output.image_generation_time_ms, content_image_ms) + if content_image_pixels > 0: + output.image_pixels = max(output.image_pixels, content_image_pixels) + + ( + metrics_image_count, + metrics_image_ms, + metrics_image_pixels, + metrics_denoise_step_ms, + ) = _image_metrics_from_stage_metrics(data.get("metrics")) + if metrics_image_count > output.image_count: + output.image_count = metrics_image_count + if metrics_image_ms > output.image_generation_time_ms: + output.image_generation_time_ms = metrics_image_ms + if metrics_image_pixels > output.image_pixels: + output.image_pixels = metrics_image_pixels + if metrics_denoise_step_ms > output.denoise_step_latency_ms: + output.denoise_step_latency_ms = metrics_denoise_step_ms + return payload_image_count + + +_VIDEO_FORM_FIELDS = ( + "seconds", + "num_frames", + "fps", + "num_inference_steps", + "seed", + "negative_prompt", + "guidance_scale", + "guidance_scale_2", + "boundary_ratio", + "flow_shift", + "true_cfg_scale", + "generate_sound", + "sound_duration", + "enable_frame_interpolation", + "frame_interpolation_exp", + "frame_interpolation_scale", + "frame_interpolation_model_path", + "lora", + "extra_params", +) + + +def _video_generation_ms_from_stage_durations(stage_durations: object) -> float: + if not isinstance(stage_durations, dict): + return 0.0 + gen_values = [ + float(value) + for key, value in stage_durations.items() + if str(key).endswith("_gen_ms") and isinstance(value, (int, float)) + ] + return max(gen_values) if gen_values else 0.0 + + +def _video_duration_from_payload(data: Mapping[str, object], request_body: Mapping[str, object]) -> float: + duration_s = coerce_positive_float_scalar(data.get("duration_s")) + if duration_s is not None and duration_s > 0: + return duration_s + + num_frames = coerce_positive_float_scalar(data.get("num_frames")) + fps = coerce_positive_float_scalar(data.get("fps")) + if num_frames is not None and fps is not None and num_frames > 0 and fps > 0: + return num_frames / fps + + seconds = coerce_positive_float_scalar(request_body.get("seconds")) + if seconds is not None and seconds > 0: + return seconds + + num_frames = coerce_positive_float_scalar(request_body.get("num_frames")) + fps = coerce_positive_float_scalar(request_body.get("fps")) + if num_frames is not None and fps is not None and num_frames > 0 and fps > 0: + return num_frames / fps + + return 0.0 + + +def _video_frames_from_payload(data: Mapping[str, object], request_body: Mapping[str, object]) -> int: + for key in ("num_frames", "video_frames", "frames"): + value = data.get(key) + if isinstance(value, list): + return len(value) + num_frames = coerce_positive_int_scalar(value) + if num_frames is not None: + return num_frames + + duration_s = coerce_positive_float_scalar(data.get("duration_s")) + fps = coerce_positive_float_scalar(data.get("fps")) + if duration_s is not None and fps is not None and duration_s > 0 and fps > 0: + return int(round(duration_s * fps)) + + num_frames = coerce_positive_int_scalar(request_body.get("num_frames")) + if num_frames is not None: + return num_frames + + seconds = coerce_positive_float_scalar(request_body.get("seconds")) + fps = coerce_positive_float_scalar(request_body.get("fps")) + if seconds is not None and fps is not None and seconds > 0 and fps > 0: + return int(round(seconds * fps)) + + return 0 + + +def _is_structured_image_reference(reference: Mapping[str, object]) -> bool: + """True for API image_reference objects ({image_url}/{file_id}).""" + image_url = reference.get("image_url") + file_id = reference.get("file_id") + has_url = isinstance(image_url, str) and bool(image_url) + has_file_id = isinstance(file_id, str) and bool(file_id) + return has_url or has_file_id + + +def _add_video_reference_to_form(form: aiohttp.FormData, reference: object) -> bool: + if isinstance(reference, dict) and "bytes" in reference: + form.add_field( + "input_reference", + reference["bytes"], + filename="benchmark-reference", + content_type=reference.get("content_type", "application/octet-stream"), + ) + return True + + if isinstance(reference, Mapping) and _is_structured_image_reference(reference): + form.add_field("image_reference", json.dumps(dict(reference))) + return True + + if isinstance(reference, list): + if reference and all(isinstance(item, Mapping) and _is_structured_image_reference(item) for item in reference): + form.add_field("image_reference", json.dumps([dict(item) for item in reference])) + return True + raise ValueError( + "Unsupported image_reference list; expected non-empty list of " + '{"image_url": "..."} and/or {"file_id": "..."} objects.' + ) + + if isinstance(reference, str): + if reference.startswith(("data:image", "http://", "https://")): + form.add_field("image_reference", json.dumps({"image_url": reference})) + return True + local_path = reference.removeprefix("file://") + if os.path.exists(local_path): + with open(local_path, "rb") as f: + reference_bytes = f.read() + form.add_field( + "input_reference", + reference_bytes, + filename=os.path.basename(local_path), + content_type=_guess_mime_type(local_path), + ) + return True + raise ValueError(f"Unsupported image_reference path or URL: {reference!r}") + + raise ValueError( + "Unsupported image_reference; expected upload bytes, local path/URL string, " + 'or {"image_url": "..."} / {"file_id": "..."} object ' + f"(got {type(reference).__name__})." + ) + + +def _add_video_extra_body_to_form( + form: aiohttp.FormData, + extra_body: Mapping[str, object], + request_body: Mapping[str, object], +) -> None: + for key in _VIDEO_FORM_FIELDS: + value = request_body.get(key) + if value is None: + continue + if isinstance(value, (dict, list)): + form.add_field(key, json.dumps(value)) + else: + form.add_field(key, str(value)) + + reserved = { + "model", + "prompt", + "size", + "width", + "height", + "poll_interval_s", + "poll_timeout_s", + # Handled only by _add_video_reference_to_form (upload / JSON image_url). + "image_reference", + "input_reference", + *_VIDEO_FORM_FIELDS, + } + for key, value in extra_body.items(): + if key in reserved or value is None: + continue + if isinstance(value, (dict, list)): + form.add_field(key, json.dumps(value)) + else: + form.add_field(key, str(value)) + + +def _apply_video_metrics_from_payload( + output: MixRequestFuncOutput, + data: Mapping[str, object], + request_body: Mapping[str, object], +) -> None: + output.video_duration = _video_duration_from_payload(data, request_body) + output.video_frames = _video_frames_from_payload(data, request_body) + _update_output_stage_metrics_from_payload(output, data, update_output_tokens=False) + _update_output_peak_memory_from_payload(output, data) + + stage_durations = data.get("stage_durations") + stage_gen_ms = _video_generation_ms_from_stage_durations(stage_durations) + if stage_gen_ms <= 0: + inference_time_s = coerce_positive_float_scalar(data.get("inference_time_s")) + if inference_time_s is not None and inference_time_s > 0: + stage_gen_ms = inference_time_s * 1000.0 + output.video_generation_time_ms = max(output.video_generation_time_ms, stage_gen_ms) + if output.video_duration <= 0: + return + # Prefer server-reported generation time so RTF is independent of client + # poll_interval_s sleep/overshoot baked into output.latency. + generation_s = output.video_generation_time_ms / 1000.0 + if generation_s > 0: + output.video_rtf = generation_s / output.video_duration + elif output.latency > 0: + output.video_rtf = output.latency / output.video_duration + + async def async_request_openai_chat_omni_completions( request_func_input: RequestFuncInput, session: aiohttp.ClientSession, @@ -1036,6 +1408,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 + output.peak_memory_mb = 0.0 completion_tokens_seen = 0 try: async with session.post(url=api_url, json=payload, headers=headers) as response: @@ -1065,6 +1438,7 @@ async def async_request_openai_chat_omni_completions( timestamp = time.perf_counter() data = json.loads(chunk) _update_output_stage_metrics_from_payload(output, data) + _update_output_peak_memory_from_payload(output, data) usage = data.get("usage") completion_tokens = None if isinstance(usage, dict): @@ -1281,16 +1655,243 @@ async def async_request_openai_chat_omni_completions( return output +def _finalize_image_json_http_response( + output: MixRequestFuncOutput, + *, + start_time: float, + status: int, + data: Mapping[str, object] | None, + error_text: str | None, +) -> None: + """Set e2el after the image JSON body has been fully read and validated.""" + output.latency = time.perf_counter() - start_time + if status != 200: + output.error = f"HTTP {status}: {error_text or ''}" + output.success = False + return + if not isinstance(data, Mapping): + output.error = "HTTP 200 response did not contain a JSON object" + output.success = False + return + payload_image_count = _apply_image_metrics_from_payload(output, data) + if payload_image_count <= 0: + output.error = "HTTP 200 response did not contain a valid image payload" + output.success = False + return + output.success = True + + +async def async_request_openai_image_generations_omni( + request_func_input: RequestFuncInput, + session: aiohttp.ClientSession, + pbar: tqdm | None = None, +) -> MixRequestFuncOutput: + """JSON request to /v1/images/generations for image generation benchmarks.""" + api_url = request_func_input.api_url + _validate_api_url(api_url, "OpenAI Image Generations API", "images/generations") + + extra_body = dict(request_func_input.extra_body or {}) + model = request_func_input.model_name if request_func_input.model_name else request_func_input.model + output = MixRequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + output.itl = [] + output.stage_metrics = {} + output.output_tokens = 0 + output.image_count = 0 + output.image_generation_time_ms = 0.0 + output.image_pixels = 0 + output.denoise_step_latency_ms = 0.0 + + size = extra_body.get("size") + if size is None: + width, height = extra_body.get("width"), extra_body.get("height") + if width is not None and height is not None: + size = f"{width}x{height}" + + payload: dict[str, object] = { + "model": model, + "prompt": request_func_input.prompt, + "n": int(extra_body.pop("n", extra_body.pop("num_outputs_per_prompt", 1)) or 1), + "response_format": "b64_json", + } + if size is not None: + payload["size"] = str(size) + + for key, value in extra_body.items(): + if key in {"height", "width"}: + continue + payload.setdefault(key, value) + + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + } + _update_headers_common(headers, request_func_input) + + st = time.perf_counter() + output.start_time = st + try: + async with session.post(url=api_url, json=payload, headers=headers) as response: + if response.status == 200: + data = await response.json() + _finalize_image_json_http_response( + output, + start_time=st, + status=response.status, + data=data if isinstance(data, Mapping) else None, + error_text=None, + ) + else: + error_text = await response.text() + _finalize_image_json_http_response( + output, + start_time=st, + status=response.status, + data=None, + error_text=error_text, + ) + except Exception: + output.latency = time.perf_counter() - st + output.success = False + output.error = traceback.format_exc() + logger.error(f"ERROR: send image generation request failed, reason is: {output.error}") + + if pbar: + pbar.update(1) + return output + + +async def async_request_openai_videos_omni( + request_func_input: RequestFuncInput, + session: aiohttp.ClientSession, + pbar: tqdm | None = None, +) -> MixRequestFuncOutput: + """Multipart request to async /v1/videos, polling metadata until completion.""" + api_url = request_func_input.api_url + _validate_api_url(api_url, "OpenAI Videos API", "videos") + + extra_body = dict(request_func_input.extra_body or {}) + model = request_func_input.model_name if request_func_input.model_name else request_func_input.model + output = MixRequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + output.itl = [] + output.stage_metrics = {} + output.output_tokens = 0 + output.video_duration = 0.0 + output.video_generation_time_ms = 0.0 + + request_body: dict[str, object] = { + "model": model, + "prompt": request_func_input.prompt, + } + request_body.update(extra_body) + + size = request_body.get("size") + if size is None: + width, height = request_body.get("width"), request_body.get("height") + if width is not None and height is not None: + size = f"{width}x{height}" + if size is not None: + request_body["size"] = str(size) + + form = aiohttp.FormData() + form.add_field("model", str(model)) + form.add_field("prompt", str(request_func_input.prompt)) + if request_body.get("size") is not None: + form.add_field("size", str(request_body["size"])) + _add_video_extra_body_to_form(form, extra_body, request_body) + + reference_added = False + for reference in _iter_image_edit_inputs(request_func_input.multi_modal_content): + if _add_video_reference_to_form(form, reference): + reference_added = True + break + if not reference_added: + image_reference = extra_body.get("image_reference") + if image_reference is not None: + _add_video_reference_to_form(form, image_reference) + + headers = { + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + } + _update_headers_common(headers, request_func_input) + + poll_interval_s = float(extra_body.get("poll_interval_s", 2.0) or 2.0) + timeout_s = float(extra_body.get("poll_timeout_s", 6 * 60 * 60) or (6 * 60 * 60)) + st = time.perf_counter() + output.start_time = st + try: + async with session.post(url=api_url, data=form, headers=headers) as response: + if response.status != 200: + output.latency = time.perf_counter() - st + output.error = f"HTTP {response.status}: {await response.text()}" + output.success = False + return output + create_payload = await response.json() + + job_id = create_payload.get("id") + job_status = create_payload.get("status") + if not isinstance(job_id, str) or not job_id: + output.latency = time.perf_counter() - st + output.error = "Video creation response missing job id." + output.success = False + return output + + job_url = f"{api_url.rstrip('/')}/{job_id}" + poll_payload = create_payload + deadline = time.perf_counter() + timeout_s + while job_status not in {"completed", "failed"}: + if time.perf_counter() >= deadline: + output.latency = time.perf_counter() - st + output.error = f"Timed out waiting for video job {job_id} to complete." + output.success = False + return output + await asyncio.sleep(poll_interval_s) + async with session.get(job_url, headers=headers) as poll_response: + if poll_response.status != 200: + output.latency = time.perf_counter() - st + output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}" + output.success = False + return output + poll_payload = await poll_response.json() + job_status = poll_payload.get("status") + + output.latency = time.perf_counter() - st + if job_status == "failed": + output.error = f"Video job failed: {poll_payload}" + output.success = False + return output + + _apply_video_metrics_from_payload(output, poll_payload, request_body) + output.success = True + except Exception: + output.latency = time.perf_counter() - st + output.success = False + output.error = traceback.format_exc() + logger.error(f"ERROR: send video request failed, reason is: {output.error}") + finally: + if pbar: + pbar.update(1) + return output + + async def async_request_openai_image_edits_omni( request_func_input: RequestFuncInput, session: aiohttp.ClientSession, pbar: tqdm | None = None, ) -> MixRequestFuncOutput: - """Streaming request to /v1/images/edits for multi-stage image-edit benchmarks.""" + """Multipart request to /v1/images/edits. + + Defaults to non-streaming JSON so single-stage edit models work. The server + rejects ``stream=true`` when ``len(stage_configs) <= 1``. Pass + ``stream: true`` in ``--extra-body`` for multi-stage SSE (AR TTFT / image + chunks). + """ api_url = request_func_input.api_url _validate_api_url(api_url, "OpenAI Image Edits API", "images/edits") extra_body = dict(request_func_input.extra_body or {}) + want_stream = coerce_bool(extra_body.pop("stream", None), default=False) model = request_func_input.model_name if request_func_input.model_name else request_func_input.model output = MixRequestFuncOutput() output.prompt_len = request_func_input.prompt_len @@ -1307,7 +1908,7 @@ async def async_request_openai_image_edits_omni( form.add_field("prompt", request_func_input.prompt) form.add_field("response_format", "b64_json") form.add_field("output_format", str(extra_body.get("output_format", "png"))) - form.add_field("stream", "true") + form.add_field("stream", "true" if want_stream else "false") size = extra_body.get("size") if size is None: @@ -1340,12 +1941,21 @@ async def async_request_openai_image_edits_omni( st = time.perf_counter() output.start_time = st - timestamp = st - most_recent_text_timestamp = st - generated_text = "" try: async with session.post(url=api_url, data=form, headers=headers) as response: - if response.status == 200: + if response.status != 200: + error_text = await response.text() + _finalize_image_json_http_response( + output, + start_time=st, + status=response.status, + data=None, + error_text=error_text, + ) + elif want_stream: + timestamp = st + most_recent_text_timestamp = st + generated_text = "" handler = StreamedResponseHandler() async for chunk_bytes in response.content.iter_any(): if not chunk_bytes: @@ -1366,6 +1976,7 @@ async def async_request_openai_image_edits_omni( data, update_output_tokens=(data.get("type") == "ar_delta"), ) + _update_output_peak_memory_from_payload(output, data) chunk_type = data.get("type") if chunk_type == "ar_delta": @@ -1401,9 +2012,16 @@ async def async_request_openai_image_edits_omni( output.generated_text = generated_text output.success = True else: - output.error = f"HTTP {response.status}: {await response.text()}" - output.success = False + data = await response.json() + _finalize_image_json_http_response( + output, + start_time=st, + status=response.status, + data=data if isinstance(data, Mapping) else None, + error_text=None, + ) except Exception: + output.latency = time.perf_counter() - st output.success = False output.error = traceback.format_exc() logger.error(f"ERROR: send image edit request failed, reason is: {output.error}") @@ -1951,6 +2569,18 @@ ASYNC_REQUEST_FUNCS["openai-chat-omni"] = async_request_openai_chat_omni_complet if "openai-chat-omni" not in OPENAI_COMPATIBLE_BACKENDS: OPENAI_COMPATIBLE_BACKENDS.append("openai-chat-omni") +ASYNC_REQUEST_FUNCS["/v1/images/edits"] = async_request_openai_image_edits_omni +if "/v1/images/edits" not in OPENAI_COMPATIBLE_BACKENDS: + OPENAI_COMPATIBLE_BACKENDS.append("/v1/images/edits") + +ASYNC_REQUEST_FUNCS["/v1/images/generations"] = async_request_openai_image_generations_omni +if "/v1/images/generations" not in OPENAI_COMPATIBLE_BACKENDS: + OPENAI_COMPATIBLE_BACKENDS.append("/v1/images/generations") + +ASYNC_REQUEST_FUNCS["/v1/videos"] = async_request_openai_videos_omni +if "/v1/videos" not in OPENAI_COMPATIBLE_BACKENDS: + OPENAI_COMPATIBLE_BACKENDS.append("/v1/videos") + ASYNC_REQUEST_FUNCS["openai-audio-speech"] = async_request_openai_audio_speech if "openai-audio-speech" not in OPENAI_COMPATIBLE_BACKENDS: OPENAI_COMPATIBLE_BACKENDS.append("openai-audio-speech") @@ -2041,21 +2671,8 @@ async def benchmark( # Reuses connections across requests to reduce TLS handshake overhead. ssl_setting = ssl_context if ssl_context is not None else ("https://" in api_url) - connector = aiohttp.TCPConnector( - limit=max_concurrency or 0, - limit_per_host=max_concurrency or 0, - ttl_dns_cache=300, - use_dns_cache=True, - enable_cleanup_closed=True, - force_close=True, - ssl=ssl_setting, - ) - - session = aiohttp.ClientSession( - connector=connector, - trust_env=True, - timeout=aiohttp.ClientTimeout(total=6 * 60 * 60), - ) + session = _build_benchmark_session(max_concurrency, ssl_setting) + print(f"Per-request timeout: {_omni_request_timeout_s():g}s") print("Starting initial single prompt test run...") test_prompt, test_prompt_len, test_output_len, test_mm_content = ( @@ -2330,6 +2947,18 @@ async def benchmark( defs.IMAGE_THROUGHPUT: getattr(metrics, defs.IMAGE_THROUGHPUT), defs.AVERAGE_PIXELS_PER_IMAGE: getattr(metrics, defs.AVERAGE_PIXELS_PER_IMAGE), defs.MEAN_DENOISE_STEP_LATENCY_MS: getattr(metrics, defs.MEAN_DENOISE_STEP_LATENCY_MS), + defs.TOTAL_VIDEO_DURATION_S: getattr(metrics, defs.TOTAL_VIDEO_DURATION_S), + defs.TOTAL_VIDEO_FRAMES: getattr(metrics, defs.TOTAL_VIDEO_FRAMES), + defs.VIDEO_THROUGHPUT: getattr(metrics, defs.VIDEO_THROUGHPUT), + defs.MEAN_VIDEO_RTF: getattr(metrics, defs.MEAN_VIDEO_RTF), + defs.MEDIAN_VIDEO_RTF: getattr(metrics, defs.MEDIAN_VIDEO_RTF), + defs.PERCENTILES_VIDEO_RTF: getattr(metrics, defs.PERCENTILES_VIDEO_RTF), + defs.MEAN_VIDEO_GENERATION_MS: getattr(metrics, defs.MEAN_VIDEO_GENERATION_MS), + defs.MEDIAN_VIDEO_GENERATION_MS: getattr(metrics, defs.MEDIAN_VIDEO_GENERATION_MS), + defs.PERCENTILES_VIDEO_GENERATION_MS: getattr(metrics, defs.PERCENTILES_VIDEO_GENERATION_MS), + defs.MEAN_PEAK_MEMORY_MB: getattr(metrics, defs.MEAN_PEAK_MEMORY_MB), + defs.MEDIAN_PEAK_MEMORY_MB: getattr(metrics, defs.MEDIAN_PEAK_MEMORY_MB), + defs.PERCENTILES_PEAK_MEMORY_MB: getattr(metrics, defs.PERCENTILES_PEAK_MEMORY_MB), "input_lens": [output.prompt_len for output in outputs], "start_times": [output.start_time for output in outputs], "output_lens": actual_output_lens, diff --git a/vllm_omni/benchmarks/serve.py b/vllm_omni/benchmarks/serve.py index 2c4f8d130..7d05eb24f 100644 --- a/vllm_omni/benchmarks/serve.py +++ b/vllm_omni/benchmarks/serve.py @@ -20,9 +20,41 @@ from vllm_omni.benchmarks.omniinteract import omniinteract_output_lock from vllm_omni.benchmarks.patch.patch import ( maybe_enable_stage_metrics, set_print_stage, + set_request_timeout_s, should_request_stage_metrics, ) +_ENDPOINT_BACKEND_KEYS = { + "/v1/images/generations", + "/v1/images/edits", + "/v1/videos", +} + + +def _normalize_endpoint(endpoint: str | None) -> str | None: + if endpoint is None: + return None + endpoint = str(endpoint).strip() + if not endpoint: + return None + if not endpoint.startswith("/"): + endpoint = f"/{endpoint}" + return endpoint + + +def _use_endpoint_backend_when_implicit(args: argparse.Namespace) -> None: + explicit_keys = getattr(args, "explicit_keys", frozenset()) + if "backend" in explicit_keys: + return + # Upstream vLLM defaults --backend to "openai". Treat that non-explicit + # default the same as an empty backend so endpoint-driven omni handlers work. + backend = getattr(args, "backend", None) + if backend not in (None, "", "openai"): + return + endpoint = _normalize_endpoint(getattr(args, "endpoint", None)) + if endpoint in _ENDPOINT_BACKEND_KEYS: + args.backend = endpoint + def main(args: argparse.Namespace) -> dict[str, Any]: if getattr(args, "seed_tts_wer_eval", False): @@ -31,6 +63,10 @@ def main(args: argparse.Namespace) -> dict[str, Any]: os.environ["SEED_TTS_WER_SAVE_ITEMS"] = "1" if getattr(args, "daily_omni_save_eval_items", False): os.environ["DAILY_OMNI_SAVE_EVAL_ITEMS"] = "1" + request_timeout_s = getattr(args, "omni_request_timeout_s", None) + if request_timeout_s is not None: + set_request_timeout_s(request_timeout_s) + _use_endpoint_backend_when_implicit(args) set_print_stage(getattr(args, "print_stage", False)) args.extra_body = maybe_enable_stage_metrics( getattr(args, "extra_body", None), diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index f53a1ef74..073cf4f59 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -45,8 +45,10 @@ from vllm_omni.config.yaml_util import ( # `from vllm_omni.config.lora import LoRAConfig` (used while data.py is still # loading) cannot close a circular import through DiffusionOutput. _LAZY_ATTRS = { + "OmniConfigResolution": ("vllm_omni.config.resolver", "OmniConfigResolution"), "StageConfigFactory": ("vllm_omni.config.config_factory", "StageConfigFactory"), "register_pipeline": ("vllm_omni.config.pipeline_registry", "register_pipeline"), + "resolve_omni_config": ("vllm_omni.config.resolver", "resolve_omni_config"), } __all__ = [ @@ -60,6 +62,8 @@ __all__ = [ "VllmOmniGenerationStageConfig", "VllmOmniDiffusionStageConfig", "StageConfigType", + "OmniConfigResolution", + "resolve_omni_config", # Structured Omni sub-configs. "OmniStageCacheConfig", "OmniStageConnectorConfig", diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index b2eb64bca..90108e7df 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -4,10 +4,10 @@ from __future__ import annotations -import dataclasses import functools +import json from collections.abc import Mapping -from dataclasses import asdict +from dataclasses import asdict, dataclass from pathlib import Path from typing import Any @@ -31,7 +31,7 @@ from vllm_omni.config.stage_config import ( merge_pipeline_deploy, normalize_pipeline_cli_overrides, ) -from vllm_omni.config.yaml_util import create_config +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.io_support import get_diffusion_output_type from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero @@ -48,6 +48,15 @@ logger = init_logger(__name__) _DEFAULT_PARALLEL_DEGREE = 1 +@dataclass(frozen=True) +class _LegacyConfigResolution: + """Factory-owned result for the temporary legacy runtime bridge.""" + + stage_configs: list[StageConfig] + pipeline_config: PipelineConfig + omni_lb_policy: str | None + + @functools.cache def _materialize_object_storage_configs(model: str) -> str: """Materialize an object-storage model URI's config files locally. @@ -104,6 +113,11 @@ def with_trust_remote_code_override( class StageConfigFactory: """Factory that loads pipeline YAML and merges CLI overrides. + Production startup source selection is owned by + :func:`vllm_omni.config.resolver.resolve_omni_config`. Factory methods are + lower-level construction primitives for that resolver, config internals, + and focused tests; entrypoints and engines must not call them directly. + Handles both single-stage and multi-stage models. Pipelines are declared in ``vllm_omni/config/pipeline_registry.py`` and @@ -394,11 +408,13 @@ class StageConfigFactory: 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. + """Build the migration-only runtime ABI 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. + The engine still consumes the legacy ``StageConfig``/OmegaConf shape. + This method is the resolver's temporary compatibility bridge, not an + alternative production source-selection entrypoint and not a stable + public contract. RFC #4021 will remove it as runtime consumers move to + ``VllmOmniConfig``. """ user_deploy_config = cls._load_user_deploy_config(deploy_config_path) pipeline_cfg = cls.get_pipeline_config( @@ -429,14 +445,29 @@ class StageConfigFactory: 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. + """Return the existing two-value legacy runtime ABI.""" + resolved = cls._resolve_legacy_from_registry( + pipeline_cfg, + cli_overrides, + deploy_config_path, + user_deploy_config, + strategy_specs, + ) + return resolved.stage_configs, resolved.omni_lb_policy + + @classmethod + def _resolve_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, + ) -> _LegacyConfigResolution: + """Create runtime stages and retain their effective topology. 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. """ cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides) deploy_cfg: DeployConfig | None @@ -475,7 +506,11 @@ class StageConfigFactory: 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 + return _LegacyConfigResolution( + stage_configs=stages, + pipeline_config=pipeline_cfg, + omni_lb_policy=omni_lb_policy, + ) @staticmethod def _apply_strategy_specs( @@ -587,65 +622,64 @@ class StageConfigFactory: @classmethod def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]: - """Single-stage diffusion - no YAML needed. + """Build the temporary runtime ABI for a generic diffusion stage. - 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. + The terminal diffusion config owns engine defaults and normalization; + this compatibility builder only adds Omni stage topology, request + defaults, and device placement. """ - # 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}" + raw_sampling_params = kwargs.get("default_sampling_params") + if isinstance(raw_sampling_params, str): + try: + raw_sampling_params = json.loads(raw_sampling_params) + except json.JSONDecodeError: + logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.") + raw_sampling_params = None + if not isinstance(raw_sampling_params, Mapping): + raw_sampling_params = None + default_sampling_params = dict(raw_sampling_params.get("0", {})) if raw_sampling_params else {} - 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") + parallel_config = DiffusionParallelConfig.from_stage_overrides(kwargs) + if kwargs.get("num_gpus") is not None: + parallel_config.resolve_data_parallel_size(int(kwargs["num_gpus"])) + engine_args = OmniDiffusionConfig.normalize_init_kwargs(kwargs) + engine_args["parallel_config"] = asdict(parallel_config) engine_args["model_stage"] = "diffusion" - # Convert dtype to string for OmegaConf - if "dtype" in engine_args: - engine_args["dtype"] = str(engine_args["dtype"]) + extras = dict(engine_args.get("extras") or {}) + for key, default in ( + ("auxiliary_text_encoder", None), + ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), + ): + value = kwargs.get(key) + if value is not None: + extras[key] = value + else: + extras.setdefault(key, default) + engine_args["extras"] = extras - engine_args.setdefault("max_num_seqs", 1) model_class_name = engine_args.get("model_class_name") + final_output_type = get_diffusion_output_type(model_class_name) + logger.info( + "Resolved generic diffusion final_output_type=%r for model_class_name=%r.", + final_output_type, + model_class_name, + ) - 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": get_diffusion_output_type(model_class_name), - } - - return [config_dict] + return [ + { + "stage_id": 0, + "stage_type": StageType.DIFFUSION.value, + "runtime": { + "process": True, + "devices": ",".join(str(i) for i in range(parallel_config.world_size)), + }, + "engine_args": engine_args, + "default_sampling_params": default_sampling_params, + "final_output": True, + "final_output_type": final_output_type, + } + ] @classmethod def _merge_cli_overrides( diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 1a1d40141..c1e380a22 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -175,7 +175,6 @@ _MODEL_PROMOTE = ( _MODEL_REQUEST_SCOPE = ( "DOTS_TTS_DIT_NUM_STEPS", - "GR00T_NOISE_SEED", "MAGI2_NEGATIVE_PROMPT", "STEP_AUDIO2_DEFAULT_PROMPT_WAV", "VLLM_GEPARD_END_FADE_MS", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 1a1c0c3ae..aaa9c3512 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -63,6 +63,7 @@ _PIPELINE_DEPLOY_CLI_FIELDS = PIPELINE_WIDE_ENGINE_FIELDS _NON_STAGE_ENGINE_CLI_FIELDS = frozenset( { "async_chunk", + "disable_log_stats", "model", "omni", "output_modalities", @@ -196,6 +197,7 @@ class _SchedulerEngineOverrides(TypedDict, total=False): class _RuntimeEngineOverrides(TypedDict, total=False): + additional_config: dict[str, Any] distributed_executor_backend: Any worker_cls: str devices: str @@ -524,6 +526,8 @@ class OmniStageConnectorConfig: class OmniStageRuntimeConfig: """Per-stage process placement and backend runtime behavior.""" + # LLM backend extensions; diffusion owns these in its config projection. + additional_config: dict[str, Any] | None = None distributed_executor_backend: Any = None worker_cls: str | None = None devices: str | None = None @@ -1271,10 +1275,12 @@ def _stage_engine_values( load_engine_fields = _LLM_LOAD_ENGINE_FIELDS cache_engine_fields = _LLM_CACHE_ENGINE_FIELDS scheduler_engine_fields = _LLM_SCHEDULER_ENGINE_FIELDS + runtime_engine_fields = _RUNTIME_ENGINE_FIELDS else: load_engine_fields = _LOAD_ENGINE_FIELDS cache_engine_fields = _CACHE_ENGINE_FIELDS scheduler_engine_fields = _SCHEDULER_ENGINE_FIELDS + runtime_engine_fields = _RUNTIME_ENGINE_FIELDS - {"additional_config"} return _StageEngineValues( quantization=cast( _QuantizationEngineOverrides, @@ -1291,7 +1297,7 @@ def _stage_engine_values( _ConnectorEngineOverrides, _select_engine_overrides(engine, _CONNECTOR_ENGINE_FIELDS), ), - runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, _RUNTIME_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), compilation_config=_copy_value(engine.get("compilation_config")), diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 5ac47dddc..b711f0cfb 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -17,9 +17,8 @@ To add a new pipeline: Out of tree pipeline configs or resolvers can also be registered with register_pipeline. -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. +NOTE: Generic single-stage diffusion is selected by ``config.resolver`` when +no registered Omni pipeline matches. """ from __future__ import annotations @@ -40,6 +39,7 @@ from vllm_omni.model_executor.models.audex.pipeline import ( AUDEX_TTA_PIPELINE, AUDEX_TTS_PIPELINE, ) +from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_PIPELINE, @@ -68,6 +68,9 @@ from vllm_omni.model_executor.models.indextts2.pipeline import ( INDEXTTS2_PIPELINE, INDEXTTS25_PIPELINE, ) +from vllm_omni.model_executor.models.joyai_vl_interaction.pipeline import ( + JOYAI_VL_INTERACTION_PIPELINE, +) from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE from vllm_omni.model_executor.models.lingbot_world.pipeline import LINGBOT_WORLD_PIPELINE from vllm_omni.model_executor.models.mammoth_moda2.pipeline import ( @@ -123,6 +126,7 @@ PipelineResolverFunc: TypeAlias = Callable[[PretrainedConfig | None], PipelineCo # --- Multi-stage omni pipelines (LLM-centric; audio / video I/O) --- OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "aura_omni": AURA_OMNI_PIPELINE, + "joyai_vl_interaction": JOYAI_VL_INTERACTION_PIPELINE, "qwen2_5_omni": QWEN2_5_OMNI_PIPELINE, "qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE, "personaplex": PERSONAPLEX_PIPELINE, @@ -169,6 +173,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "voxtral_tts": VOXTRAL_TTS_PIPELINE, "glm_tts": GLM_TTS_PIPELINE, "fish_qwen3_omni": FISH_SPEECH_PIPELINE, + "arktts": AUDIO8_TTS_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, diff --git a/vllm_omni/config/resolver.py b/vllm_omni/config/resolver.py new file mode 100644 index 000000000..4f5ee8a6c --- /dev/null +++ b/vllm_omni/config/resolver.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, fields, is_dataclass +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs +from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override +from vllm_omni.config.endpoint_policy import EndpointRestriction +from vllm_omni.config.omni_config import VllmOmniConfig +from vllm_omni.config.stage_config import PipelineConfig +from vllm_omni.config.yaml_util import create_config +from vllm_omni.diffusion.data import resolve_model_class_name +from vllm_omni.diffusion.registry import DiffusionModelRegistry +from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class OmniConfigResolution: + """Migration envelope returned by the production config resolver. + + ``stage_configs`` intentionally carries the current OmegaConf-compatible + runtime ABI only until stage startup consumes ``VllmOmniConfig`` directly. + ``pipeline_config`` is the effective topology after runtime stage injection, + so both fields always describe the same set of stages. + It is not a stable authoring or extension API; new production callers + should resolve through :func:`resolve_omni_config` and must not construct or + merge this compatibility shape themselves. + """ + + config_path: str | None + stage_configs: tuple[Any, ...] # Temporary StageConfig/OmegaConf bridge. + pipeline_config: PipelineConfig | None = None + omni_lb_policy: str | None = None + + @property + def endpoint_restrictions(self) -> tuple[EndpointRestriction, ...]: + if self.pipeline_config is None: + return () + return tuple(self.pipeline_config.endpoint_restrictions) + + def stage_by_id(self, stage_id: int) -> Any: + for stage in self.stage_configs: + if stage.stage_id == stage_id: + return stage + available = [stage.stage_id for stage in self.stage_configs] + raise KeyError(f"no stage {stage_id}; resolved stages: {available}") + + +def _filter_dict_like_object(obj: dict | Any) -> dict: + """Convert a dict-like object while dropping OmegaConf-incompatible callables.""" + result = {} + filtered_keys = [] + for key, value in obj.items(): + # Preserve class objects as import paths for consumers such as + # custom_pipeline_args.pipeline_class. + if isinstance(value, type): + module = getattr(value, "__module__", None) + qualname = getattr(value, "__qualname__", getattr(value, "__name__", None)) + result[key] = f"{module}.{qualname}" if module and qualname and module != "builtins" else qualname + elif callable(value): + filtered_keys.append(str(key)) + else: + result[key] = _convert_dataclasses_to_dict(value) + if filtered_keys: + logger.warning( + "Filtered out %d callable object(s) from base_engine_args that are not compatible with OmegaConf: %s.", + len(filtered_keys), + filtered_keys, + ) + return result + + +def _convert_dataclasses_to_dict(obj: Any) -> Any: + """Recursively convert caller values to OmegaConf-compatible types.""" + # Check by class name before dict to cover both collections.Counter and + # vllm.utils.Counter without importing either implementation. + if hasattr(obj, "__class__") and obj.__class__.__name__ == "Counter": + try: + return dict(obj) + except (TypeError, ValueError): + return {} + if isinstance(obj, set): + return list(obj) + if is_dataclass(obj) and not isinstance(obj, type): + result = {} + for config_field in fields(obj): + if not config_field.init: + continue + value = getattr(obj, config_field.name) + # At the CLI/deploy boundary, None means "unset" rather than + # "clear an inherited value". Preserve the dataclass default by + # omitting fields whose declared default is already None. + if value is None and config_field.default is None: + continue + result[config_field.name] = _convert_dataclasses_to_dict(value) + return result + if isinstance(obj, dict): + return _filter_dict_like_object(obj) + if isinstance(obj, type): + module = getattr(obj, "__module__", None) + qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", None)) + return f"{module}.{qualname}" if module and qualname and module != "builtins" else qualname + if callable(obj): + logger.warning( + "Cannot convert callable %r to an OmegaConf-compatible value.", + obj, + ) + raise TypeError(f"callable {obj!r} is not an OmegaConf-compatible value") + if isinstance(obj, (list, tuple)): + converted = [] + for item in obj: + if callable(item): + logger.warning( + "Filtered callable %r from an OmegaConf-compatible sequence.", + item, + ) + continue + converted.append(_convert_dataclasses_to_dict(item)) + return type(obj)(converted) + if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, (str, bytes)): + try: + return _filter_dict_like_object(obj) + except (TypeError, ValueError, AttributeError) as exc: + logger.warning( + "Failed to convert dict-like %s to an OmegaConf-compatible mapping.", + type(obj).__name__, + ) + raise TypeError(f"cannot convert dict-like {type(obj).__name__}") from exc + return obj + + +def _flatten_stage_overrides( + cli_overrides: dict[str, Any], + stage_overrides: Mapping[str, Mapping[str, Any]] | None, +) -> None: + """Translate per-stage mappings to the factory's flat CLI override ABI.""" + if not stage_overrides: + return + for stage_id, overrides in stage_overrides.items(): + if not isinstance(overrides, Mapping): + raise TypeError(f"stage override {stage_id!r} must be a mapping, got {type(overrides).__name__}") + for key, value in overrides.items(): + cli_overrides[f"stage_{stage_id}_{key}"] = value + + +def _apply_generic_stage_overrides( + cli_overrides: dict[str, Any], + stage_overrides: Mapping[str, Mapping[str, Any]] | None, +) -> None: + """Apply stage-zero overrides to the generic single-stage fallback.""" + if not stage_overrides or not (stage_zero := stage_overrides.get("0")): + return + for key, value in stage_zero.items(): + if key == "extras": + if not isinstance(value, Mapping): + raise TypeError(f"stage override '0'.extras must be a mapping, got {type(value).__name__}") + cli_overrides[key] = {**dict(cli_overrides.get(key) or {}), **value} + else: + cli_overrides[key] = value + + +def _load_strategy_specs(strategy_config_path: str | None) -> Mapping[Any, Any] | None: + if strategy_config_path is None: + return None + + return load_strategy_specs(strategy_config_path) + + +def _build_registered_resolution( + structured_config: VllmOmniConfig, + *, + cli_overrides: dict[str, Any], + strategy_config_path: str | None, +) -> OmniConfigResolution: + """Build the temporary runtime view for an already-resolved pipeline.""" + # Runtime consumers have not yet moved to typed per-stage configs. Use the + # factory-owned compatibility bridge instead of reimplementing legacy YAML + # discovery and merging in this resolver. + effective_deploy_path = structured_config.orchestrator_config.deploy_config_path + legacy_resolution = StageConfigFactory._resolve_legacy_from_registry( + structured_config.pipeline_config, + cli_overrides, + effective_deploy_path, + strategy_specs=_load_strategy_specs(strategy_config_path), + ) + + return OmniConfigResolution( + config_path=effective_deploy_path, + stage_configs=tuple(stage.to_omegaconf() for stage in legacy_resolution.stage_configs), + pipeline_config=legacy_resolution.pipeline_config, + omni_lb_policy=legacy_resolution.omni_lb_policy, + ) + + +def _resolve_generic_diffusion_model_class( + model: str, + cli_overrides: Mapping[str, Any], +) -> tuple[bool, str | None]: + """Detect generic diffusion support and its serving pipeline class.""" + model_class_name = cli_overrides.get("model_class_name") or resolve_model_class_name( + model, + str(cli_overrides.get("diffusion_load_format") or "default"), + cli_overrides.get("revision"), + ) + supported = bool(model_class_name and model_class_name in DiffusionModelRegistry.get_supported_archs()) + if not supported: + supported = is_diffusion_model(model) + return supported, str(model_class_name) if model_class_name else None + + +def resolve_omni_config( + model: str, + *, + trust_remote_code: bool | None, + deploy_config_path: str | None, + cli_overrides: Mapping[str, Any] | None, + stage_overrides: Mapping[str, Mapping[str, Any]] | None, + strategy_config_path: str | None, +) -> OmniConfigResolution: + """Resolve registry/deploy inputs through the single public entrypoint.""" + normalized_overrides = _convert_dataclasses_to_dict(dict(cli_overrides or {})) + normalized_overrides = with_trust_remote_code_override(normalized_overrides, trust_remote_code) + registry_overrides = dict(normalized_overrides) + _flatten_stage_overrides(registry_overrides, stage_overrides) + + structured_config = StageConfigFactory.create_from_model( + model, + trust_remote_code=trust_remote_code, + cli_overrides=registry_overrides, + deploy_config_path=deploy_config_path, + ) + if structured_config is not None: + return _build_registered_resolution( + structured_config, + cli_overrides=registry_overrides, + strategy_config_path=strategy_config_path, + ) + + _apply_generic_stage_overrides(normalized_overrides, stage_overrides) + supported, model_class_name = _resolve_generic_diffusion_model_class(model, normalized_overrides) + if not supported: + raise ValueError( + f"Model {model!r} did not resolve to a registered Omni pipeline or a supported diffusion model." + ) + if model_class_name is not None: + normalized_overrides.setdefault("model_class_name", model_class_name) + default_stages = StageConfigFactory.create_default_diffusion(normalized_overrides) + + return OmniConfigResolution( + config_path=deploy_config_path, + stage_configs=tuple(create_config(_convert_dataclasses_to_dict(default_stages))), + ) + + +__all__ = [ + "OmniConfigResolution", + "resolve_omni_config", +] diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 06840e55a..2ac1f99dd 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -509,6 +509,7 @@ class DuplexSessionRuntimeConfig: max_pending_turns_per_session: int = 4 max_sessions: int = 1 completed_append_cache_size: int = 256 + server_vad_model_path: str | None = None # Startup warmup: run this many silent 80 ms-style frames through a # throwaway realtime session before real clients are admitted, so # one-time costs (kernel JIT, first prefill/decode paths, codec caches) @@ -528,6 +529,10 @@ class DuplexSessionRuntimeConfig: } if self.idle_ttl_s is not None and self.idle_ttl_s <= 0: raise ValueError("duplex_session.idle_ttl_s must be positive or null") + if self.server_vad_model_path is not None and ( + not isinstance(self.server_vad_model_path, str) or not self.server_vad_model_path.strip() + ): + raise ValueError("duplex_session.server_vad_model_path must be a non-empty string or null") for name, value in positive.items(): if value <= 0: raise ValueError(f"duplex_session.{name} must be positive") diff --git a/vllm_omni/deploy/audio8_tts.yaml b/vllm_omni/deploy/audio8_tts.yaml new file mode 100644 index 000000000..644d77745 --- /dev/null +++ b/vllm_omni/deploy/audio8_tts.yaml @@ -0,0 +1,80 @@ +# Audio8 TTS Preview 0.6B deploy: async-chunk streaming slow_ar -> codec_decoder. +# Model: Audio8/Audio8-TTS-Preview-0.6b (model_type = "arktts"). +# +# trust_remote_code is deliberately OFF: the checkpoint ships an `auto_map`, and +# transformers prefers remote code over registered classes when it is enabled, +# which would bypass vllm-omni's Audio8TTSConfig (the Qwen2-shaped view that +# Qwen2Model needs). The tokenizer is a plain PreTrainedTokenizerFast and needs +# no remote code either. +async_chunk: true +enable_chunked_prefill: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + # ~21.5 Hz codec; 25 frames ~= 1.16 s of audio. + codec_chunk_frames: 25 + codec_left_context_frames: 25 + initial_codec_chunk_frames: 4 + # float32 keeps the Snake activations / weight-normed convs clean. + audio8_tts_codec_dtype: float32 + +stages: + - stage_id: 0 + # >= 4 so concurrent requests pipeline; a default of 1 starves the codec + # window and produces audible gaps under load. + # + # Measured on 1x H20 (20 prompts, mean audio 4.5 s): at client concurrency 8, + # raising this to 8 buys +27% throughput and 4.5x lower first-packet latency + # (1124 ms -> 247 ms) at the cost of a 0.22 s worst-case buffer deficit in + # 1 of 20 requests. 4 is the default because every measured point here had + # zero underrun; see examples/online_serving/text_to_speech/README.md. + max_num_seqs: 4 + # 0.6B weights + a 2048-token KV cache need very little: sized for a small + # slice of the GPU so the codec stage and other tenants still fit. + gpu_memory_utilization: 0.2 + enforce_eager: false + trust_remote_code: false + enable_prefix_caching: false + async_scheduling: false + # The checkpoint's max_seq_len is 2048 packed text+audio positions; with + # chunked prefill off, max_num_batched_tokens must be >= max_model_len. + max_model_len: 2048 + max_num_batched_tokens: 2048 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + # generation_config.json of the released checkpoint. + temperature: 0.7 + top_k: 50 + top_p: 0.9 + max_tokens: 1024 + seed: 42 + repetition_penalty: 1.0 + + - stage_id: 1 + max_num_seqs: 1 + # The codec is ~1.4 GB of weights plus per-chunk activations. + gpu_memory_utilization: 0.1 + enforce_eager: true + trust_remote_code: false + enable_prefix_caching: false + async_scheduling: false + max_model_len: 2048 + max_num_batched_tokens: 2048 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 diff --git a/vllm_omni/deploy/cosyvoice3.yaml b/vllm_omni/deploy/cosyvoice3.yaml index 551b42af6..09fd83290 100644 --- a/vllm_omni/deploy/cosyvoice3.yaml +++ b/vllm_omni/deploy/cosyvoice3.yaml @@ -23,6 +23,7 @@ # On a slower GPU class, lower max_num_seqs (e.g. 4) and/or codec_chunk_frames # first. Raising connector_get_sleep_s also relaxes the consumer's polling # cadence, but trades against TTFP — keep it as a last resort. +pipeline: cosyvoice3 async_chunk: true connectors: connector_of_shared_memory: diff --git a/vllm_omni/deploy/joyai_vl_interaction.yaml b/vllm_omni/deploy/joyai_vl_interaction.yaml new file mode 100644 index 000000000..59eaa8f6c --- /dev/null +++ b/vllm_omni/deploy/joyai_vl_interaction.yaml @@ -0,0 +1,79 @@ +# Single-GPU JoyAI-VL-Interaction -> Qwen3-TTS all-sync profile. +pipeline: joyai_vl_interaction +async_chunk: false +trust_remote_code: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + 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 + # Qwen3-TTS Code2Wav buckets: 25, 72+1, 72+25, 72+72+25, and 300+25 frames. + decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325] + +stages: + # GPU 0: JoyAI action model. + - stage_id: 0 + model: jdopensource/JoyAI-VL-Interaction-Preview + devices: "0" + gpu_memory_utilization: 0.5 + max_num_seqs: 4 + # Some OmniInteract video prompts contain more than 12K input tokens. + max_model_len: 16384 + max_num_batched_tokens: 16384 + limit_mm_per_prompt: {image: 256, video: 1} + enable_prefix_caching: true + default_sampling_params: + temperature: 0.8 + top_p: 0.9 + top_k: 40 + max_tokens: 256 + repetition_penalty: 1.1 + presence_penalty: 1.5 + + # GPU 0: Qwen3-TTS Talker. + - stage_id: 1 + model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice + devices: "0" + gpu_memory_utilization: 0.2 + max_num_seqs: 8 + max_model_len: 4096 + max_num_batched_tokens: 512 + enable_prefix_caching: false + output_connectors: + to_stage_2: connector_of_shared_memory + default_sampling_params: + temperature: 0.9 + top_k: 50 + min_tokens: 2 + max_tokens: 4096 + repetition_penalty: 1.05 + subtalker_sampling_params: + do_sample: true + temperature: 0.9 + top_k: 50 + top_p: 1.0 + + # GPU 0: Qwen3-TTS Code2Wav. + - stage_id: 2 + model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice + devices: "0" + gpu_memory_utilization: 0.2 + max_num_seqs: 8 + max_model_len: 65536 + max_num_batched_tokens: 65536 + enable_prefix_caching: false + 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 + repetition_penalty: 1.0 diff --git a/vllm_omni/deploy/omnivoice.yaml b/vllm_omni/deploy/omnivoice.yaml index 6769dd5e9..df9e6a5da 100644 --- a/vllm_omni/deploy/omnivoice.yaml +++ b/vllm_omni/deploy/omnivoice.yaml @@ -11,3 +11,5 @@ stages: trust_remote_code: true distributed_executor_backend: "mp" dtype: "float32" + max_num_seqs: 8 + request_batch_max_wait_ms: 10 diff --git a/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml index 2aecdbf9b..b600ff1e1 100644 --- a/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml +++ b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml @@ -5,8 +5,9 @@ # --forced-aligner sets just the model path. # # The stage runs the model with runner="pooling" + pooling_task="token_classify". -# Prompt tokens, chat template and word segmentation are Qwen-specific and live -# in qwen3_force_align_processor (not configurable here). +# Prompt tokens and word segmentation are Qwen-specific and live in +# qwen3_force_align_processor (not configurable here); the prompt is the same +# string qwen_asr's encode_timestamp builds, with no chat template. forced_aligner: model: Qwen/Qwen3-ForcedAligner-0.6B diff --git a/vllm_omni/deploy/step_audio2_ci.yaml b/vllm_omni/deploy/step_audio2_ci.yaml index 4373dfaa8..060ae6c54 100644 --- a/vllm_omni/deploy/step_audio2_ci.yaml +++ b/vllm_omni/deploy/step_audio2_ci.yaml @@ -15,7 +15,7 @@ stages: max_model_len: 1024 max_num_batched_tokens: 1024 max_num_seqs: 1 - gpu_memory_utilization: 0.7 + gpu_memory_utilization: 0.8 skip_mm_profiling: true enforce_eager: true async_scheduling: false diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py index f75833831..20327b7e9 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py @@ -7,6 +7,22 @@ from __future__ import annotations import torch from vllm.triton_utils import tl, triton +# Ascend CANN limits the launch coreDim to uint16. +_MAX_1D_GRID_SIZE = 65535 + + +def _iter_row_chunks(rows: int): + for row_offset in range(0, rows, _MAX_1D_GRID_SIZE): + yield row_offset, min(rows - row_offset, _MAX_1D_GRID_SIZE) + + +def _launch_row_chunks(kernel, rows: int, device_type: str, *args, **kwargs) -> None: + if device_type != "npu": + kernel[(rows,)](*args, 0, **kwargs) + return + for row_offset, chunk_rows in _iter_row_chunks(rows): + kernel[(chunk_rows,)](*args, row_offset, **kwargs) + @triton.jit def _indexed_scale_shift_kernel( @@ -20,9 +36,10 @@ def _indexed_scale_shift_kernel( stride_shift_row, stride_scale_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -51,9 +68,10 @@ def _indexed_gate_kernel( stride_gate_row, stride_other_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -83,9 +101,10 @@ def _rms_norm_indexed_scale_shift_kernel( stride_shift_row, stride_scale_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -123,9 +142,10 @@ def _indexed_gate_rms_norm_scale_shift_kernel( stride_shift_row, stride_scale_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -162,7 +182,10 @@ def indexed_scale_shift_( rows, hidden_size = x.shape if rows == 0: return x - _indexed_scale_shift_kernel[(rows,)]( + _launch_row_chunks( + _indexed_scale_shift_kernel, + rows, + x.device.type, x, x, shift, @@ -192,7 +215,10 @@ def indexed_gate( rows, hidden_size = x.shape if rows == 0: return output - _indexed_gate_kernel[(rows,)]( + _launch_row_chunks( + _indexed_gate_kernel, + rows, + x.device.type, output, x, gate, @@ -229,7 +255,10 @@ def rms_norm_indexed_scale_shift( output = torch.empty_like(x) rows, hidden_size = x.shape if rows: - _rms_norm_indexed_scale_shift_kernel[(rows,)]( + _launch_row_chunks( + _rms_norm_indexed_scale_shift_kernel, + rows, + x.device.type, output, x, weight, @@ -274,7 +303,10 @@ def indexed_gate_rms_norm_scale_shift( modulated_out = torch.empty_like(residual) rows, hidden_size = residual.shape if rows: - _indexed_gate_rms_norm_scale_shift_kernel[(rows,)]( + _launch_row_chunks( + _indexed_gate_rms_norm_scale_shift_kernel, + rows, + residual.device.type, residual_out, modulated_out, residual, diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index cbaa8ba62..28298c0a2 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -608,15 +608,15 @@ class UlyssesParallelAttention: if attn_metadata.attn_mask is None: assert attn_metadata.joint_attn_mask is not None attn_metadata.attn_mask = torch.ones( - [query.shape[0], query.shape[1] - attn_metadata.joint_attn_mask.shape[1]], + [key.shape[0], key.shape[1] - attn_metadata.joint_attn_mask.shape[1]], dtype=torch.bool, - device=query.device, + device=key.device, ) elif attn_metadata.joint_attn_mask is None: attn_metadata.joint_attn_mask = torch.ones( - [query.shape[0], query.shape[1] - attn_metadata.attn_mask.shape[1]], + [key.shape[0], key.shape[1] - attn_metadata.attn_mask.shape[1]], dtype=torch.bool, - device=query.device, + device=key.device, ) attn_metadata.attn_mask = ( torch.cat([attn_metadata.joint_attn_mask, attn_metadata.attn_mask], dim=1) @@ -625,9 +625,10 @@ class UlyssesParallelAttention: ) if attn_metadata.attn_mask is not None: - # the final attn_mask is ready, the length should be aligedn with query length - assert attn_metadata.attn_mask.shape[1] == query.shape[1], ( - f"attn_mask length: {attn_metadata.attn_mask.shape[1]} != query length: {query.shape[1]}" + # A 2D mask describes keys. Joint attention may provide K/V-only + # context, so key and query lengths are not necessarily equal. + assert attn_metadata.attn_mask.shape[1] == key.shape[1], ( + f"attn_mask length: {attn_metadata.attn_mask.shape[1]} != key length: {key.shape[1]}" ) attn_metadata.attn_mask = attn_metadata.attn_mask.bool().contiguous() return query, key, value, attn_metadata, ctx diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index fd6f0237d..c30529520 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import copy +import json import math import os import random @@ -42,45 +43,57 @@ if TYPE_CHECKING: logger = init_logger(__name__) -def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]: +def normalize_omni_diffusion_kwargs(raw_kwargs: Mapping[str, Any]) -> dict[str, Any]: """Normalize legacy diffusion kwargs before config construction.""" - normalized = dict(kwargs) + config_kwargs = dict(raw_kwargs) + + dtype = config_kwargs.get("dtype") + if dtype is None: + config_kwargs["dtype"] = "auto" + elif isinstance(dtype, torch.dtype): + config_kwargs["dtype"] = str(dtype).removeprefix("torch.") + elif not isinstance(dtype, str): + raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(dtype).__name__}") # 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) + if "static_lora_scale" in config_kwargs: + if "lora_scale" not in config_kwargs: + config_kwargs["lora_scale"] = config_kwargs["static_lora_scale"] + config_kwargs.pop("static_lora_scale", None) + + diffusion_quantization = config_kwargs.pop("diffusion_quantization_config", None) + if config_kwargs.get("quantization_config") is None and diffusion_quantization is not None: + config_kwargs["quantization_config"] = diffusion_quantization # 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") + if "quantization" in config_kwargs and config_kwargs.get("quantization_config", None) is None: + config_kwargs["quantization_config"] = config_kwargs.pop("quantization") else: - normalized.pop("quantization", None) + config_kwargs.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") + if config_kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in config_kwargs: + config_kwargs["diffusion_kv_cache_dtype"] = config_kwargs.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") + config_kwargs.pop("kv_cache_dtype", None) + if config_kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in config_kwargs: + config_kwargs["diffusion_kv_cache_skip_steps"] = config_kwargs.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") + config_kwargs.pop("kv_cache_skip_steps", None) + if config_kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in config_kwargs: + config_kwargs["diffusion_kv_cache_skip_layers"] = config_kwargs.pop("kv_cache_skip_layers") else: - normalized.pop("kv_cache_skip_layers", None) + config_kwargs.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) - fastvideo_vsa_topk = normalized.pop("fastvideo_vsa_topk", None) + diffusion_attn_backend = config_kwargs.pop("diffusion_attention_backend", None) + fastvideo_vsa_topk = config_kwargs.pop("fastvideo_vsa_topk", None) if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None: - existing = normalized.get("diffusion_attention_config") - normalized["diffusion_attention_config"] = parse_attention_config( + existing = config_kwargs.get("diffusion_attention_config") + config_kwargs["diffusion_attention_config"] = parse_attention_config( existing, attention_backend=diffusion_attn_backend, fastvideo_vsa_topk=fastvideo_vsa_topk, @@ -88,21 +101,33 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any] # 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: + if "cache_backend" not in config_kwargs: 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" - elif normalized["cache_backend"] is None: + config_kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none" + elif config_kwargs["cache_backend"] is None: # Callers (e.g. example CLIs with `default=None`) pass an explicit # None for "no cache"; canonicalize it so every consumer sees the # declared `str` value instead of relying on per-model None handling. - normalized["cache_backend"] = "none" + config_kwargs["cache_backend"] = "none" + + cache_config = config_kwargs.get("cache_config") + if isinstance(cache_config, str): + try: + config_kwargs["cache_config"] = json.loads(cache_config) + except json.JSONDecodeError: + logger.warning("Invalid cache_config JSON, using backend defaults.") + config_kwargs.pop("cache_config", None) + + if config_kwargs.get("streaming_output") is None and config_kwargs.get("diffusion_streaming_output") is not None: + config_kwargs["streaming_output"] = config_kwargs["diffusion_streaming_output"] + config_kwargs.pop("diffusion_streaming_output", 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] = {} + if key in config_kwargs and config_kwargs[key] is None: + config_kwargs[key] = {} - return normalized + return config_kwargs def validate_host_weight_runtime_options(*, mode: object, root: object) -> None: @@ -439,6 +464,27 @@ class DiffusionParallelConfig: raise TypeError(f"Expected parallel config dict, got {type(data)!r}") return cls(**data) + @classmethod + def from_stage_overrides(cls, kwargs: Mapping[str, Any]) -> "DiffusionParallelConfig": + """Resolve nested or flat stage inputs through the config's own defaults.""" + parallel_config = kwargs.get("parallel_config") + if isinstance(parallel_config, cls): + return parallel_config + + valid_fields = {item.name for item in fields(cls) if item.init} + if isinstance(parallel_config, Mapping): + return cls.from_dict(dict(parallel_config)) + elif parallel_config is not None and hasattr(parallel_config, "__dict__"): + return cls.from_dict(dict(vars(parallel_config))) + elif parallel_config is not None: + raise TypeError( + f"parallel_config must be a mapping or DiffusionParallelConfig, got {type(parallel_config).__name__}" + ) + else: + values = {key: value for key, value in kwargs.items() if key in valid_fields} + + return cls(**{key: value for key, value in values.items() if value is not None}) + @dataclass class TransformerConfig: @@ -1550,15 +1596,18 @@ class OmniDiffusionConfig: raise @classmethod - def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig": - kwargs = normalize_omni_diffusion_kwargs(kwargs) - - # Filter kwargs to only include valid fields + def normalize_init_kwargs(cls, raw_kwargs: Mapping[str, Any]) -> dict[str, Any]: + config_kwargs = normalize_omni_diffusion_kwargs(raw_kwargs) valid_fields = {f.name for f in fields(cls)} - filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields} + # Remaining ``None`` values mean "unset" at the CLI/deploy boundary. + # Drop them so non-optional dataclass defaults are not overwritten. + # Fields where ``None`` has normalization semantics (for example dtype + # and nullable container inputs) are handled above before this filter. + return {key: value for key, value in config_kwargs.items() if key in valid_fields and value is not None} - instance = cls(**filtered_kwargs) - return instance + @classmethod + def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig": + return cls(**cls.normalize_init_kwargs(kwargs)) @dataclass diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 44c6efe66..eddcb1ad6 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from contextlib import nullcontext from typing import Any @@ -9,6 +9,7 @@ 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 diffusers.utils.accelerate_utils import apply_forward_hook from vllm.logger import init_logger from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard @@ -18,6 +19,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor GridSpec, TileTask, ) +from vllm_omni.diffusion.models.interface import DecodedChunkConsumer from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -45,9 +47,133 @@ class OmniAutoencoderKLWan(AutoencoderKLWan): return super().encode(x, return_dict=return_dict) def decode(self, z: torch.Tensor, return_dict: bool = True): + """Decode a Wan latent using the Diffusers-compatible full-tensor API.""" with self._execution_context(): return super().decode(z, return_dict=return_dict) + def _decode(self, z: torch.Tensor, return_dict: bool = True): + """Decode one latent batch, preserving the parent Diffusers contract.""" + _, _, _, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + + if getattr(self, "use_tiling", False) and (width > tile_latent_min_width or height > tile_latent_min_height): + return self.tiled_decode(z, return_dict=return_dict) + + return self._decode_temporal(z, return_dict=return_dict) + + @apply_forward_hook + def decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: + """Decode ``z`` while synchronously delivering temporal chunks. + + The ordinary :meth:`decode` contract remains the Diffusers contract. + Chunk callbacks deliberately reject a sliced batch: Diffusers slices + each sample before decoding, while a streaming consumer needs a stable + batch ownership contract. Callers can disable ``use_slicing`` or + invoke this capability once per sample. + + Every rank participating in distributed VAE execution must invoke this + method. Spatial-sharded execution invokes ``on_chunk`` only on the + output-owning rank. Chunks have shape ``(B, C, T, H, W)`` and + values clamped to ``[-1, 1]``. After a callback failure, remaining + chunks are decoded and discarded before the exception is re-raised. + """ + if not callable(on_chunk): + raise TypeError("on_chunk must be callable") + if getattr(self, "use_slicing", False) and z.shape[0] > 1: + raise ValueError( + "Wan chunked decode does not support use_slicing for batch size > 1; " + "disable VAE slicing or decode one sample at a time" + ) + + with self._execution_context(): + self._decode_with_chunks(z, on_chunk=on_chunk) + + def _decode_temporal( + self, + z: torch.Tensor, + *, + on_chunk: DecodedChunkConsumer | None = None, + return_dict: bool = True, + ): + """Decode the non-tiled temporal path, optionally streaming chunks.""" + self.clear_cache() + callback_error: BaseException | None = None + try: + x = self.post_quant_conv(z) + decoded_chunks: list[torch.Tensor] = [] + for i in range(z.shape[2]): + self._conv_idx = [0] + if i == 0: + chunk = self.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + else: + chunk = self.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + ) + + if on_chunk is None: + decoded_chunks.append(chunk) + continue + + if callback_error is None: + try: + if self.config.patch_size is not None: + chunk = unpatchify(chunk, patch_size=self.config.patch_size) + on_chunk(torch.clamp(chunk, min=-1.0, max=1.0)) + except BaseException as exc: + # A previously installed spatial-shard decoder still + # runs collectives, even for latents below the tiling + # threshold. Drain before surfacing consumer failures. + callback_error = exc + del chunk + + if callback_error is not None: + raise callback_error + if on_chunk is not None: + return None + + out = torch.cat(decoded_chunks, dim=2) + if self.config.patch_size is not None: + out = unpatchify(out, patch_size=self.config.patch_size) + out = torch.clamp(out, min=-1.0, max=1.0) + if not return_dict: + return (out,) + return DecoderOutput(sample=out) + finally: + self.clear_cache() + + @staticmethod + def _discard_decoded_chunk(_chunk: torch.Tensor) -> None: + """Consume a non-owner rank's chunk without retaining its tensor.""" + + def _decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: + """Run the temporal streaming primitive without API or hook handling.""" + distributed_enabled = False + if hasattr(self, "is_distributed_enabled"): + distributed_enabled = self.is_distributed_enabled() + + if getattr(self, "use_tiling", False): + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + exceeds_tile = z.shape[-1] > tile_latent_min_width or z.shape[-2] > tile_latent_min_height + if exceeds_tile: + if distributed_enabled: + self._tiled_decode_with_chunks(z, on_chunk=on_chunk) + return + raise ValueError("Wan chunk callbacks are unsupported for spatial tiling") + + callback = on_chunk + if distributed_enabled and dist.get_rank(group=self.distributed_executor.group) != 0: + callback = self._discard_decoded_chunk + self._decode_temporal(z, on_chunk=callback) + class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): @classmethod @@ -78,7 +204,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): blend_height = self.tile_sample_min_height - tile_sample_stride_height blend_width = self.tile_sample_min_width - tile_sample_stride_width - tiletask_list = [] + tiletask_list: list[TileTask] = [] for i in range(0, height, tile_latent_stride_height): for j in range(0, width, tile_latent_stride_width): time_list = [] @@ -151,7 +277,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): blend_height = tile_latent_min_height - tile_latent_stride_height blend_width = tile_latent_min_width - tile_latent_stride_width - tiletask_list = [] + tiletask_list: list[TileTask] = [] temporal_compression = self.config.scale_factor_temporal for i in range(0, height, tile_sample_stride_height): for j in range(0, width, tile_sample_stride_width): @@ -307,8 +433,10 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): if not self.is_distributed_enabled(): return super().tiled_decode(z, return_dict=return_dict) - if self._spatial_shard_decode_enabled(z): + spatial_shard_enabled = self._spatial_shard_decode_enabled(z) + if spatial_shard_enabled: split_dim = self._spatial_shard_decode_split_dim() + assert split_dim is not None logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim) return wan_spatial_shard.spatial_shard_decode( self, @@ -329,6 +457,27 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): return DecoderOutput(sample=result) + def _tiled_decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: + """Decode tiled latents and deliver temporal chunks to ``on_chunk``.""" + if not self.is_distributed_enabled(): + raise ValueError("Wan chunk callbacks are unsupported for spatial tiling") + + spatial_shard_enabled = self._spatial_shard_decode_enabled(z) + if not spatial_shard_enabled: + raise ValueError("Wan chunk callbacks require spatial-shard decode when tiling is enabled") + + split_dim = self._spatial_shard_decode_split_dim() + assert split_dim is not None + logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim) + wan_spatial_shard.spatial_shard_decode( + self, + z, + group=self.distributed_executor.group, + return_dict=True, + split_dim=split_dim, + on_chunk=on_chunk, + ) + def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: """ Encode using distributed VAE executor. diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py index 07edad694..ac5ea472b 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # # The halo-exchange spatial-parallel decode here is adapted from SGLang's # spatial-parallel VAE decode @@ -24,7 +24,7 @@ from contextlib import nullcontext from contextvars import ContextVar from dataclasses import dataclass from types import MethodType -from typing import Any +from typing import Any, cast import torch import torch.distributed as dist @@ -34,6 +34,8 @@ 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.models.interface import DecodedChunkConsumer + logger = init_logger(__name__) @@ -644,7 +646,7 @@ def _replace_child( ) return if isinstance(child, nn.ZeroPad2d): - padding = tuple(int(p) for p in child.padding) + padding = cast(tuple[int, int, int, int], 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 @@ -794,7 +796,8 @@ def spatial_shard_decode( group: dist.ProcessGroup, return_dict: bool = True, split_dim: str = "height", -) -> DecoderOutput | tuple[torch.Tensor]: + on_chunk: DecodedChunkConsumer | None = None, +) -> DecoderOutput | tuple[torch.Tensor] | None: install_wan_spatial_shard_decode(vae, group, split_dim=split_dim) if z.shape[2] == 0: @@ -806,6 +809,7 @@ def spatial_shard_decode( produce_output = world_size <= 1 or rank == 0 vae.clear_cache() + callback_error: BaseException | None = None try: context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext() with context: @@ -820,18 +824,33 @@ def spatial_shard_decode( first_chunk=(i == 0), ) if produce_output: - decoded_chunks.append(chunk) + if on_chunk is not None and callback_error is None: + try: + if vae.config.patch_size is not None: + chunk = unpatchify(chunk, patch_size=vae.config.patch_size) + on_chunk(torch.clamp(chunk, min=-1.0, max=1.0)) + except BaseException as exc: + # Keep all ranks in the temporal collective loop; + # surface the callback failure only after decode. + callback_error = exc + elif on_chunk is None: + decoded_chunks.append(chunk) + del chunk - if produce_output: + if produce_output and on_chunk is None: 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 callback_error is not None: + raise callback_error + if on_chunk is not None: + return None + if not produce_output: + out = z.new_zeros(0) if not return_dict: return (out,) return DecoderOutput(sample=out) diff --git a/vllm_omni/diffusion/lora/layers/base_linear.py b/vllm_omni/diffusion/lora/layers/base_linear.py index c1f941a10..1a78c6ad4 100644 --- a/vllm_omni/diffusion/lora/layers/base_linear.py +++ b/vllm_omni/diffusion/lora/layers/base_linear.py @@ -25,6 +25,9 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): lora_a_stacked: tuple[torch.Tensor, ...] lora_b_stacked: tuple[torch.Tensor, ...] + # Mask saved while the adapter is suspended; None means "not suspended". + _diffusion_lora_suspended_slices: tuple[bool, ...] | None = None + def _move_lora_buffers(self, device: torch.device) -> None: self.lora_a_stacked = tuple(tensor.to(device=device) for tensor in self.lora_a_stacked) self.lora_b_stacked = tuple(tensor.to(device=device) for tensor in self.lora_b_stacked) @@ -54,11 +57,14 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): object.__setattr__(self, "_diffusion_base_layer_ref", base_layer) n_slices = getattr(self, "n_slices", 1) self._diffusion_lora_active_slices = (False,) * int(n_slices) + self._diffusion_lora_suspended_slices = None def reset_lora(self, index: int): super().reset_lora(index) n_slices = getattr(self, "n_slices", 1) self._diffusion_lora_active_slices = (False,) * int(n_slices) + # The upload is gone, so a saved mask no longer describes the buffers. + self._diffusion_lora_suspended_slices = None def set_lora( self, @@ -82,6 +88,30 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): # Single-slice layer. self._diffusion_lora_active_slices = (True,) + def suspend_lora(self) -> None: + """Gate every slice off without dropping the uploaded weights. + + `apply()` reads the stacked buffers only for slices this mask marks + active, so a suspended adapter costs nothing while staying resident. + A saved mask means the layer is already suspended; keep the original + rather than overwriting it with the all-False mask now in place. + """ + if self._diffusion_lora_suspended_slices is not None: + return + self._diffusion_lora_suspended_slices = self._diffusion_lora_active_slices + self._diffusion_lora_active_slices = (False,) * len(self._diffusion_lora_active_slices) + + def resume_lora(self) -> None: + """Restore the mask captured by `suspend_lora()`. + + A layer registered after the suspend, or one whose buffers were torn + down meanwhile, has no mask to restore and keeps its own inactive one. + """ + saved = self._diffusion_lora_suspended_slices + if saved is not None: + self._diffusion_lora_active_slices = saved + self._diffusion_lora_suspended_slices = None + def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: """ override: Use simple matmul instead of punica_wrapper.add_lora_linear(). diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py index 3f74b4853..a13363c7a 100644 --- a/vllm_omni/diffusion/lora/loader.py +++ b/vllm_omni/diffusion/lora/loader.py @@ -10,11 +10,11 @@ from diffusers.loaders.lora_conversion_utils import ( _convert_non_diffusers_qwen_lora_to_diffusers, _convert_non_diffusers_wan_lora_to_diffusers, ) -from huggingface_hub import hf_hub_download from safetensors.torch import load_file from vllm.logger import init_logger from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline +from vllm_omni.transformers_utils.repo_utils import hf_api logger = init_logger(__name__) @@ -131,7 +131,7 @@ def _load_lora_state_dict( # finally, we try to load it from the internet try: - model_file = hf_hub_download( + model_file = hf_api().hf_hub_download( pretrained_model_name_or_path, filename=weights_name, subfolder=subfolder, diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 1b63692fc..2f6d47780 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -48,6 +48,11 @@ class DiffusionLoRAManager: # Valid max allowed ranks for LoRA in vLLM _VALID_MAX_RANKS: list[int] = sorted(get_args(MaxLoRARanks)) + # Adapter whose weights are still uploaded while gated off. Class-level so + # a manager built without __init__ (some tests use object.__new__) reads as + # "nothing suspended" instead of raising. + _suspended_adapter_id: int | None = None + def __init__( self, pipeline: nn.Module, @@ -493,9 +498,11 @@ class DiffusionLoRAManager: fully_sharded_loras=False, ) - # Recreate per-layer buffers with the new maximum rank. + # Recreate per-layer buffers with the new maximum rank. The previous + # upload is gone, so nothing may be re-armed afterwards. for lora_layer in self._lora_modules.values(): lora_layer.create_lora_weights(max_loras=1, lora_config=lora_config, model_config=None) + self._suspended_adapter_id = None # Re-apply active adapter if needed (buffers were reset). if self._active_adapter_id is not None: @@ -679,18 +686,32 @@ class DiffusionLoRAManager: def _reset_lora_layers(self) -> None: for lora_layer in self._lora_modules.values(): lora_layer.reset_lora(0) + self._suspended_adapter_id = None def _activate_adapter(self, adapter_id: int, scale: float) -> None: if self._is_active_at_scale(adapter_id, scale): logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale) return + if self._suspended_adapter_id == adapter_id and self._adapter_scales.get( + adapter_id + ) == DiffusionLoRAManager._get_rounded_scale(scale): + # Weights are still uploaded from an earlier activation; re-arming + # the masks avoids rebuilding and re-uploading every layer. + for lora_layer in self._lora_modules.values(): + lora_layer.resume_lora() + self._suspended_adapter_id = None + self._active_adapter_id = adapter_id + logger.debug("Re-armed suspended adapter %d", adapter_id) + return + logger.info("Activating adapter: id=%d", adapter_id) lora_model = self._registered_adapters[adapter_id] # Binding overwrites slot 0 incrementally. Invalidate the fast-path # state before the first mutation and leave every wrapper inactive if # any set_lora() call or model validator fails. self._active_adapter_id = None + self._suspended_adapter_id = None try: self._bind_adapter_weights(lora_model, scale) except Exception: @@ -704,8 +725,10 @@ class DiffusionLoRAManager: if self._active_adapter_id is None: logger.debug("All adapters already inactive") return - logger.info("Deactivating all adapters: %d layers", len(self._lora_modules)) - self._reset_lora_layers() + logger.info("Suspending all adapters: %d layers", len(self._lora_modules)) + for lora_layer in self._lora_modules.values(): + lora_layer.suspend_lora() + self._suspended_adapter_id = self._active_adapter_id self._active_adapter_id = None logger.debug("All adapters deactivated") @@ -772,6 +795,11 @@ class DiffusionLoRAManager: if self._active_adapter_id == adapter_id: self._deactivate_all_adapters() + if self._suspended_adapter_id == adapter_id: + # The adapter is going away, so the upload can never be resumed. + # Tear it down instead of leaving it in the stacked buffers. + self._reset_lora_layers() + del self._registered_adapters[adapter_id] self._adapter_scales.pop(adapter_id, None) self._adapter_access_order.pop(adapter_id, None) diff --git a/vllm_omni/diffusion/model_loader/hub_prefetch.py b/vllm_omni/diffusion/model_loader/hub_prefetch.py index 3594aa4e1..a1e893b3b 100644 --- a/vllm_omni/diffusion/model_loader/hub_prefetch.py +++ b/vllm_omni/diffusion/model_loader/hub_prefetch.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Best-effort HuggingFace Hub prefetch for multi-subfolder pipelines. @@ -297,7 +297,7 @@ def prefetch_subfolders( logger.info("Prefetching %s subfolders: %s", model, subfolders) try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except ImportError: # pragma: no cover - huggingface_hub is a hard dep logger.debug("huggingface_hub unavailable; skipping prefetch of %s", model) return @@ -334,7 +334,7 @@ def prefetch_subfolders( for attempt in range(1, _PREFETCH_MAX_ATTEMPTS + 1): try: with _repo_prefetch_lock(model): - snapshot_download( + hf_api().snapshot_download( repo_id=model, revision=revision, allow_patterns=allow_patterns, diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index 38df83846..c91d2ead5 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -11,6 +11,10 @@ class DiffusionModelMetadata: supports_multimodal_inputs: bool = False max_multimodal_image_inputs: int | None = None supports_mixed_reference_inputs: bool = False + # Multipart controls are exposed as ``control_reference`` plus + # ``control_type`` by the video API. A pipeline that opts in receives the + # persisted upload through ``extra_args[control_type]["control_path"]``. + supported_control_upload_types: tuple[str, ...] = () attention_mask_free: bool = False final_output_type: str | None = None @@ -94,8 +98,14 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { "LongCatVideoAvatarPipeline": DiffusionModelMetadata(final_output_type="video"), "MagiHumanPipeline": DiffusionModelMetadata(final_output_type="video"), "DreamIDOmniPipeline": DiffusionModelMetadata(final_output_type="video"), - "Cosmos3OmniDiffusersPipeline": DiffusionModelMetadata(final_output_type="video"), - "Cosmos3OmniPipeline": DiffusionModelMetadata(final_output_type="video"), + "Cosmos3OmniDiffusersPipeline": DiffusionModelMetadata( + supported_control_upload_types=("edge", "blur", "depth", "seg", "wsm"), + final_output_type="video", + ), + "Cosmos3OmniPipeline": DiffusionModelMetadata( + supported_control_upload_types=("edge", "blur", "depth", "seg", "wsm"), + final_output_type="video", + ), "SanaVideoPipeline": DiffusionModelMetadata(final_output_type="video"), "SanaImageToVideoPipeline": DiffusionModelMetadata(final_output_type="video"), "SanaWmPipeline": DiffusionModelMetadata( @@ -111,6 +121,7 @@ _DIFFUSION_MODEL_METADATA_ALIASES = { "LTX2DistilledOneStagePipeline": "LTX2DistilledPipeline", "LTX2DistilledTwoStagePipeline": "LTX2DistilledPipeline", "LingBotWorldCausalDMDPipeline": "LingBotVideoPipeline", + "BooguImageTurboPipeline": "BooguImagePipeline", } diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py index 2025859ed..66031ca3a 100644 --- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py +++ b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py @@ -35,6 +35,7 @@ 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 from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.models.utils import make_attention_mask from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: @@ -1554,17 +1555,11 @@ class BooguImageTransformer2DModel(nn.Module): img_hidden_states = combined_img_hidden_states # Joint mask for [instruct + image]. - max_seq_len = max(seq_lengths) - joint_attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool) - for i, seq_len in enumerate(seq_lengths): - joint_attention_mask[i, :seq_len] = True + joint_attention_mask = make_attention_mask(hidden_states, seq_lengths) # Dual-stream (double-stream) stage. if self.num_double_stream_layers > 0: - max_img_len = max(combined_img_seq_lengths) - img_attention_mask = hidden_states.new_zeros(batch_size, max_img_len, dtype=torch.bool) - for i, img_seq_len in enumerate(combined_img_seq_lengths): - img_attention_mask[i, :img_seq_len] = True + img_attention_mask = make_attention_mask(hidden_states, combined_img_seq_lengths) for layer in self.double_stream_layers: img_hidden_states, instruct_hidden_states = layer( diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index a171e2d3b..49a4e4027 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -3522,7 +3522,12 @@ class Cosmos3OmniDiffusersPipeline( control_chunks_per_hint[key].append(control[:, :, current_conditional_frames:]) if not is_output_rank: - return DiffusionOutput(output={"video": output_video}, custom_output={"fps": frame_rate}) + return DiffusionOutput( + output={ + "payload": {"video": output_video}, + "metadata": {"video": {"fps": frame_rate}}, + }, + ) full_output = torch.cat(output_chunks, dim=2)[:, :, :total_frames] full_controls = { diff --git a/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py b/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py index 5d6ff8dd9..8c1337ec9 100644 --- a/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py +++ b/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Cosmos3 sound tokenizer integration. The tokenizer decodes model-generated sound latents for video+audio output. It @@ -338,9 +338,9 @@ class Cosmos3SoundTokenizer: model_root = str(model_path) if model_path and os.path.isdir(model_path) else None if model_root is None and model_path and not explicit_avae_path: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_root = snapshot_download( + model_root = hf_api().snapshot_download( repo_id=str(model_path), revision=getattr(od_config, "revision", None), allow_patterns=[ diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index c946eb4ca..fe7a34c8b 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -705,7 +705,21 @@ class Cosmos3CrossAttention(nn.Module): # standard reverse-all-to-all path (no joint-output splitting). joint_q = q.new_empty(B, 0, self.num_heads_local, self.head_dim) + gen_mask = None + joint_mask = None + if is_forward_context_available(): + ctx = get_forward_context() + if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size + gen_mask = torch.ones(B, padded_seq_len, dtype=torch.bool, device=q.device) + gen_mask[:, ctx.sp_original_seq_len :] = False + # UND K/V are trimmed to the common real text length before GEN + # attention, so every remaining joint key is valid. + joint_mask = torch.ones(B, k_und.shape[1], dtype=torch.bool, device=q.device) + attn_metadata = AttentionMetadata( + attn_mask=gen_mask, + joint_attn_mask=joint_mask, joint_query=joint_q, joint_key=k_und, joint_value=v_und, @@ -1055,6 +1069,13 @@ class Cosmos3GenSPPrepare(nn.Module): freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # A transfer denoising step may run a full-control branch followed by a + # shorter no-control branch in one ForwardContext. Auto-padding metadata + # belongs to one branch only and must not truncate a later gather. + if is_forward_context_available(): + ctx = get_forward_context() + ctx.sp_padding_size = 0 + ctx.sp_original_seq_len = None return hidden_gen, freqs_cos, freqs_sin @@ -1112,9 +1133,9 @@ class Cosmos3VFMTransformer(nn.Module): _sp_plan = { "gen_sp_prepare": { - 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), - 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), - 2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 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), + 2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), }, "gen_sp_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3), } @@ -1586,47 +1607,6 @@ class Cosmos3VFMTransformer(nn.Module): self.cached_kv = None self.cached_freqs_gen = None - @staticmethod - 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 = [] - 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: - detail_parts.append(f"sound tokens {s_sound}") - detail = " = " + " + ".join(detail_parts) if len(detail_parts) > 1 else "" - adjust_detail = ( - "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_control or has_action or has_sound - else ( - "Adjust the spatial resolution so that " - "t * ceil(h/patch) * ceil(w/patch) is a multiple " - "of ulysses_degree." - ) - ) - raise ValueError( - f"GEN sequence length ({s_gen}{detail}) must be divisible by " - f"ulysses_degree ({ulysses_size}). {adjust_detail}" - ) - def sound_latent_frames_for_sequence_parallel( self, *, @@ -1907,17 +1887,6 @@ class Cosmos3VFMTransformer(nn.Module): # 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=1 if use_multi_control_attention else ulysses_size, - ) freqs_cos, freqs_sin = self.cached_freqs_gen if not use_multi_control_attention: hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index bf19f0a0f..09bd536cd 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -21,7 +21,6 @@ from contextlib import contextmanager import numpy as np import torch import torch.nn as nn -from huggingface_hub import hf_hub_download from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel from vllm.model_executor.model_loader.weight_utils import default_weight_loader @@ -67,6 +66,7 @@ from vllm_omni.experimental.world_models.session_state import ( resolve_session_state_config, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.transformers_utils.repo_utils import hf_api logger = logging.getLogger(__name__) MAX_DREAMZERO_SESSIONS = 64 @@ -581,7 +581,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): return json.load(f) try: - json_path = hf_hub_download(model_path, relative_path) + json_path = hf_api().hf_hub_download(model_path, relative_path) with open(json_path) as f: return json.load(f) except Exception: diff --git a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py index 3fcbce203..6116572ee 100644 --- a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py +++ b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import logging -import os from typing import Any import torch @@ -138,26 +137,22 @@ class Gr00tN1d7ActionHead(nn.Module): state_features: [B, state_horizon, input_embedding_dim] embodiment_id: [B] (embodiment IDs) backbone_output: Output from the backbone model + options: ``generator`` (a ``torch.Generator`` on the model device) draws + the initial action noise; ``None`` leaves ``torch.randn`` on the + global RNG, as upstream Isaac-GR00T does. The ``rtc_*`` keys drive + real-time chunking when ``action_input`` carries a previous chunk. """ 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, - ) + generator = options.get("generator") if options else None + actions = torch.randn( + size=(batch_size, self.config.action_horizon, self.action_dim), + dtype=vl_embeds.dtype, + device=device, + generator=generator, + ) dt = 1.0 / self.num_inference_timesteps vel_strength = torch.ones_like(actions) diff --git a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py index ade5eac83..d30d1c829 100644 --- a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py +++ b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import json import os import re import warnings +from collections.abc import Callable from copy import deepcopy from pathlib import Path from typing import Any @@ -46,20 +47,119 @@ class LetterBoxTransform: return transforms.functional.pad(img, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0) +def _smallest_edge_resize(image: np.ndarray, max_size: int) -> np.ndarray: + """albumentations 1.4.18 ``SmallestMaxSize(max_size, interpolation=cv2.INTER_AREA)``. + + Scale off the shortest edge, skip entirely at scale 1.0, round each dimension + independently; ``cv2.resize`` takes ``dsize`` as ``(width, height)``. + """ + import cv2 + + height, width = image.shape[:2] + scale = max_size / float(min(height, width)) + if scale == 1.0: + return image + new_height, new_width = (round(dim * scale) for dim in (height, width)) + if (new_height, new_width) == (height, width): + return image + return cv2.resize(image, dsize=(new_width, new_height), interpolation=cv2.INTER_AREA) + + +def _fractional_center_crop(image: np.ndarray, fraction: float) -> np.ndarray: + """Isaac-GR00T ``FractionalCenterCrop``: ``int()`` truncation (not rounding), floor 1, centred.""" + height, width = image.shape[:2] + crop_height = max(1, int(height * fraction)) + crop_width = max(1, int(width * fraction)) + y_min = (height - crop_height) // 2 + x_min = (width - crop_width) // 2 + return image[y_min : y_min + crop_height, x_min : x_min + crop_width] + + +def _letter_box_pad(image: np.ndarray) -> np.ndarray: + """Isaac-GR00T ``LetterBoxPad``: zero-pad to square, centred, remainder on the far side.""" + import cv2 + + height, width = image.shape[:2] + if height == width: + return image + longest = max(height, width) + pad_h, pad_w = longest - height, longest - width + top, left = pad_h // 2, pad_w // 2 + return cv2.copyMakeBorder(image, top, pad_h - top, left, pad_w - left, cv2.BORDER_CONSTANT, value=0) + + +class AlbumentationsEvalImageTransform: + """Eval transform of the Isaac-GR00T ``use_albumentations=True`` recipe, without albumentations. + + ``[LetterBoxPad] -> SmallestMaxSize(INTER_AREA) -> FractionalCenterCrop -> SmallestMaxSize(INTER_AREA)``, + transcribed from the eval branch of Isaac-GR00T ``build_image_transformations_albumentations`` and + albumentations 1.4.18, so a checkpoint trained with that recipe sees the same pixels at inference. + Per frame: HWC uint8 RGB array in, ``(C, H, W)`` uint8 tensor out (the ``_get_vlm_inputs`` contract). + """ + + def __init__(self, max_size: int, fraction: float, letter_box_transform: bool): + self.max_size = int(max_size) + self.fraction = float(fraction) + self.letter_box_transform = letter_box_transform + + def __call__(self, image: np.ndarray) -> torch.Tensor: + array = np.asarray(image) + if array.dtype != np.uint8 or array.ndim != 3 or array.shape[2] != 3: + raise ValueError( + "GR00T eval image transform expects an HWC uint8 RGB frame, " + f"got shape={tuple(array.shape)} dtype={array.dtype}" + ) + if self.letter_box_transform: + array = _letter_box_pad(array) + array = _smallest_edge_resize(array, self.max_size) + array = _fractional_center_crop(array, self.fraction) + array = _smallest_edge_resize(array, self.max_size) + # The crop is a non-contiguous view; the caller stacks frames right after this. + return torch.from_numpy(np.ascontiguousarray(array)).permute(2, 0, 1) + + 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( + *, + shortest_image_edge: int | None = None, + crop_fraction: float | None = None, + use_albumentations: bool = False, + letter_box_transform: bool = False, +) -> Callable[[np.ndarray], torch.Tensor]: + """Deterministic eval/inference image transform, selected the way Isaac-GR00T's processor does. + + ``use_albumentations=True``: ``[letterbox] -> resize (shortest edge) -> fractional centre crop -> resize`` + (:class:`AlbumentationsEvalImageTransform`). Otherwise the torchvision recipe + ``[letterbox] -> resize -> centre crop -> resize``. Letterboxing is gated by ``letter_box_transform`` + on both paths. + """ + if use_albumentations: + if crop_fraction is None: + if image_crop_size is None or image_target_size is None: + raise ValueError("image_crop_size and image_target_size are required when crop_fraction is None") + fraction = image_crop_size[0] / image_target_size[0] + else: + fraction = crop_fraction + if shortest_image_edge is None: + if image_target_size is None: + raise ValueError("image_target_size is required when shortest_image_edge is None") + max_size = image_target_size[0] + else: + max_size = shortest_image_edge + return AlbumentationsEvalImageTransform(max_size, fraction, letter_box_transform) + + steps: list[Callable[..., torch.Tensor]] = [transforms.ToImage()] + if letter_box_transform: + steps.append(LetterBoxTransform()) + steps.extend( [ - transforms.ToImage(), - LetterBoxTransform(), transforms.Resize(size=image_target_size), transforms.CenterCrop(size=image_crop_size), transforms.Resize(size=image_target_size), ] ) + return transforms.Compose(steps) logger = init_logger(__name__) @@ -191,6 +291,9 @@ class Gr00tN1d7Processor(ProcessorMixin): exclude_state: bool = False, # Normalization use_mean_std: bool = False, + # Image pipeline selection; must match the recipe the checkpoint was trained with + use_albumentations: bool = False, + letter_box_transform: bool = False, **kwargs, # absorb deprecated training-only keys from saved processor_config.json ): if transformers_loading_kwargs is None: @@ -237,6 +340,8 @@ class Gr00tN1d7Processor(ProcessorMixin): self.embodiment_id_mapping[k] = v self.shortest_image_edge = shortest_image_edge self.crop_fraction = crop_fraction + self.use_albumentations = use_albumentations + self.letter_box_transform = letter_box_transform self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {} @@ -244,6 +349,10 @@ class Gr00tN1d7Processor(ProcessorMixin): self.eval_image_transform = _build_eval_image_transform( image_target_size, image_crop_size, + shortest_image_edge=shortest_image_edge, + crop_fraction=crop_fraction, + use_albumentations=use_albumentations, + letter_box_transform=letter_box_transform, ) self._collator = self.data_collator_class( model_name=model_name, @@ -428,7 +537,7 @@ class Gr00tN1d7Processor(ProcessorMixin): image_keys: list[str], images: list[Image.Image], masks: dict[str, list[np.ndarray]] | None, - image_transform: transforms.Compose, + image_transform: Callable[[np.ndarray], torch.Tensor], language: str, ): temporal_stacked_images = {} @@ -470,6 +579,8 @@ class Gr00tN1d7Processor(ProcessorMixin): "color_jitter_params": self.color_jitter_params, "shortest_image_edge": self.shortest_image_edge, "crop_fraction": self.crop_fraction, + "use_albumentations": self.use_albumentations, + "letter_box_transform": self.letter_box_transform, "model_name": self.model_name, "model_type": self.model_type, "formalize_language": self.formalize_language, diff --git a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py index 1154dcf71..87108522b 100644 --- a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py +++ b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections.abc import Iterable, Mapping from typing import Any @@ -13,6 +13,7 @@ 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 +from vllm_omni.inputs.data import OmniDiffusionSamplingParams logger = init_logger(__name__) @@ -30,6 +31,10 @@ class Gr00tN1d7Pipeline(nn.Module): 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"]`. + + The request's `seed` (materialised by the runner into `sampling_params.generator`) + draws the flow-matching noise, so repeating a request with the same seed reproduces + its action chunk. """ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: @@ -123,7 +128,8 @@ class Gr00tN1d7Pipeline(nn.Module): self.reset() policy_obs = _normalize_observation(robot_obs, language_key=self.policy.language_key) - result = self.policy.get_action(policy_obs) + generator = _noise_generator(req.sampling_params, self.device) + result = self.policy.get_action(policy_obs, {"generator": generator} if generator is not None else None) 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.") @@ -132,6 +138,35 @@ class Gr00tN1d7Pipeline(nn.Module): return DiffusionOutput(output={"actions": _to_float32_action_dict(actions)}) +def _noise_generator(sampling_params: OmniDiffusionSamplingParams, device: str) -> torch.Generator | None: + """Pick the generator that draws the action head's initial noise. + + The runner materialises ``sampling_params.seed`` into ``sampling_params.generator`` on + the model device before every forward, and ``OmniDiffusionRequest`` assigns a random + seed when the caller omits one, so in serving the generator is always present. The + seed fallback serves callers that bypass the runner (direct use, unit tests). ``None`` + leaves the action head on the global RNG, as upstream Isaac-GR00T runs. GR00T draws one + noise tensor per request, so a generator in any other form (for example a list of several + generators) is ignored with a warning and the seed fallback applies. + """ + generator = sampling_params.generator + if isinstance(generator, list) and len(generator) == 1: + generator = generator[0] + if isinstance(generator, torch.Generator): + return generator + if generator is not None: + received = f"a list of {len(generator)} generators" if isinstance(generator, list) else type(generator).__name__ + logger.warning( + "GR00T expects a single torch.Generator per request but got %s; ignoring it. The initial noise " + "falls back to sampling_params.seed=%s, or to the global RNG when that is None.", + received, + sampling_params.seed, + ) + if sampling_params.seed is not None: + return torch.Generator(device=device).manual_seed(int(sampling_params.seed)) + return None + + def _normalize_observation(robot_obs: Mapping[str, Any], *, language_key: str) -> dict[str, Any]: obs: dict[str, Any] = {} if "video" in robot_obs: diff --git a/vllm_omni/diffusion/models/gr00t/policy.py b/vllm_omni/diffusion/models/gr00t/policy.py index c16b34463..118f57545 100644 --- a/vllm_omni/diffusion/models/gr00t/policy.py +++ b/vllm_omni/diffusion/models/gr00t/policy.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from pathlib import Path from typing import Any @@ -328,7 +328,8 @@ class Gr00tPolicy: Args: observation: Batched observation dictionary - options: Optional parameters (currently unused) + options: Optional per-request model parameters, forwarded verbatim to + the action head (e.g. ``generator``) Returns: Tuple of (actions_dict, info_dict) @@ -351,7 +352,7 @@ class Gr00tPolicy: # Step 4: Run model inference to predict actions with torch.inference_mode(): - model_pred = self.model.get_action(**collated_inputs) + model_pred = self.model.get_action(**collated_inputs, options=options) normalized_action = model_pred["action_pred"].float() # Step 5: Decode actions from normalized space back to physical units diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index 45c6d0693..610fb2b8f 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -76,9 +76,9 @@ def load_json_config(model_path: str, subfolder: str, filename: str, local_files return json.load(f) else: try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - config_path = hf_hub_download( + config_path = hf_api().hf_hub_download( repo_id=model_path, filename=f"{subfolder}/{filename}", ) @@ -1358,6 +1358,8 @@ class HeliosPipeline( """Single-stage denoising loop for one chunk.""" batch_size = latents.shape[0] do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None + # Distilled (DMD) needs the chunk-start noise for renoise; mirror step_scheduler. + stage1_start_latents = latents with self.progress_bar(total=len(timesteps)) as pbar: for i, t in enumerate(timesteps): @@ -1420,7 +1422,20 @@ class HeliosPipeline( cfg_normalize=False, ) - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + if self.is_distilled: + latents = self.scheduler.step( + noise_pred, + t, + latents, + return_dict=False, + cur_sampling_step=i, + dmd_noisy_tensor=stage1_start_latents, + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=timesteps, + )[0] + else: + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) pbar.update() diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index 75cc5a0d2..dac9a0b1d 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from typing import ( TYPE_CHECKING, @@ -9,6 +10,7 @@ from typing import ( ClassVar, Literal, Protocol, + TypeGuard, runtime_checkable, ) @@ -46,6 +48,34 @@ class SupportAudioOutput(Protocol): support_audio_output: ClassVar[bool] = True +DecodedChunkConsumer = Callable[["torch.Tensor"], None] + + +@runtime_checkable +class SupportsChunkedVAEDecode(Protocol): + """Optional capability for VAEs that stream decoded temporal chunks. + + Implementations must drain decoding before surfacing callback errors. In a + distributed VAE, every rank invokes the method so collectives stay in + lockstep, while ``on_chunk`` is called only on the rank that owns output. + """ + + def decode_with_chunks( + self, + z: torch.Tensor, + *, + on_chunk: DecodedChunkConsumer, + ) -> None: + """Decode ``z`` and synchronously deliver committed chunks in order.""" + ... + + +def supports_chunked_vae_decode(vae: object) -> TypeGuard[SupportsChunkedVAEDecode]: + """Return whether ``vae`` exposes the optional chunked decode capability.""" + + return isinstance(vae, SupportsChunkedVAEDecode) + + @runtime_checkable class SupportsStepExecution(Protocol): """State-driven step-level execution protocol for diffusion pipelines. diff --git a/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py b/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py index 33b0e2a45..6cb2dd7fb 100644 --- a/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py +++ b/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -376,9 +376,9 @@ def _ensure_local_dir(model: str | os.PathLike[str], allow_patterns: list[str] | model_path = Path(model) if model_path.exists(): return model_path - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return Path(snapshot_download(str(model), allow_patterns=allow_patterns)) + return Path(hf_api().snapshot_download(str(model), allow_patterns=allow_patterns)) def _ensure_avatar_model_index(model_dir: Path, use_int8: bool) -> None: diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index c00fdd835..72e63a22f 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -20,7 +20,6 @@ from diffusers.pipelines.ltx2 import LTX2TextConnectors from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder from diffusers.video_processor import VideoProcessor -from huggingface_hub import hf_hub_download from safetensors import safe_open from safetensors.torch import load_file from transformers import AutoModelForImageTextToText, AutoTokenizer, Gemma3ForConditionalGeneration @@ -32,6 +31,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.offloader.module_collector import ModuleDiscovery +from vllm_omni.transformers_utils.repo_utils import hf_api if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -279,7 +279,7 @@ def resolve_ltx_artifact( # independently pinned artifact revision. revision = model_revision if model == repo_id else artifact_revision try: - return hf_hub_download( + return hf_api().hf_hub_download( repo_id=repo_id, filename=filename, revision=revision, @@ -351,7 +351,7 @@ def _load_ltx_metadata_json(model: str, filename: str, revision: str | None = No return {} else: try: - path = hf_hub_download(repo_id=model, filename=filename, revision=revision) + path = hf_api().hf_hub_download(repo_id=model, filename=filename, revision=revision) except Exception: return {} try: @@ -529,7 +529,7 @@ def _detect_vocoder_output_sample_rate(model: str, revision: str | None = 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", revision=revision) + vocoder_config_path = hf_api().hf_hub_download(model, "vocoder/config.json", revision=revision) except Exception: return None try: @@ -839,7 +839,7 @@ def load_transformer_config( if not os.path.exists(config_path): raise FileNotFoundError(f"LTX transformer config not found: {config_path}") else: - config_path = hf_hub_download( + config_path = hf_api().hf_hub_download( repo_id=model_path, filename=f"{subfolder}/config.json", revision=revision, diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py index 7ebd94143..f35ce25c5 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared denoise execution primitives for LTX pipelines.""" @@ -578,12 +578,6 @@ class LTXPhaseExecutor: video_audio_step_adapter=video_audio_step_adapter, ) video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents) - ring_degree = getattr(pipeline.od_config.parallel_config, "ring_degree", 1) or 1 - if padded_audio_num_frames > original_audio_num_frames and ring_degree > 1: - raise ValueError( - "LTX audio padding requires an attention mask, which Ring sequence parallelism does not support. " - "Use Ulysses-only SP or choose a request whose audio latent length is divisible by the SP size." - ) denoise_ctx = LTXDenoiseContext( latents=latents, audio_latents=audio_latents, diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py index 22230e486..f42c43709 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -497,7 +497,6 @@ class LTXRequestMixin: image_crf = _normalize_image_crf(first_image_crf if first_image_crf is not None else image_crf) sampling = sampling_params_list[0] - is_dummy_run = req.is_dummy_run() prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt negative_prompt = _resolve_negative_prompts( req.prompts, @@ -511,8 +510,6 @@ class LTXRequestMixin: num_inference_steps = ( sampling.num_inference_steps or num_inference_steps or self.pipeline_recipe.num_inference_steps ) - if is_dummy_run and self.pipeline_recipe.fixed_num_inference_steps: - num_inference_steps = self.pipeline_recipe.num_inference_steps num_inference_steps = max(int(num_inference_steps), 2) num_videos_per_prompt = ( @@ -528,8 +525,6 @@ class LTXRequestMixin: for item in sampling_params_list ] guidance = guidance_specs[0] - if is_dummy_run and not self.pipeline_recipe.phases[0].allow_guidance_override: - guidance = self.pipeline_recipe.request_guidance if any(item != guidance for item in guidance_specs[1:]): raise ValueError("Batched LTX requests must use identical video/audio guidance parameters.") diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index 0ffa09dcb..2c51b17d9 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -5,14 +5,16 @@ from __future__ import annotations -from collections.abc import Iterable -from contextlib import nullcontext +import time +from collections.abc import Iterable, Iterator +from contextlib import contextmanager, nullcontext from dataclasses import replace from typing import Any, ClassVar import torch from diffusers.utils.torch_utils import randn_tensor from torch import nn +from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -60,6 +62,8 @@ from .ltx2_request import ( validate_pipeline_request, ) +logger = init_logger(__name__) + def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.Tensor: """Run the BWE vocoder in FP32, matching the official LTX pipeline.""" @@ -195,7 +199,7 @@ class LTXRuntime( connector_batches_cfg = False distributed_video_decode = True support_image_input = False - dummy_run_num_frames = 1 + dummy_run_num_frames: ClassVar[int] = 0 preserve_sp_padded_audio_duration = False reports_stage_durations = False @@ -207,6 +211,14 @@ class LTXRuntime( f"{self.__class__.__name__} does not support ulysses_mode='advanced_uaa'. " "Use the default ulysses_mode='strict' for LTX sequence parallelism." ) + ring_degree = getattr(parallel_config, "ring_degree", 1) + allgather_degree = getattr(parallel_config, "allgather_degree", 1) + if ring_degree != 1 or allgather_degree != 1: + raise ValueError( + f"{self.__class__.__name__} supports pure Ulysses sequence parallelism only. " + "Set ring_degree=1 and allgather_degree=1; " + f"got {ring_degree=} and {allgather_degree=}." + ) self.model_version = detect_ltx_model_version(od_config.model, revision=getattr(od_config, "revision", None)) self.use_diffusion_decoder = _ltx2_use_diffusion_decoder(od_config, self.model_version) self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version) @@ -354,27 +366,29 @@ class LTXRuntime( ) ) self._enter_phase(phase_recipe) - phase_inputs = self._build_phase_inputs( - request_inputs, - phase_recipe, - phase_results[-1] if phase_results else None, - ) + with self._profile_phase_method("_build_phase_inputs", phase_recipe.name): + phase_inputs = self._build_phase_inputs( + request_inputs, + phase_recipe, + phase_results[-1] if phase_results else None, + ) if phase_sigmas is not None: phase_inputs = replace(phase_inputs, num_inference_steps=len(phase_sigmas) - 1) noise_scale = phase_recipe.noise_scale if override_sigmas is not None and phase_recipe.input_transform == "spatial_upsample": noise_scale = float(override_sigmas[0]) - phase_result = self.run_phase( - req, - phase_inputs, - noise_scale=noise_scale, - sigmas=phase_sigmas, - timesteps=None, - attention_kwargs=None, - phase_recipe=phase_recipe, - image=image, - prompt_context=prompt_context, - ) + with self._profile_phase_method("run_phase", phase_recipe.name): + phase_result = self.run_phase( + req, + phase_inputs, + noise_scale=noise_scale, + sigmas=phase_sigmas, + timesteps=None, + attention_kwargs=None, + phase_recipe=phase_recipe, + image=image, + prompt_context=prompt_context, + ) phase_results.append(phase_result) prompt_context = phase_result.forward_context.prompt_context @@ -384,7 +398,45 @@ class LTXRuntime( video=phase_results[self.pipeline_recipe.video_output_phase].video, audio=phase_results[self.pipeline_recipe.audio_output_phase].audio, ) - return self.decode_phase(output_phase) + with self._profile_phase_method("decode_phase"): + output = self.decode_phase(output_phase) + return self._attach_profiled_stage_durations(output) + + @contextmanager + def _profile_phase_method(self, method_name: str, phase_name: str | None = None) -> Iterator[None]: + if not getattr(self, "enable_diffusion_pipeline_profiler", False): + yield + return + + metric_parts = [self.__class__.__name__] + if phase_name is not None: + metric_parts.append(phase_name) + metric_parts.append(method_name) + metric_name = ".".join(metric_parts) + + if current_omni_platform.is_available(): + current_omni_platform.synchronize() + start_time = time.perf_counter() + try: + yield + finally: + if current_omni_platform.is_available(): + current_omni_platform.synchronize() + duration = time.perf_counter() - start_time + logger.info("[DiffusionPipelineProfiler] %s took %.6fs", metric_name, duration) + with self._profiler_lock: + self._stage_durations[metric_name] = self._stage_durations.get(metric_name, 0.0) + duration + + def _attach_profiled_stage_durations( + self, output: DiffusionOutput | list[DiffusionOutput] + ) -> DiffusionOutput | list[DiffusionOutput]: + if not getattr(self, "enable_diffusion_pipeline_profiler", False): + return output + stage_durations = self.stage_durations + outputs = output if isinstance(output, list) else [output] + for item in outputs: + item.stage_durations = dict(stage_durations) + return output def _enter_phase(self, phase: LTXPhaseRecipe) -> None: self._active_phase_name = phase.name diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_sequence_parallel.py b/vllm_omni/diffusion/models/ltx2/ltx2_sequence_parallel.py new file mode 100644 index 000000000..874df90e3 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_sequence_parallel.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.parallel.base import ( + ParallelAttentionContext, +) +from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D +from vllm_omni.diffusion.distributed.group_coordinator import ( + SequenceParallelGroupCoordinator, +) + + +@dataclass(frozen=True, slots=True) +class _VideoToAudioContext(ParallelAttentionContext): + ulysses_pg: dist.ProcessGroup + use_sync: bool + + +class LTX2VideoToAudioParallelAttention: + """LTX video-to-audio SP with replicated audio queries. + + Video K/V arrive sequence-sharded, while the much shorter audio query is + replicated. Redistribute only K/V across sequence/head dimensions, slice + audio query heads locally, then gather output heads. This uses three + collectives instead of standard Ulysses' four and keeps audio replicated. + """ + + def __init__( + self, + sp_group: SequenceParallelGroupCoordinator, + *, + use_sync: bool = False, + ) -> None: + if sp_group.ring_world_size != 1: + raise NotImplementedError( + "LTX video-to-audio sequence parallelism currently supports pure Ulysses only (ring_degree must be 1)." + ) + self._sp_group = sp_group + self._ulysses_pg = sp_group.ulysses_group + self._use_sync = use_sync + + @property + def enabled(self) -> bool: + return True + + @property + def name(self) -> str: + return "ltx2_video_to_audio" + + def pre_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ): + if query.ndim != 4 or key.ndim != 4 or value.ndim != 4: + raise ValueError( + "LTX video-to-audio SP expects Q/K/V in (B, S, H, D) layout, " + f"got {query.shape=}, {key.shape=}, {value.shape=}." + ) + if key.shape != value.shape: + raise ValueError( + "LTX video-to-audio SP requires matching K/V shapes, " + f"got key={tuple(key.shape)} and value={tuple(value.shape)}." + ) + + world_size = self._sp_group.ulysses_world_size + rank = self._sp_group.ulysses_rank + query_heads = int(query.shape[2]) + if query_heads % world_size != 0: + raise ValueError( + "LTX video-to-audio SP requires query heads divisible by " + f"ulysses_degree, got heads={query_heads}, ulysses_degree={world_size}." + ) + if int(key.shape[2]) % world_size != 0: + raise ValueError( + "LTX video-to-audio SP requires K/V heads divisible by " + f"ulysses_degree, got heads={key.shape[2]}, ulysses_degree={world_size}." + ) + + heads_per_rank = query_heads // world_size + head_start = rank * heads_per_rank + query = query[:, :, head_start : head_start + heads_per_rank, :].contiguous() + key = SeqAllToAll4D.apply(self._ulysses_pg, key, 2, 1, self._use_sync) + value = SeqAllToAll4D.apply(self._ulysses_pg, value, 2, 1, self._use_sync) + + ctx = _VideoToAudioContext( + name=self.name, + ulysses_pg=self._ulysses_pg, + use_sync=self._use_sync, + ) + return query, key, value, attn_metadata, ctx + + def post_attention( + self, + attn_output: torch.Tensor, + ctx: ParallelAttentionContext | None, + ) -> torch.Tensor: + if not isinstance(ctx, _VideoToAudioContext): + raise TypeError(f"Unexpected LTX video-to-audio context: {type(ctx)!r}.") + + # Gather along heads by moving H to the leading dimension required by + # all_gather_into_tensor, then restore the backend's (B, S, H, D) layout. + local_heads_first = attn_output.permute(2, 0, 1, 3).contiguous() + world_size = dist.get_world_size(ctx.ulysses_pg) + gathered = torch.empty( + (local_heads_first.shape[0] * world_size, *local_heads_first.shape[1:]), + dtype=attn_output.dtype, + device=attn_output.device, + ) + dist.all_gather_into_tensor(gathered, local_heads_first, group=ctx.ulysses_pg) + if ctx.use_sync: + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.synchronize() + return gathered.permute(1, 2, 0, 3).contiguous() diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 0b0e7fc08..1d95fc3df 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2025 The Lightricks team and The HuggingFace Team. # All rights reserved. # @@ -47,9 +50,12 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module +from vllm_omni.diffusion.distributed.parallel_state import get_sp_group 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 .ltx2_sequence_parallel import LTX2VideoToAudioParallelAttention + logger = init_logger(__name__) _RMSNORM_INIT_PARAMS = inspect.signature(RMSNorm.__init__).parameters @@ -359,6 +365,34 @@ def to_ltx_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor: class _LTX2ParallelAttention(Attention): """Preserve LTX SP collectives while applying model-specific padding masks.""" + _SP_MODES = frozenset({"local", "video_self", "video_to_audio"}) + + def __init__(self, *args, ltx_sp_mode: str = "local", **kwargs): + if ltx_sp_mode not in self._SP_MODES: + raise ValueError( + f"Unsupported LTX sequence-parallel mode {ltx_sp_mode!r}; expected one of {sorted(self._SP_MODES)}." + ) + super().__init__(*args, **kwargs) + self.ltx_sp_mode = ltx_sp_mode + self._video_to_audio_strategy = None + + def _get_active_parallel_strategy(self): + base_strategy = super()._get_active_parallel_strategy() + if self.ltx_sp_mode == "local" or not base_strategy.enabled: + return self._no_parallel_strategy + if self.ltx_sp_mode == "video_self": + return base_strategy + if base_strategy.name != "ulysses": + raise NotImplementedError( + f"LTX video-to-audio sequence parallelism requires pure Ulysses; got strategy={base_strategy.name!r}." + ) + if self._video_to_audio_strategy is None: + self._video_to_audio_strategy = LTX2VideoToAudioParallelAttention( + get_sp_group(), + use_sync=self.use_sync, + ) + return self._video_to_audio_strategy + def _run_local_attention(self, query, key, value, attn_metadata): attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None if attention_mask is not None: @@ -597,6 +631,7 @@ class LTX2Attention(torch.nn.Module): quant_config: "QuantizationConfig | None" = None, prefix: str = "", disable_kv_quant: bool = False, + sequence_parallel_mode: str = "local", ): super().__init__() # LTX-2 uses "rms_norm_across_heads", LTX-2.3 uses "rms_norm" -- both @@ -721,6 +756,7 @@ class LTX2Attention(torch.nn.Module): causal=False, prefix=prefix, disable_kv_quant=disable_kv_quant, + ltx_sp_mode=sequence_parallel_mode, ) # LTX-2.3: per-head gated attention @@ -876,6 +912,7 @@ class LTX2VideoTransformerBlock(nn.Module): apply_gated_attention=video_gated_attn, quant_config=quant_config, prefix=f"{prefix}.attn1" if prefix else "attn1", + sequence_parallel_mode="video_self", ) self.audio_norm1 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) @@ -965,6 +1002,7 @@ class LTX2VideoTransformerBlock(nn.Module): apply_gated_attention=audio_gated_attn, quant_config=quant_config, prefix=f"{prefix}.video_to_audio_attn" if prefix else "video_to_audio_attn", + sequence_parallel_mode="video_to_audio", ) # 4. Feedforward layers @@ -1568,38 +1606,25 @@ class LTX2VideoTransformer3DModel(nn.Module): return { "": { - # Shard video/audio latents across sequence + # Keep only the long video stream sequence-sharded. Audio and + # prompt streams remain replicated on every SP rank. "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), - "audio_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), "keyframes_mask": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), - # Shard prompt embeds across sequence - "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), - "audio_encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), - # Keep per-token video/audio timestep embeddings aligned with - # their corresponding latent sequence shards. Official LTX - # multi-guidance supplies both tensors as (B, seq_len). + # Keep per-token video timestep embeddings aligned with the + # video latent shard. LTX guidance supplies a separate, + # replicated audio_timestep under SP. "timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False), - "audio_timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False), }, "rope": { 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), }, - "audio_rope": { - 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), - 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), - }, "cross_attn_rope": { 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), }, - "cross_attn_audio_rope": { - 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), - 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), - }, - # Gather outputs before returning + # Video output is sharded; audio output is already replicated. "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), - "audio_proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), } def __init__( diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index 2e830fa04..e98659409 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -38,6 +38,27 @@ MINIMAX_H3_AUDIO_CHANNELS = 2 logger = init_logger(__name__) +@contextmanager +def _minimax_h3_keyframe_encode_context( + device: torch.device, +) -> Iterator[None]: + if device.type != "cuda": + yield + return + + # The official keyframe latent uses cuDNN's TF32 convolution path. The + # default non-deterministic algorithm can select numerically different + # reductions on H100s, and the difference is amplified by the denoiser. + # Pin both the algorithm and math mode for this sensitive encode only. + with torch.backends.cudnn.flags( + enabled=True, + benchmark=False, + deterministic=True, + allow_tf32=True, + ): + yield + + def _load_component_config(component_path: str) -> dict[str, Any]: config_path = Path(component_path) / "config.json" config = json.loads(config_path.read_text(encoding="utf-8")) @@ -280,10 +301,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): for device in devices: with self.device_module.device(device): self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) - latent = self.model.encode_images( - image, - use_fp16_latent=True, - )[0] + with _minimax_h3_keyframe_encode_context(parameter.device): + latent = self.model.encode_images( + image, + use_fp16_latent=True, + )[0] finally: self.model.parallel_tiling = previous_parallel if previous_dtype != torch.float32: diff --git a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py index 62ac027ca..bad1c6154 100644 --- a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py +++ b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ OmniVoice TTS Pipeline for vLLM-Omni diffusion engine. @@ -12,10 +12,13 @@ Uses request-mode execution (all steps in one forward() call). from __future__ import annotations import json +import math import os +import random import re -from collections.abc import Iterable -from typing import ClassVar +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any, ClassVar import numpy as np import torch @@ -26,10 +29,17 @@ 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 SupportAudioOutput +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 StepRequestState +from vllm_omni.errors import OmniClientError 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 +from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( + OmniVoiceGenerator, + _build_cu_seqs, + _get_time_steps, +) from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -43,6 +53,15 @@ import torchaudio logger = init_logger(__name__) +@dataclass +class _PreparedOmniVoiceRequest: + input_ids: torch.Tensor + audio_mask: torch.Tensor + cond_len: int + target_len: int + seed: int | None + + def get_omnivoice_post_process_func(od_config: OmniDiffusionConfig): """Post-processing: convert audio tensor to numpy for WAV encoding.""" @@ -132,6 +151,8 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput): """ support_audio_output: ClassVar[bool] = True + supports_request_batch: ClassVar[bool] = True + supports_step_execution: ClassVar[bool] = True def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__() @@ -141,9 +162,9 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput): # Resolve model path (HF hub ID → local cache) if not os.path.isdir(self.model_path): - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - self.model_path = snapshot_download(self.model_path) + self.model_path = hf_api().snapshot_download(self.model_path) # Load OmniVoice config config_path = os.path.join(self.model_path, "config.json") @@ -152,7 +173,7 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput): self.config = OmniVoiceConfig(**hf_config) # Build generator and decoder - self.generator = OmniVoiceGenerator(self.config) + self.generator = OmniVoiceGenerator(self.config, od_config) self.decoder = OmniVoiceDecoder(self.config) # Tokenizer (low-level, avoids HF tokenizer extra_special_tokens issue) @@ -205,50 +226,36 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput): tokens = tokens.squeeze(0) # [8, T_ref] return tokens - @torch.inference_mode() - 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: - {"text": "...", "ref_audio": (samples, sr), "ref_text": "...", - "lang": "...", "instruct": "..."} - """ - prompt = req.prompts[0] if req.prompts else "" + def _prepare_request_input( + self, + prompt: Any, + extra: dict[str, Any], + ) -> _PreparedOmniVoiceRequest | DiffusionOutput: + """Build one request's conditional/unconditional model inputs.""" ref_audio = None ref_text = None lang = "None" instruct = "None" - extra = req.sampling_params.extra_args or {} + voice_name = None seed = extra.get("seed", None) - voice_name = None if isinstance(prompt, dict): - # Top-level keys (used by serving_speech.py /v1/audio/speech path) text = prompt.get("input") or prompt.get("text") or prompt.get("prompt") ref_audio = prompt.get("ref_audio") ref_text = prompt.get("ref_text") voice_name = prompt.get("voice_name") lang = prompt.get("lang") instruct = prompt.get("instruct") - # OmniTextPrompt format (used by offline Omni.generate path): - # ref_audio comes via multi_modal_data["audio"] and the rest via - # mm_processor_kwargs. Fall back to those when top-level keys are - # absent so both invocation styles work. mm_data = prompt.get("multi_modal_data") or {} mm_kwargs = prompt.get("mm_processor_kwargs") or {} if ref_audio is None: audio_field = mm_data.get("audio") - # Standard multimodal shape allows a list of audios; OmniVoice - # voice cloning conditions on a single reference clip, so - # unwrap a length-1 list and reject multi-reference prompts up - # front (otherwise a list would later crash inside - # ``_encode_ref_audio`` when it calls ``audio.dim()``). if isinstance(audio_field, list): if len(audio_field) == 1: 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)}" # noqa: E501 + error=f"OmniVoice voice cloning supports a single reference audio; got {len(audio_field)}" ) else: audio_field = None @@ -264,7 +271,6 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput): lang = mm_kwargs.get("lang") if instruct is None: instruct = mm_kwargs.get("instruct") - if not text: return DiffusionOutput(error="Empty text prompt") lang = lang or "None" @@ -274,105 +280,325 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput): if not text: return DiffusionOutput(error="Empty text prompt") - device = self.device - num_cb = self.config.num_audio_codebook - mask_id = self.config.audio_mask_id - - # Estimate target duration target_len = self.duration_estimator.estimate_duration(text, "Nice to meet you.", 25) target_len = max(1, int(target_len)) - # Build text prompt with control tokens style_text = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>" full_text = _combine_text(ref_text=ref_text, text=text) wrapped_text = f"<|text_start|>{full_text}<|text_end|>" style_tokens = self.tokenizer.encode(style_text).ids text_tokens = _tokenize_with_nonverbal_tags(wrapped_text, self.tokenizer) encoding_ids = style_tokens + text_tokens - text_tokens = torch.tensor(encoding_ids, dtype=torch.long, device=device) - text_len = text_tokens.shape[0] + text_tokens_tensor = torch.tensor(encoding_ids, dtype=torch.long, device=self.device) + text_len = text_tokens_tensor.shape[0] - # Encode reference audio tokens if provided (with voice caching) ref_audio_tokens = None if ref_audio is not None: if self.audio_tokenizer is None: raise RuntimeError( "Voice cloning requires transformers>=5.3.0. Try: uv pip install 'transformers>=5.3.0'" ) - # Check speaker cache first - _cache_key = None + cache_key = None if voice_name: - _cache_key = self._speaker_cache.make_cache_key( + cache_key = self._speaker_cache.make_cache_key( voice_name, model_type="omnivoice", created_at=int(prompt.get("voice_created_at") or 0), ) - cached = self._speaker_cache.get(_cache_key) + cached = self._speaker_cache.get(cache_key) if cached is not None: - ref_audio_tokens = cached["ref_audio_tokens"].to(device) - _cache_key = None # hit → don't store again + ref_audio_tokens = cached["ref_audio_tokens"].to(self.device) + cache_key = None logger.debug("Speaker cache HIT for OmniVoice speaker '%s'", voice_name) if ref_audio_tokens is None: audio_signal, sr = ref_audio if isinstance(audio_signal, np.ndarray): audio_signal = torch.from_numpy(audio_signal).float() - ref_audio_tokens = self._encode_ref_audio(audio_signal, int(sr)).to(device) - - # Store in cache for next request - if _cache_key is not None: - self._speaker_cache.put(_cache_key, {"ref_audio_tokens": ref_audio_tokens.cpu()}) + ref_audio_tokens = self._encode_ref_audio(audio_signal, int(sr)).to(self.device) + if cache_key is not None: + self._speaker_cache.put(cache_key, {"ref_audio_tokens": ref_audio_tokens.cpu()}) logger.debug("Speaker cache STORE for OmniVoice speaker '%s'", voice_name) - # Build conditional + unconditional batches [2, 8, max_len] - text_ids = text_tokens.unsqueeze(0).repeat(num_cb, 1) - target_ids = torch.full((num_cb, target_len), mask_id, dtype=torch.long, device=device) - - if ref_audio_tokens is not None: - cond_ids = torch.cat([text_ids, ref_audio_tokens, target_ids], dim=1) - else: - cond_ids = torch.cat([text_ids, target_ids], dim=1) + num_cb = self.config.num_audio_codebook + mask_id = self.config.audio_mask_id + text_ids = text_tokens_tensor.unsqueeze(0).repeat(num_cb, 1) + target_ids = torch.full((num_cb, target_len), mask_id, dtype=torch.long, device=self.device) + cond_ids = ( + torch.cat([text_ids, ref_audio_tokens, target_ids], dim=1) + if ref_audio_tokens is not None + else torch.cat([text_ids, target_ids], dim=1) + ) cond_len = cond_ids.shape[1] uncond_ids = target_ids.clone() - uncond_len = target_len - max_len = max(cond_len, uncond_len) - if uncond_len < max_len: - pad = torch.full( - (num_cb, max_len - uncond_len), - mask_id, - dtype=torch.long, - device=device, + input_ids = torch.cat([cond_ids, uncond_ids], dim=1).transpose(0, 1).contiguous() + + max_len = input_ids.shape[0] + audio_mask = torch.zeros(max_len, dtype=torch.bool, device=self.device) + audio_mask[text_len:] = True + + return _PreparedOmniVoiceRequest( + input_ids=input_ids, + audio_mask=audio_mask, + cond_len=cond_len, + target_len=target_len, + seed=seed, + ) + + def _collate_request_inputs( + self, + prepared_requests: Sequence[_PreparedOmniVoiceRequest], + ) -> tuple[torch.Tensor, torch.Tensor, list[int]]: + """Pack request-major [cond, uncond] token sequences.""" + input_ids: list[torch.Tensor] = [] + audio_masks: list[torch.Tensor] = [] + cond_lens: list[int] = [] + + for request in prepared_requests: + input_id = request.input_ids + audio_mask = request.audio_mask + cond_len = request.cond_len + input_ids.append(input_id) + audio_masks.append(audio_mask) + cond_lens.append(cond_len) + + input_ids = torch.cat(input_ids, dim=0) + audio_masks = torch.cat(audio_masks, dim=0) + + return input_ids, audio_masks, cond_lens + + def prepare_encode(self, state: StepRequestState) -> StepRequestState: + prompt = state.prompt if state.prompt else "" + extra = state.sampling.extra_args or {} + prepared = self._prepare_request_input(prompt, extra) + if isinstance(prepared, DiffusionOutput): + raise OmniClientError(prepared.error or "OmniVoice request preparation failed") + + prepared_request = prepared + cond_len = prepared_request.cond_len + target_len = prepared_request.target_len + input_ids = prepared_request.input_ids + audio_mask = prepared_request.audio_mask + seed = prepared_request.seed + device = self.device + mask_id = self.config.audio_mask_id + num_codebooks = self.config.num_audio_codebook + if seed is None: + seed = random.randint(0, 2**63 - 1) + num_step = ( + state.sampling.num_inference_steps if state.sampling.num_inference_steps is not None else self.num_step + ) + + t_shift = self.t_shift + + # Initialize all target tokens as [MASK] + tokens = torch.full((1, num_codebooks, target_len), mask_id, dtype=torch.long, device=device) + + timesteps = _get_time_steps(0.0, 1.0, num_step + 1, t_shift) + + # Compute unmasking schedule + schedules = [] + total_mask = target_len * num_codebooks + rem = total_mask + sched = [] + for step in range(num_step): + num = ( + rem + if step == num_step - 1 + else min( + math.ceil(total_mask * (timesteps[step + 1] - timesteps[step])), + rem, + ) ) - uncond_ids = torch.cat([uncond_ids, pad], dim=1) + sched.append(int(num)) + rem -= int(num) + schedules = torch.tensor(sched, dtype=torch.long, device=device) - batch_input_ids = torch.stack([cond_ids, uncond_ids]) + layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1) + generator = torch.Generator(device=device).manual_seed(seed) - batch_audio_mask = torch.zeros(2, max_len, dtype=torch.bool, device=device) - batch_audio_mask[0, text_len:cond_len] = True - batch_audio_mask[1, :uncond_len] = True + guidance_scale = ( + state.sampling.guidance_scale if state.sampling.guidance_scale is not None else self.guidance_scale + ) + state.latents = input_ids + state.timesteps = schedules + state.guidance = guidance_scale + state.extra["schedules"] = schedules + state.extra["layer_ids"] = layer_ids + state.extra["generator"] = generator + state.extra["t_shift"] = t_shift + state.extra["cond_len"] = cond_len + state.extra["target_len"] = target_len + state.extra["audio_mask"] = audio_mask + state.extra["tokens"] = tokens + return state - batch_attn_mask = torch.zeros(2, 1, max_len, max_len, dtype=torch.bool, device=device) - batch_attn_mask[0, :, :cond_len, :cond_len] = True - batch_attn_mask[1, :, :uncond_len, :uncond_len] = True + def denoise_step(self, input_batch: InputBatch, *, states: Sequence[StepRequestState] | None = None, **kwargs: Any): + input_ids = input_batch.latents + use_cuda_graph = self.generator._cuda_graph_fwd is not None and input_ids.is_cuda + layer_ids = states[0].extra["layer_ids"] + audio_masks: list[torch.Tensor] = [] + target_lens: list[int] = [] + batch_tokens: list[torch.Tensor] = [] + cond_lens: list[int] = [] + + steps: list[int] = [] + schedules: list[torch.Tensor] = [] + generators: list[torch.Generator] = [] + guidance_scales: list[float] = [] + + for state in states: + audio_masks.append(state.extra.get("audio_mask", None)) + cond_lens.append(state.extra["cond_len"]) + target_lens.append(state.extra["target_len"]) + batch_tokens.append(state.extra["tokens"]) + guidance_scales.append(state.guidance) + generators.append(state.extra.get("generator", None)) + schedules.append(state.extra["schedules"]) + steps.append(state.step_index) + + audio_masks = torch.cat(audio_masks, dim=0) + + B = len(target_lens) + cu_seqs = _build_cu_seqs(cond_lens, target_lens, input_ids.device) + + position_temperature = self.position_temperature + class_temperature = self.class_temperature + layer_penalty_factor = self.layer_penalty_factor + if use_cuda_graph: + # Replay a fixed packed-token bucket with dynamic varlen metadata. + batch_logits = self.generator._cuda_graph_fwd(input_ids, audio_masks, cu_seqs, B) + else: + # Run packed eager attention for the current active requests. + inputs_embeds = self.generator._prepare_embeddings(input_ids, audio_masks) + hidden_states = self.generator._transformer_forward( + inputs_embeds, + cu_seqs, + max_seqlen=max(cond_lens), + ) + # fp32 cast deferred to the per-item slices below. + batch_logits = self.generator._get_logits(hidden_states) + # batch_logits: [8, total_seq_len, 1025] + + target_offsets: list[int] = [] + target_offset = 0 + for target_len in target_lens: + target_offsets.append(target_offset) + target_offset += target_len + + sequence_offsets: list[int] = [] + sequence_offset = 0 + for cond_len, target_len in zip(cond_lens, target_lens): + sequence_offsets.append(sequence_offset) + sequence_offset += cond_len + target_len + + for i in range(B): + k = schedules[i][steps[i]] + if k <= 0: + continue + + c_len = cond_lens[i] + t_len = target_lens[i] + + # Extract logits for target region; upcast only the slices we actually consume. + request_start = sequence_offsets[i] + cond_end = request_start + c_len + uncond_start = cond_end + + # Extract logits for target region; upcast only the slices we actually consume. + c_logits = batch_logits[:, cond_end - t_len : cond_end, :].unsqueeze(0).to(torch.float32) + u_logits = batch_logits[:, uncond_start : uncond_start + t_len, :].unsqueeze(0).to(torch.float32) + sample = batch_tokens[i] + sample_tokens = sample[..., :t_len] + self.generator._unmask_one_request( + c_logits, + u_logits, + sample_tokens, + num_to_unmask=k, + guidance_scale=guidance_scales[i], + generator=generators[i], + class_temperature=class_temperature, + position_temperature=position_temperature, + layer_penalty_factor=layer_penalty_factor, + layer_ids=layer_ids, + ) + + # Mirror update into both cond and uncond input_ids halves for the next step. + packed_sample_tokens = sample_tokens.squeeze(0).transpose(0, 1) + input_ids[cond_end - t_len : cond_end] = packed_sample_tokens + input_ids[uncond_start : uncond_start + t_len] = packed_sample_tokens + states[i].extra["tokens"] = sample_tokens + + # InputBatch reuses its latents buffer across steps. Returning that + # same storage would make the Runner persist per-request views into the + # cached destination; the next make_batch() would then copy overlapping + # source/destination slices. Break the alias at the lifecycle boundary. + return input_ids.clone() + + def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any): + state.latents = noise_pred + state.step_index += 1 + + def post_decode(self, state: StepRequestState, **kwargs: Any): + tokens = state.extra["tokens"] + if tokens.dim() == 2: + tokens = tokens.unsqueeze(0) + audio = self.decoder(tokens) + return DiffusionOutput(output=audio) + + @torch.inference_mode() + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + """Generate speech audio from text, optionally with voice cloning. + + Accepts either a plain text prompt or a structured dict: + {"text": "...", "ref_audio": (samples, sr), "ref_text": "...", + "lang": "...", "instruct": "..."} + """ + prepared_requests: list[_PreparedOmniVoiceRequest] = [] + outputs = [None] * len(req.requests) + prepared_indices: list[int] = [] + for i, request in enumerate(req.requests): + prompt = request.prompt if request.prompt else "" + extra = request.sampling_params.extra_args or {} + prepared = self._prepare_request_input(prompt, extra) + if isinstance(prepared, DiffusionOutput): + outputs[i] = prepared + continue + prepared_indices.append(i) + prepared_requests.append(prepared) + + if not prepared_requests: + return outputs + + batch_target_len = [request.target_len for request in prepared_requests] + batch_seeds = [request.seed for request in prepared_requests] + batch_input_ids, batch_audio_mask, batch_cond_lens = self._collate_request_inputs(prepared_requests) # Run 32-step iterative unmasking + sampling = req.requests[0].sampling_params + num_step = sampling.num_inference_steps if sampling.num_inference_steps is not None else self.num_step + guidance_scale = sampling.guidance_scale if sampling.guidance_scale is not None else self.guidance_scale tokens = self.generator( input_ids=batch_input_ids, audio_mask=batch_audio_mask, - attention_mask=batch_attn_mask, - target_lens=[target_len], - num_step=self.num_step, - guidance_scale=self.guidance_scale, + cond_lens=batch_cond_lens, + target_lens=batch_target_len, + num_step=num_step, + guidance_scale=guidance_scale, t_shift=self.t_shift, layer_penalty_factor=self.layer_penalty_factor, position_temperature=self.position_temperature, class_temperature=self.class_temperature, - seed=seed, + seed=batch_seeds, ) - # Decode tokens to audio - audio = self.decoder(tokens) # [1, 1, samples] - return DiffusionOutput(output=audio) + target_offset = 0 + for i, target_len in enumerate(batch_target_len): + request_tokens = tokens[:, :, target_offset : target_offset + target_len] + audio = self.decoder(request_tokens) + outputs[prepared_indices[i]] = DiffusionOutput(output=audio) + target_offset += target_len + return outputs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights from model directory (not from the iterator). diff --git a/vllm_omni/diffusion/models/pi0/pipeline_pi0.py b/vllm_omni/diffusion/models/pi0/pipeline_pi0.py index 1eeca2969..a03801c8d 100644 --- a/vllm_omni/diffusion/models/pi0/pipeline_pi0.py +++ b/vllm_omni/diffusion/models/pi0/pipeline_pi0.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """π0 (Pi-Zero) VLA pipeline for vllm-omni. Entry point for ``DiffusionEngine.step() → pipeline.forward(req)``. Mirrors the @@ -94,9 +94,9 @@ class Pi0Pipeline(nn.Module): return None if os.path.isdir(model): return model - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download( + return hf_api().snapshot_download( repo_id=model, allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"], ) 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 de636fb6d..2be1a9c7a 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -27,7 +27,6 @@ import numpy as np import torch import torch.nn as nn import torchvision.transforms as T -from huggingface_hub import snapshot_download from PIL import Image from transformers import AutoTokenizer from vllm.logger import init_logger @@ -109,7 +108,9 @@ def _resolve_model_path(model_path: str) -> str: """Resolve a HuggingFace model ID or local path to a local directory.""" if os.path.isdir(model_path): return model_path - return snapshot_download(model_path) + from vllm_omni.transformers_utils.repo_utils import hf_api + + return hf_api().snapshot_download(model_path) def _round_by_factor(number: float, factor: int) -> int: diff --git a/vllm_omni/diffusion/models/utils.py b/vllm_omni/diffusion/models/utils.py index 122646219..429c34b9d 100644 --- a/vllm_omni/diffusion/models/utils.py +++ b/vllm_omni/diffusion/models/utils.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -31,6 +31,34 @@ if TYPE_CHECKING: Style = Literal["colwise", "colwise_rep", "rowwise", "rowwise_rep", "replicate"] +def make_attention_mask( + hidden_states: torch.Tensor, + seq_lengths: list[int], +) -> torch.Tensor | None: + """Build a padding mask only when the batch contains padding. + + ``seq_lengths`` must be non-empty and contain valid sequence lengths. + Returns a boolean mask of shape ``(len(seq_lengths), max(seq_lengths))`` + on the device of ``hidden_states``, with True marking valid tokens, + or None when all sequence lengths are equal. + + Passing an all-true CUDA mask makes attention backends inspect a device + scalar to select the dense path, which introduces a host-device sync. + """ + max_seq_len = max(seq_lengths) + if all(seq_len == max_seq_len for seq_len in seq_lengths): + return None + + attention_mask = hidden_states.new_zeros( + len(seq_lengths), + max_seq_len, + dtype=torch.bool, + ) + for i, seq_len in enumerate(seq_lengths): + attention_mask[i, :seq_len] = True + return attention_mask + + def replace_linear_class( linear: nn.Linear, style: Style = "replicate", @@ -140,8 +168,8 @@ def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> with open(path) as f: return json.load(f) else: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - cached = hf_hub_download(repo_id=model_path, filename=filename) + cached = hf_api().hf_hub_download(repo_id=model_path, filename=filename) with open(cached) as f: return json.load(f) 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 2383fc6fd..c9ef325c4 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -156,9 +156,9 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc else: # Try to download config from HF Hub try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - config_path = hf_hub_download( + config_path = hf_api().hf_hub_download( repo_id=model_path, filename=f"{subfolder}/config.json", ) @@ -364,9 +364,9 @@ class Wan22Pipeline( else: # For remote models, download and read model_index.json try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_index_path = hf_hub_download(repo_id=model, filename="model_index.json") + model_index_path = hf_api().hf_hub_download(repo_id=model, filename="model_index.json") with open(model_index_path) as f: model_index = json.load(f) self.expand_timesteps = model_index.get("expand_timesteps", False) diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 0df4039eb..42ba45cc9 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -20,6 +20,12 @@ from .distributed_layerwise_backend import ( from .layerwise_backend import LayerWiseOffloadBackend from .module_residency import BoundedAllocatorCache, PinnedModuleStager from .offload_plan import OffloadPlan, get_offload_plan +from .plan_resolver import ( + BlockStack, + ResolvedComponent, + ResolvedOffloadPlan, + resolve_offload_plan, +) from .sequential_backend import ( ModelLevelOffloadBackend, apply_sequential_offload, @@ -42,6 +48,10 @@ __all__ = [ "OffloadConfig", "OffloadPlan", "OffloadStrategy", + "BlockStack", + "ResolvedComponent", + "ResolvedOffloadPlan", + "resolve_offload_plan", "SupportsModelCpuOffload", "LayerWiseOffloadBackend", "DistributedLayerwiseOffloadBackend", diff --git a/vllm_omni/diffusion/offloader/component_utils.py b/vllm_omni/diffusion/offloader/component_utils.py index 1000b29d6..cbc63159d 100644 --- a/vllm_omni/diffusion/offloader/component_utils.py +++ b/vllm_omni/diffusion/offloader/component_utils.py @@ -4,7 +4,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Sequence from itertools import chain from operator import attrgetter from typing import TYPE_CHECKING, Any @@ -14,14 +14,12 @@ from torch import nn from torch.distributed.tensor import DTensor from vllm.logger import init_logger -from .block_discovery import get_blocks_from_dit -from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT +from .config import TEXT_ENCODER_COMPONENT from .offload_plan import OffloadPlan from .tensor_utils import set_tensor_storage if TYPE_CHECKING: - from .base import OffloadConfig - from .module_collector import PipelineModules + from .plan_resolver import BlockStack, ResolvedComponent, ResolvedOffloadPlan logger = init_logger(__name__) @@ -76,30 +74,25 @@ def get_encoder_block_groups( def iter_streamable_dits( - modules: PipelineModules, - config: OffloadConfig, + resolved: ResolvedOffloadPlan, device: torch.device, - plan: OffloadPlan | None, -) -> Iterator[tuple[str, nn.Module, list[str], list[nn.Module]]]: - """Yield selected DiTs whose block metadata resolves successfully.""" - if not config.offloads(DIT_COMPONENT): - return - for name, module in zip(modules.dit_names, modules.dits): - logger.info("Applying hooks on %s (%s)", name, module.__class__.__name__) - planned_attrs = None if plan is None else plan.block_attrs.get(name) - block_attrs, blocks = get_blocks_from_dit(module, planned_attrs) - if blocks: - yield name, module, block_attrs, blocks +) -> Iterator[tuple[ResolvedComponent, BlockStack]]: + """Yield selected DiTs that resolved to a streamable block stack.""" + for component in resolved.dits: + if not component.selected: continue - if config.components is not None: - raise ValueError(f"Selected DiT {name!r} has no streamable layerwise-offload blocks") - logger.warning("Target layers (blocks) not found. Skipping offloading on %s (%s)", name, type(module).__name__) - module.to(device) + logger.info("Applying hooks on %s (%s)", component.path, component.module.__class__.__name__) + if not component.stacks: + # The resolver already warned; a legacy selection keeps the + # unstreamable DiT resident instead of failing the run. + component.module.to(device) + continue + yield component, component.stacks[0] def move_non_block_state_to_device( module: nn.Module, - block_groups: list[nn.ModuleList], + block_groups: Sequence[Sequence[nn.Module]], device: torch.device, ) -> None: """Keep component state outside streamed block lists resident on device.""" @@ -120,7 +113,7 @@ def move_non_block_state_to_device( def set_encoder_layerwise_state( module: nn.Module, hooks: list[Any], - block_groups: list[nn.ModuleList], + block_groups: Sequence[Sequence[nn.Module]], ) -> None: """Publish the backend-neutral state used by encoder stage lifecycles.""" module._omni_layerwise_hooks = hooks @@ -163,59 +156,38 @@ def prepare_component( def prepare_pipeline_components( - modules: PipelineModules, - config: OffloadConfig, - plan: OffloadPlan | None, + resolved: ResolvedOffloadPlan, *, device: torch.device, staged_components: list[nn.Module], - enable_encoder_blocks: Callable[[nn.Module, str, OffloadPlan | None, bool], bool], + enable_encoder_blocks: Callable[[ResolvedComponent], bool], ) -> None: """Apply the shared encoder/VAE/resident placement policy.""" - if config.components is not None and config.offloads(TEXT_ENCODER_COMPONENT): - selected_encoder_names = [name for name in modules.encoder_names if config.offloads_encoder(name, plan)] - if not selected_encoder_names: - raise ValueError("No text encoder modules found for selected text_encoder offload") - - if plan is not None: - for encoder, name in zip(modules.encoders, modules.encoder_names): - if config.should_offload_encoder(name, plan) and name in plan.on_demand_component_paths: - validate_on_demand_component(encoder, name) - - for encoder, name in zip(modules.encoders, modules.encoder_names): - selected = config.should_offload_encoder(name, plan) - stage_on_demand = bool(selected and plan is not None and name in plan.on_demand_component_paths) - blockwise = selected and enable_encoder_blocks(encoder, name, plan, stage_on_demand) - if stage_on_demand and not blockwise and config.uses_allgather(TEXT_ENCODER_COMPONENT): - raise ValueError( - f"Text encoder {name!r} cannot use AllGather without a model-declared streamable block plan" - ) - if selected and config.components is not None and not (blockwise or stage_on_demand): - raise ValueError(f"Selected text encoder {name!r} requires a model-declared streamable or on-demand plan") + for component in resolved.encoders: + blockwise = bool(component.stacks) and enable_encoder_blocks(component) prepare_component( - encoder, - name, + component.module, + component.path, device=device, - stage_on_demand=stage_on_demand, + stage_on_demand=component.on_demand, blockwise=blockwise, staged_components=staged_components, ) - for vae, name in zip(modules.vaes, modules.vae_names): - legacy_staged = config.components is None and plan is not None and name in plan.on_demand_component_paths + for component in resolved.vaes: prepare_component( - vae, - name, + component.module, + component.path, device=device, - stage_on_demand=legacy_staged, + stage_on_demand=component.on_demand, blockwise=False, staged_components=staged_components, ) - for name, module in zip(modules.resident_names, modules.resident_modules): - module.to(device) - logger.debug("Moved resident module %s to %s", name, device) + for component in resolved.residents: + component.module.to(device) + logger.debug("Moved resident module %s to %s", component.path, device) - if not config.offloads(DIT_COMPONENT): - for dit in modules.dits: - dit.to(device) + for component in resolved.dits: + if not component.selected: + component.module.to(device) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index e626d87a5..299ec3d3a 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -21,6 +21,7 @@ from __future__ import annotations import threading import time import weakref +from collections.abc import Sequence from itertools import chain from typing import Any @@ -39,7 +40,6 @@ from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, run_cleanup_steps from .component_utils import ( clear_encoder_layerwise_state, - get_encoder_block_groups, iter_streamable_dits, move_non_block_state_to_device, prepare_component, @@ -53,11 +53,8 @@ from .host_registration import ( HostRegistrationError, register_host_mappings, ) -from .module_collector import ModuleDiscovery -from .offload_plan import ( - OffloadPlan, - get_offload_plan, -) +from .offload_plan import OffloadPlan +from .plan_resolver import ResolvedComponent, resolve_offload_plan from .tensor_utils import ( clear_block_storage, clear_tensor_storage, @@ -1437,7 +1434,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): def _install_hook_group( self, - blocks: list[nn.Module] | nn.ModuleList, + blocks: Sequence[nn.Module] | nn.ModuleList, component: str, *, use_dit_mmap: bool = False, @@ -1483,48 +1480,27 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): hooks[1]._is_group_first = True return hooks - def _try_layerwise_offload_encoder( - self, - module: nn.Module, - name: str, - plan: OffloadPlan | None, - stage_on_demand: bool = False, - ) -> bool: - """Apply DLO hooks to a plan-declared text encoder. + def _try_layerwise_offload_encoder(self, component: ResolvedComponent) -> bool: + """Apply DLO hooks to a text encoder's resolved block stacks. Rank-local transfer is always safe because it preserves the tensors produced by the model loader. Multi-rank AllGather is enabled only when the model declares that those tensors are replicated across the - selected DLO group; encoder TP groups contain different shards and - are deliberately not used for this transport. + selected DLO group; the resolver rejects every other layout before + this method installs anything. """ - if plan is None: - return False - - group_size = self._component_transport(TEXT_ENCODER_COMPONENT)[1] - if group_size > 1 and name not in plan.encoder_dlo_weight_replication: - raise ValueError( - f"Text encoder {name!r} cannot use DLO AllGather across the DiT offload group: " - "its loader-produced weights are not declared replicated across that group. " - "Set layer_options.text_encoder.weight_transfer='rank-local' in diffusion_offload_config " - "for encoder-TP or rank-specific layouts." - ) - block_groups = get_encoder_block_groups( - module, - name, - plan, - strict=self.config.components is not None or group_size > 1, - ) + block_groups = [stack.blocks for stack in component.stacks] if not block_groups: return False + module = component.module encoder_hooks: list[DistributedLayerwiseOffloadHook] = [] for blocks in block_groups: encoder_hooks.extend(self._install_hook_group(blocks, TEXT_ENCODER_COMPONENT)) # Track the module before placement, which may fail, so outer rollback # also removes its partially-installed block hooks and marker state. self._encoder_modules.append(module) - if not stage_on_demand: + if not component.on_demand: move_non_block_state_to_device(module, block_groups, self.device) set_encoder_layerwise_state( module, @@ -1534,10 +1510,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): logger.info( "Enabled %s DLO transfer for text encoder %s (%d blocks across %d stacks, group_size=%d)", self.config.transfer_for(TEXT_ENCODER_COMPONENT).value, - name, + component.path, sum(len(blocks) for blocks in block_groups), len(block_groups), - group_size, + self._component_transport(TEXT_ENCODER_COMPONENT)[1], ) return True @@ -1700,31 +1676,15 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if self.dp_group is None and self._has_multirank_allgather(): self._init_dp_group() - modules = ModuleDiscovery.discover(pipeline) - if not modules.dits and self.config.offloads(DIT_COMPONENT): + # One resolver call validates the whole topology; an explicit selection + # has already failed by the time any of the checks below run. + resolved = resolve_offload_plan(pipeline, self.config) + if not resolved.dits and self.config.offloads(DIT_COMPONENT): if self.host_weight_plan is not None: raise RuntimeError( "DLO received a loader-owned host-weight plan, but no DiT modules were discovered to consume it" ) - message = "No DiT/transformer modules found for selected DiT offload" - if self.config.components is not None: - raise ValueError(message) - logger.warning(message) - - # Retrieve optional declarative OffloadPlan from the pipeline. - # When present, replaces heuristic block discovery. - plan = get_offload_plan(pipeline) - - if self.config.dlo_resident_layers: - resident_paths = frozenset() if plan is None else plan.resident_dit_paths - if not resident_paths.intersection(modules.dit_names): - message = ( - f"resident_layers={self.config.dlo_resident_layers} was requested, but this model declares " - "no matching resident_dit_paths" - ) - if self.config.components is not None: - raise ValueError(message) - logger.warning("%s; all blocks will be streamed.", message) + logger.warning("No DiT/transformer modules found for selected DiT offload") # Storage selection belongs to the loader. DLO consumes the exact # prevalidated plan that caused the loader to skip materialization; @@ -1755,7 +1715,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): elif host_weight_plan.backing_kind == "checkpoint_mmap": self._load_weights_via_mmap( pipeline, - modules, + resolved.modules, host_weight_plan, ) else: @@ -1768,10 +1728,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): else: remaining_meta = [ name - for dit_name, dit_module in zip(modules.dit_names, modules.dits) + for component in resolved.dits for name, tensor in chain( - dit_module.named_parameters(), - dit_module.named_buffers(), + component.module.named_parameters(), + component.module.named_buffers(), ) if getattr(tensor, "is_meta", False) ] @@ -1784,66 +1744,46 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # Apply each selected encoder transfer while keeping explicit VAEs and # unselected components resident. prepare_pipeline_components( - modules, - self.config, - plan, + resolved, device=self.device, staged_components=self._staged_components, enable_encoder_blocks=self._try_layerwise_offload_encoder, ) if self.config.offloads(DIT_COMPONENT): - logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names) + logger.info( + "Applying distributed layer-wise offloading on %s", + [component.path for component in resolved.dits], + ) # Collect all DiT module objects to detect submodules that are # already handled as a separate DiT module (avoids duplicate hooks). - all_dit_modules = set(id(m) for m in modules.dits) + all_dit_modules = set(id(component.module) for component in resolved.dits) # Apply hooks for each DiT module - for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits( - modules, self.config, self.device, plan - ): - resident_count = 0 - if plan is not None and dit_name in plan.resident_dit_paths: - resident_count = min(self.config.dlo_resident_layers, len(blocks)) - if resident_count: - resident_blocks = blocks[:resident_count] - self._resident_blocks.extend(resident_blocks) - blocks = blocks[resident_count:] + for component, stack in iter_streamable_dits(resolved, self.device): + streaming = list(stack.streaming) + if stack.resident_head: + self._resident_blocks.extend(stack.resident) logger.info( "Keeping %d leading blocks resident on %s; streaming %d tail blocks", - resident_count, - dit_name, - len(blocks), + stack.resident_head, + component.path, + len(streaming), ) self._prepare_dit_non_block_modules( - dit_module, - blocks_attr_names, + component.module, + list(stack.attrs), all_dit_modules, - plan, + resolved.declaration, ) - num_blocks = len(blocks) - if num_blocks == 0: - logger.info("All blocks for %s are resident; no streaming hooks required", dit_name) - continue - if num_blocks <= 1: - if self.config.components is not None: - raise ValueError( - f"Selected DiT {dit_name!r} leaves only one streaming block after " - f"resident_layers={resident_count}; choose a resident count that " - "leaves zero or at least two streaming blocks" - ) - logger.warning( - "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)", - dit_name, - dit_module.__class__.__name__, - ) - self._resident_blocks.extend(blocks) + if not streaming: + logger.info("All blocks for %s are resident; no streaming hooks required", component.path) continue - self._install_hook_group(blocks, DIT_COMPONENT, use_dit_mmap=True) + self._install_hook_group(streaming, DIT_COMPONENT, use_dit_mmap=True) if self._resident_blocks: self._resident_layer_group = PinnedResidentLayerGroup( self._resident_blocks, diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index 9b63f9162..a1e594f8b 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations +from collections.abc import Sequence from itertools import chain from typing import Any @@ -20,15 +21,13 @@ from .block_discovery import ( ) from .component_utils import ( clear_encoder_layerwise_state, - get_encoder_block_groups, iter_streamable_dits, move_non_block_state_to_device, prepare_pipeline_components, set_encoder_layerwise_state, ) from .config import DIT_COMPONENT -from .module_collector import ModuleDiscovery -from .offload_plan import OffloadPlan, get_offload_plan +from .plan_resolver import ResolvedComponent, resolve_offload_plan from .tensor_utils import ( clear_block_storage, clear_tensor_storage, @@ -279,7 +278,7 @@ def remove_block_hook(module: nn.Module) -> None: def _install_layerwise_hook_group( - blocks: list[nn.Module] | nn.ModuleList, + blocks: Sequence[nn.Module] | nn.ModuleList, device: torch.device, stream: Any, pin_memory: bool, @@ -327,36 +326,29 @@ def _install_layerwise_hook_group( def enable_plan_encoder_layerwise_offload( - module: nn.Module, - name: str, - plan: OffloadPlan | None, + component: ResolvedComponent, *, device: torch.device, stream: current_omni_platform.Stream, pin_memory: bool, - stage_on_demand: bool = False, - strict: bool = False, ) -> bool: - """Apply rank-local layerwise hooks to plan-declared encoder stacks.""" + """Apply rank-local layerwise hooks to the encoder's resolved stacks.""" + module = component.module if getattr(module, "_omni_layerwise_enabled", False): return True - hooks: list[LayerwiseOffloadHook] = [] - hooked_blocks: list[nn.Module] = [] - block_groups = get_encoder_block_groups( - module, - name, - plan, - strict=strict, - ) + block_groups = [stack.blocks for stack in component.stacks] if not block_groups: return False + + hooks: list[LayerwiseOffloadHook] = [] + hooked_blocks: list[nn.Module] = [] try: for blocks in block_groups: group_hooks = _install_layerwise_hook_group(blocks, device, stream, pin_memory) hooks.extend(group_hooks) hooked_blocks.extend(blocks) - if not stage_on_demand: + if not component.on_demand: move_non_block_state_to_device(module, block_groups, device) except BaseException: run_cleanup_steps( @@ -376,7 +368,7 @@ def enable_plan_encoder_layerwise_offload( ) logger.info( "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)", - name, + component.path, sum(len(blocks) for blocks in block_groups), len(block_groups), ) @@ -440,41 +432,27 @@ class LayerWiseOffloadBackend(OffloadBackend): logger.warning("LayerWiseOffloadBackend already enabled") return - modules = ModuleDiscovery.discover(pipeline) - plan = get_offload_plan(pipeline) - if not modules.dits and self.config.offloads(DIT_COMPONENT): - message = "No DiT/transformer modules found for selected DiT layerwise offload" - if self.config.components is not None: - raise ValueError(message) - logger.warning(message) + resolved = resolve_offload_plan(pipeline, self.config) + if not resolved.dits and self.config.offloads(DIT_COMPONENT): + # An explicit selection already failed while resolving. + logger.warning("No DiT/transformer modules found for selected DiT layerwise offload") return - def enable_encoder_blocks( - module: nn.Module, - name: str, - component_plan: OffloadPlan | None, - stage_on_demand: bool, - ) -> bool: + def enable_encoder_blocks(component: ResolvedComponent) -> bool: enabled = enable_plan_encoder_layerwise_offload( - module, - name, - component_plan, + component, device=self.device, stream=self.copy_stream, pin_memory=self.config.pin_cpu_memory, - stage_on_demand=stage_on_demand, - strict=self.config.components is not None, ) if enabled: # Record each successful installation immediately so a later # component failure can remove these hooks transactionally. - self._encoder_modules.append(module) + self._encoder_modules.append(component.module) return enabled prepare_pipeline_components( - modules, - self.config, - plan, + resolved, device=self.device, staged_components=self._staged_components, enable_encoder_blocks=enable_encoder_blocks, @@ -489,30 +467,17 @@ class LayerWiseOffloadBackend(OffloadBackend): ) return - logger.info("Applying layer-wise offloading on %s", modules.dit_names) + logger.info("Applying layer-wise offloading on %s", [component.path for component in resolved.dits]) # Apply block-wise offloading hook for each of the blocks in DiT model(s) # Note that there might exist multiple DiT models in specific pipelines - for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits( - modules, self.config, self.device, plan - ): - num_blocks = len(blocks) - if num_blocks <= 1: - if self.config.components is not None: - raise ValueError( - f"Selected DiT {dit_name!r} requires at least two streamable layerwise-offload blocks" - ) - logger.warning( - "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", - dit_name, - dit_module.__class__.__name__, - ) - dit_module.to(self.device) - continue + for component, stack in iter_streamable_dits(resolved, self.device): + dit_module = component.module + blocks = list(stack.blocks) # Move non-block modules to GPU (they stay resident) for name, m in dit_module.named_children(): - if name not in blocks_attr_names: + if name not in stack.attrs: m.to(self.device) logger.debug(f"Moved {name} to device {self.device}") else: @@ -540,7 +505,7 @@ class LayerWiseOffloadBackend(OffloadBackend): # zero once; later denoising iterations prefetch it from the ring. block_hooks[0].prefetch_layer(non_blocking=False) - logger.info(f"Layer-wise offloading enabled on {num_blocks} layers (blocks)") + logger.info(f"Layer-wise offloading enabled on {len(blocks)} layers (blocks)") # Track hooked blocks for cleanup self._blocks.append(blocks) diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index 4aa37ae8f..5cc0c24c7 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Declarative model capabilities for component layerwise offload. Models declare an ``_offload_plan`` class attribute on the pipeline. Both the @@ -22,9 +22,11 @@ class OffloadPlan: pipeline class. When present, both layerwise backends use it instead of model-specific backend branches, making new integrations data-driven. - If not declared, the offloader falls back to: - 1. ``_layerwise_offload_blocks_attrs`` on each DiT module class. - 2. Heuristic search for ``layers`` / ``blocks`` / ``h`` attributes. + If not declared, the offloader falls back to + ``_layerwise_offload_blocks_attrs`` on each DiT module class. + + :func:`~vllm_omni.diffusion.offloader.plan_resolver.resolve_offload_plan` + is the only consumer; backends read the resolved artifact it returns. Attributes: block_attrs: Maps DiT path → tuple of block-list attribute names. diff --git a/vllm_omni/diffusion/offloader/plan_resolver.py b/vllm_omni/diffusion/offloader/plan_resolver.py new file mode 100644 index 000000000..b2e951b4e --- /dev/null +++ b/vllm_omni/diffusion/offloader/plan_resolver.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""One topology resolver shared by every diffusion offload backend. + +Model topology is declared through several mechanisms that grew independently: +:class:`~vllm_omni.diffusion.models.interface.SupportsComponentDiscovery` class +variables, the pipeline-level :class:`OffloadPlan`, per-DiT +``_layerwise_offload_blocks_attrs``, and a leaf-name rule for encoders. This +module runs them in one defined order and returns one backend-neutral artifact. + +The resolver is pure: it moves no tensor, installs no hook, writes no module +attribute, and touches no process group. Everything it can reject is rejected +here, so plan-dependent failures happen before a backend mutates the model. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import cache +from typing import TYPE_CHECKING + +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery + +from .block_discovery import get_blocks_from_dit +from .component_utils import get_encoder_block_groups, validate_on_demand_component +from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT, OffloadStrategy +from .module_collector import ModuleDiscovery, PipelineModules +from .offload_plan import OffloadPlan, get_offload_plan + +if TYPE_CHECKING: + from .base import OffloadConfig + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class BlockStack: + """One ring of repeated blocks that a backend streams together. + + ``attrs`` names the owner attributes that contributed blocks; backends use + it to skip those children when they place the non-streamed remainder. It is + filled for DiT stacks only — encoder stacks are hooked as whole stacks and + their remainder is placed by tensor identity. + """ + + blocks: tuple[nn.Module, ...] + attrs: tuple[str, ...] = () + resident_head: int = 0 + + @property + def resident(self) -> tuple[nn.Module, ...]: + """Leading blocks that stay on the device instead of streaming.""" + return self.blocks[: self.resident_head] + + @property + def streaming(self) -> tuple[nn.Module, ...]: + """Blocks a hook ring transfers on demand.""" + return self.blocks[self.resident_head :] + + +@dataclass(frozen=True) +class ResolvedComponent: + """One pipeline component with its resolved, validated offload topology. + + ``selected`` means the active selector covers this component. ``on_demand`` + means the pipeline owns its residency through ``load_to_device`` / + ``offload_to_cpu``; a VAE staged by a legacy model plan is ``on_demand`` + without being selectable through the public component grammar. + """ + + path: str + module: nn.Module + selected: bool + stacks: tuple[BlockStack, ...] = () + on_demand: bool = False + + +@dataclass(frozen=True) +class ResolvedOffloadPlan: + """Backend-neutral topology for one pipeline under one offload config.""" + + dits: tuple[ResolvedComponent, ...] + encoders: tuple[ResolvedComponent, ...] + vaes: tuple[ResolvedComponent, ...] + residents: tuple[ResolvedComponent, ...] + # Raw discovery output and the model declaration remain reachable for the + # loader-owned mmap path and for DLO's nested-submodule staging, which + # still interpret topology themselves. + modules: PipelineModules + declaration: OffloadPlan | None = None + + @property + def components(self) -> tuple[ResolvedComponent, ...]: + return self.dits + self.encoders + self.vaes + self.residents + + +@cache +def _warn_legacy_discovery(pipeline_class: str) -> None: + """Warn once per pipeline class that component discovery guessed. + + The key is the fully qualified class name so two same-named pipeline + classes in different modules each get their own warning. + """ + logger.warning( + "%s does not declare SupportsComponentDiscovery; offload components were " + "discovered by scanning well-known attribute names. Declare _dit_modules, " + "_encoder_modules, and _vae_modules; the attribute scan is deprecated.", + pipeline_class, + ) + + +def _resolve_dit_stacks( + module: nn.Module, + path: str, + declaration: OffloadPlan | None, + config: OffloadConfig, + *, + explicit: bool, +) -> tuple[BlockStack, ...]: + """Resolve one DiT's streamed blocks and its device-resident head.""" + planned_attrs = None if declaration is None else declaration.block_attrs.get(path) + attrs, blocks = get_blocks_from_dit(module, planned_attrs) + if not blocks: + if explicit: + raise ValueError(f"Selected DiT {path!r} has no streamable layerwise-offload blocks") + logger.warning( + "Target layers (blocks) not found. Skipping offloading on %s (%s)", + path, + type(module).__name__, + ) + return () + + distributed = config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + resident_head = 0 + if ( + distributed + and config.dlo_resident_layers + and declaration is not None + and path in declaration.resident_dit_paths + ): + resident_head = min(config.dlo_resident_layers, len(blocks)) + + if len(blocks) - resident_head == 1: + if explicit and distributed: + raise ValueError( + f"Selected DiT {path!r} leaves only one streaming block after " + f"resident_layers={resident_head}; choose a resident count that " + "leaves zero or at least two streaming blocks" + ) + if explicit: + raise ValueError(f"Selected DiT {path!r} requires at least two streamable layerwise-offload blocks") + if distributed: + # One streamed block cannot form a prefetch ring; keeping it + # resident preserves the model instead of skipping placement. + logger.warning( + "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)", + path, + type(module).__name__, + ) + resident_head = len(blocks) + else: + logger.warning( + "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", + path, + type(module).__name__, + ) + return () + + return (BlockStack(blocks=tuple(blocks), attrs=tuple(attrs), resident_head=resident_head),) + + +def _resolve_encoder_stacks( + module: nn.Module, + path: str, + declaration: OffloadPlan | None, + config: OffloadConfig, + *, + explicit: bool, +) -> tuple[BlockStack, ...]: + """Resolve one encoder's streamable stacks and check transfer safety.""" + if declaration is None: + return () + + group_size = 1 + if config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and config.uses_allgather(TEXT_ENCODER_COMPONENT): + group_size = config.dp_size + if group_size > 1 and path not in declaration.encoder_dlo_weight_replication: + # Every rank must reject an unsafe group, including ranks whose local + # encoder is an unloaded stub, so the check precedes block discovery. + raise ValueError( + f"Text encoder {path!r} cannot use DLO AllGather across the DiT offload group: " + "its loader-produced weights are not declared replicated across that group. " + "Set layer_options.text_encoder.weight_transfer='rank-local' in diffusion_offload_config " + "for encoder-TP or rank-specific layouts." + ) + + groups = get_encoder_block_groups(module, path, declaration, strict=explicit or group_size > 1) + return tuple(BlockStack(blocks=tuple(blocks)) for blocks in groups) + + +def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> ResolvedOffloadPlan: + """Resolve and validate one pipeline's offload topology. + + Raises ``ValueError`` for every topology that the requested configuration + cannot serve, before the caller places a module or installs a hook. + """ + modules = ModuleDiscovery.discover(pipeline) + declaration = get_offload_plan(pipeline) + if not isinstance(pipeline, SupportsComponentDiscovery) and (modules.dits or modules.encoders or modules.vaes): + pipeline_class = type(pipeline) + _warn_legacy_discovery(f"{pipeline_class.__module__}.{pipeline_class.__qualname__}") + + explicit = config.components is not None + layerwise = config.strategy in (OffloadStrategy.LAYER_WISE, OffloadStrategy.DISTRIBUTED_LAYER_WISE) + dit_selected = config.offloads(DIT_COMPONENT) + + if explicit: + if config.strategy is OffloadStrategy.MODEL_LEVEL: + # Model-level offload swaps the DiT against an encoder stage, so + # both sides must exist regardless of which side was selected. + if not modules.dits: + raise ValueError("Component-selective model offload requires a DiT/transformer module") + if not modules.encoders: + raise ValueError("Component-selective model offload requires an encoder execution stage") + elif dit_selected and not modules.dits: + raise ValueError("No DiT/transformer modules found for selected DiT offload") + if config.offloads(TEXT_ENCODER_COMPONENT) and not any( + config.offloads_encoder(path, declaration) for path in modules.encoder_names + ): + raise ValueError("No text encoder modules found for selected text_encoder offload") + + if config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and config.dlo_resident_layers: + resident_paths = frozenset() if declaration is None else declaration.resident_dit_paths + if not resident_paths.intersection(modules.dit_names): + message = ( + f"resident_layers={config.dlo_resident_layers} was requested, but this model declares " + "no matching resident_dit_paths" + ) + if explicit: + raise ValueError(message) + logger.warning("%s; all blocks will be streamed.", message) + + dits: list[ResolvedComponent] = [] + encoders: list[ResolvedComponent] = [] + vaes: list[ResolvedComponent] = [] + residents: list[ResolvedComponent] = [] + + for path, module in zip(modules.dit_names, modules.dits): + stacks: tuple[BlockStack, ...] = () + if dit_selected and layerwise: + stacks = _resolve_dit_stacks(module, path, declaration, config, explicit=explicit) + dits.append(ResolvedComponent(path=path, module=module, selected=dit_selected, stacks=stacks)) + + for path, module in zip(modules.encoder_names, modules.encoders): + selected = config.should_offload_encoder(path, declaration) + # Staged residency belongs to the layer backends; model-level offload + # swaps whole modules through its own hooks. + on_demand = bool( + selected and layerwise and declaration is not None and path in declaration.on_demand_component_paths + ) + if on_demand: + validate_on_demand_component(module, path) + stacks = () + if selected and layerwise: + stacks = _resolve_encoder_stacks(module, path, declaration, config, explicit=explicit) + if on_demand and not stacks and config.uses_allgather(TEXT_ENCODER_COMPONENT): + raise ValueError( + f"Text encoder {path!r} cannot use AllGather without a model-declared streamable block plan" + ) + if explicit and not (stacks or on_demand): + raise ValueError( + f"Selected text encoder {path!r} requires a model-declared streamable or on-demand plan" + ) + encoders.append( + ResolvedComponent( + path=path, + module=module, + selected=selected, + stacks=stacks, + on_demand=on_demand, + ) + ) + + for path, module in zip(modules.vae_names, modules.vaes): + # VAEs are not part of the public selector. A model plan may still own + # their residency, which the compatibility topology preserves. + legacy_staged = ( + layerwise and not explicit and declaration is not None and path in declaration.on_demand_component_paths + ) + if legacy_staged: + validate_on_demand_component(module, path) + vaes.append(ResolvedComponent(path=path, module=module, selected=False, on_demand=legacy_staged)) + + for path, module in zip(modules.resident_names, modules.resident_modules): + residents.append(ResolvedComponent(path=path, module=module, selected=False)) + + resolved = ResolvedOffloadPlan( + dits=tuple(dits), + encoders=tuple(encoders), + vaes=tuple(vaes), + residents=tuple(residents), + modules=modules, + declaration=declaration, + ) + _validate_unique_ownership(resolved) + return resolved + + +def _validate_unique_ownership(resolved: ResolvedOffloadPlan) -> None: + """Reject topologies where two components would hook the same blocks.""" + owner_by_block: dict[int, str] = {} + for component in resolved.components: + for block in (block for stack in component.stacks for block in stack.blocks): + owner = owner_by_block.setdefault(id(block), component.path) + if owner != component.path: + raise ValueError( + f"Offload block is claimed by both {owner!r} and {component.path!r}; " + "declare one owner for each block list" + ) + + +__all__ = [ + "BlockStack", + "ResolvedComponent", + "ResolvedOffloadPlan", + "resolve_offload_plan", +] diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index b6b596070..1187c8cee 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections.abc import Collection, Iterator from contextlib import contextmanager @@ -13,9 +13,8 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload -from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT -from .module_collector import ModuleDiscovery -from .offload_plan import get_offload_plan +from .config import DIT_COMPONENT +from .plan_resolver import resolve_offload_plan logger = init_logger(__name__) @@ -329,59 +328,44 @@ class ModelLevelOffloadBackend(OffloadBackend): ) return - modules = ModuleDiscovery.discover(pipeline) - plan = get_offload_plan(pipeline) - selected_encoders = [ - encoder - for encoder, name in zip(modules.encoders, modules.encoder_names) - if self.config.should_offload_encoder(name, plan) - ] - if self.config.components is not None: - if not modules.dits: - raise ValueError("Component-selective model offload requires a DiT/transformer module") - if not modules.encoders: - raise ValueError("Component-selective model offload requires an encoder execution stage") - if self.config.offloads(TEXT_ENCODER_COMPONENT) and not selected_encoders: - raise ValueError("No text encoder modules found for selected text_encoder module offload") + resolved = resolve_offload_plan(pipeline, self.config) + dits = [component.module for component in resolved.dits] + encoders = [component.module for component in resolved.encoders] + vaes = [component.module for component in resolved.vaes] + residents = [component.module for component in resolved.residents] + selected_encoders = [component.module for component in resolved.encoders if component.selected] - all_modules = [ - *modules.dits, - *modules.encoders, - *modules.vaes, - *modules.resident_modules, - ] + all_modules = [*dits, *encoders, *vaes, *residents] initial_devices = _capture_tensor_devices(all_modules) try: - for encoder in modules.encoders: + for encoder in encoders: encoder.to(self.device) - for vae in modules.vaes: + for vae in vaes: vae.to(self.device, non_blocking=True) - for resident in modules.resident_modules: + for resident in residents: resident.to(self.device) - if not modules.dits: + if not dits: logger.warning("No DiT/transformer modules found, skipping model-level offloading") return - if not modules.encoders: - for dit in modules.dits: + if not encoders: + for dit in dits: dit.to(self.device) logger.warning("No encoder modules found, skipping model-level offloading") return apply_sequential_offload( - dit_modules=modules.dits, - encoder_modules=modules.encoders, + dit_modules=dits, + encoder_modules=encoders, device=self.device, pin_memory=self.config.pin_cpu_memory, use_hsdp=self.config.use_hsdp, - offload_dit_modules=( - modules.dits if self.config.components is None or self.config.offloads(DIT_COMPONENT) else () - ), + offload_dit_modules=(dits if self.config.offloads(DIT_COMPONENT) else ()), offload_encoder_modules=selected_encoders, ) except BaseException: try: - remove_sequential_offload([*modules.dits, *modules.encoders]) + remove_sequential_offload([*dits, *encoders]) except BaseException: logger.exception("Failed to remove every model-level hook during rollback") try: @@ -391,15 +375,19 @@ class ModelLevelOffloadBackend(OffloadBackend): raise # Track modules for cleanup - self._offload_modules = [*modules.dits, *modules.encoders] + self._offload_modules = [*dits, *encoders] self.enabled = True logger.info( "Model-level offloading enabled: %s <-> %s (mutual exclusion)%s", - ", ".join(modules.dit_names), - ", ".join(modules.encoder_names), - f"; resident on GPU: {', '.join(modules.resident_names)}" if modules.resident_names else "", + ", ".join(component.path for component in resolved.dits), + ", ".join(component.path for component in resolved.encoders), + ( + f"; resident on GPU: {', '.join(component.path for component in resolved.residents)}" + if resolved.residents + else "" + ), ) def disable(self) -> None: diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 46ddaad1d..39fff7df2 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -166,6 +166,11 @@ _DIFFUSION_MODELS = { "pipeline_boogu_image", "BooguImagePipeline", ), + "BooguImageTurboPipeline": ( + "boogu_image", + "pipeline_boogu_image", + "BooguImageTurboPipeline", + ), "LancePipeline": ( "lance", "pipeline_lance", @@ -549,6 +554,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "ZImagePipeline": "get_post_process_func", "OvisImagePipeline": "get_ovis_image_post_process_func", "BooguImagePipeline": "get_boogu_image_post_process_func", + "BooguImageTurboPipeline": "get_boogu_image_post_process_func", "WanPipeline": "get_wan22_post_process_func", "WanDMDPipeline": "get_wan22_post_process_func", "WanVACEPipeline": "get_wan22_vace_post_process_func", @@ -620,6 +626,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "BagelPipeline": "get_bagel_pre_process_func", "GlmImagePipeline": "get_glm_image_pre_process_func", "BooguImagePipeline": "get_boogu_image_pre_process_func", + "BooguImageTurboPipeline": "get_boogu_image_pre_process_func", "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func", "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func", "LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func", diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index df67f9d99..615df090f 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Video/audio muxing utilities using PyAV (no ffmpeg binary dependency).""" from __future__ import annotations import io +import queue +import threading from collections.abc import Iterable from fractions import Fraction from typing import Any, cast @@ -12,6 +14,153 @@ from typing import Any, cast import av import numpy as np +_CHUNKED_MP4_DONE = object() + + +def _validate_video_chunk(chunk: np.ndarray, *, width: int, height: int) -> None: + """Validate a ``(T, H, W, 3)`` uint8 RGB chunk against a session's frame size.""" + if chunk.ndim != 4 or chunk.shape[-1] != 3: + raise ValueError("video chunk must have shape (T, H, W, 3)") + if chunk.dtype != np.uint8: + raise ValueError("video chunk must have dtype uint8") + if chunk.shape[1] != height or chunk.shape[2] != width: + raise ValueError("video chunks in a session must use a consistent frame size") + + +class ChunkedMP4Encoder: + """Encode temporal video chunks while the producer is still decoding. + + A bounded queue and one muxing worker provide ordered backpressure while + keeping host memory bounded by ``max_pending`` chunks. Chunks use the same + ``(T, H, W, 3)`` uint8 contract as :func:`mux_video_audio_bytes`. + """ + + def __init__( + self, + *, + width: int, + height: int, + fps: float, + audio_waveform: np.ndarray | None = None, + audio_sample_rate: int | None = None, + max_pending: int = 2, + video_codec: str = "h264", + audio_codec: str = "aac", + crf: str = "18", + video_codec_options: dict[str, str] | None = None, + ) -> None: + if max_pending <= 0: + raise ValueError("max_pending must be positive") + if width <= 0 or height <= 0: + raise ValueError("width and height must be positive") + self.width = int(width) + self.height = int(height) + self._queue: queue.Queue[object] = queue.Queue(maxsize=max_pending) + self._result: bytes | None = None + self._error: BaseException | None = None + self._closed = False + self._aborted = False + self._input_done = False + self._state_lock = threading.Lock() + + def run() -> None: + try: + self._result = mux_av_video_audio_bytes( + self._frames(), + width=self.width, + height=self.height, + audio_waveform=audio_waveform, + fps=fps, + audio_sample_rate=audio_sample_rate, + video_codec=video_codec, + audio_codec=audio_codec, + crf=crf, + video_codec_options=video_codec_options, + ) + except BaseException as exc: + self._error = exc + self._drain_until_done() + + self._thread = threading.Thread(target=run, name="chunked-mp4", daemon=True) + self._thread.start() + + def _drain_until_done(self) -> None: + # A flush/mux failure can occur after _frames consumed the sentinel. + if not self._input_done: + while self._queue.get() is not _CHUNKED_MP4_DONE: + pass + + def _frames(self): + while True: + chunk = self._queue.get() + if chunk is _CHUNKED_MP4_DONE: + self._input_done = True + return + assert isinstance(chunk, np.ndarray) + for frame_data in chunk: + yield av.VideoFrame.from_ndarray(frame_data, format="rgb24") + + def _validate_chunk(self, chunk: np.ndarray) -> None: + _validate_video_chunk(chunk, width=self.width, height=self.height) + + def _raise_if_failed(self) -> None: + if self._error is not None: + raise self._error + if self._closed: + raise RuntimeError("ChunkedMP4Encoder is already closed") + + def push(self, chunk: np.ndarray) -> None: + """Queue a chunk; asynchronous failures surface here or at finish().""" + self._validate_chunk(chunk) + self._raise_if_failed() + self._queue.put(chunk) + self._raise_if_failed() + + def _send_done(self) -> None: + self._queue.put(_CHUNKED_MP4_DONE) + + def finish(self) -> bytes: + """Flush the muxer and return complete progressive MP4 bytes.""" + with self._state_lock: + if not self._closed: + self._closed = True + self._send_done() + self._thread.join() + if self._error is not None: + raise self._error + if self._aborted: + raise RuntimeError("ChunkedMP4Encoder was aborted") + assert self._result is not None + return self._result + + def abort(self) -> None: + """Discard pending chunks and join the worker after cancellation. + + Any chunk already being encoded must finish before the worker exits. + """ + with self._state_lock: + if not self._closed: + self._closed = True + self._aborted = True + while True: + try: + self._queue.get_nowait() + except queue.Empty: + break + self._send_done() + self._thread.join() + + close = finish + + def __enter__(self) -> ChunkedMP4Encoder: + return self + + def __exit__(self, exc_type, exc, tb) -> None: + if exc_type is None: + self.finish() + else: + self.abort() + class FragmentedMP4Muxer: """Incrementally mux video frames into one fragmented MP4 byte stream.""" @@ -57,12 +206,7 @@ class FragmentedMP4Muxer: """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.") + _validate_video_chunk(video_frames, width=self._stream.width, height=self._stream.height) for frame_data in video_frames: frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24") diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index 5158dbfbf..efc935a5d 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/vllm_omni/diffusion/worker/utils.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Per-request mutable state for step-wise diffusion execution.""" from __future__ import annotations diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 812f4c3e7..bd6a71562 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -381,9 +381,9 @@ class OmniEngineArgs(EngineArgs): if tokenizer_subfolder: # Download just the tokenizer files from the subfolder try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - local_dir = snapshot_download( + local_dir = hf_api().snapshot_download( model_path, allow_patterns=[ f"{tokenizer_subfolder}/tokenizer*", @@ -518,7 +518,7 @@ class OrchestratorArgs: # === Config Files === deploy_config: str | None = None - stage_overrides: str | None = None # raw JSON string; parsed downstream + stage_overrides: dict[str, dict[str, Any]] | None = None # 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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 4cec16382..6850e6b29 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -13,7 +13,6 @@ from __future__ import annotations import asyncio import concurrent.futures import copy -import json import queue import shutil import threading @@ -21,30 +20,23 @@ import time import uuid import weakref from collections.abc import Mapping, Sequence -from dataclasses import asdict from typing import Any, Literal, cast import janus -import torch -from omegaconf import OmegaConf from vllm import envs as vllm_envs from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.input_processor import InputProcessor -from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override +from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config from vllm_omni.config.stage_config import ( DuplexSessionRuntimeConfig, + PipelineConfig, load_deploy_config, ) from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS -from vllm_omni.diffusion.data import ( - DiffusionParallelConfig, - parse_attention_config, - resolve_model_class_name, -) -from vllm_omni.diffusion.io_support import get_diffusion_output_type from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.async_engine_utils import ( SHUTDOWN_ENQUEUE_TIMEOUT_S, @@ -88,10 +80,7 @@ from vllm_omni.engine.stage_runtime import ( create_stage_runtime, ) from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin -from vllm_omni.entrypoints.utils import ( - load_and_resolve_stage_configs, - parse_stage_overrides, -) +from vllm_omni.entrypoints.utils import parse_stage_overrides from vllm_omni.inputs.data import OmniInteractionPrompt, OmniSamplingParams from vllm_omni.metrics.prometheus import OmniRequestCounter @@ -99,6 +88,27 @@ logger = init_logger(__name__) _STARTUP_POLL_INTERVAL_S = 1.0 _REQUEST_QUEUE_MAXSIZE = 256 +_ConfigResolutionResult = OmniConfigResolution | tuple[str | None, list[Any], str | None] + + +def load_and_resolve_stage_configs( + model: str, + kwargs: dict[str, Any], + *, + trust_remote_code: bool | None, + deploy_config_path: str | None, + stage_overrides: Mapping[str, Mapping[str, Any]] | None, + strategy_config_path: str | None, +) -> OmniConfigResolution: + """Compatibility seam delegating to the single config resolver.""" + return resolve_omni_config( + model, + trust_remote_code=trust_remote_code, + cli_overrides=kwargs, + deploy_config_path=deploy_config_path, + stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, + ) class AsyncOmniEngine: @@ -198,41 +208,28 @@ class AsyncOmniEngine: self._omni_master_port, ) - # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand. + # Keep the historical tuple return from _resolve_stage_configs while + # retaining the richer resolver result for pipeline-wide settings. + # Overrides of that private seam fall back to the factory below. deploy_config_path = kwargs.get("deploy_config") - # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not - # specified" so stage-config resolution can defer to the deploy yaml's - # per-stage value (see ``with_trust_remote_code_override``). The - # restriction path below loads the top-level HF config via vLLM's - # ``get_config``, which needs a real bool, so collapse ``None`` to the - # default ``False`` here (#5495). - pipeline_config = StageConfigFactory.get_pipeline_config( - model=model, - trust_remote_code=bool(trust_remote_code), - deploy_config_path=deploy_config_path, - ) - self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () - self._duplex_runtime_extension_path = ( - pipeline_config.duplex_runtime_extension if pipeline_config is not None else None - ) - self.duplex_serving_adapter_path = ( - pipeline_config.duplex_serving_adapter if pipeline_config is not None else None - ) - self._duplex_control_enabled = bool(pipeline_config and pipeline_config.duplex_control_enabled) - self.duplex_session_config = DuplexSessionRuntimeConfig() - if deploy_config_path is not None: - self.duplex_session_config = load_deploy_config(deploy_config_path).duplex_session - - # Tri-state: None means "not specified" — the deploy yaml's per-stage - # trust_remote_code stays in effect. An explicit True/False here is a - # global override (precedence: caller > deploy yaml > default False); - # the merge rule lives in with_trust_remote_code_override. - kwargs = with_trust_remote_code_override(kwargs, trust_remote_code) + self._config_resolution: OmniConfigResolution | None = None self.config_path, self.stage_configs = self._resolve_stage_configs( model, kwargs, trust_remote_code=trust_remote_code, ) + if self._config_resolution is None: + pipeline_config = StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=bool(trust_remote_code), + deploy_config_path=deploy_config_path, + ) + self._set_pipeline_runtime_config(pipeline_config, self.config_path) + else: + self._set_pipeline_runtime_config( + self._config_resolution.pipeline_config, + self._config_resolution.config_path, + ) self.num_stages = len(self.stage_configs) stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None @@ -921,52 +918,6 @@ class AsyncOmniEngine: len(companions), ) - @staticmethod - def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | None: - if cache_backend == "cache_dit": - return { - "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": None, - "scm_steps_policy": "dynamic", - } - if cache_backend == "tea_cache": - return { - "rel_l1_thresh": 0.2, - } - if cache_backend == "mag_cache": - return { - "mag_threshold": 0.24, - "mag_max_skip_steps": 5, - "mag_retention_ratio": 0.1, - } - if cache_backend == "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 - def _normalize_cache_config(cache_backend: str | None, cache_config: Any | None) -> Any | None: - if isinstance(cache_config, str): - try: - cache_config = json.loads(cache_config) - except json.JSONDecodeError: - logger.warning("Invalid cache_config JSON, using defaults.") - cache_config = None - if cache_config is None and cache_backend not in (None, "", "none"): - cache_config = AsyncOmniEngine._get_default_cache_config(cache_backend) - return cache_config - def _detect_pd_config(self) -> dict[str, Any] | None: """Detect PD (Prefill-Decode) disaggregation config from stage_configs. Returns a dict with 'pd_pair' and 'bootstrap_addr', or None. @@ -1009,219 +960,6 @@ class AsyncOmniEngine: "prefill_engine_id": prefill_engine_id, } - @staticmethod - def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: - """Create a default single-stage diffusion config from kwargs.""" - # We temporally create a default config for diffusion stage. - # In the future, we should merge the default config with the user-provided config. - normalized_kwargs = dict(kwargs) - default_sampling_params = normalized_kwargs.get("default_sampling_params") - if isinstance(default_sampling_params, str): - try: - default_sampling_params = json.loads(default_sampling_params) - except json.JSONDecodeError: - logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.") - default_sampling_params = None - if not isinstance(default_sampling_params, dict): - default_sampling_params = None - stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {} - if normalized_kwargs.get("dtype") is None: - normalized_kwargs["dtype"] = "auto" - - # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. - if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str): - if not isinstance(normalized_kwargs["dtype"], torch.dtype): - raise TypeError( - f"Provided dtype must be a string or torch.dtype, got {type(normalized_kwargs['dtype']).__name__}" - ) - normalized_kwargs["dtype"] = str(normalized_kwargs["dtype"]).removeprefix("torch.") - - cache_backend = normalized_kwargs.get("cache_backend", "none") - cache_config = AsyncOmniEngine._normalize_cache_config( - cache_backend, - normalized_kwargs.get("cache_config", None), - ) - - parallel_config = normalized_kwargs.get("parallel_config") - if isinstance(parallel_config, dict): - parallel_config = DiffusionParallelConfig.from_dict(parallel_config) - if parallel_config is None: - ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1 - ring_degree = normalized_kwargs.get("ring_degree") or 1 - allgather_degree = normalized_kwargs.get("allgather_degree") or 1 - ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" - ulysses_a2a_permute = bool(normalized_kwargs.get("ulysses_a2a_permute", False)) - sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") - pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 - data_parallel_size = normalized_kwargs.get("data_parallel_size") - tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 - 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" - text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1 - enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False - 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 = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree - - parallel_config = DiffusionParallelConfig( - pipeline_parallel_size=pipeline_parallel_size, - data_parallel_size=data_parallel_size, - tensor_parallel_size=tensor_parallel_size, - enable_expert_parallel=enable_expert_parallel, - sequence_parallel_size=sequence_parallel_size, - ulysses_degree=ulysses_degree, - ring_degree=ring_degree, - allgather_degree=allgather_degree, - ulysses_mode=ulysses_mode, - ulysses_a2a_permute=ulysses_a2a_permute, - cfg_parallel_size=cfg_parallel_size, - vae_patch_parallel_size=vae_patch_parallel_size, - vae_parallel_mode=vae_parallel_mode, - text_encoder_tp_size=text_encoder_tp_size, - use_hsdp=use_hsdp, - hsdp_shard_size=hsdp_shard_size, - hsdp_replicate_size=hsdp_replicate_size, - ) - - num_gpus = normalized_kwargs.get("num_gpus") - if num_gpus is not None: - num_gpus = int(num_gpus) - parallel_config.resolve_data_parallel_size(num_gpus) - - num_devices = max(1, int(parallel_config.world_size)) - devices = ",".join(str(i) for i in range(num_devices)) - model_class_name = kwargs.get("model_class_name", None) - final_output_type = get_diffusion_output_type(model_class_name) - - attention_config = None - if ( - kwargs.get("diffusion_attention_config") is not None - or kwargs.get("diffusion_attention_backend") is not None - or kwargs.get("fastvideo_vsa_topk") is not None - ): - attention_config = parse_attention_config( - kwargs.get("diffusion_attention_config"), - attention_backend=kwargs.get("diffusion_attention_backend"), - fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"), - ) - - extras = dict(kwargs.get("extras") or {}) - for key, default in ( - ("auxiliary_text_encoder", None), - ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), - ): - top_level_value = kwargs.get(key) - if top_level_value is not None: - extras[key] = top_level_value - else: - extras.setdefault(key, default) - - stage_engine_args = { - "max_num_seqs": kwargs.get("max_num_seqs") or 1, - "parallel_config": parallel_config, - # Default-stage construction bypasses the structured projection. - # Runner selection remains owned by the selected engine/platform. - "engine_backend": kwargs.get("engine_backend", "default"), - "model_class_name": kwargs.get("model_class_name", None), - "task_type": kwargs.get("task_type", None), - "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, - "cache_config": cache_config, - "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), - "diffusion_offload_config": kwargs.get("diffusion_offload_config", None), - "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), - "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), - "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), - "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), - "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), - "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"), - "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"), - "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0), - "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), - "diffusion_compile_granularity": ( - "regional" - if kwargs.get("diffusion_compile_granularity") is None - else kwargs["diffusion_compile_granularity"] - ), - "diffusion_compile_dynamic": ( - True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"] - ), - "fa_deterministic": bool(kwargs.get("fa_deterministic", False)), - "boundary_ratio": kwargs.get("boundary_ratio", None), - "flow_shift": kwargs.get("flow_shift", None), - "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), - "lora_path": kwargs.get("lora_path", None), - "lora_scale": kwargs.get("lora_scale", 1.0), - "lora_backend": kwargs.get("lora_backend", "peft"), - "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), - "worker_extension_cls": kwargs.get("worker_extension_cls", None), - "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]), - "distributed_executor_backend": kwargs.get("distributed_executor_backend"), - "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), - "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False), - "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32), - "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": extras, - **( - { - "profiler_config": asdict(kwargs["profiler_config"]) - if hasattr(kwargs["profiler_config"], "__dataclass_fields__") - else kwargs["profiler_config"] - } - if kwargs.get("profiler_config") is not None - else {} - ), - } - if num_gpus is not None: - stage_engine_args["num_gpus"] = num_gpus - # Only set dtype if it was already explicitly passed and normalized - if "dtype" in normalized_kwargs: - stage_engine_args["dtype"] = normalized_kwargs["dtype"] - - # New split fields for diffusers adapter kwargs. - if kwargs.get("diffusers_load_kwargs") is not None: - stage_engine_args["diffusers_load_kwargs"] = kwargs["diffusers_load_kwargs"] - if kwargs.get("diffusers_call_kwargs") is not None: - stage_engine_args["diffusers_call_kwargs"] = kwargs["diffusers_call_kwargs"] - - default_stage_cfg = [ - { - "stage_id": 0, - "stage_type": "diffusion", - "runtime": { - "process": True, - "devices": devices, - }, - "engine_args": stage_engine_args, - "engine_input_source": [], - "default_sampling_params": stage_default_sampling_params, - "final_output": True, - "final_output_type": final_output_type, - } - ] - default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" - return default_stage_cfg - def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: """Apply a strategy-derived ``omni_lb_policy`` to the engine. @@ -1251,6 +989,29 @@ class AsyncOmniEngine: ) self._omni_lb_policy = str(derived) + @staticmethod + def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list[dict[str, Any]]: + """Compatibility seam for the factory-owned diffusion fallback.""" + return StageConfigFactory.create_default_diffusion(kwargs) + + def _set_pipeline_runtime_config( + self, + pipeline_config: PipelineConfig | None, + config_path: str | None, + ) -> None: + """Initialize engine-wide settings resolved from pipeline metadata.""" + self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () + self._duplex_runtime_extension_path = ( + pipeline_config.duplex_runtime_extension if pipeline_config is not None else None + ) + self.duplex_serving_adapter_path = ( + pipeline_config.duplex_serving_adapter if pipeline_config is not None else None + ) + self._duplex_control_enabled = bool(pipeline_config and pipeline_config.duplex_control_enabled) + self.duplex_session_config = DuplexSessionRuntimeConfig() + if config_path is not None: + self.duplex_session_config = load_deploy_config(config_path).duplex_session + def _resolve_stage_configs( self, model: str, @@ -1264,161 +1025,48 @@ class AsyncOmniEngine: if legacy_arg in kwargs: raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.") + # log_stats is captured by __init__; its CLI-only negative alias must + # not cross into per-stage structured config ownership validation. + kwargs.pop("disable_log_stats", 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) + # CLI callers arrive pre-parsed; offline Python callers may use the + # JSON-string form documented in recipes. + stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None)) # ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg; - # stage configs know the field as ``streaming_output``. The unregistered - # single-stage fallback translates it in - # ``_create_default_diffusion_stage_cfg``, but a registered pipeline - # resolves through StageConfigFactory, which only passes through keys the - # stage schema recognizes — so mirror the translation here. Only a truthy - # value is mirrored: the serve CLI always carries the flag's ``False`` - # default, which must not override a deploy YAML's ``streaming_output``. + # stage configs know the field as ``streaming_output``. Mirror only a + # truthy value so the CLI's False default does not override deploy YAML. if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None: kwargs["streaming_output"] = True - # Parse --stage-overrides JSON string if provided - stage_overrides = parse_stage_overrides(stage_overrides_json) - - # Unregistered diffusion checkpoints use the single-stage fallback - # below instead of StageConfigFactory, so fold stage-0 model extras - # into the fallback's input as well. Registered pipelines still get - # the complete override mapping through load_and_resolve_stage_configs. - default_stage_kwargs = kwargs - stage_zero_overrides = (stage_overrides or {}).get("0", {}) - if "extras" in stage_zero_overrides: - override_extras = stage_zero_overrides["extras"] - if not isinstance(override_extras, Mapping): - raise TypeError("stage 0 extras must be a mapping") - default_stage_kwargs = { - **kwargs, - "extras": { - **(kwargs.get("extras") or {}), - **override_extras, - }, - } - - def create_default_stage_config() -> list: - fallback_kwargs = dict(default_stage_kwargs) - if not fallback_kwargs.get("model_class_name"): - model_class_name = resolve_model_class_name( - model, - fallback_kwargs.get("diffusion_load_format", "default"), - ) - if model_class_name is not None: - fallback_kwargs["model_class_name"] = model_class_name - return self._create_default_diffusion_stage_cfg(fallback_kwargs) - - config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs( - model, - kwargs, - trust_remote_code=trust_remote_code, - default_stage_cfg_factory=create_default_stage_config, - deploy_config_path=deploy_config_path, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, + resolution = cast( + _ConfigResolutionResult, + load_and_resolve_stage_configs( + model, + kwargs, + trust_remote_code=trust_remote_code, + deploy_config_path=deploy_config_path, + stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, + ), ) + if isinstance(resolution, OmniConfigResolution): + self._config_resolution = resolution + config_path = resolution.config_path + stage_configs = list(resolution.stage_configs) + strategy_lb_policy = resolution.omni_lb_policy + else: + # Compatibility for overrides of the historical tuple-returning + # seam. Production always receives OmniConfigResolution above. + config_path, stage_configs, strategy_lb_policy = resolution # 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: - if not hasattr(cfg, "engine_args") or cfg.engine_args is None: - cfg.engine_args = OmegaConf.create({}) - global_sleep_mode = kwargs.get("enable_sleep_mode") - if global_sleep_mode is not None: - if not hasattr(cfg.engine_args, "enable_sleep_mode") or cfg.engine_args.enable_sleep_mode is None: - cfg.engine_args.enable_sleep_mode = global_sleep_mode - if getattr(cfg, "stage_type", None) != "diffusion": - continue - if not hasattr(cfg, "engine_args") or cfg.engine_args is None: - cfg.engine_args = OmegaConf.create({}) - additional_config = kwargs.get("additional_config") - if additional_config is not None: - current_additional_config = getattr(cfg.engine_args, "additional_config", None) - if current_additional_config in (None, {}): - cfg.engine_args.additional_config = additional_config - if kwargs.get("lora_path") is not None: - if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None: - cfg.engine_args.lora_path = kwargs["lora_path"] - lora_scale = kwargs.get("lora_scale") - if lora_scale is None: - # Backwards compatibility for older callers. - lora_scale = kwargs.get("static_lora_scale") - if lora_scale is not None: - if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None: - cfg.engine_args.lora_scale = lora_scale - if kwargs.get("lora_backend") is not None: - if not hasattr(cfg.engine_args, "lora_backend") or cfg.engine_args.lora_backend is None: - cfg.engine_args.lora_backend = kwargs["lora_backend"] - if ( - kwargs.get("diffusion_attention_config") is not None - or kwargs.get("diffusion_attention_backend") is not None - or kwargs.get("fastvideo_vsa_topk") is not None - ): - has_stage_attention = ( - getattr(cfg.engine_args, "diffusion_attention_config", None) is not None - or getattr(cfg.engine_args, "diffusion_attention_backend", None) is not None - ) - if not has_stage_attention: - cfg.engine_args.diffusion_attention_config = parse_attention_config( - kwargs.get("diffusion_attention_config"), - attention_backend=kwargs.get("diffusion_attention_backend"), - fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"), - ) - 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") - or cfg.engine_args.quantization_config is None - ): - cfg.engine_args.quantization_config = quantization_config - # Inject profiler flags for diffusion stages - for profiler_key in ( - "enable_diffusion_pipeline_profiler", - "enable_ar_profiler", - ): - val = kwargs.get(profiler_key) - if val: - if not hasattr(cfg.engine_args, profiler_key) or not getattr( - cfg.engine_args, profiler_key, False - ): - setattr(cfg.engine_args, profiler_key, val) - quantization = kwargs.get("quantization") - if quantization is not None: - if not hasattr(cfg.engine_args, "quantization") or cfg.engine_args.quantization is None: - cfg.engine_args.quantization = quantization - diffusion_kv_cache_dtype = kwargs.get("diffusion_kv_cache_dtype") - if diffusion_kv_cache_dtype is not None: - if ( - not hasattr(cfg.engine_args, "diffusion_kv_cache_dtype") - or cfg.engine_args.diffusion_kv_cache_dtype is None - ): - cfg.engine_args.diffusion_kv_cache_dtype = diffusion_kv_cache_dtype - diffusion_kv_cache_skip_steps = kwargs.get("diffusion_kv_cache_skip_steps") - if diffusion_kv_cache_skip_steps is not None: - if ( - not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_steps") - or cfg.engine_args.diffusion_kv_cache_skip_steps is None - ): - cfg.engine_args.diffusion_kv_cache_skip_steps = diffusion_kv_cache_skip_steps - diffusion_kv_cache_skip_layers = kwargs.get("diffusion_kv_cache_skip_layers") - if diffusion_kv_cache_skip_layers is not None: - if ( - not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_layers") - or cfg.engine_args.diffusion_kv_cache_skip_layers is None - ): - cfg.engine_args.diffusion_kv_cache_skip_layers = diffusion_kv_cache_skip_layers - except Exception as e: - logger.warning("Failed to inject LoRA config for stage: %s", e) - - return config_path, stage_configs + return cast(str, config_path), stage_configs # ==================== Public API ==================== @@ -1944,7 +1592,7 @@ class AsyncOmniEngine: tokens generated before abort (empty for diffusion / no OP state). """ if not request_ids or getattr(self, "_shutdown_called", False): - return + return [] if self.request_queue is None: raise RuntimeError("request_queue is not initialized") transport = self._correlated_rpc_client @@ -1971,7 +1619,7 @@ class AsyncOmniEngine: result_msg = await loop.run_in_executor(None, _wait) except Exception as exc: if getattr(self, "_shutdown_called", False) and is_abort_transport_shutdown(exc): - return + return [] raise if not result_msg.success: raise RuntimeError(result_msg.error or "abort failed") diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index e86dc7673..299e1dd07 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -75,7 +75,7 @@ from vllm_omni.engine.messages import ( 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, StageUnavailableError -from vllm_omni.errors import DEFAULT_CLIENT_ERROR_TYPE +from vllm_omni.errors import DEFAULT_CLIENT_ERROR_TYPE, OmniClientError from vllm_omni.metrics import definitions as metric_defs from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger @@ -1647,9 +1647,11 @@ class Orchestrator: stage_id: int, error: str, *, + status_code: int = HTTPStatus.BAD_REQUEST.value, + error_type: str = DEFAULT_CLIENT_ERROR_TYPE, close_duplex_sessions: bool = False, ) -> None: - """Fail one request with a non-fatal 400 (bad input, engine survives). + """Fail one request with a non-fatal client error (400 by default). The non-fatal counterpart of `_fail_request_dead_stage`: emits a client-error ErrorMessage (default `fatal=False`) so the engine keeps @@ -1658,8 +1660,8 @@ class Orchestrator: await self.output_async_queue.put( ErrorMessage( error=error, - status_code=HTTPStatus.BAD_REQUEST.value, - error_type=DEFAULT_CLIENT_ERROR_TYPE, + status_code=status_code, + error_type=error_type, request_id=req_id, stage_id=stage_id, ) @@ -2073,6 +2075,8 @@ class Orchestrator: ) if ( req_state.streaming.enabled + # A failed forward may already have cleaned up this request. + and self.request_states.get(req_id) is req_state and finished and not final_only_finished and not self._is_duplex_session_request(req_state) @@ -2734,6 +2738,16 @@ class Orchestrator: req_state.prompt, streaming_context=req_state.streaming, ) + except OmniClientError as exc: + await self._fail_request_client_error( + req_id, + next_logical, + str(exc), + status_code=exc.status_code, + error_type=exc.error_type, + close_duplex_sessions=self._is_duplex_session_request(req_state), + ) + return except Exception as exc: logger.exception( "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s", diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index f96cbd1d3..87baaed5f 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -396,36 +396,54 @@ class StageEngineCoreClientBase(StageClientBase): and the original prompt. """ if self.custom_process_input_func is not None: - signature = inspect.signature(self.custom_process_input_func) - extra_kwargs: dict[str, Any] = {} - if "next_stage_hf_config" in signature.parameters: - # Let a processor size the next stage's prompt from that - # stage's model config (e.g. a talker whose engine positions - # must cover a speaker-prompt prefill). - extra_kwargs["next_stage_hf_config"] = self._stage_hf_config - # Match the context parameter by name, including the - # underscore-prefixed spelling some processors use (e.g. - # MiniCPM-o's ``llm2tts(..., _streaming_context)``), so bridge - # state keeps flowing to them. - if "streaming_context" in signature.parameters or "_streaming_context" in signature.parameters: - return self.custom_process_input_func( - source_outputs, - prompt, - self.requires_multimodal_data, - streaming_context, - **extra_kwargs, - ) - return self.custom_process_input_func( - source_outputs, - prompt, - self.requires_multimodal_data, - **extra_kwargs, - ) + return self._call_custom_process_input(source_outputs, prompt, streaming_context) if not self.engine_input_source: raise ValueError(f"engine_input_source empty for stage {self.stage_id}") return _default_process_engine_inputs(source_outputs, prompt, self.requires_multimodal_data) + def _call_custom_process_input( + self, + source_outputs: list[Any], + prompt: Any, + streaming_context: Any | None, + ) -> list[OmniTokensPrompt]: + """Call a stage input processor with its explicitly requested context.""" + processor = self.custom_process_input_func + assert processor is not None + signature = inspect.signature(processor) + extra_kwargs: dict[str, Any] = {} + if "next_stage_hf_config" in signature.parameters: + # Let a processor size the next stage's prompt from that + # stage's model config (e.g. a talker whose engine positions + # must cover a speaker-prompt prefill). + extra_kwargs["next_stage_hf_config"] = self._stage_hf_config + target_model_config = signature.parameters.get("target_model_config") + if target_model_config is not None: + # The JoyAI bridge needs the Talker tokenizer and model config + # to calculate the exact prompt length. + if target_model_config.kind is not inspect.Parameter.KEYWORD_ONLY: + raise TypeError("target_model_config must be a keyword-only parameter") + extra_kwargs["target_model_config"] = self.vllm_config.model_config + # Match the context parameter by name, including the + # underscore-prefixed spelling some processors use (e.g. + # MiniCPM-o's ``llm2tts(..., _streaming_context)``), so bridge + # state keeps flowing to them. + if "streaming_context" in signature.parameters or "_streaming_context" in signature.parameters: + return processor( + source_outputs, + prompt, + self.requires_multimodal_data, + streaming_context, + **extra_kwargs, + ) + return processor( + source_outputs, + prompt, + self.requires_multimodal_data, + **extra_kwargs, + ) + async def collective_rpc_async( self, method: str, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index f2bc42ab1..1e57036a0 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -31,7 +31,6 @@ from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer from vllm.sampling_params import SamplingParams from vllm.tokenizers import cached_tokenizer_from_config -from vllm.transformers_utils.repo_utils import hf_api from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.input_processor import InputProcessor @@ -53,12 +52,13 @@ from vllm_omni.config.stage_config import StageType from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs 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.entrypoints.utils import filter_dataclass_kwargs 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 +from vllm_omni.transformers_utils.repo_utils import hf_api logger = init_logger(__name__) @@ -1087,7 +1087,8 @@ def _project_omni_stage_engine_args( ), ( stage_config.runtime_config, - frozenset({"devices", "num_replicas", "env", "num_gpus"}), + frozenset({"devices", "num_replicas", "env", "num_gpus"}) + | (frozenset({"additional_config"}) if is_diffusion else frozenset()), ), ): engine_args.update( @@ -1811,7 +1812,7 @@ def release_device_locks(lock_fds: list[int]) -> None: 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. + """Load omni transfer config from the resolver-selected deploy config. Resolves ``base_config`` inheritance (CI overlay → base deploy YAML) so that connectors defined in the base config are visible to the transfer @@ -1820,12 +1821,11 @@ def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> from vllm_omni.distributed.omni_connectors import load_omni_transfer_config try: - resolved_config_path = config_path or resolve_model_config_path(model) - if resolved_config_path is None: + if config_path is None: return None from vllm_omni.config.stage_config import resolve_deploy_yaml - resolved_dict = resolve_deploy_yaml(resolved_config_path) + resolved_dict = resolve_deploy_yaml(config_path) return load_omni_transfer_config(config_dict=resolved_dict) except Exception as e: logger.warning("[stage_init] Failed to load transfer config: %s", e) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index e7de5f51c..0ab6c1da7 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -739,10 +739,16 @@ class StagePool: def _infer_output_unit_type(self, request_outputs: list[Any], *, token_count: int) -> str: final_output_type = getattr(self.stage_client, "final_output_type", None) - if self._has_image_output(request_outputs) or final_output_type in {"image", "images"}: - return "image" - if self._has_video_output(request_outputs) or final_output_type in {"video", "videos"}: + # Prefer declared modality over payload heuristics: video diffusion often + # stores frames in ``images`` (see serving_video / output_formatter). + if final_output_type in {"video", "videos"}: return "video" + if final_output_type in {"image", "images"}: + return "image" + if self._has_video_output(request_outputs): + return "video" + if self._has_image_output(request_outputs): + return "image" if self._has_audio_output(request_outputs) or final_output_type == "audio": return "audio" if self._has_trajectory_latent_output(request_outputs) or self._has_latent_output(request_outputs): @@ -778,6 +784,10 @@ class StagePool: total_videos = sum(self._count_videos(ro) for ro in request_outputs) if total_videos > 0: return total_videos + # Video payloads are commonly carried on ``images`` for diffusion. + total_images = sum(self._count_images(ro) for ro in request_outputs) + if total_images > 0: + return total_images if unit_type == "latent": total_latents = sum( self._count_value_units(getattr(ro, "trajectory_latents", None)) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 4026907e1..948f7e53a 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -1416,19 +1416,9 @@ class AsyncOmni(EngineClient, OmniBase): ``EngineCore.sleep(level>=1)`` already clears the P1 receiver cache. Clearing P0 avoids hash-only follow-up requests after that reset. """ - processor = getattr(self, "input_processor", None) - if processor is None: - processor = getattr(self.engine, "input_processor", None) - cache = getattr(processor, "mm_processor_cache", None) - if cache is None: - logger.debug("[AsyncOmni] reset_mm_cache: no frontend mm_processor_cache") - return - for name in ("clear", "reset", "clear_cache"): - fn = getattr(cache, name, None) - if callable(fn): - fn() - return - logger.debug("[AsyncOmni] reset_mm_cache: cache has no clear/reset method") + renderer = self.renderer + if renderer is not None: + await renderer.clear_mm_cache_async() async def reset_encoder_cache(self) -> None: """Reset the encoder cache for all stages. diff --git a/vllm_omni/entrypoints/cli/benchmark/cli_args.py b/vllm_omni/entrypoints/cli/benchmark/cli_args.py index 2cf668769..21f528954 100644 --- a/vllm_omni/entrypoints/cli/benchmark/cli_args.py +++ b/vllm_omni/entrypoints/cli/benchmark/cli_args.py @@ -92,6 +92,16 @@ def add_multi_stage_cli_args(parser: argparse.ArgumentParser) -> None: "and tpop controls both text TPOP and internal stream TPOP." ), ) + group.add_argument( + "--omni-request-timeout-s", + type=float, + default=None, + help=( + "Total per-request timeout in seconds for benchmark HTTP requests. Timed-out requests are " + "recorded as failed instead of stalling the run on a hung server. Defaults to 900 when " + "omitted; set <= 0 to restore the legacy 6 h cap." + ), + ) def add_diffusion_cli_args(parser: argparse.ArgumentParser) -> None: @@ -103,8 +113,8 @@ def add_diffusion_cli_args(parser: argparse.ArgumentParser) -> None: type=str, default="think", help=( - "Default bot_task form field for --backend openai-image-edits-omni " - "(/v1/images/edits). " + "Default bot_task form field for image edits " + "(--backend openai-image-edits-omni or --endpoint /v1/images/edits). " 'Use --extra-body \'{"bot_task":"..."}\' to override per run.' ), ) @@ -352,6 +362,12 @@ def preprocess_serve_args(args: argparse.Namespace) -> None: raise ValueError("OmniInteract requires --max-concurrency to be positive") extra_body = dict(getattr(args, "extra_body", None) or {}) bot_task = getattr(args, "bot_task", None) - if getattr(args, "backend", None) == "openai-image-edits-omni" and bot_task is not None: + backend = getattr(args, "backend", None) + endpoint = getattr(args, "endpoint", None) + # serve.py remaps implicit backend to the endpoint path for image edits; + # inject bot_task for both the named backend and /v1/images/edits. + if bot_task is not None and ( + backend in ("openai-image-edits-omni", "/v1/images/edits") or endpoint == "/v1/images/edits" + ): extra_body.setdefault("bot_task", bot_task) args.extra_body = extra_body diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index ea56acb1a..78b212f98 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -14,6 +14,7 @@ import math import os import signal from types import FrameType +from typing import Any, cast import uvloop from vllm.entrypoints.cli.types import CLISubcommand @@ -23,6 +24,7 @@ from vllm.logger import init_logger from vllm_omni.entrypoints.cli.logo import log_logo from vllm_omni.entrypoints.openai.api_server import omni_run_server +from vllm_omni.entrypoints.utils import parse_stage_overrides from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace logger = init_logger(__name__) @@ -47,6 +49,17 @@ Search by using: `--help=` to explore options by section (e.g., """ +def _parse_stage_overrides(value: str) -> dict[str, dict[str, Any]]: + """Adapt shared stage-override validation to argparse's error type.""" + try: + parsed = parse_stage_overrides(value) + except ValueError as exc: + raise argparse.ArgumentTypeError(str(exc)) from exc + if parsed is None: + raise argparse.ArgumentTypeError("--stage-overrides requires a JSON object") + return parsed + + def _nonneg_finite_float(value: str) -> float: """Argparse type for finite, non-negative floats (rejects nan/inf).""" try: @@ -286,7 +299,7 @@ class OmniServeCommand(CLISubcommand): ) omni_config_group.add_argument( "--stage-overrides", - type=str, + type=_parse_stage_overrides, default=None, help="Per-stage JSON overrides. Example: " '\'{"0": {"gpu_memory_utilization": 0.8}, "2": {"enforce_eager": true}}\'', @@ -945,6 +958,7 @@ def run_headless(args: TrackingNamespace) -> None: from vllm.v1.executor.multiproc_executor import MultiprocExecutor from vllm.version import __version__ as VLLM_VERSION + from vllm_omni.config.resolver import resolve_omni_config from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_omni_kv_config_for_stage from vllm_omni.engine.stage_engine_startup import ( get_headless_replica_devices, @@ -959,10 +973,6 @@ def run_headless(args: TrackingNamespace) -> None: load_omni_transfer_config_for_model, prepare_engine_environment, ) - from vllm_omni.entrypoints.utils import ( - load_and_resolve_stage_configs, - parse_stage_overrides, - ) model = args.model stage_id: int | None = args.stage_id @@ -983,6 +993,13 @@ def run_headless(args: TrackingNamespace) -> None: # Filter down to a dict of things explicitly requested by the user args_dict = args.get_explicit_kwargs_dict() + # This CLI-only negative alias is consumed below when selecting the + # launcher log_stats value; it is not a per-stage config override. + args_dict.pop("disable_log_stats", None) + + deploy_config_path = args_dict.pop("deploy_config", None) + strategy_config_path = args_dict.pop("strategy_config", None) + stage_overrides = args_dict.pop("stage_overrides", None) # ``--replica-id`` is deprecated and ignored — replica ids are # auto-assigned by ``OmniMasterServer`` so headless processes carry @@ -996,33 +1013,27 @@ def run_headless(args: TrackingNamespace) -> None: "master server.", args.replica_id, ) + args_dict.pop("replica_id") - # 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( + resolved = resolve_omni_config( model, - args_dict, # store_true cannot express an explicit False: absent maps to None # ("not specified") so the deploy yaml's per-stage value applies. trust_remote_code=getattr(args, "trust_remote_code", None) or None, - deploy_config_path=args_dict.get("deploy_config"), + cli_overrides=args_dict, + deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, - strategy_config_path=args_dict.get("strategy_config"), + strategy_config_path=strategy_config_path, ) + config_path = resolved.config_path + stage_configs = list(resolved.stage_configs) - # Locate the stage config that matches stage_id. - stage_cfg = None - for cfg in stage_configs: - if cfg.stage_id == stage_id: - stage_cfg = cfg - break - if stage_cfg is None: + try: + stage_cfg = resolved.stage_by_id(stage_id) + except KeyError: raise ValueError( f"No stage config found for stage_id={stage_id}. Available stage ids: {[c.stage_id for c in stage_configs]}" - ) + ) from None prepare_engine_environment() per_replica_devices = get_headless_replica_devices(stage_cfg, stage_id, omni_dp_size_local) @@ -1037,7 +1048,7 @@ def run_headless(args: TrackingNamespace) -> None: omni_master_port=omni_master_port, omni_dp_size_local=omni_dp_size_local, per_replica_devices=per_replica_devices, - config_path=config_path, + config_path=cast(str, config_path), replica_bind_address=omni_replica_address, ) return diff --git a/vllm_omni/entrypoints/duplex/audio.py b/vllm_omni/entrypoints/duplex/audio.py index 644a7a297..747293a1a 100644 --- a/vllm_omni/entrypoints/duplex/audio.py +++ b/vllm_omni/entrypoints/duplex/audio.py @@ -21,8 +21,12 @@ MIN_INPUT_SAMPLE_RATE_HZ = 8_000 MAX_INPUT_SAMPLE_RATE_HZ = 192_000 -def validate_input_sample_rate_hz(sample_rate_hz: int | float) -> int: - if isinstance(sample_rate_hz, bool) or (isinstance(sample_rate_hz, float) and not math.isfinite(sample_rate_hz)): +def validate_input_sample_rate_hz(sample_rate_hz: object) -> int: + if ( + isinstance(sample_rate_hz, bool) + or not isinstance(sample_rate_hz, int | float) + or (isinstance(sample_rate_hz, float) and not math.isfinite(sample_rate_hz)) + ): raise ValueError("sample_rate_hz must be a finite integer") rate = int(sample_rate_hz) if rate != sample_rate_hz: @@ -35,14 +39,27 @@ def validate_input_sample_rate_hz(sample_rate_hz: int | float) -> int: def resample_pcm16_mono(raw: bytes, *, source_rate_hz: int, target_rate_hz: int) -> bytes: if source_rate_hz <= 0 or target_rate_hz <= 0 or source_rate_hz == target_rate_hz: return raw - samples = np.frombuffer(raw, dtype=" StreamingAudioResampler._max_polyphase_factor: + # The FIR size scales with the reduced ratio. Keep unusual client + # rates bounded while common audio rates use the high-quality path. + if samples.size <= 1: + return raw + target_size = max(1, int(round(samples.size * target_rate_hz / source_rate_hz))) + source_x = np.linspace(0.0, 1.0, num=samples.size, endpoint=True) + target_x = np.linspace(0.0, 1.0, num=target_size, endpoint=True) + resampled = np.interp(target_x, source_x, samples) + return np.clip(resampled, -32_768, 32_767).astype(" bytes: @@ -116,6 +133,24 @@ def wav_payload_to_pcm16(raw: bytes) -> tuple[bytes | None, int | None]: return None, None +def encode_float32_mono_wav_base64( + samples: np.ndarray, + *, + sample_rate_hz: int, +) -> str: + """Package normalized mono float32 samples as a PCM16 WAV data payload.""" + rate = validate_input_sample_rate_hz(sample_rate_hz) + normalized = np.ascontiguousarray(samples, dtype=np.float32).reshape(-1) + pcm16 = np.clip(np.rint(normalized * 32768.0), -32768, 32767).astype(" bool: - """Enable duplex routes only for deployments that explicitly opt in.""" + """Enable the realtime session handler for explicitly configured deployments.""" if stage_configs: for stage in stage_configs: session_mode = ( @@ -21,14 +21,13 @@ def should_enable_duplex_endpoint( return True if config_path: try: - from omegaconf import OmegaConf + from vllm_omni.config.stage_config import resolve_deploy_yaml - raw_config = OmegaConf.load(config_path) - session_mode = raw_config.get("session_mode") if hasattr(raw_config, "get") else None - if session_mode == "duplex": + raw_config = resolve_deploy_yaml(config_path) + if raw_config.get("session_mode") == "duplex" or isinstance(raw_config.get("duplex_session"), dict): return True except Exception as exc: - logger.warning("Failed to inspect duplex session_mode from %s: %s", config_path, exc) + logger.warning("Failed to inspect realtime session configuration from %s: %s", config_path, exc) return False diff --git a/vllm_omni/entrypoints/duplex/chat_fallback.py b/vllm_omni/entrypoints/duplex/chat_fallback.py index 31df63833..6e4d0a0e0 100644 --- a/vllm_omni/entrypoints/duplex/chat_fallback.py +++ b/vllm_omni/entrypoints/duplex/chat_fallback.py @@ -38,7 +38,14 @@ class ChatFallbackProjectorMixin: request = self._build_chat_request(session, request_id) result = await self._chat_service.create_chat_completion(request, raw_request=None) if isinstance(result, ErrorResponse): - await send_json({"type": "error", "error": result.message, "code": result.type or "chat_error"}) + error = result.error + await send_json( + { + "type": "error", + "error": error.message if error else "Chat request failed", + "code": error.type if error else "chat_error", + } + ) session.end_response(commit_text=False) return if hasattr(result, "__aiter__"): diff --git a/vllm_omni/entrypoints/duplex/protocol.py b/vllm_omni/entrypoints/duplex/protocol.py index 9fd3b0ef5..ff91e0db5 100644 --- a/vllm_omni/entrypoints/duplex/protocol.py +++ b/vllm_omni/entrypoints/duplex/protocol.py @@ -10,6 +10,16 @@ from enum import Enum from types import MappingProxyType from uuid import uuid4 +import numpy as np + +from vllm_omni.entrypoints.duplex.audio import ( + encode_float32_mono_wav_base64, +) +from vllm_omni.entrypoints.duplex.server_vad import ( + ServerVADConfig, + parse_session_turn_detection, +) + class DuplexOverlapPolicy(str, Enum): LISTEN_ONLY = "listen_only" @@ -231,10 +241,12 @@ class DuplexSessionConfig: overlap_barge_in_ms: int = 1200 overlap_silence_rms: float = 0.003 playback_commit_policy: str = DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value + turn_detection_configured: bool = False + server_vad: ServerVADConfig | None = None extra_body: dict[str, object] = field(default_factory=dict) def as_dict(self) -> dict[str, object]: - return { + payload: dict[str, object] = { "model": self.model, "modalities": list(self.modalities), "instructions": self.instructions, @@ -253,6 +265,9 @@ class DuplexSessionConfig: "playback_commit_policy": self.playback_commit_policy, "extra_body": dict(self.extra_body), } + if self.turn_detection_configured: + payload["turn_detection"] = self.server_vad.as_dict() if self.server_vad is not None else None + return payload @classmethod def from_event(cls, event: dict[str, object]) -> DuplexSessionConfig: @@ -295,9 +310,13 @@ class DuplexSessionConfig: config.playback_commit_policy = cls._normalize_playback_commit_policy(source["playback_commit_policy"]) if isinstance(source.get("modalities"), list) and all(isinstance(x, str) for x in source["modalities"]): config.modalities = list(source["modalities"]) + turn_detection_configured, server_vad = parse_session_turn_detection(source) if isinstance(source.get("extra_body"), dict): config.extra_body = normalize_native_duplex_key(dict(source["extra_body"])) extra = config.extra_body + raw_realtime_session = extra.get("realtime_session_payload") + if not turn_detection_configured and isinstance(raw_realtime_session, dict): + turn_detection_configured, server_vad = parse_session_turn_detection(raw_realtime_session) if isinstance(extra.get("overlap_policy"), str): config.overlap_policy = cls._normalize_overlap_policy(extra["overlap_policy"]) if isinstance(extra.get("overlap_short_ack_ms"), int | float): @@ -308,6 +327,9 @@ class DuplexSessionConfig: config.overlap_silence_rms = max(0.0, float(extra["overlap_silence_rms"])) if isinstance(extra.get("playback_commit_policy"), str): config.playback_commit_policy = cls._normalize_playback_commit_policy(extra["playback_commit_policy"]) + if turn_detection_configured: + config.turn_detection_configured = True + config.server_vad = server_vad return config @staticmethod @@ -354,6 +376,13 @@ class DuplexCommittedInput: input_commit_seq: int +@dataclass +class PendingServerVADTurn: + committed: DuplexCommittedInput + create_response: bool + item_id: str + + @dataclass class DuplexAssistantAudioTextMark: text_chars: int @@ -368,6 +397,13 @@ class InputBufferState: overlap_speech_ms: int = 0 reserved_input_bytes: int = 0 pending_turns: int = 0 + turn_detection_config_locked: bool = False + vad_prefix_audio: list[np.ndarray] = field(default_factory=list) + vad_prefix_samples: int = 0 + vad_utterance_audio: list[np.ndarray] = field(default_factory=list) + vad_speech_active: bool = False + active_server_vad_item_id: str | None = None + pending_server_vad_turn: PendingServerVADTurn | None = None @dataclass @@ -627,6 +663,9 @@ class DuplexSession: return active_turn_id is None or int(turn_id) == active_turn_id def append_history_message(self, message: dict[str, object]) -> None: + response_id = self.active_response_id + if response_id is not None: + self.reserve_history_item(f"item_{response_id}") self._conversation.messages.append(message) def stage_pending_history_item(self, item_id: str, message: dict[str, object]) -> None: @@ -653,6 +692,13 @@ class DuplexSession: self._input.reserved_input_bytes += size return True + @property + def turn_detection_config_locked(self) -> bool: + return self._input.turn_detection_config_locked + + def lock_turn_detection_config(self) -> None: + self._input.turn_detection_config_locked = True + def release_input_bytes(self, size: int) -> None: self._input.reserved_input_bytes = max(0, self._input.reserved_input_bytes - max(0, int(size))) @@ -677,6 +723,92 @@ class DuplexSession: self._input.pending_audio.append(DuplexAudioChunk(data=data, format=fmt, sample_rate_hz=sample_rate_hz)) self.turn_state = DuplexTurnState.USER_SPEAKING + def append_server_vad_frame( + self, + samples: np.ndarray, + *, + speech_started: bool, + speech_stopped: bool, + prefix_samples: int, + ) -> int: + """Store one detector frame while keeping commit-eligible audio session-owned. + + Returns the number of raw float32 bytes released when the rolling + pre-speech prefix is trimmed. + """ + # Session-owned audio must not retain the batch backing a detector frame. + frame = np.array(samples, dtype=np.float32, copy=True).reshape(-1) + released_samples = 0 + if not self._input.vad_speech_active: + if speech_started: + self._input.vad_utterance_audio = [*self._input.vad_prefix_audio, frame] + self._input.vad_prefix_audio = [] + self._input.vad_prefix_samples = 0 + self._input.vad_speech_active = True + else: + self._input.vad_prefix_audio.append(frame) + self._input.vad_prefix_samples += frame.size + while self._input.vad_prefix_samples > prefix_samples and self._input.vad_prefix_audio: + overflow = self._input.vad_prefix_samples - prefix_samples + first = self._input.vad_prefix_audio[0] + if first.size <= overflow: + self._input.vad_prefix_audio.pop(0) + self._input.vad_prefix_samples -= first.size + released_samples += first.size + else: + self._input.vad_prefix_audio[0] = first[overflow:].copy() + self._input.vad_prefix_samples -= overflow + released_samples += overflow + else: + self._input.vad_utterance_audio.append(frame) + + if speech_stopped: + self._input.vad_speech_active = False + if self._input.vad_speech_active or speech_started: + self.turn_state = DuplexTurnState.USER_SPEAKING + return released_samples * np.dtype(np.float32).itemsize + + def stage_server_vad_audio_for_commit(self) -> bool: + if not self._input.vad_utterance_audio: + return False + samples = np.concatenate(self._input.vad_utterance_audio) + encoded = encode_float32_mono_wav_base64(samples, sample_rate_hz=16_000) + self._input.pending_audio.append(DuplexAudioChunk(data=encoded, format="wav", sample_rate_hz=16_000)) + self._input.vad_utterance_audio.clear() + self._input.vad_speech_active = False + return True + + def clear_server_vad_audio(self) -> None: + self._input.vad_prefix_audio.clear() + self._input.vad_prefix_samples = 0 + self._input.vad_utterance_audio.clear() + self._input.vad_speech_active = False + self._input.active_server_vad_item_id = None + + @property + def active_server_vad_item_id(self) -> str | None: + return self._input.active_server_vad_item_id + + @property + def server_vad_utterance_bytes(self) -> int: + return sum(frame.nbytes for frame in self._input.vad_utterance_audio) + + def discard_uncommitted_server_vad_utterance(self) -> int: + released = self.server_vad_utterance_bytes + self._input.vad_utterance_audio.clear() + self._input.vad_speech_active = False + self._input.active_server_vad_item_id = None + self._input.pending_text.clear() + return released + + def begin_server_vad_speech(self, item_id: str) -> None: + self._input.active_server_vad_item_id = item_id + + def finish_server_vad_speech(self) -> str | None: + item_id = self._input.active_server_vad_item_id + self._input.active_server_vad_item_id = None + return item_id + def mark_user_input_activity(self) -> None: self.turn_state = DuplexTurnState.USER_SPEAKING @@ -692,10 +824,16 @@ class DuplexSession: self._input.pending_audio.clear() self._input.reserved_input_bytes = 0 self._input.pending_turns = 0 + self.clear_server_vad_audio() self.turn_state = DuplexTurnState.IDLE return cancelled - def commit_user_input(self) -> DuplexCommittedInput | None: + def commit_user_input( + self, + *, + append_history: bool = True, + retain_input_bytes: int = 0, + ) -> DuplexCommittedInput | None: text = "".join(self._input.pending_text).strip() audio_chunks = list(self._input.pending_audio) if not text and not audio_chunks: @@ -722,10 +860,11 @@ class DuplexSession: self._bind_active_response_to_input_commit(self._input.commit_seq + 1) self._input.commit_seq += 1 message = {"role": "user", "content": content} - self._conversation.messages.append(message) + if append_history: + self._conversation.messages.append(message) self._input.pending_text.clear() self._input.pending_audio.clear() - self._input.reserved_input_bytes = 0 + self._input.reserved_input_bytes = max(0, int(retain_input_bytes)) self.turn_state = DuplexTurnState.USER_COMMITTED return DuplexCommittedInput( message=message, @@ -734,6 +873,40 @@ class DuplexSession: input_commit_seq=self._input.commit_seq, ) + @property + def pending_server_vad_turn(self) -> PendingServerVADTurn | None: + return self._input.pending_server_vad_turn + + def stage_server_vad_turn( + self, + committed: DuplexCommittedInput, + *, + create_response: bool, + item_id: str, + ) -> None: + if self._input.pending_server_vad_turn is not None: + raise RuntimeError("A Server VAD turn is already pending") + response_id = self.active_response_id + if response_id is not None: + # A pending user turn advances input_commit_seq before the active + # response completes. Reserve its current history position first. + self.reserve_history_item(f"item_{response_id}") + self._input.pending_server_vad_turn = PendingServerVADTurn( + committed=committed, + create_response=create_response, + item_id=item_id, + ) + + def activate_server_vad_turn(self) -> PendingServerVADTurn | None: + pending = self._input.pending_server_vad_turn + if pending is None: + return None + self._input.pending_server_vad_turn = None + self._conversation.messages.append(pending.committed.message) + self.register_history_item(pending.item_id, pending.committed.message) + self.turn_state = DuplexTurnState.USER_COMMITTED + return pending + def commit_native_audio_input( self, *, @@ -1020,7 +1193,14 @@ class DuplexSession: response_input_commit_seq = self._response.active_response_input_commit_seq if response_input_commit_seq is None: response_input_commit_seq = self.input_commit_seq - response_history_is_late = self.input_commit_seq > response_input_commit_seq + response_item_id = f"item_{response_id}" if response_id is not None else None + response_history_is_reserved = response_item_id is not None and ( + response_item_id in self._conversation.item_ids + or response_item_id in self._conversation.history_item_placeholders + ) + response_history_is_late = ( + self.input_commit_seq > response_input_commit_seq and not response_history_is_reserved + ) assistant_text = "".join(self._response.assistant_text_buffer).strip() message = None if assistant_text: @@ -1041,9 +1221,8 @@ class DuplexSession: committed_text = "" if commit_text and committed_text: message = {"role": "assistant", "content": committed_text} - item_id = f"item_{response_id}" if response_id is not None else None - if item_id is not None and item_id in self._conversation.history_item_placeholders: - self._store_history_item_message(item_id, message) + if response_item_id is not None and response_history_is_reserved: + self._store_history_item_message(response_item_id, message) else: self._conversation.messages.append(message) effective_playback_policy = playback_commit_policy or self.config.playback_commit_policy @@ -1054,7 +1233,7 @@ class DuplexSession: and effective_playback_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value ): self.register_history_item(f"item_{response_id}", None) - elif response_id is not None and not assistant_text: + elif response_id is not None and (not assistant_text or not commit_text): self._discard_history_item_placeholder(f"item_{response_id}") self._response.assistant_text_buffer.clear() if not preserve_request: @@ -1113,6 +1292,11 @@ class DuplexSession: ) def delete_history_item(self, item_id: str) -> bool: + removed_pending_turn = False + pending_server_vad_turn = self._input.pending_server_vad_turn + if pending_server_vad_turn is not None and pending_server_vad_turn.item_id == item_id: + self._input.pending_server_vad_turn = None + removed_pending_turn = True response_id = item_id.removeprefix("item_") if item_id.startswith("item_") else None message = self._conversation.item_ids.pop(item_id, None) self._conversation.item_audio_text_marks.pop(item_id, None) @@ -1125,7 +1309,7 @@ class DuplexSession: if response_id is not None: self._conversation.assistant_response_snapshots.pop(response_id, None) if message is None: - return pending is not None or removed_placeholder + return removed_pending_turn or pending is not None or removed_placeholder try: self._conversation.messages.remove(message) except ValueError: @@ -1390,6 +1574,8 @@ class DuplexSession: self._response.active_response_turn_id = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False + self.clear_server_vad_audio() + self._input.pending_server_vad_turn = None self._clear_response_metrics() self._restore_response_config() @@ -1420,6 +1606,8 @@ class DuplexSession: "playback": self.playback.as_dict(), "capabilities": self.capabilities.as_dict(), } + if self.config.turn_detection_configured: + payload["turn_detection"] = self.config.server_vad.as_dict() if self.config.server_vad is not None else None if isinstance(self.config.extra_body.get("realtime_tools"), list): payload["tools"] = self.config.extra_body["realtime_tools"] if isinstance(self.config.extra_body.get("realtime_tool_choice"), str | dict): diff --git a/vllm_omni/entrypoints/duplex/realtime_input.py b/vllm_omni/entrypoints/duplex/realtime_input.py index 03710b43b..7fdadbcb2 100644 --- a/vllm_omni/entrypoints/duplex/realtime_input.py +++ b/vllm_omni/entrypoints/duplex/realtime_input.py @@ -7,27 +7,23 @@ import asyncio import base64 import binascii from collections.abc import Callable -from typing import Any from uuid import uuid4 import numpy as np from vllm_omni.entrypoints.duplex.audio import ( convert_input_audio_with_rate, + encode_float32_mono_wav_base64, + validate_input_sample_rate_hz, ) from vllm_omni.entrypoints.duplex.protocol import normalize_native_duplex_key from vllm_omni.entrypoints.duplex.realtime_state import ( REALTIME_INPUT_AUDIO_FORMATS, REALTIME_OUTPUT_AUDIO_FORMATS, RealtimeStateOwner, + realtime_default_sample_rate_hz, ) -from vllm_omni.entrypoints.duplex.vad import ( - SILERO_VAD_MIN_THRESHOLD, - ServerVADUnavailableError, - SileroStreamingVAD, - SileroVADConfig, - StreamingVADResult, -) +from vllm_omni.entrypoints.duplex.server_vad import ServerVADConfig, parse_session_turn_detection class RealtimeInputTranslator(RealtimeStateOwner): @@ -49,7 +45,6 @@ class RealtimeInputTranslator(RealtimeStateOwner): _conversation_items: dict[str, dict[str, object]] _item_truncation_cursors: dict[str, tuple[int, int]] _pending_outbound: asyncio.Queue[dict[str, object]] - _pending_commit_item_ids: asyncio.Queue[str] async def _send_realtime_payload(self, payload: dict[str, object]) -> None: raise NotImplementedError @@ -66,6 +61,17 @@ class RealtimeInputTranslator(RealtimeStateOwner): _response_is_done: Callable[[object], bool] + @property + def _turn_based_server_vad(self) -> bool: + return self._turn_detection is not None and not self._native_input_append + + def _reset_realtime_input_buffer_state(self) -> None: + self._input_speech_started = False + self._active_input_item_id = None + self._input_audio_buffer_has_audio = False + self._input_audio_buffer_had_non_speech = False + self._input_audio_buffer_transcript_parts.clear() + async def discard_pending_input_audio( self, *, @@ -88,11 +94,7 @@ class RealtimeInputTranslator(RealtimeStateOwner): "item_id": self._active_input_item_id, } ) - self._input_speech_started = False - self._active_input_item_id = None - self._input_audio_buffer_has_audio = False - self._input_audio_buffer_had_non_speech = False - self._input_audio_buffer_transcript_parts.clear() + self._reset_realtime_input_buffer_state() async def _send_realtime_input_ack(self, event: dict[str, object]) -> None: if event.get("type") != "conversation.item.create": @@ -100,7 +102,6 @@ class RealtimeInputTranslator(RealtimeStateOwner): item = event.get("item") if not isinstance(item, dict): return - item = self._normalize_conversation_item(item) previous_item_id = event.get("previous_item_id") if isinstance(previous_item_id, str): item["_previous_item_id"] = previous_item_id @@ -115,18 +116,40 @@ class RealtimeInputTranslator(RealtimeStateOwner): if event_type == "session.update": session = event.get("session") session_payload: dict[str, object] = session if isinstance(session, dict) else event - format_error = self._validate_realtime_session_audio_formats(session_payload) - if format_error is not None: + try: + turn_detection_configured, server_vad = parse_session_turn_detection(session_payload) + except ValueError as exc: await self._send_realtime_payload( self._realtime_error_payload( - "unsupported_audio_format", - format_error, + "unsupported_turn_detection", + str(exc), event_id=event.get("event_id"), + param="turn_detection", ) ) return None - turn_detection_error = self._validate_realtime_turn_detection(session_payload) + normalized_turn_detection = server_vad.as_dict() if server_vad is not None else None + previous_protocol_config = self._realtime_protocol_config() + try: + self._apply_realtime_session_defaults(session_payload) + except ValueError as exc: + self._set_realtime_protocol_config(previous_protocol_config) + await self._send_realtime_payload( + self._realtime_error_payload( + "unsupported_audio_format", + str(exc), + event_id=event.get("event_id"), + param="audio.input.format", + ) + ) + return None + turn_detection_error = self._validate_realtime_turn_detection( + session_payload, + configured=turn_detection_configured, + normalized=normalized_turn_detection, + ) if turn_detection_error is not None: + self._set_realtime_protocol_config(previous_protocol_config) await self._send_realtime_payload( self._realtime_error_payload( "unsupported_turn_detection", @@ -136,31 +159,38 @@ class RealtimeInputTranslator(RealtimeStateOwner): ) ) return None - self._prepare_realtime_turn_detection_update(session_payload) - self._apply_realtime_session_defaults(session_payload) - session_payload.update(self._realtime_overlap_fields(session_payload)) - if not self._opened: - self._opened = True - self._initial_session_update = True - return self._session_create_from_realtime(session_payload) - return { - "type": "turn.signal", - "event": "session.update", - "payload": session_payload, - "realtime_event_id": event.get("event_id"), - } - if event_type == "conversation.item.create": - item = event.get("item") - format_error = self._validate_conversation_item_audio_formats(item) - if format_error is not None: - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - format_error, - event_id=event.get("event_id"), - ) + turn_detection_candidate = self._prepare_realtime_turn_detection_update( + session_payload, + configured=turn_detection_configured, + normalized=normalized_turn_detection, + ) + try: + session_payload.update(self._realtime_overlap_fields(session_payload)) + if not self._opened: + self._opened = True + self._initial_session_update = True + translated = self._session_create_from_realtime(session_payload) + event_id = event.get("event_id") + if isinstance(event_id, str) and event_id: + translated["_realtime_event_id"] = event_id + else: + translated = { + "type": "turn.signal", + "event": "session.update", + "payload": session_payload, + "realtime_event_id": event.get("event_id"), + } + candidate_protocol_config = self._realtime_protocol_config() + finally: + self._set_realtime_protocol_config(previous_protocol_config) + if turn_detection_candidate[0] or candidate_protocol_config != previous_protocol_config: + self._pending_turn_detection_update = ( + *turn_detection_candidate, + asyncio.Event(), + candidate_protocol_config, ) - return None + return translated + if event_type == "conversation.item.create": return await self._conversation_item_to_duplex(event) if event_type == "conversation.item.delete": item_id = event.get("item_id") @@ -293,27 +323,20 @@ class RealtimeInputTranslator(RealtimeStateOwner): return None if event_type == "input_audio_buffer.append": audio = event.get("audio") or event.get("delta") - fmt, format_rate = self._parse_realtime_audio_format( - event.get("format") or event.get("input_audio_format") or self._input_audio_format - ) - sample_rate_hz = ( - event.get("sample_rate_hz") or event.get("sample_rate") or format_rate or self._input_sample_rate_hz - ) - if not self._is_supported_realtime_input_format(fmt): - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - f"Unsupported input_audio_format: {fmt}", - event_id=event.get("event_id"), - ) - ) - return None + event_format = event.get("format") or event.get("input_audio_format") + native_server_vad = self._turn_detection is not None and self._native_input_append + turn_based_server_vad = self._turn_based_server_vad try: - audio, fmt, sample_rate_hz = convert_input_audio_with_rate( - audio, - fmt, - sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, - ) + fmt, format_rate = self._parse_realtime_audio_format(event_format or self._input_audio_format) + sample_rate_hz = event.get("sample_rate_hz") or event.get("sample_rate") or format_rate + if sample_rate_hz is None: + sample_rate_hz = ( + realtime_default_sample_rate_hz(fmt) + if event_format is not None and fmt != self._input_audio_format + else self._input_sample_rate_hz + ) + if turn_based_server_vad: + sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz) except ValueError as exc: await self._send_realtime_payload( self._realtime_error_payload( @@ -324,63 +347,65 @@ class RealtimeInputTranslator(RealtimeStateOwner): ) ) return None - detector = self._server_vad - vad_result: StreamingVADResult | None = None - if isinstance(detector, SileroStreamingVAD): + channels = event.get("channels", 1) + if not self._is_supported_realtime_input_format(fmt): + await self._send_realtime_payload( + self._realtime_error_payload( + "unsupported_audio_format", + f"Unsupported input_audio_format: {fmt}", + event_id=event.get("event_id"), + ) + ) + return None + if turn_based_server_vad and (fmt != "pcm16" or channels != 1 or sample_rate_hz not in {16_000, 24_000}): + await self._send_realtime_payload( + self._realtime_error_payload( + "unsupported_audio_format", + "server_vad requires mono PCM16 input at 16 kHz or 24 kHz", + event_id=event.get("event_id"), + param="audio.input.format", + ) + ) + return None + looks_like_speech: bool | None = None + if not turn_based_server_vad: try: - vad_result = detector.process_base64(audio, fmt=fmt, sample_rate_hz=sample_rate_hz) - except (ServerVADUnavailableError, ValueError) as exc: - unavailable = isinstance(exc, ServerVADUnavailableError) + audio, fmt, sample_rate_hz = convert_input_audio_with_rate( + audio, + fmt, + sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, + ) + except ValueError as exc: await self._send_realtime_payload( self._realtime_error_payload( - "server_vad_unavailable" if unavailable else "bad_audio", + "bad_event", str(exc), event_id=event.get("event_id"), - param="turn_detection" if unavailable else "audio", + param="sample_rate_hz", ) ) return None - looks_like_speech = ( - vad_result.is_speech - if vad_result is not None - else self._input_looks_like_speech(event, audio=audio, fmt=fmt) - ) - self._input_audio_buffer_has_audio = self._input_audio_buffer_has_audio or ( - looks_like_speech and isinstance(audio, str) and bool(audio) - ) - self._input_audio_buffer_had_non_speech = self._input_audio_buffer_had_non_speech or ( - not looks_like_speech and isinstance(audio, str) and bool(audio) - ) - stopped_event: dict[str, object] | None = None - if vad_result is not None and vad_result.speech_stopped: - stopped_event = dict(event) - if vad_result.speech_end_ms is not None: - stopped_event["audio_end_ms"] = vad_result.speech_end_ms - if stopped_event is not None and vad_result is not None and vad_result.speech_active: - await self._emit_input_speech_stopped( - stopped_event, - item_id=self._active_input_item_id or f"item_{uuid4().hex}", + if not turn_based_server_vad and not native_server_vad: + looks_like_speech = self._input_looks_like_speech(event, audio=audio, fmt=fmt) + has_audio_payload = isinstance(audio, str) and bool(audio) + self._input_audio_buffer_has_audio = self._input_audio_buffer_has_audio or ( + looks_like_speech and has_audio_payload ) - if vad_result is not None and vad_result.speech_started: - started_event = dict(event) - if vad_result.speech_start_ms is not None: - started_event["audio_start_ms"] = vad_result.speech_start_ms - await self._emit_input_speech_started(started_event) - elif vad_result is None and looks_like_speech: - await self._emit_input_speech_started(event) - if stopped_event is not None and vad_result is not None and not vad_result.speech_active: - await self._emit_input_speech_stopped( - stopped_event, - item_id=self._active_input_item_id or f"item_{uuid4().hex}", + self._input_audio_buffer_had_non_speech = self._input_audio_buffer_had_non_speech or ( + not looks_like_speech and has_audio_payload ) - if looks_like_speech: - self._remember_input_transcript_hint(event) + if looks_like_speech: + await self._emit_input_speech_started(event) + self._remember_input_transcript_hint(event) payload = { "type": "input_audio_buffer.append", "audio": audio, "format": fmt, "sample_rate_hz": sample_rate_hz, } + event_id = event.get("event_id") + if self._turn_detection is not None and isinstance(event_id, str) and event_id: + payload["_realtime_event_id"] = event_id video_frames = event.get("video_frames") if video_frames is not None: frames_error = self._validate_realtime_video_frames(video_frames, event.get("max_slice_nums")) @@ -397,25 +422,27 @@ class RealtimeInputTranslator(RealtimeStateOwner): if isinstance(video_frames, list): payload["video_frames"] = [frame for frame in video_frames if isinstance(frame, str) and frame] self._copy_realtime_input_hints(event, payload) - payload["is_speech"] = looks_like_speech - if vad_result is not None: - payload["vad"] = { - "backend": "silero", - "is_speech": vad_result.is_speech, - "speech_active": vad_result.speech_active, - "speech_started": vad_result.speech_started, - "speech_stopped": vad_result.speech_stopped, - "speech_probability": vad_result.speech_probability, - } - if vad_result.speech_active: - payload["force_listen"] = True + if looks_like_speech is not None: + payload["is_speech"] = looks_like_speech return payload if event_type == "input_audio_buffer.commit": + if self._turn_based_server_vad: + await self._send_realtime_payload( + self._realtime_error_payload( + "server_vad_manual_commit_unsupported", + ( + "input_audio_buffer.commit is not supported while server_vad is enabled; " + "wait for server endpointing, or start a session with turn_detection=null " + "to use manual commit" + ), + event_id=event.get("event_id"), + ) + ) + return None if not self._input_audio_buffer_has_audio: if self._input_audio_buffer_had_non_speech: self._input_audio_buffer_had_non_speech = False self._active_input_item_id = None - self._reset_server_vad() return { "type": "input_audio_buffer.commit", "final": event.get("final", True), @@ -431,29 +458,22 @@ class RealtimeInputTranslator(RealtimeStateOwner): ) return None item_id = self._active_input_item_id or f"item_{uuid4().hex}" - await self._pending_commit_item_ids.put(item_id) await self._emit_input_speech_stopped(event, item_id=item_id) transcript = self._consume_input_transcript_hint() self._active_input_item_id = None self._input_audio_buffer_has_audio = False self._input_audio_buffer_had_non_speech = False - self._reset_server_vad() payload = { "type": "input_audio_buffer.commit", "final": event.get("final", True), - "realtime_item_id": item_id, + "item_id": item_id, "response_create": bool(event.get("response_create", False)), } if transcript: payload["transcript"] = transcript return payload if event_type == "input_audio_buffer.clear": - self._input_speech_started = False - self._active_input_item_id = None - self._input_audio_buffer_has_audio = False - self._input_audio_buffer_had_non_speech = False - self._input_audio_buffer_transcript_parts.clear() - self._reset_server_vad() + self._reset_realtime_input_buffer_state() return {"type": "input_audio_buffer.clear", "reason": event_type} if event_type == "output_audio_buffer.clear": payload = {"type": "output_audio_buffer.clear", "reason": event_type} @@ -643,19 +663,19 @@ class RealtimeInputTranslator(RealtimeStateOwner): "previous_item_id": previous_item_id, "item": item, }, - { - "type": "conversation.item.created", - "previous_item_id": previous_item_id, - "item": item, - }, ] def _conversation_item_done_event(self, item: dict[str, object]) -> dict[str, object]: item_id = item.get("id") previous_item_id = self._previous_item_id(item_id) if isinstance(item_id, str) else None if isinstance(item_id, str) and item_id: + already_known = item_id in self._conversation_items self._conversation_items[item_id] = item - self._last_conversation_item_id = item_id + # Completing an item that was already added must not move it to the + # end of the conversation. A later user item may have been added + # while this response was still in progress. + if not already_known: + self._last_conversation_item_id = item_id return { "type": "conversation.item.done", "previous_item_id": previous_item_id, @@ -737,17 +757,14 @@ class RealtimeInputTranslator(RealtimeStateOwner): self, *, item_id: str, - event: dict[str, Any], ) -> dict[str, object]: return { "type": "input_audio_buffer.committed", "previous_item_id": self._previous_item_id(item_id), "item_id": item_id, - "event": event, } def _session_create_from_realtime(self, session_payload: dict[str, object]) -> dict[str, object]: - self._apply_realtime_session_defaults(session_payload) model = session_payload.get("model") audio_config = session_payload.get("audio") audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None @@ -825,8 +842,13 @@ class RealtimeInputTranslator(RealtimeStateOwner): audio_input = audio_config.get("input") if isinstance(audio_input, dict): input_format = audio_input.get("format") + input_format_configured = input_format is not None input_format, input_rate = self._parse_realtime_audio_format(input_format) - if isinstance(input_format, str) and input_format.lower() in REALTIME_INPUT_AUDIO_FORMATS: + if input_format is not None and not ( + isinstance(input_format, str) and input_format.lower() in REALTIME_INPUT_AUDIO_FORMATS + ): + raise ValueError(f"Unsupported input_audio_format: {input_format}") + if isinstance(input_format, str): self._input_audio_format = input_format output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") output_rate_raw: object | None = None @@ -836,9 +858,13 @@ class RealtimeInputTranslator(RealtimeStateOwner): output_format = audio_output.get("format") output_rate_raw = audio_output.get("sample_rate_hz") or audio_output.get("sample_rate") output_format, output_rate = self._parse_realtime_audio_format(output_format) - if output_rate is None and isinstance(output_rate_raw, int | float) and output_rate_raw > 0: - output_rate = int(output_rate_raw) - if isinstance(output_format, str) and output_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: + if output_rate is None and output_rate_raw is not None: + output_rate = validate_input_sample_rate_hz(output_rate_raw) + if output_format is not None and not ( + isinstance(output_format, str) and output_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS + ): + raise ValueError(f"Unsupported output_audio_format: {output_format}") + if isinstance(output_format, str): self._output_audio_format = self._realtime_output_format(output_format) sample_rate = session_payload.get("sample_rate_hz") or session_payload.get("sample_rate") if sample_rate is None and isinstance(audio_config, dict): @@ -847,41 +873,19 @@ class RealtimeInputTranslator(RealtimeStateOwner): sample_rate = audio_input.get("sample_rate_hz") or audio_input.get("sample_rate") if sample_rate is None: sample_rate = input_rate - if isinstance(sample_rate, int | float) and sample_rate > 0: - self._input_sample_rate_hz = int(sample_rate) - if isinstance(output_rate, int | float) and output_rate > 0: - self._output_sample_rate_hz = int(output_rate) + if sample_rate is not None: + self._input_sample_rate_hz = validate_input_sample_rate_hz(sample_rate) + elif input_format_configured: + default_input_rate = realtime_default_sample_rate_hz(input_format) + if default_input_rate is not None: + self._input_sample_rate_hz = default_input_rate + if output_rate is not None: + self._output_sample_rate_hz = output_rate overlap_fields = self._realtime_overlap_fields(session_payload) overlap_silence_rms = overlap_fields.get("overlap_silence_rms") if isinstance(overlap_silence_rms, int | float): self._overlap_silence_rms = max(0.0, float(overlap_silence_rms)) - @classmethod - def _validate_realtime_session_audio_formats(cls, session_payload: dict[str, object]) -> str | None: - audio_config = session_payload.get("audio") - input_format: object = session_payload.get("input_audio_format") - if input_format is None and isinstance(audio_config, dict): - audio_input = audio_config.get("input") - if isinstance(audio_input, dict): - input_format = audio_input.get("format") - parsed_input, _ = cls._parse_realtime_audio_format(input_format) - if input_format is not None and not ( - isinstance(parsed_input, str) and parsed_input.lower() in REALTIME_INPUT_AUDIO_FORMATS - ): - return f"Unsupported input_audio_format: {input_format}" - - output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") - if output_format is None and isinstance(audio_config, dict): - audio_output = audio_config.get("output") - if isinstance(audio_output, dict): - output_format = audio_output.get("format") - parsed_output, _ = cls._parse_realtime_audio_format(output_format) - if output_format is not None and not ( - isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS - ): - return f"Unsupported output_audio_format: {output_format}" - return None - @classmethod def _validate_realtime_response_audio_formats(cls, response_payload: dict[str, object]) -> str | None: output_format: object = response_payload.get("output_audio_format") or response_payload.get("response_format") @@ -890,33 +894,16 @@ class RealtimeInputTranslator(RealtimeStateOwner): audio_output = audio_config.get("output") if isinstance(audio_output, dict): output_format = audio_output.get("format") - parsed_output, _ = cls._parse_realtime_audio_format(output_format) + try: + parsed_output, _ = cls._parse_realtime_audio_format(output_format) + except ValueError as exc: + return str(exc) if output_format is not None and not ( isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS ): return f"Unsupported output_audio_format: {output_format}" return None - @classmethod - def _validate_conversation_item_audio_formats(cls, item: object) -> str | None: - if not isinstance(item, dict): - return None - content = item.get("content") - if not isinstance(content, list): - return None - for part in content: - if not isinstance(part, dict): - continue - if part.get("type") not in {"input_audio", "audio"}: - continue - raw_format = part.get("format") - parsed_format, _ = cls._parse_realtime_audio_format(raw_format) - if raw_format is not None and not ( - isinstance(parsed_format, str) and parsed_format.lower() in REALTIME_INPUT_AUDIO_FORMATS - ): - return f"Unsupported input_audio format in conversation.item.create: {raw_format}" - return None - @staticmethod def _json_safe_realtime_payload(payload: dict[str, object]) -> dict[str, object]: clean: dict[str, object] = {} @@ -958,7 +945,7 @@ class RealtimeInputTranslator(RealtimeStateOwner): if not isinstance(raw_format, dict): return raw_format, None rate = raw_format.get("rate") or raw_format.get("sample_rate_hz") or raw_format.get("sample_rate") - sample_rate_hz = int(rate) if isinstance(rate, int | float) and rate > 0 else None + sample_rate_hz = validate_input_sample_rate_hz(rate) if rate is not None else None fmt = raw_format.get("type") or raw_format.get("format") if not isinstance(fmt, str): return raw_format, sample_rate_hz @@ -1107,124 +1094,122 @@ class RealtimeInputTranslator(RealtimeStateOwner): return "video_frames entries must be JPEG or PNG images" return None - @staticmethod - def _configured_realtime_turn_detection(session_payload: dict[str, object]) -> tuple[str | None, object]: - field = "turn_detection" if "turn_detection" in session_payload else None - selected = session_payload.get("turn_detection") - audio_config = session_payload.get("audio") - if isinstance(audio_config, dict): - audio_input = audio_config.get("input") - if isinstance(audio_input, dict) and "turn_detection" in audio_input: - field, selected = "audio.input.turn_detection", audio_input["turn_detection"] - return field, selected - - @classmethod - def _validate_realtime_turn_detection(cls, session_payload: dict[str, object]) -> str | None: - field, turn_detection = cls._configured_realtime_turn_detection(session_payload) - if field is None: + def _validate_realtime_turn_detection( + self, + session_payload: dict[str, object], + *, + configured: bool, + normalized: dict[str, object] | None, + ) -> str | None: + effective_turn_detection = normalized if configured else self._turn_detection + if effective_turn_detection is None: if session_payload.get("overlap_policy") == "barge_in_on_speech": - return ( - "overlap_policy='barge_in_on_speech' requires " - "turn_detection.type='server_vad' on the Realtime endpoint" - ) + return "overlap_policy='barge_in_on_speech' requires turn_detection.type='server_vad'" return None - if turn_detection is not None and not isinstance(turn_detection, dict): - return f"{field} must be null or an object with type='server_vad'" - if isinstance(turn_detection, dict): - if turn_detection.get("type") != "server_vad": - return f"{field}.type must be 'server_vad'" - threshold = turn_detection.get("threshold", 0.5) - if not cls._valid_vad_number(threshold, minimum=SILERO_VAD_MIN_THRESHOLD, strict=True, maximum=1): - return f"{field}.threshold must be greater than {SILERO_VAD_MIN_THRESHOLD} and at most 1" - for name in ("prefix_padding_ms", "silence_duration_ms", "min_speech_duration_ms"): - value = turn_detection.get(name) - if value is not None and not cls._valid_vad_number(value, minimum=0): - return f"{field}.{name} must be a non-negative number" - if turn_detection.get("interrupt_response", True) is not True: - return ( - f"{field}.interrupt_response=false is unsupported; " - "use turn_detection=null for model-owned listen/speak" - ) - desired_policy = "barge_in_on_speech" if turn_detection is not None else "listen_only" + + interrupt_response = effective_turn_detection.get("interrupt_response") + if interrupt_response is None: + # Serving resolves the default and validates overlap policy once + # the session's runtime capabilities are known. + return None + desired_policy = "barge_in_on_speech" if interrupt_response else "listen_only" overlap_policy = session_payload.get("overlap_policy") if isinstance(overlap_policy, str) and overlap_policy != desired_policy: return f"overlap_policy={overlap_policy!r} conflicts with turn_detection; expected {desired_policy!r}" return None - @staticmethod - def _valid_vad_number(value: object, *, minimum: float, maximum: float | None = None, strict: bool = False) -> bool: + def validate_server_vad_input_audio(self, session_payload: dict[str, object]) -> str | None: + """Validate turn-based VAD audio after serving resolves runtime capabilities.""" + candidate = self._pending_turn_detection_update + input_format, input_rate, *_ = candidate[3] if candidate is not None else self._realtime_protocol_config() + audio_config = session_payload.get("audio") + audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None + channels = audio_input.get("channels") if isinstance(audio_input, dict) else None + raw_format = session_payload.get("input_audio_format") + if raw_format is None and isinstance(audio_input, dict): + raw_format = audio_input.get("format") + if channels is None and isinstance(raw_format, dict): + channels = raw_format.get("channels") + if channels is None: + channels = 1 + if input_format != "pcm16" or channels != 1 or input_rate not in {16_000, 24_000}: + return "server_vad requires mono PCM16 input at 16 kHz or 24 kHz" + return None + + def _prepare_realtime_turn_detection_update( + self, + session_payload: dict[str, object], + *, + configured: bool, + normalized: dict[str, object] | None, + ) -> tuple[bool, dict[str, object] | None]: + if self._pending_turn_detection_update is not None: + raise RuntimeError("A Realtime session update is already pending") + + turn_detection = dict(normalized) if configured and normalized is not None else None + if configured: + if turn_detection is None: + session_payload.setdefault("overlap_policy", "listen_only") + elif "interrupt_response" in turn_detection: + session_payload.setdefault( + "overlap_policy", "barge_in_on_speech" if turn_detection["interrupt_response"] else "listen_only" + ) + session_payload["turn_detection"] = turn_detection + return configured, turn_detection + + def _realtime_protocol_config(self) -> tuple[str, int, str, int | None, float]: return ( - isinstance(value, int | float) - and not isinstance(value, bool) - and np.isfinite(float(value)) - and (float(value) > minimum if strict else float(value) >= minimum) - and (maximum is None or float(value) <= maximum) + self._input_audio_format, + self._input_sample_rate_hz, + self._output_audio_format, + self._output_sample_rate_hz, + self._overlap_silence_rms, ) - def _prepare_realtime_turn_detection_update(self, session_payload: dict[str, object]) -> None: - field, configured = self._configured_realtime_turn_detection(session_payload) - if field is None: - return - if self._pending_turn_detection_update is not None: - raise RuntimeError("A Realtime turn-detection update is already pending") + def _set_realtime_protocol_config( + self, + config: tuple[str, int, str, int | None, float], + ) -> None: + ( + self._input_audio_format, + self._input_sample_rate_hz, + self._output_audio_format, + self._output_sample_rate_hz, + self._overlap_silence_rms, + ) = config - turn_detection: dict[str, object] | None = None - detector: SileroStreamingVAD | None = None - if configured is None: - session_payload["overlap_policy"] = "listen_only" - else: - assert isinstance(configured, dict) - turn_detection = { - "type": "server_vad", - "interrupt_response": True, - "threshold": 0.5, - "prefix_padding_ms": 300, - "silence_duration_ms": 500, - "min_speech_duration_ms": 96, - **configured, - } - detector = SileroStreamingVAD( - SileroVADConfig( - threshold=float(turn_detection["threshold"]), - prefix_padding_ms=int(turn_detection["prefix_padding_ms"]), - silence_duration_ms=int(turn_detection["silence_duration_ms"]), - min_speech_duration_ms=max(32, int(turn_detection["min_speech_duration_ms"])), - ) - ) - session_payload["turn_detection"] = dict(turn_detection) - session_payload["overlap_policy"] = "barge_in_on_speech" - - self._pending_turn_detection_update = (turn_detection, detector, asyncio.Event()) + def pending_realtime_input_audio_config_changed(self) -> bool: + candidate = self._pending_turn_detection_update + if candidate is None: + return False + candidate_protocol_config = candidate[3] + return candidate_protocol_config[:2] != self._realtime_protocol_config()[:2] async def wait_for_realtime_turn_detection_update(self) -> None: candidate = self._pending_turn_detection_update if candidate is not None: await candidate[2].wait() - def commit_realtime_turn_detection_update(self) -> None: - self._resolve_realtime_turn_detection_update(commit=True) + def commit_realtime_turn_detection_update(self, server_vad: ServerVADConfig | None = None) -> None: + self._resolve_realtime_turn_detection_update(commit=True, server_vad=server_vad) def reject_realtime_turn_detection_update(self) -> None: self._resolve_realtime_turn_detection_update(commit=False) - def _resolve_realtime_turn_detection_update(self, *, commit: bool) -> None: + def _resolve_realtime_turn_detection_update( + self, *, commit: bool, server_vad: ServerVADConfig | None = None + ) -> None: candidate = self._pending_turn_detection_update if candidate is None: return - turn_detection, detector, resolved = candidate - discarded = self._server_vad if commit else detector + configured, turn_detection, resolved, protocol_config = candidate if commit: - self._turn_detection, self._server_vad = turn_detection, detector + self._set_realtime_protocol_config(protocol_config) + if configured: + self._turn_detection = server_vad.as_dict() if server_vad is not None else turn_detection + self._turn_detection_configured = True self._pending_turn_detection_update = None resolved.set() - reset = getattr(discarded, "reset", None) - if callable(reset): - reset() - - def _reset_server_vad(self) -> None: - reset = getattr(self._server_vad, "reset", None) - if callable(reset): - reset() @staticmethod def _input_audio_transcription_config(session_payload: dict[str, object]) -> dict[str, object] | None: @@ -1257,19 +1242,19 @@ class RealtimeInputTranslator(RealtimeStateOwner): return None async def _conversation_item_to_duplex(self, event: dict[str, object]) -> dict[str, object] | None: - item = event.get("item") - if not isinstance(item, dict): + raw_item = event.get("item") + if not isinstance(raw_item, dict): return None - item = self._normalize_conversation_item(item) + public_item = self._normalize_conversation_item(raw_item) previous_item_id = event.get("previous_item_id") if isinstance(previous_item_id, str): - item["_previous_item_id"] = previous_item_id - item_id = str(item["id"]) - item_type = item.get("type") - role = item.get("role") + public_item["_previous_item_id"] = previous_item_id + item_id = str(public_item["id"]) + item_type = public_item.get("type") + role = public_item.get("role") if item_type == "function_call_output": - call_id = item.get("call_id") - output = item.get("output") + call_id = public_item.get("call_id") + output = public_item.get("output") matching_call = any( known.get("type") == "function_call" and known.get("call_id") == call_id for known in self._conversation_items.values() @@ -1307,30 +1292,110 @@ class RealtimeInputTranslator(RealtimeStateOwner): # ``conversation.item.created`` event records it in the shared # input/output protocol state; on failure the client may retry the # same call_id instead of being rejected by provisional history. + content = public_item.get("content") + if not isinstance(content, list): + return None + parsed_audio_formats: dict[int, tuple[object, int | None]] = {} + for index, part in enumerate(content): + if not isinstance(part, dict) or part.get("type") not in {"input_audio", "audio"}: + continue + raw_format = part.get("format") + try: + parsed_format = self._parse_realtime_audio_format(raw_format or self._input_audio_format) + except ValueError as exc: + await self._send_realtime_payload( + self._realtime_error_payload( + "unsupported_audio_format", + str(exc), + event_id=event.get("event_id"), + param="item.content.audio", + ) + ) + return None + if not self._is_supported_realtime_input_format(parsed_format[0]): + await self._send_realtime_payload( + self._realtime_error_payload( + "unsupported_audio_format", + f"Unsupported input_audio format in conversation.item.create: {raw_format}", + event_id=event.get("event_id"), + param="item.content.audio", + ) + ) + return None + parsed_audio_formats[index] = parsed_format if item_type != "message" or role in {"assistant", "system"}: + event["item"] = public_item return { "type": "turn.signal", "event": "conversation.item.create", - "payload": {"item": item}, + "payload": {"item": public_item}, } - self._conversation_items[item_id] = item - content = item.get("content") - if not isinstance(content, list): - return None + server_vad_full_audio_item = self._turn_based_server_vad text_chunks: list[str] = [] audio_events: list[dict[str, object]] = [] - for part in content: + normalized_content: list[object] = [] + normalized_full_audio = False + for index, part in enumerate(content): if not isinstance(part, dict): + normalized_content.append(part) continue if part.get("type") in {"input_text", "text"} and isinstance(part.get("text"), str): text_chunks.append(str(part["text"])) if part.get("type") in {"input_audio", "audio"}: audio = part.get("audio") or part.get("data") - fmt, format_rate = self._parse_realtime_audio_format(part.get("format") or self._input_audio_format) - sample_rate_hz = ( - part.get("sample_rate_hz") or part.get("sample_rate") or format_rate or self._input_sample_rate_hz - ) - if not self._is_supported_realtime_input_format(fmt): + if not isinstance(audio, str) or not audio: + await self._send_realtime_payload( + self._realtime_error_payload( + "bad_audio", + "conversation.item.create input_audio must not be empty", + event_id=event.get("event_id"), + param="item.content.audio", + ) + ) + return None + raw_part_format = part.get("format") + fmt, format_rate = parsed_audio_formats[index] + explicit_sample_rate_hz = part.get("sample_rate_hz") + if explicit_sample_rate_hz is None: + explicit_sample_rate_hz = part.get("sample_rate") + if explicit_sample_rate_hz is None and isinstance(raw_part_format, dict): + explicit_sample_rate_hz = ( + raw_part_format.get("rate") + if "rate" in raw_part_format + else raw_part_format.get("sample_rate_hz", raw_part_format.get("sample_rate")) + ) + sample_rate_hz = explicit_sample_rate_hz or format_rate or self._input_sample_rate_hz + if server_vad_full_audio_item and ( + fmt != self._input_audio_format + or (explicit_sample_rate_hz is not None and explicit_sample_rate_hz != self._input_sample_rate_hz) + ): + await self._send_realtime_payload( + self._realtime_error_payload( + "bad_audio", + "conversation.item.create input_audio must use the Session input format and sample rate", + event_id=event.get("event_id"), + param="item.content.audio", + ) + ) + return None + if server_vad_full_audio_item: + try: + normalized_audio = self._normalize_complete_audio_item_part(audio, self._input_sample_rate_hz) + except ValueError as exc: + await self._send_realtime_payload( + self._realtime_error_payload( + "bad_audio", + str(exc), + event_id=event.get("event_id"), + param="item.content.audio", + ) + ) + return None + normalized_part = dict(part) + normalized_part.pop("data", None) + normalized_part.update(audio=normalized_audio, format="wav", sample_rate_hz=16_000) + normalized_content.append(normalized_part) + normalized_full_audio = True continue try: audio, fmt, sample_rate_hz = convert_input_audio_with_rate( @@ -1341,15 +1406,13 @@ class RealtimeInputTranslator(RealtimeStateOwner): except ValueError as exc: await self._send_realtime_payload( self._realtime_error_payload( - "bad_event", + "bad_audio", str(exc), event_id=event.get("event_id"), - param="sample_rate_hz", + param="item.content.audio", ) ) return None - if not isinstance(audio, str) or not audio: - continue speech_hints = dict(event) speech_hints.update(part) if not self._input_looks_like_speech(speech_hints, audio=audio, fmt=fmt): @@ -1364,31 +1427,69 @@ class RealtimeInputTranslator(RealtimeStateOwner): self._copy_realtime_input_hints(part, payload) self._copy_realtime_input_hints(event, payload) audio_events.append(payload) + continue + normalized_content.append(part) + if server_vad_full_audio_item and normalized_full_audio: + history_item = dict(public_item) + history_item["content"] = normalized_content + history_item["status"] = "completed" + event["item"] = public_item + return { + "type": "turn.signal", + "event": "conversation.item.create", + "payload": { + "item": public_item, + "history_item": history_item, + }, + } if audio_events: - self._input_audio_buffer_has_audio = True - transcript = self._input_transcript_from_item(item) + transcript = self._input_transcript_from_item(public_item) for extra_event in audio_events[1:]: self._pending_outbound.put_nowait(extra_event) commit_payload: dict[str, object] = { "type": "input_audio_buffer.commit", "final": True, - "realtime_item_id": item_id, + "item_id": item_id, "response_create": False, } if transcript: commit_payload["transcript"] = transcript self._pending_outbound.put_nowait(commit_payload) + # This complete item is already committed locally; a later ack + # must not be needed to start accepting the next input turn. + self._reset_realtime_input_buffer_state() + event["item"] = public_item return audio_events[0] if not text_chunks: return None - text_item = dict(item) + text_item = dict(public_item) text_item["status"] = "completed" + event["item"] = text_item return { "type": "turn.signal", "event": "conversation.item.create", "payload": {"item": text_item}, } + @staticmethod + def _normalize_complete_audio_item_part( + audio: str, + sample_rate_hz: int, + ) -> str: + """Package a complete Realtime audio item for the chat fallback path.""" + try: + raw = base64.b64decode(audio.strip(), validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError("conversation.item.create input_audio is not valid Base64") from exc + if not raw or len(raw) % np.dtype(" None: transcript = event.get("transcript") if not isinstance(transcript, str): @@ -1454,12 +1555,6 @@ class RealtimeInputTranslator(RealtimeStateOwner): } return None - def _pop_pending_commit_item_id(self) -> str: - try: - return self._pending_commit_item_ids.get_nowait() - except asyncio.QueueEmpty: - return f"item_{uuid4().hex}" - async def _emit_input_speech_started(self, event: dict[str, object]) -> None: if self._input_speech_started: return diff --git a/vllm_omni/entrypoints/duplex/realtime_output.py b/vllm_omni/entrypoints/duplex/realtime_output.py index e8c99f218..82ce9761e 100644 --- a/vllm_omni/entrypoints/duplex/realtime_output.py +++ b/vllm_omni/entrypoints/duplex/realtime_output.py @@ -39,6 +39,28 @@ class RealtimeOutputProjector: if event_type == "session.updated": session = self._realtime_session_payload(event.get("session")) return [{"type": "session.updated", "session": session}] + if event_type == "input_audio_buffer.speech_started": + item_id = event.get("item_id") + if isinstance(item_id, str): + self._active_input_item_id = item_id + self._input_speech_started = True + self._input_audio_buffer_has_audio = True + return [ + { + "type": "input_audio_buffer.speech_started", + "audio_start_ms": max(0, int(event.get("audio_start_ms", 0) or 0)), + "item_id": item_id, + } + ] + if event_type == "input_audio_buffer.speech_stopped": + self._input_speech_started = False + return [ + { + "type": "input_audio_buffer.speech_stopped", + "audio_end_ms": max(0, int(event.get("audio_end_ms", 0) or 0)), + "item_id": event.get("item_id") or self._active_input_item_id, + } + ] if event_type in { "session.resumed", "session.heartbeat_ack", @@ -201,14 +223,15 @@ class RealtimeOutputProjector: ), ] if event_type == "input.committed": - event_item_id = event.get("realtime_item_id") - item_id = ( - event_item_id - if isinstance(event_item_id, str) and event_item_id - else self._pop_pending_commit_item_id() - ) + event_item_id = event.get("item_id") + # Acknowledgements carry their own identity; untracked silence + # commits must not consume another turn's item ID. + item_id = event_item_id if isinstance(event_item_id, str) and event_item_id else f"item_{uuid4().hex}" + # The reader may already have accepted audio for the next turn. + if self._active_input_item_id == item_id: + self._reset_realtime_input_buffer_state() item = self._conversation_items.get(item_id) - created_payload: list[dict[str, object]] = [] + payloads = [self._input_audio_buffer_committed_event(item_id=item_id)] if item is None: message = event.get("message") no_response = event.get("no_response") is True @@ -226,19 +249,18 @@ class RealtimeOutputProjector: ), } self._conversation_items[item_id] = item - created_payload.extend(self._conversation_item_added_events(item)) + payloads.extend(self._conversation_item_added_events(item)) item["status"] = "completed" - payloads = created_payload + [ - self._input_audio_buffer_committed_event(item_id=item_id, event=event), - ] transcription_event = self._input_audio_transcription_completed_event(item_id, item) if transcription_event is not None: payloads.append(transcription_event) payloads.append(self._conversation_item_done_event(item)) return payloads if event_type == "input.cancelled": + self._reset_realtime_input_buffer_state() return [{"type": "input_audio_buffer.cleared"}] if event_type == "input_audio_buffer.cleared": + self._reset_realtime_input_buffer_state() return [{"type": "input_audio_buffer.cleared"}] if event_type == "audio.cancelled": response_id = event.get("response_id") @@ -416,7 +438,14 @@ class RealtimeOutputProjector: return [event] message = str(raw_error or event.get("message") or "Duplex runtime error") code = str(event.get("code") or "duplex_error") - return [self._realtime_error_payload(code, message, event_id=event.get("realtime_event_id"))] + return [ + self._realtime_error_payload( + code, + message, + event_id=event.get("realtime_event_id") or event.get("event_id"), + param=event.get("param"), + ) + ] if event_type == "session.closed": return [{"type": "session.closed", "event": event}] return [{"type": f"duplex.{event_type}", "event": event}] @@ -490,7 +519,14 @@ class RealtimeOutputProjector: }, }, ) - payload.setdefault("turn_detection", payload.get("turn_detection")) + if self._turn_detection_configured or "turn_detection" in payload: + turn_detection = payload.get("turn_detection", self._turn_detection) + audio = dict(payload["audio"]) if isinstance(payload.get("audio"), dict) else {} + audio_input = dict(audio["input"]) if isinstance(audio.get("input"), dict) else {} + audio_input["turn_detection"] = turn_detection + audio["input"] = audio_input + payload["audio"] = audio + payload["turn_detection"] = turn_detection payload.setdefault("input_audio_transcription", payload.get("input_audio_transcription")) payload.setdefault("tracing", payload.get("tracing")) return payload diff --git a/vllm_omni/entrypoints/duplex/realtime_session.py b/vllm_omni/entrypoints/duplex/realtime_session.py index c4604ff39..122a7ab52 100644 --- a/vllm_omni/entrypoints/duplex/realtime_session.py +++ b/vllm_omni/entrypoints/duplex/realtime_session.py @@ -42,6 +42,10 @@ class NativeRealtimeSessionProtocol( def bind_sender(self, send_realtime_json) -> None: self._send_realtime_json = send_realtime_json + def bind_native_input_append(self, enabled: bool) -> None: + """Use the serving-selected input path, independent of interruption policy.""" + self._native_input_append = enabled + def _default_session_payload(self) -> dict[str, object]: payload: dict[str, object] = { "model": self._default_model, diff --git a/vllm_omni/entrypoints/duplex/realtime_state.py b/vllm_omni/entrypoints/duplex/realtime_state.py index 4a1e7bdba..babb470a5 100644 --- a/vllm_omni/entrypoints/duplex/realtime_state.py +++ b/vllm_omni/entrypoints/duplex/realtime_state.py @@ -24,6 +24,32 @@ REALTIME_OUTPUT_AUDIO_FORMATS = { "g711_ulaw", "g711_alaw", } +REALTIME_PCM_DEFAULT_SAMPLE_RATE_HZ = 24_000 +REALTIME_PCM_F32_DEFAULT_SAMPLE_RATE_HZ = 16_000 +REALTIME_G711_DEFAULT_SAMPLE_RATE_HZ = 8_000 + +_RealtimeProtocolConfig = tuple[str, int, str, int | None, float] +_PendingTurnDetectionUpdate = tuple[ + bool, + dict[str, object] | None, + asyncio.Event, + _RealtimeProtocolConfig, +] + + +def realtime_default_sample_rate_hz(fmt: object) -> int | None: + if not isinstance(fmt, str): + return None + normalized = fmt.lower() + if normalized in {"pcm16", "pcm_s16le", "s16le", "pcm", "audio/pcm"}: + return REALTIME_PCM_DEFAULT_SAMPLE_RATE_HZ + if normalized in {"pcm_f32le", "audio/pcm_f32le"}: + return REALTIME_PCM_F32_DEFAULT_SAMPLE_RATE_HZ + if normalized in {"g711_ulaw", "g711_alaw", "audio/g711_ulaw", "audio/g711_alaw"}: + return REALTIME_G711_DEFAULT_SAMPLE_RATE_HZ + return None + + REALTIME_ERROR_TYPES_BY_CODE = { "bad_event": "invalid_request_error", "bad_audio": "invalid_request_error", @@ -43,12 +69,16 @@ REALTIME_ERROR_TYPES_BY_CODE = { "runtime_data_plane_text_without_audio": "server_error", "response_error": "server_error", "chat_error": "server_error", + "server_vad_initialization_failed": "server_error", + "server_vad_inference_failed": "server_error", "duplex_session_busy": "rate_limit_error", "resource_exhausted": "rate_limit_error", + "input_backpressure": "rate_limit_error", "response_already_active": "invalid_request_error", "response_not_active": "invalid_request_error", "response_create_without_input": "invalid_request_error", "input_audio_buffer_empty": "invalid_request_error", + "server_vad_manual_commit_unsupported": "invalid_request_error", "missing_item_id": "invalid_request_error", "item_not_found": "invalid_request_error", "playback_item_mismatch": "invalid_request_error", @@ -63,7 +93,6 @@ REALTIME_ERROR_TYPES_BY_CODE = { "native_text_append_unsupported": "invalid_request_error", "runtime_native_stage_role_required": "server_error", "runtime_native_runner_kv_required": "server_error", - "server_vad_unavailable": "server_error", } @@ -109,12 +138,12 @@ class RealtimeSessionState: _default_session_id: object | None = None _default_extra_body: dict[str, object] = field(default_factory=dict) _input_audio_format: str = "pcm16" - _input_sample_rate_hz: int = 16000 + _input_sample_rate_hz: int = REALTIME_PCM_DEFAULT_SAMPLE_RATE_HZ _output_audio_format: str = "pcm16" _overlap_silence_rms: float = 0.003 _turn_detection: dict[str, object] | None = None - _server_vad: object | None = None - _pending_turn_detection_update: tuple[dict[str, object] | None, object | None, asyncio.Event] | None = None + _native_input_append: bool = False + _pending_turn_detection_update: _PendingTurnDetectionUpdate | None = None _send_realtime_json: Any = None _initial_session_update: bool = False _input_speech_started: bool = False @@ -123,13 +152,13 @@ class RealtimeSessionState: _active_response_id: str | None = None _last_response_id: str | None = None _conversation_items: dict[str, dict[str, object]] = field(default_factory=dict) - _pending_commit_item_ids: asyncio.Queue[str] = field(default_factory=asyncio.Queue) _last_conversation_item_id: str | None = None _output_sample_rate_hz: int | None = None _active_input_item_id: str | None = None _input_audio_buffer_has_audio: bool = False _input_audio_buffer_had_non_speech: bool = False _input_audio_buffer_transcript_parts: list[str] = field(default_factory=list) + _turn_detection_configured: bool = False @classmethod def from_query_params(cls, query_params: Any) -> RealtimeSessionState: @@ -197,10 +226,8 @@ class RealtimeStateOwner: _output_audio_format: str = _RealtimeStateField() _overlap_silence_rms: float = _RealtimeStateField() _turn_detection: dict[str, object] | None = _RealtimeStateField() - _server_vad: object | None = _RealtimeStateField() - _pending_turn_detection_update: tuple[dict[str, object] | None, object | None, asyncio.Event] | None = ( - _RealtimeStateField() - ) + _native_input_append: bool = _RealtimeStateField() + _pending_turn_detection_update: _PendingTurnDetectionUpdate | None = _RealtimeStateField() _send_realtime_json: Any = _RealtimeStateField() _initial_session_update: bool = _RealtimeStateField() _input_speech_started: bool = _RealtimeStateField() @@ -209,10 +236,10 @@ class RealtimeStateOwner: _active_response_id: str | None = _RealtimeStateField() _last_response_id: str | None = _RealtimeStateField() _conversation_items: dict[str, dict[str, object]] = _RealtimeStateField() - _pending_commit_item_ids: asyncio.Queue[str] = _RealtimeStateField() _last_conversation_item_id: str | None = _RealtimeStateField() _output_sample_rate_hz: int | None = _RealtimeStateField() _active_input_item_id: str | None = _RealtimeStateField() _input_audio_buffer_has_audio: bool = _RealtimeStateField() _input_audio_buffer_had_non_speech: bool = _RealtimeStateField() _input_audio_buffer_transcript_parts: list[str] = _RealtimeStateField() + _turn_detection_configured: bool = _RealtimeStateField() diff --git a/vllm_omni/entrypoints/duplex/runtime_adapter.py b/vllm_omni/entrypoints/duplex/runtime_adapter.py index 8bf23e1ed..a873f33a8 100644 --- a/vllm_omni/entrypoints/duplex/runtime_adapter.py +++ b/vllm_omni/entrypoints/duplex/runtime_adapter.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio from collections.abc import Awaitable, Callable, Iterable, Mapping, MutableMapping +from dataclasses import dataclass, field from importlib import import_module from typing import Any, Protocol @@ -106,6 +107,104 @@ class ServingRuntimeSessionState(Protocol): def clear_continuation(self) -> None: ... +class _TurnBasedPcmAppendBuffer: + """Empty native-audio buffer for chat-fallback sessions. + + Turn-based sessions keep input audio in ``DuplexSession`` and never use + the model-native append path. The runner still shares lifecycle and + cancellation code with native sessions, so it needs an inert buffer for + those common operations. + """ + + @property + def pending_byte_count(self) -> int: + return 0 + + def clear(self) -> None: + return None + + def clear_force_listen(self) -> None: + return None + + def has_pending(self) -> bool: + return False + + def has_reserved(self) -> bool: + return False + + def prepare_append( + self, + payload: dict[str, object], + *, + operation_id: str, + chunk_period_ms: int, + allow_emit: bool, + ) -> PcmAppendReservation | None: + del payload, operation_id, chunk_period_ms, allow_emit + raise RuntimeError("Model-native audio append is not enabled for this session") + + def prepare_commit( + self, + *, + operation_id: str, + chunk_period_ms: int, + ) -> PcmAppendReservation: + del operation_id, chunk_period_ms + raise RuntimeError("Model-native audio append is not enabled for this session") + + def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None: + del chunk_period_ms + return None + + +@dataclass(slots=True) +class TurnBasedServingSessionState: + """Runner-local state for sessions served through chat fallback.""" + + audio_buffer: PcmAppendBuffer = field(default_factory=_TurnBasedPcmAppendBuffer) + input_since_commit: bool = False + speech_since_commit: bool = False + native_context_locked: bool = False + committed_audio_payload: dict[str, object] | None = None + committed_audio_operation_id: str | None = None + committed_audio_reserved_bytes: int = 0 + deferred_response_create: bool = False + deferred_precreate_response: bool = False + data_plane_task: asyncio.Task[None] | None = None + data_plane_restart_requested: bool = False + continuation_owner_id: str | None = None + continuation_units: int = 0 + pending_silence_task: asyncio.Task[bool] | None = None + pending_silence_owner_id: str | None = None + silence_continuation_scheduler: Callable[..., Awaitable[bool]] | None = None + + def retain_committed_audio( + self, + payload: dict[str, object], + *, + operation_id: str | None, + reserved_bytes: int = 0, + ) -> None: + self.committed_audio_payload = payload + self.committed_audio_operation_id = operation_id + self.committed_audio_reserved_bytes += max(0, int(reserved_bytes)) + + def clear_committed_audio(self) -> int: + reserved_bytes = self.committed_audio_reserved_bytes + self.committed_audio_payload = None + self.committed_audio_operation_id = None + self.committed_audio_reserved_bytes = 0 + self.deferred_response_create = False + self.deferred_precreate_response = False + return reserved_bytes + + def clear_continuation(self) -> None: + self.continuation_owner_id = None + self.continuation_units = 0 + self.pending_silence_task = None + self.pending_silence_owner_id = None + + class RuntimeDataPlane(Protocol): def begin_request(self, request_id: str) -> None: ... diff --git a/vllm_omni/entrypoints/duplex/runtime_bridge.py b/vllm_omni/entrypoints/duplex/runtime_bridge.py index 6db5be236..bf839523c 100644 --- a/vllm_omni/entrypoints/duplex/runtime_bridge.py +++ b/vllm_omni/entrypoints/duplex/runtime_bridge.py @@ -39,6 +39,11 @@ class NativeRuntimeBridgeMixin: } async def _open_runtime_session(self, session: DuplexSession, send_json) -> dict[str, object] | bool: + # Turn-based sessions use the chat-fallback path and do not own a + # model-native duplex runtime. Do not send control messages merely + # because the shared engine client exposes the duplex RPC methods. + if not self._uses_serving_runtime_adapter(session.config): + return True contract_error = self._native_runtime_contract_error(session) if contract_error is not None: await send_json( @@ -535,6 +540,8 @@ class NativeRuntimeBridgeMixin: session_config: dict[str, object] | None = None, runtime_config: dict[str, object] | None = None, ) -> bool: + if not self._uses_serving_runtime_adapter(session.config): + return True signal_turn = getattr(self._chat_service.engine_client, "signal_duplex_turn_async", None) if not callable(signal_turn): return True @@ -582,6 +589,8 @@ class NativeRuntimeBridgeMixin: return True async def _close_runtime_session(self, session: DuplexSession, *, reason: str, send_json=None) -> bool: + if not self._uses_serving_runtime_adapter(session.config): + return True close_session = getattr(self._chat_service.engine_client, "close_duplex_session_async", None) if not callable(close_session): return True @@ -1301,11 +1310,19 @@ class NativeRuntimeBridgeMixin: def _cleanup_duplex_session_state(self, session: DuplexSession) -> None: session_id = session.session_id - self._serving_runtime_adapter.remove_session_state(session_id) - self._serving_runtime_adapter.data_plane.close_session( - session_id, - active_request_id=session.active_request_id, - ) + pipeline = self._server_vad_pipelines.pop(session_id, None) + if pipeline is not None: + pipeline.reset() + self._realtime_vad_metrics.session_finished() + adapter = self._serving_runtime_adapter + if adapter is None: + self._serving_session_states.pop(session_id, None) + else: + adapter.remove_session_state(session_id) + adapter.data_plane.close_session( + session_id, + active_request_id=session.active_request_id, + ) def _encode_native_data_plane_audio( self, diff --git a/vllm_omni/entrypoints/duplex/server_vad.py b/vllm_omni/entrypoints/duplex/server_vad.py new file mode 100644 index 000000000..610e51c8c --- /dev/null +++ b/vllm_omni/entrypoints/duplex/server_vad.py @@ -0,0 +1,528 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import asyncio +import hashlib +import threading +import time +from collections.abc import Mapping +from dataclasses import dataclass, replace +from pathlib import Path +from typing import TYPE_CHECKING, Literal, Protocol + +import numpy as np +from vllm.transformers_utils.repo_utils import try_get_local_file + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.audio_utils_mixin import StreamingAudioResampler + +SILERO_VAD_REPO_ID = "istupakov/silero-vad-onnx" +SILERO_VAD_REVISION = "8b14476858ef240c50b3884bb38cc67290c1cc70" +SILERO_VAD_FILENAME = "silero_vad.onnx" +SILERO_VAD_SHA256 = "1a153a22f4509e292a94e67d6f9b85e8deb25b4988682b7e174c65279d8788e3" +SILERO_VAD_MIN_THRESHOLD = 0.15 +SILERO_VAD_DEFAULT_MIN_SPEECH_DURATION_MS = 96 + +_SERVER_VAD_FIELDS = { + "type", + "threshold", + "prefix_padding_ms", + "silence_duration_ms", + "create_response", + "interrupt_response", + "min_speech_duration_ms", +} + + +@dataclass(frozen=True, slots=True) +class ServerVADConfig: + """Validated OpenAI-compatible ``server_vad`` session configuration.""" + + type: Literal["server_vad"] = "server_vad" + threshold: float = 0.5 + prefix_padding_ms: int = 300 + silence_duration_ms: int = 500 + create_response: bool = True + # Omission is resolved from runtime capabilities before the session is opened. + interrupt_response: bool | None = None + min_speech_duration_ms: int | None = None + + @classmethod + def from_value(cls, value: object) -> ServerVADConfig: + if not isinstance(value, dict): + raise ValueError("turn_detection must be null or an object") + unknown = sorted(set(value) - _SERVER_VAD_FIELDS) + if unknown: + raise ValueError(f"Unknown server_vad field(s): {', '.join(unknown)}") + vad_type = value.get("type") + if vad_type != "server_vad": + raise ValueError("turn_detection.type must be 'server_vad'") + + threshold = value.get("threshold", 0.5) + if isinstance(threshold, bool) or not isinstance(threshold, int | float) or not 0 <= threshold <= 1: + raise ValueError("server_vad.threshold must be a number between 0 and 1") + + prefix_padding_ms = value.get("prefix_padding_ms", 300) + if isinstance(prefix_padding_ms, bool) or not isinstance(prefix_padding_ms, int) or prefix_padding_ms < 0: + raise ValueError("server_vad.prefix_padding_ms must be a non-negative integer") + + silence_duration_ms = value.get("silence_duration_ms", 500) + if ( + isinstance(silence_duration_ms, bool) + or not isinstance(silence_duration_ms, int) + or silence_duration_ms <= 0 + ): + raise ValueError("server_vad.silence_duration_ms must be a positive integer") + + create_response = value.get("create_response", True) + if not isinstance(create_response, bool): + raise ValueError("server_vad.create_response must be a boolean") + + interrupt_response = value.get("interrupt_response") + if "interrupt_response" in value and not isinstance(interrupt_response, bool): + raise ValueError("server_vad.interrupt_response must be a boolean") + + min_speech_duration_ms = value.get("min_speech_duration_ms") + if min_speech_duration_ms is not None and ( + isinstance(min_speech_duration_ms, bool) + or not isinstance(min_speech_duration_ms, int | float) + or not np.isfinite(float(min_speech_duration_ms)) + or min_speech_duration_ms < 0 + ): + raise ValueError("server_vad.min_speech_duration_ms must be a non-negative number") + + return cls( + type=vad_type, + threshold=float(threshold), + prefix_padding_ms=prefix_padding_ms, + silence_duration_ms=silence_duration_ms, + create_response=create_response, + interrupt_response=interrupt_response, + min_speech_duration_ms=(int(min_speech_duration_ms) if min_speech_duration_ms is not None else None), + ) + + def as_dict(self) -> dict[str, object]: + result: dict[str, object] = { + "type": self.type, + "threshold": self.threshold, + "prefix_padding_ms": self.prefix_padding_ms, + "silence_duration_ms": self.silence_duration_ms, + "create_response": self.create_response, + } + if self.interrupt_response is not None: + result["interrupt_response"] = self.interrupt_response + if self.min_speech_duration_ms is not None: + result["min_speech_duration_ms"] = self.min_speech_duration_ms + return result + + +def parse_session_turn_detection( + payload: Mapping[str, object], +) -> tuple[bool, ServerVADConfig | None]: + """Extract and validate OpenAI-compatible turn detection aliases.""" + configured_values: list[tuple[str, object]] = [] + if "turn_detection" in payload: + configured_values.append(("turn_detection", payload["turn_detection"])) + audio = payload.get("audio") + if isinstance(audio, Mapping): + audio_input = audio.get("input") + if isinstance(audio_input, Mapping) and "turn_detection" in audio_input: + configured_values.append(("audio.input.turn_detection", audio_input["turn_detection"])) + if not configured_values: + return False, None + + first_path, first_value = configured_values[0] + first_config = None if first_value is None else ServerVADConfig.from_value(first_value) + for field_path, value in configured_values[1:]: + config = None if value is None else ServerVADConfig.from_value(value) + if first_config is not None and config is not None: + # An omitted option must not conflict with an explicit alias value. + if first_config.interrupt_response is None: + first_config = replace(first_config, interrupt_response=config.interrupt_response) + if config.interrupt_response is None: + config = replace(config, interrupt_response=first_config.interrupt_response) + if config != first_config: + raise ValueError(f"{first_path} and {field_path} must not specify conflicting values") + return True, first_config + + +class SpeechDetectorBackend(Protocol): + """Detector contract; per-stream state is owned by the pipeline.""" + + frame_samples: int + + def new_state(self) -> object: ... + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: ... + + +class SpeechDetectorBackendProvider(Protocol): + def get(self) -> SpeechDetectorBackend: ... + + +@dataclass(frozen=True, slots=True) +class SpeechEndpointDecision: + speech_started: bool = False + speech_stopped: bool = False + audio_start_ms: int | None = None + audio_end_ms: int | None = None + endpoint_delay_ms: int | None = None + + +class SpeechEndpointPolicy(Protocol): + @property + def speech_active(self) -> bool: ... + + def update( + self, + probability: float, + *, + frame_start_sample: int, + frame_samples: int, + ) -> SpeechEndpointDecision: ... + + def reset(self) -> None: ... + + +class ThresholdEndpointPolicy: + """Apply threshold and trailing-silence endpoint rules.""" + + def __init__(self, config: ServerVADConfig, *, sample_rate_hz: int) -> None: + self.config = config + self.sample_rate_hz = sample_rate_hz + self._speech_active = False + self._candidate_samples = 0 + self._candidate_start_sample: int | None = None + self._silence_samples = 0 + + @property + def speech_active(self) -> bool: + return self._speech_active + + def update( + self, + probability: float, + *, + frame_start_sample: int, + frame_samples: int, + ) -> SpeechEndpointDecision: + if probability >= self.config.threshold: + self._silence_samples = 0 + if not self._speech_active: + if self._candidate_samples == 0: + self._candidate_start_sample = frame_start_sample + self._candidate_samples += frame_samples + min_speech_duration_ms = self.config.min_speech_duration_ms or 0 + min_speech_samples = max(1, round(min_speech_duration_ms * self.sample_rate_hz / 1000)) + if self._candidate_samples < min_speech_samples: + return SpeechEndpointDecision() + self._speech_active = True + detected_start_sample = ( + self._candidate_start_sample if self._candidate_start_sample is not None else frame_start_sample + ) + detected_start_ms = round(detected_start_sample * 1000 / self.sample_rate_hz) + self._candidate_samples = 0 + self._candidate_start_sample = None + return SpeechEndpointDecision( + speech_started=True, + audio_start_ms=max(0, detected_start_ms - self.config.prefix_padding_ms), + ) + return SpeechEndpointDecision() + + if not self._speech_active: + self._candidate_samples = 0 + self._candidate_start_sample = None + return SpeechEndpointDecision() + + # Match Silero v6.2's streaming VAD hysteresis: activation uses the + # configured threshold, while potential speech end starts only below + # ``max(threshold - 0.15, 0.01)``. Once a silence candidate exists, + # intermediate frames keep elapsed time moving but cannot themselves + # close the turn. + negative_threshold = max(self.config.threshold - SILERO_VAD_MIN_THRESHOLD, 0.01) + below_negative_threshold = probability < negative_threshold + if not below_negative_threshold and self._silence_samples == 0: + return SpeechEndpointDecision() + + self._silence_samples += frame_samples + if not below_negative_threshold: + return SpeechEndpointDecision() + silence_limit = max(1, round(self.config.silence_duration_ms * self.sample_rate_hz / 1000)) + if self._silence_samples < silence_limit: + return SpeechEndpointDecision() + + # OpenAI defines ``audio_end_ms`` as the end of the audio sent to the + # model, including the trailing silence used for endpoint detection. + audio_end_sample = frame_start_sample + frame_samples + endpoint_delay_ms = round(self._silence_samples * 1000 / self.sample_rate_hz) + self.reset() + return SpeechEndpointDecision( + speech_stopped=True, + audio_end_ms=max(0, round(audio_end_sample * 1000 / self.sample_rate_hz)), + endpoint_delay_ms=endpoint_delay_ms, + ) + + def reset(self) -> None: + self._speech_active = False + self._candidate_samples = 0 + self._candidate_start_sample = None + self._silence_samples = 0 + + +@dataclass(frozen=True, slots=True) +class ServerVADFrame: + samples: np.ndarray + decision: SpeechEndpointDecision + + +@dataclass(frozen=True, slots=True) +class ServerVADBatch: + frames: tuple[ServerVADFrame, ...] + inference_ms: float + + +class ServerVADPipeline: + """Frame continuous 16 kHz audio and apply acoustic endpointing. + + The pipeline retains only input-resampling residuals, a partial frame, + detector state, and timing counters. Commit-eligible prefix and utterance + audio remain session-owned. + """ + + sample_rate_hz = 16_000 + + def __init__( + self, + backend: SpeechDetectorBackend, + config: ServerVADConfig, + *, + endpoint_policy: SpeechEndpointPolicy | None = None, + ) -> None: + self.backend = backend + self.config = config + self.endpoint_policy = endpoint_policy or ThresholdEndpointPolicy( + config, + sample_rate_hz=self.sample_rate_hz, + ) + self._input_resampler: StreamingAudioResampler | None = None + self._source_sample_rate_hz: int | None = None + self._scratch = np.empty(0, dtype=np.float32) + self._detector_state = backend.new_state() + self._processed_samples = 0 + self._stream_start_ms = 0 + + @property + def speech_active(self) -> bool: + return self.endpoint_policy.speech_active + + @property + def source_sample_rate_hz(self) -> int | None: + return self._source_sample_rate_hz + + @property + def scratch_bytes(self) -> int: + resampler_bytes = self._input_resampler.scratch_bytes if self._input_resampler is not None else 0 + return int(self._scratch.nbytes) + resampler_bytes + + async def push_pcm16( + self, + samples: bytes, + *, + source_sample_rate_hz: int, + ) -> ServerVADBatch: + """Normalize chunked mono PCM16 input and run endpoint detection.""" + if isinstance(source_sample_rate_hz, bool) or source_sample_rate_hz not in {16_000, 24_000}: + raise ValueError("server_vad PCM16 input sample rate must be 16000 or 24000 Hz") + + if len(samples) % np.dtype(" ServerVADBatch: + normalized = np.ascontiguousarray(samples, dtype=np.float32).reshape(-1) + return await asyncio.to_thread(self._push_sync, normalized) + + def _push_sync(self, samples: np.ndarray) -> ServerVADBatch: + started_at = time.perf_counter() + if self._scratch.size: + samples = np.concatenate((self._scratch, samples)) + frame_samples = int(self.backend.frame_samples) + complete_samples = samples.size - samples.size % frame_samples + # Retained state must not keep the whole input chunk alive through a view. + self._scratch = samples[complete_samples:].copy() + + results: list[ServerVADFrame] = [] + for offset in range(0, complete_samples, frame_samples): + frame = np.ascontiguousarray(samples[offset : offset + frame_samples], dtype=np.float32) + frame_start = self._processed_samples + probability, self._detector_state = self.backend.infer(frame, self._detector_state) + probability = min(1.0, max(0.0, float(probability))) + decision = self.endpoint_policy.update( + probability, + frame_start_sample=frame_start, + frame_samples=frame_samples, + ) + if decision.speech_started and decision.audio_start_ms is not None: + # Prefix audio before the most recent reset is no longer available. + decision = replace(decision, audio_start_ms=max(self._stream_start_ms, decision.audio_start_ms)) + + self._processed_samples += frame_samples + results.append( + ServerVADFrame( + samples=frame, + decision=decision, + ) + ) + return ServerVADBatch( + frames=tuple(results), + inference_ms=(time.perf_counter() - started_at) * 1000, + ) + + def reset(self) -> None: + """Discard stream state while preserving the session-relative audio clock.""" + self._stream_start_ms = round(self._processed_samples * 1000 / self.sample_rate_hz) + self._input_resampler = None + self._source_sample_rate_hz = None + self._scratch = np.empty(0, dtype=np.float32) + self._detector_state = self.backend.new_state() + self.endpoint_policy.reset() + + +@dataclass(frozen=True, slots=True) +class _SileroVADState: + model_state: np.ndarray + context: np.ndarray + + +class SileroVADBackend: + """Shared ONNX Runtime Silero v6.2 detector running on CPU.""" + + sample_rate_hz = 16_000 + frame_samples = 512 + context_samples = 64 + model_state_shape = (2, 1, 128) + + def __init__(self, model_path: str | Path) -> None: + try: + import onnxruntime as ort + except ImportError as exc: # pragma: no cover - platform packaging supplies ORT. + raise RuntimeError("server_vad requires ONNX Runtime") from exc + + session_options = ort.SessionOptions() + session_options.inter_op_num_threads = 1 + session_options.intra_op_num_threads = 1 + + self.model_path = Path(model_path) + self._session = ort.InferenceSession( + str(self.model_path), + providers=["CPUExecutionProvider"], + sess_options=session_options, + ) + input_names = {item.name for item in self._session.get_inputs()} + if "input" not in input_names or "sr" not in input_names or "state" not in input_names: + raise RuntimeError(f"Unsupported Silero ONNX input contract: {sorted(input_names)}") + self._warm_up() + + def _warm_up(self) -> None: + self.infer(np.zeros(self.frame_samples, dtype=np.float32), self.new_state()) + + def new_state(self) -> _SileroVADState: + return _SileroVADState( + model_state=np.zeros(self.model_state_shape, dtype=np.float32), + context=np.zeros((1, self.context_samples), dtype=np.float32), + ) + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: + if not isinstance(state, _SileroVADState): + raise TypeError("Silero detector state must be created by SileroVADBackend.new_state()") + model_state = np.ascontiguousarray(state.model_state, dtype=np.float32) + context = np.ascontiguousarray(state.context, dtype=np.float32) + if model_state.shape != self.model_state_shape: + raise ValueError(f"Silero model state must have shape {self.model_state_shape}, got {model_state.shape}") + expected_context_shape = (1, self.context_samples) + if context.shape != expected_context_shape: + raise ValueError(f"Silero context must have shape {expected_context_shape}, got {context.shape}") + + audio = np.ascontiguousarray(frame, dtype=np.float32).reshape(1, -1) + if audio.shape[1] != self.frame_samples: + raise ValueError( + f"Silero detector frame must contain exactly {self.frame_samples} samples, got {audio.shape[1]}" + ) + model_input = np.concatenate((context, audio), axis=1) + inputs = { + "input": model_input, + "state": model_state, + "sr": np.asarray(self.sample_rate_hz, dtype=np.int64), + } + # ONNX Runtime permits concurrent Run calls on one CPU session. Stream + # state is explicit, so independent pipelines can safely share the session. + output = self._session.run(None, inputs) + if len(output) < 2: + raise RuntimeError("Silero ONNX model did not return probability and model state") + probability = float(np.asarray(output[0]).reshape(-1)[0]) + next_state = _SileroVADState( + model_state=np.ascontiguousarray(output[1], dtype=np.float32), + context=np.ascontiguousarray(model_input[:, -self.context_samples :], dtype=np.float32), + ) + return probability, next_state + + +class SileroVADBackendProvider: + """Resolve, verify, and load one Silero model instance per process.""" + + def __init__(self, *, model_path: str | None = None) -> None: + self.model_path = model_path + self._backend: SileroVADBackend | None = None + self._lock = threading.Lock() + + def get(self) -> SileroVADBackend: + if self._backend is not None: + return self._backend + with self._lock: + if self._backend is None: + path = self._resolve_local_artifact() + self._verify_checksum(path) + self._backend = SileroVADBackend(path) + return self._backend + + def _resolve_local_artifact(self) -> Path: + if self.model_path: + path = Path(self.model_path).expanduser() + if path.is_file(): + return path + raise RuntimeError(f"Configured Silero VAD model does not exist: {path}") + cached = try_get_local_file( + model=SILERO_VAD_REPO_ID, + file_name=SILERO_VAD_FILENAME, + revision=SILERO_VAD_REVISION, + ) + if isinstance(cached, Path) and cached.is_file(): + return cached + raise RuntimeError( + "Silero VAD artifact is not available locally. Pre-download the pinned artifact " + "or configure duplex_session.server_vad_model_path." + ) + + def _verify_checksum(self, path: Path) -> None: + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != SILERO_VAD_SHA256: + raise RuntimeError(f"Silero VAD checksum mismatch for {path}: expected {SILERO_VAD_SHA256}, got {digest}") diff --git a/vllm_omni/entrypoints/duplex/serving.py b/vllm_omni/entrypoints/duplex/serving.py index d1a028bd4..7bd690841 100644 --- a/vllm_omni/entrypoints/duplex/serving.py +++ b/vllm_omni/entrypoints/duplex/serving.py @@ -8,9 +8,9 @@ import base64 import binascii import inspect import json -from collections.abc import Mapping +from collections.abc import Mapping, MutableMapping from contextlib import suppress -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING import numpy as np @@ -49,12 +49,21 @@ from vllm_omni.entrypoints.duplex.runtime_adapter import ( ServingRuntimeAdapter, ServingRuntimeConfigError, ServingRuntimeSessionState, + TurnBasedServingSessionState, load_serving_runtime_adapter, validate_serving_runtime_adapter, ) from vllm_omni.entrypoints.duplex.runtime_bridge import ( NativeRuntimeBridgeMixin, ) +from vllm_omni.entrypoints.duplex.server_vad import ( + SILERO_VAD_DEFAULT_MIN_SPEECH_DURATION_MS, + SILERO_VAD_MIN_THRESHOLD, + ServerVADPipeline, + SileroVADBackendProvider, + SpeechDetectorBackendProvider, + parse_session_turn_detection, +) from vllm_omni.entrypoints.duplex.session_attachment import ( DuplexJournalGapError, DuplexSessionAttachmentRegistry, @@ -68,6 +77,7 @@ from vllm_omni.entrypoints.duplex.websocket import ( DuplexSessionTasks, DuplexWebSocketActor, ) +from vllm_omni.metrics.realtime import RealtimeVADMetrics if TYPE_CHECKING: from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat @@ -85,6 +95,7 @@ class _DuplexSessionHandshake: session: DuplexSession resumed: bool = False attachment_generation: int | None = None + event_id: str | None = None class OmniDuplexSessionHandler( @@ -105,16 +116,27 @@ class OmniDuplexSessionHandler( self, *, chat_service: OmniOpenAIServingChat, + served_model_name: str | None = None, config_timeout_s: float = _DEFAULT_CONFIG_TIMEOUT_S, idle_timeout_s: float = _DEFAULT_IDLE_TIMEOUT_S, + log_stats: bool = True, duplex_session_config: DuplexSessionRuntimeConfig | None = None, serving_runtime_adapter: ServingRuntimeAdapter | None = None, serving_runtime_adapter_path: str | None = None, + server_vad_backend_provider: SpeechDetectorBackendProvider | None = None, ) -> None: self._chat_service = chat_service self._config_timeout_s = config_timeout_s self._idle_timeout_s = idle_timeout_s self._duplex_session_config = duplex_session_config or DuplexSessionRuntimeConfig() + self._server_vad_backend_provider = server_vad_backend_provider or SileroVADBackendProvider( + model_path=self._duplex_session_config.server_vad_model_path, + ) + self._server_vad_pipelines: dict[str, ServerVADPipeline] = {} + self._realtime_vad_metrics = RealtimeVADMetrics( + served_model_name or chat_service.model_config.model, + log_stats=log_stats, + ) self._registry = DuplexSessionRegistry( DuplexCapabilities( supports_model_native_turn_policy=False, @@ -131,6 +153,7 @@ class OmniDuplexSessionHandler( "duplex_serving_adapter_path", None, ) + self._serving_runtime_adapter: ServingRuntimeAdapter | None if serving_runtime_adapter is not None: self._serving_runtime_adapter = validate_serving_runtime_adapter(serving_runtime_adapter) elif isinstance(adapter_path, str) and adapter_path: @@ -139,7 +162,10 @@ class OmniDuplexSessionHandler( self._encode_native_data_plane_audio, ) else: - raise ValueError("A duplex serving runtime adapter must be explicitly configured") + self._serving_runtime_adapter = None + self._serving_session_states: MutableMapping[str, ServingRuntimeSessionState] = ( + self._serving_runtime_adapter.session_states if self._serving_runtime_adapter is not None else {} + ) self._session_tasks: dict[str, DuplexSessionTasks] = {} self._realtime_protocols: dict[str, NativeRealtimeSessionProtocol] = {} self._lease_generations: dict[str, int] = {} @@ -223,7 +249,7 @@ class OmniDuplexSessionHandler( active_response_task.cancel() with suppress(asyncio.CancelledError): await active_response_task - native = self._serving_runtime_adapter.session_states.get(session.session_id) + native = self._serving_session_states.get(session.session_id) if native is not None and native.data_plane_task is not None: data_plane_task = native.data_plane_task native.data_plane_task = None @@ -794,8 +820,27 @@ class OmniDuplexSessionHandler( duration_ms = cls._input_audio_duration_ms(event, payload) return 0 < duration_ms <= session.config.overlap_short_ack_ms + def _create_runtime_session_state(self) -> ServingRuntimeSessionState: + adapter = self._serving_runtime_adapter + if adapter is not None: + return adapter.create_session_state() + return TurnBasedServingSessionState() + def _runtime_session_state(self, session: DuplexSession) -> ServingRuntimeSessionState: - return self._serving_runtime_adapter.session_state(session.session_id) + adapter = self._serving_runtime_adapter + if adapter is not None: + return adapter.session_state(session.session_id) + state = self._serving_session_states.get(session.session_id) + if state is None: + state = TurnBasedServingSessionState() + self._serving_session_states[session.session_id] = state + return state + + def _require_serving_runtime_adapter(self) -> ServingRuntimeAdapter: + adapter = self._serving_runtime_adapter + if adapter is None: + raise RuntimeError("Model-native duplex runtime is not configured for this session") + return adapter def _should_force_listen_for_auto_response_overlap( self, @@ -823,21 +868,25 @@ class OmniDuplexSessionHandler( *, session: DuplexSession, ) -> bool: - for key in ("is_speech", "speech"): - value = event.get(key) - if isinstance(value, bool): - return value - vad = event.get("vad") - if isinstance(vad, dict): - value = vad.get("is_speech") - if isinstance(value, bool): - return value - probability = vad.get("speech_probability", vad.get("probability")) + # Prefer serving-derived fields in ``payload`` over optional client + # hints retained in ``event``. In server-VAD mode the pipeline, rather + # than the client or the RMS fallback below, owns speech classification. + for source in (payload, event): + for key in ("is_speech", "speech"): + value = source.get(key) + if isinstance(value, bool): + return value + vad = source.get("vad") + if isinstance(vad, dict): + value = vad.get("is_speech") + if isinstance(value, bool): + return value + probability = vad.get("speech_probability", vad.get("probability")) + if isinstance(probability, int | float): + return float(probability) >= 0.5 + probability = source.get("speech_probability") if isinstance(probability, int | float): return float(probability) >= 0.5 - probability = event.get("speech_probability") - if isinstance(probability, int | float): - return float(probability) >= 0.5 fmt = payload.get("format") audio = payload.get("audio") @@ -925,14 +974,26 @@ class OmniDuplexSessionHandler( ) return None - config = DuplexSessionConfig.from_event(event) + try: + config = DuplexSessionConfig.from_event(event) + except ValueError as exc: + error: dict[str, object] = { + "type": "error", + "error": str(exc), + "code": "invalid_request_error", + } + event_id = event.get("_realtime_event_id") + if isinstance(event_id, str) and event_id: + error["event_id"] = event_id + await send_json(error) + return None if config.idle_timeout_s == _DEFAULT_IDLE_TIMEOUT_S: config.idle_timeout_s = self._idle_timeout_s - use_native_runtime = self._uses_serving_runtime_adapter(config) + runtime_adapter = self._serving_runtime_adapter if self._uses_serving_runtime_adapter(config) else None runtime_config: dict[str, object] = {} - if use_native_runtime: + if runtime_adapter is not None: try: - runtime_config = await self._serving_runtime_adapter.prepare_runtime_config( + runtime_config = await runtime_adapter.prepare_runtime_config( config, model_config=getattr(self._chat_service, "model_config", None), ) @@ -944,14 +1005,20 @@ class OmniDuplexSessionHandler( return None session_id = event.get("session_id") if isinstance(event.get("session_id"), str) else None session = self._registry.create(config=config, session_id=session_id) - if use_native_runtime: + if runtime_adapter is not None: session.replace_capabilities( - self._serving_runtime_adapter.capabilities( - max_sessions=self._duplex_session_config.max_sessions, - ) + runtime_adapter.capabilities(max_sessions=self._duplex_session_config.max_sessions) ) session.replace_runtime_config(runtime_config) - return _DuplexSessionHandshake(session=session) + session_payload = event.get("session") + self._resolve_server_vad_defaults( + session, config, session_payload if isinstance(session_payload, dict) else event + ) + event_id = event.get("_realtime_event_id") + return _DuplexSessionHandshake( + session=session, + event_id=event_id if isinstance(event_id, str) and event_id else None, + ) async def _resume_session_handshake( self, @@ -1151,7 +1218,8 @@ class OmniDuplexSessionHandler( return None def _uses_serving_runtime_adapter(self, config: DuplexSessionConfig) -> bool: - return self._serving_runtime_adapter.is_enabled(config) + adapter = self._serving_runtime_adapter + return adapter is not None and adapter.is_enabled(config) def _runtime_session_update_error( self, @@ -1161,7 +1229,7 @@ class OmniDuplexSessionHandler( if not self._uses_native_input_append(session): return None try: - self._serving_runtime_adapter.validate_client_extra_body(payload.get("extra_body")) + self._require_serving_runtime_adapter().validate_client_extra_body(payload.get("extra_body")) except ServingRuntimeConfigError as exc: return { "type": "error", @@ -1178,24 +1246,109 @@ class OmniDuplexSessionHandler( ) -> dict[str, object]: if not self._uses_native_input_append(session): return dict(session.runtime_config) - return self._serving_runtime_adapter.runtime_config_for_update( + return self._require_serving_runtime_adapter().runtime_config_for_update( candidate_config, dict(session.runtime_config), ) + def _resolve_server_vad_defaults( + self, + session: DuplexSession, + config: DuplexSessionConfig, + payload: dict[str, object], + ) -> None: + """Resolve omitted options without overriding explicit client policies.""" + if config.server_vad is None: + return + if config.server_vad.interrupt_response is None: + config.server_vad = replace(config.server_vad, interrupt_response=self._uses_native_input_append(session)) + extra = payload.get("extra_body") + overlap_policy = extra.get("overlap_policy") if isinstance(extra, dict) else None + if not isinstance(overlap_policy, str): + overlap_policy = payload.get("overlap_policy") + config.overlap_policy = ( + overlap_policy + if isinstance(overlap_policy, str) + else "barge_in_on_speech" + if config.server_vad.interrupt_response + else "listen_only" + ) + def _runtime_session_candidate_update_error( self, session: DuplexSession, candidate_config: DuplexSessionConfig, + *, + realtime_protocol: NativeRealtimeSessionProtocol | None = None, ) -> dict[str, object] | None: + if candidate_config.server_vad is not None: + desired_policy = "barge_in_on_speech" if candidate_config.server_vad.interrupt_response else "listen_only" + if candidate_config.overlap_policy != desired_policy: + return { + "type": "error", + "code": "unsupported_turn_detection", + "error": f"overlap_policy conflicts with turn_detection; expected {desired_policy!r}", + "param": "overlap_policy", + } if not self._uses_native_input_append(session): - if candidate_config.overlap_policy != DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value: - return None - return { - "type": "error", - "code": "server_vad_requires_native_duplex", - "error": "Realtime server_vad requires a model-native duplex runtime", - } + if ( + candidate_config.server_vad is not None + and candidate_config.server_vad.min_speech_duration_ms is not None + ): + return { + "type": "error", + "code": "unsupported_turn_detection", + "error": "server_vad.min_speech_duration_ms is only supported by model-native duplex runtimes", + } + requests_interruption = ( + candidate_config.server_vad is not None and candidate_config.server_vad.interrupt_response + ) or candidate_config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + if requests_interruption: + return { + "type": "error", + "code": "server_vad_requires_native_duplex", + "error": ( + "Realtime server_vad interruption requires a model-native duplex runtime; " + "set interrupt_response=false for turn-based models" + ), + } + if candidate_config.server_vad is not None and realtime_protocol is not None: + session_payload = candidate_config.extra_body.get("realtime_session_payload") + audio_error = realtime_protocol.validate_server_vad_input_audio( + session_payload if isinstance(session_payload, dict) else {} + ) + if audio_error is not None: + return { + "type": "error", + "code": "unsupported_turn_detection", + "error": audio_error, + "param": "audio.input.format", + } + return None + if candidate_config.server_vad is not None: + if not candidate_config.server_vad.interrupt_response: + return { + "type": "error", + "session_id": session.session_id, + "code": "unsupported_turn_detection", + "error": "Model-native duplex server_vad requires interrupt_response=true", + } + if candidate_config.server_vad.threshold <= SILERO_VAD_MIN_THRESHOLD: + return { + "type": "error", + "session_id": session.session_id, + "code": "unsupported_turn_detection", + "error": ( + f"Model-native duplex server_vad.threshold must be greater than {SILERO_VAD_MIN_THRESHOLD}" + ), + } + if not candidate_config.server_vad.create_response: + return { + "type": "error", + "session_id": session.session_id, + "code": "unsupported_turn_detection", + "error": "Model-native duplex server_vad requires create_response=true", + } if ( candidate_config.instructions != session.config.instructions and self._runtime_session_state(session).native_context_locked @@ -1280,7 +1433,7 @@ class OmniDuplexSessionHandler( session: DuplexSession, committed: DuplexCommittedInput, *, - realtime_item_id: object | None = None, + item_id: object | None = None, ) -> dict[str, object]: payload: dict[str, object] = { "type": "input.committed", @@ -1290,15 +1443,15 @@ class OmniDuplexSessionHandler( "history_len": len(session.history), "message": committed.message, } - if isinstance(realtime_item_id, str) and realtime_item_id: - payload["realtime_item_id"] = realtime_item_id + if isinstance(item_id, str) and item_id: + payload["item_id"] = item_id return payload @staticmethod def _commit_native_audio_input( session: DuplexSession, *, - realtime_item_id: object | None = None, + item_id: object | None = None, transcript: object | None = None, turn_id: int | None = None, ) -> DuplexCommittedInput: @@ -1307,8 +1460,8 @@ class OmniDuplexSessionHandler( transcript=clean_transcript or None, turn_id=turn_id, ) - if isinstance(realtime_item_id, str) and realtime_item_id: - session.register_history_item(realtime_item_id, committed.message) + if isinstance(item_id, str) and item_id: + session.register_history_item(item_id, committed.message) return committed @staticmethod @@ -1316,7 +1469,7 @@ class OmniDuplexSessionHandler( session: DuplexSession, *, committed: DuplexCommittedInput | None = None, - realtime_item_id: object | None = None, + item_id: object | None = None, transcript: object | None = None, ) -> dict[str, object]: message = committed.message if committed is not None else None @@ -1342,12 +1495,15 @@ class OmniDuplexSessionHandler( } if isinstance(transcript, str) and transcript: payload["transcript"] = transcript - if isinstance(realtime_item_id, str) and realtime_item_id: - payload["realtime_item_id"] = realtime_item_id + if isinstance(item_id, str) and item_id: + payload["item_id"] = item_id return payload - @staticmethod - def _apply_session_update(session: DuplexSession, payload: dict[str, object]) -> dict[str, object] | None: + def _apply_session_update( + self, + session: DuplexSession, + payload: dict[str, object], + ) -> dict[str, object] | None: model = payload.get("model") audio_config = payload.get("audio") audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None @@ -1385,6 +1541,14 @@ class OmniDuplexSessionHandler( extra_body_payload = normalize_native_duplex_key(dict(extra_body_payload)) payload = {**payload, "extra_body": extra_body_payload} if isinstance(extra_body_payload, dict) and NATIVE_DUPLEX_KEY in extra_body_payload: + requested_native_duplex = extra_body_payload[NATIVE_DUPLEX_KEY] + if not isinstance(requested_native_duplex, bool): + return { + "type": "error", + "session_id": session.session_id, + "code": "invalid_request_error", + "error": "extra_body.native_duplex must be a boolean", + } current_native_duplex = native_duplex_opt_in(session.config.extra_body) if current_native_duplex is not None and extra_body_payload[NATIVE_DUPLEX_KEY] != current_native_duplex: return { @@ -1498,8 +1662,69 @@ class OmniDuplexSessionHandler( session.config.extra_body["realtime_session_payload"] = ( NativeRealtimeSessionProtocol._json_safe_realtime_payload(payload) ) + try: + turn_detection_configured, server_vad = parse_session_turn_detection(payload) + if turn_detection_configured: + vad_config = replace(session.config, server_vad=server_vad) + self._resolve_server_vad_defaults(session, vad_config, payload) + server_vad = vad_config.server_vad + if ( + not self._uses_native_input_append(session) + and session.turn_detection_config_locked + and (not session.config.turn_detection_configured or server_vad != session.config.server_vad) + ): + return { + "type": "error", + "session_id": session.session_id, + "code": "unsupported_turn_detection", + "error": "turn_detection cannot be changed after the first audio append", + } + session.config.turn_detection_configured = True + session.config.server_vad = server_vad + session.config.overlap_policy = vad_config.overlap_policy + except ValueError as exc: + return { + "type": "error", + "session_id": session.session_id, + "code": "invalid_request_error", + "error": str(exc), + } return None + async def _configure_server_vad(self, session: DuplexSession) -> None: + pipeline = await self._prepare_server_vad_pipeline(session, session.config) + self._install_server_vad_pipeline(session, pipeline) + + def _install_server_vad_pipeline( + self, + session: DuplexSession, + pipeline: ServerVADPipeline | None, + ) -> None: + existing = self._server_vad_pipelines.pop(session.session_id, None) + if existing is not None: + existing.reset() + self._realtime_vad_metrics.session_finished() + if pipeline is None: + return + self._server_vad_pipelines[session.session_id] = pipeline + self._realtime_vad_metrics.session_started() + + async def _prepare_server_vad_pipeline( + self, + session: DuplexSession, + config: DuplexSessionConfig, + ) -> ServerVADPipeline | None: + if config.server_vad is None: + return None + backend = await asyncio.to_thread(self._server_vad_backend_provider.get) + pipeline_config = config.server_vad + if self._uses_native_input_append(session) and pipeline_config.min_speech_duration_ms is None: + pipeline_config = replace( + pipeline_config, + min_speech_duration_ms=SILERO_VAD_DEFAULT_MIN_SPEECH_DURATION_MS, + ) + return ServerVADPipeline(backend, pipeline_config) + def _apply_response_create_options( self, session: DuplexSession, @@ -1580,7 +1805,7 @@ class OmniDuplexSessionHandler( response_extra.update( (key, value) for key, value in extra_body.items() - if key not in self._serving_runtime_adapter.private_runtime_config_keys + if key not in self._require_serving_runtime_adapter().private_runtime_config_keys ) else: response_extra.update(extra_body) @@ -1841,7 +2066,8 @@ class OmniDuplexSessionHandler( # Epoch-scoped request ids prevent state reuse, but explicitly # release projector/parser cursors so cancelled epochs do not # accumulate until the whole session closes. - self._serving_runtime_adapter.data_plane.close_stream(old_request_id) + if self._serving_runtime_adapter is not None: + self._serving_runtime_adapter.data_plane.close_stream(old_request_id) await self._abort_request_background( session, old_request_id, diff --git a/vllm_omni/entrypoints/duplex/session_runner.py b/vllm_omni/entrypoints/duplex/session_runner.py index 821e2c9c7..ed18d8591 100644 --- a/vllm_omni/entrypoints/duplex/session_runner.py +++ b/vllm_omni/entrypoints/duplex/session_runner.py @@ -4,17 +4,23 @@ from __future__ import annotations import asyncio +import binascii import json import uuid from contextlib import suppress from copy import deepcopy +import numpy as np +import pybase64 as base64 from fastapi import WebSocket, WebSocketDisconnect from vllm.logger import init_logger from vllm_omni.engine.duplex.lease import DuplexLeaseActivity from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.entrypoints.duplex.audio import convert_input_audio_with_rate +from vllm_omni.entrypoints.duplex.audio import ( + convert_input_audio_with_rate, + validate_input_sample_rate_hz, +) from vllm_omni.entrypoints.duplex.commit_policy import ( CommitAction, CommitSnapshot, @@ -36,6 +42,7 @@ from vllm_omni.entrypoints.duplex.runtime_adapter import ( ServingRuntimeSessionState, payload_turn_id, ) +from vllm_omni.entrypoints.duplex.server_vad import ServerVADPipeline from vllm_omni.entrypoints.duplex.session_attachment import ( DuplexJournalOverflowError, ) @@ -138,7 +145,7 @@ class DuplexSessionRunnerMixin: await actor.send_json(raw_payload) realtime_protocol.bind_sender(send_realtime_raw) - native: ServingRuntimeSessionState = self._serving_runtime_adapter.create_session_state() + native: ServingRuntimeSessionState = self._create_runtime_session_state() def begin_close(reason: str) -> None: actor.closing = True @@ -291,6 +298,120 @@ class DuplexSessionRunnerMixin: return True return False + def start_chat_response() -> None: + assert session is not None + + async def run_and_notify() -> None: + task = asyncio.current_task() + try: + await self._run_response(session, emit_event) + finally: + await actor.enqueue_event( + { + "type": "__chat_response_finished__", + "_response_task": task, + } + ) + + actor.active_response_task = asyncio.create_task(run_and_notify()) + + def activate_waiting_server_vad_turn() -> None: + assert session is not None + pending_turn = session.activate_server_vad_turn() + if pending_turn is not None and pending_turn.create_response: + start_chat_response() + + def correlate_realtime_error( + payload: dict[str, object], + event_id: object, + ) -> dict[str, object]: + if isinstance(event_id, str) and event_id: + payload["event_id"] = event_id + return payload + + async def reject_and_clear_server_vad_input( + pipeline: ServerVADPipeline, + *, + code: str, + error: str, + metric_reason: str, + event_id: object, + ) -> None: + assert session is not None + session.cancel_pending_input() + pipeline.reset() + self._realtime_vad_metrics.error(metric_reason) + await emit_event( + correlate_realtime_error( + { + "type": "error", + "session_id": session.session_id, + "code": code, + "error": error, + }, + event_id, + ) + ) + await emit_event( + { + "type": "input_audio_buffer.cleared", + "session_id": session.session_id, + } + ) + + async def commit_server_vad_turn( + *, + item_id: str, + create_response: bool, + retain_input_bytes: int, + event_id: object = None, + ) -> None: + assert session is not None + # Retain ownership until the actor consumes the terminal event so + # queued audio cannot overtake pending-turn activation. + active_response = actor.active_response_task is not None + if active_response and session.pending_server_vad_turn is not None: + released = session.discard_uncommitted_server_vad_utterance() + session.release_input_bytes(released) + self._realtime_vad_metrics.error("pending_turn_backpressure") + await emit_event( + correlate_realtime_error( + { + "type": "error", + "session_id": session.session_id, + "code": "input_backpressure", + "error": "Only one committed Server VAD turn may wait for the active response", + }, + event_id, + ) + ) + return + if not session.stage_server_vad_audio_for_commit(): + return + committed = session.commit_user_input( + append_history=not active_response, + retain_input_bytes=retain_input_bytes, + ) + if committed is None: + return + if active_response: + session.stage_server_vad_turn( + committed, + create_response=create_response, + item_id=item_id, + ) + else: + session.register_history_item(item_id, committed.message) + await emit_event( + self._input_committed_payload( + session, + committed, + item_id=item_id, + ) + ) + if create_response and not active_response: + start_chat_response() + def clear_completed_pending_silence() -> None: task = native.pending_silence_task if task is not None and task.done(): @@ -679,11 +800,18 @@ class DuplexSessionRunnerMixin: if handshake is None: return session = handshake.session - if initial_update_error := self._runtime_session_candidate_update_error(session, session.config): - await emit_event(initial_update_error) + if initial_update_error := self._runtime_session_candidate_update_error( + session, session.config, realtime_protocol=realtime_protocol + ): + await emit_event( + correlate_realtime_error( + initial_update_error, + handshake.event_id, + ) + ) return if handshake.resumed: - native = self._serving_runtime_adapter.session_states[session.session_id] + native = self._serving_session_states[session.session_id] actor.tasks = self._session_tasks[session.session_id] persisted_protocol = self._realtime_protocols.get(session.session_id) if persisted_protocol is None: @@ -697,18 +825,36 @@ class DuplexSessionRunnerMixin: self._ensure_lifecycle_listener() reader_task = asyncio.create_task(read_event_loop(), name="duplex-session-reader") else: - self._serving_runtime_adapter.session_states[session.session_id] = native + self._serving_session_states[session.session_id] = native self._session_tasks[session.session_id] = actor.tasks if realtime_protocol is not None: self._realtime_protocols[session.session_id] = realtime_protocol native.silence_continuation_scheduler = schedule_native_silence_continuation if realtime_protocol is not None: + realtime_protocol.bind_native_input_append(self._uses_native_input_append(session)) session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value if not handshake.resumed: open_result = await self._open_runtime_session(session, emit_event) if open_result is False: return runtime_opened = True + try: + await self._configure_server_vad(session) + except Exception as exc: + logger.exception("Server VAD initialization failed: %s", exc) + self._realtime_vad_metrics.error("initialization") + await emit_event( + correlate_realtime_error( + { + "type": "error", + "session_id": session.session_id, + "code": "server_vad_initialization_failed", + "error": str(exc), + }, + handshake.event_id, + ) + ) + return created_attachment = await self._attachment_registry.create( session.session_id, incarnation=session.incarnation, @@ -734,7 +880,7 @@ class DuplexSessionRunnerMixin: if isinstance(open_result, dict): created_payload["runtime_control"] = self._redact_runtime_control_result(open_result) if realtime_protocol is not None: - realtime_protocol.commit_realtime_turn_detection_update() + realtime_protocol.commit_realtime_turn_detection_update(session.config.server_vad) await emit_event(created_payload) resume_credential_delivered = session.capabilities.supports_session_resume reader_task = asyncio.create_task(read_event_loop(), name="duplex-session-reader") @@ -746,6 +892,14 @@ class DuplexSessionRunnerMixin: if event_type == "__replaced_attachment__": return + if event_type == "__chat_response_finished__": + completed_task = event.get("_response_task") + if completed_task is not actor.active_response_task: + continue + actor.active_response_task = None + activate_waiting_server_vad_turn() + continue + if event_type == "__timeout__": begin_close("timeout") native.audio_buffer.clear() @@ -896,6 +1050,9 @@ class DuplexSessionRunnerMixin: if event_type == "input_audio_buffer.clear": native.audio_buffer.clear() session.release_all_input_bytes() + server_vad_pipeline = self._server_vad_pipelines.get(session.session_id) + if server_vad_pipeline is not None: + server_vad_pipeline.reset() native.input_since_commit = False native.speech_since_commit = False native.clear_committed_audio() @@ -932,10 +1089,8 @@ class DuplexSessionRunnerMixin: if event_type == "response.cancel": requested_response_id = event.get("response_id") has_active_response_work = native_response_in_progress() - if ( - isinstance(requested_response_id, str) - and session.active_response_id is not None - and requested_response_id != session.active_response_id + if isinstance(requested_response_id, str) and ( + session.active_response_id is None or requested_response_id != session.active_response_id ): await emit_event( { @@ -946,6 +1101,7 @@ class DuplexSessionRunnerMixin: } ) continue + if event_type == "response.cancel": if not has_active_response_work: if realtime_protocol is not None and isinstance(requested_response_id, str): continue @@ -1056,6 +1212,7 @@ class DuplexSessionRunnerMixin: } ) actor.active_response_task = None + activate_waiting_server_vad_turn() continue if not cancelled: await self._cancel_pending_input(session, emit_event, reason="barge_in") @@ -1077,6 +1234,7 @@ class DuplexSessionRunnerMixin: ): continue actor.active_response_task = None + activate_waiting_server_vad_turn() continue if event_type == "turn.signal": @@ -1119,6 +1277,27 @@ class DuplexSessionRunnerMixin: ) reject_update() continue + if ( + realtime_protocol is not None + and not self._uses_native_input_append(session) + and session.config.server_vad is not None + and session.turn_detection_config_locked + and realtime_protocol.pending_realtime_input_audio_config_changed() + ): + await emit_update_event( + { + "type": "error", + "session_id": session.session_id, + "code": "unsupported_audio_format", + "error": ( + "session.update cannot change Server VAD input format or sample rate " + "after the first audio append" + ), + "param": "audio.input.format", + } + ) + reject_update() + continue runtime_update_error = self._runtime_session_update_error(session, payload) if runtime_update_error is not None: await emit_update_event(runtime_update_error) @@ -1131,6 +1310,7 @@ class DuplexSessionRunnerMixin: update_error = self._apply_session_update(session, payload) finally: session.replace_config(previous_config) + server_vad_changed = candidate_config.server_vad != previous_config.server_vad if update_error is not None: await emit_update_event(update_error) reject_update() @@ -1138,6 +1318,7 @@ class DuplexSessionRunnerMixin: runtime_update_error = self._runtime_session_candidate_update_error( session, candidate_config, + realtime_protocol=realtime_protocol, ) if runtime_update_error is not None: await emit_update_event(runtime_update_error) @@ -1159,19 +1340,46 @@ class DuplexSessionRunnerMixin: ) reject_update() continue - if not await self._signal_runtime_session( + candidate_server_vad_pipeline = None + if server_vad_changed: + try: + candidate_server_vad_pipeline = await self._prepare_server_vad_pipeline( + session, + candidate_config, + ) + except Exception as exc: + logger.exception("Server VAD session update failed: %s", exc) + self._realtime_vad_metrics.error("initialization") + await emit_update_event( + { + "type": "error", + "session_id": session.session_id, + "code": "server_vad_initialization_failed", + "error": str(exc), + } + ) + reject_update() + continue + runtime_updated = await self._signal_runtime_session( session, turn_event, emit_update_event, session_config=candidate_config.as_dict(), runtime_config=candidate_runtime_config, - ): + ) + if not runtime_updated: reject_update() continue session.replace_config(candidate_config) session.replace_runtime_config(candidate_runtime_config) + if server_vad_changed: + self._install_server_vad_pipeline( + session, + candidate_server_vad_pipeline, + ) if realtime_protocol is not None: - realtime_protocol.commit_realtime_turn_detection_update() + realtime_protocol.bind_native_input_append(self._uses_native_input_append(session)) + realtime_protocol.commit_realtime_turn_detection_update(session.config.server_vad) await emit_update_event( { "type": "session.updated", @@ -1229,7 +1437,10 @@ class DuplexSessionRunnerMixin: expected_model_turn_id=session.turn_id, ) continue - message = self._realtime_item_to_history_message(item) + history_item = payload.get("history_item") if isinstance(payload, dict) else None + if not isinstance(history_item, dict): + history_item = item + message = self._realtime_item_to_history_message(history_item) item_id = item.get("id") if isinstance(item, dict) else None if message is not None: session.append_history_message(message) @@ -1346,24 +1557,30 @@ class DuplexSessionRunnerMixin: default_sample_rate_hz = 16000 sr_raw = event.get("sample_rate_hz") or event.get("sample_rate") sample_rate_hz = sr_raw if isinstance(sr_raw, int | float) else default_sample_rate_hz - try: - audio, fmt, sample_rate_hz = convert_input_audio_with_rate( - audio, - fmt, - sample_rate_hz=sample_rate_hz, - ) - except ValueError as exc: - await emit_event({"type": "error", "error": str(exc), "code": "bad_event"}) - continue - if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le"}: - await emit_event( - { - "type": "error", - "error": "input_audio_buffer.append pcm16 audio could not be decoded", - "code": "bad_audio", - } - ) - continue + native_input = self._uses_native_input_append(session) + server_vad_config = session.config.server_vad + turn_based_server_vad = server_vad_config is not None and not native_input + if not turn_based_server_vad: + try: + audio, fmt, sample_rate_hz = convert_input_audio_with_rate( + audio, + fmt, + sample_rate_hz=sample_rate_hz, + ) + except ValueError as exc: + await emit_event({"type": "error", "error": str(exc), "code": "bad_event"}) + continue + if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le"}: + await emit_event( + { + "type": "error", + "error": "input_audio_buffer.append pcm16 audio could not be decoded", + "code": "bad_audio", + } + ) + continue + if server_vad_config is None: + session.lock_turn_detection_config() force_listen = bool(event.get("force_listen", False)) payload = { "type": "audio", @@ -1377,11 +1594,186 @@ class DuplexSessionRunnerMixin: frames = [frame for frame in video_frames if isinstance(frame, str) and frame] if frames: payload["video_frames"] = frames + if server_vad_config is not None: + realtime_event_id = event.get("_realtime_event_id") + pipeline = self._server_vad_pipelines.get(session.session_id) + if pipeline is None: + await emit_event( + correlate_realtime_error( + { + "type": "error", + "session_id": session.session_id, + "code": "internal_error", + "error": "server_vad is configured but its detector pipeline is unavailable", + }, + realtime_event_id, + ) + ) + continue + speech_was_active = pipeline.speech_active + reserved_bytes = 0 + previous_scratch_bytes = pipeline.scratch_bytes + try: + if native_input: + if fmt != "pcm_f32le" or sample_rate_hz != pipeline.sample_rate_hz: + raise ValueError("native server_vad input must be mono PCM float32 at 16000 Hz") + raw_audio = base64.b64decode(audio, validate=True) + if len(raw_audio) % np.dtype(" 0 + else: + if fmt != "pcm16": + raise ValueError("server_vad input must be mono PCM16") + source_sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz) + if source_sample_rate_hz not in {16_000, 24_000}: + raise ValueError("server_vad PCM16 input sample rate must be 16000 or 24000 Hz") + if ( + pipeline.source_sample_rate_hz is not None + and source_sample_rate_hz != pipeline.source_sample_rate_hz + ): + raise ValueError( + "server_vad input sample rate cannot change within a continuous audio stream" + ) + try: + raw_audio = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError( + "server_vad input audio must be valid base64-encoded PCM16" + ) from exc + if len(raw_audio) % np.dtype(" 0 + except (binascii.Error, ValueError) as exc: + await emit_event( + correlate_realtime_error( + { + "type": "error", + "session_id": session.session_id, + "code": "bad_audio", + "error": str(exc), + "param": "audio", + }, + realtime_event_id, + ) + ) + continue + if not native_input: + normalized_sample_upper_bound = ( + source_samples + if source_sample_rate_hz == pipeline.sample_rate_hz + else ((source_samples + 2) // 3) * 2 + ) + reserved_bytes = normalized_sample_upper_bound * np.dtype(np.float32).itemsize + if not session.reserve_input_bytes( + reserved_bytes, + limit=self._duplex_session_config.max_pending_input_bytes_per_session, + ): + await reject_and_clear_server_vad_input( + pipeline, + code="input_backpressure", + error="Server VAD input exceeds the per-session input limit", + metric_reason="input_backpressure", + event_id=realtime_event_id, + ) + continue + try: + vad_batch = ( + await pipeline.push(normalized_audio) + if native_input + else await pipeline.push_pcm16( + raw_audio, + source_sample_rate_hz=source_sample_rate_hz, + ) + ) + except Exception as exc: + logger.exception("Server VAD inference failed: %s", exc) + await reject_and_clear_server_vad_input( + pipeline, + code="server_vad_inference_failed", + error=str(exc), + metric_reason="inference", + event_id=realtime_event_id, + ) + continue + if not native_input: + retained_delta = ( + sum(frame.samples.nbytes for frame in vad_batch.frames) + + pipeline.scratch_bytes + - previous_scratch_bytes + ) + session.release_input_bytes(max(0, reserved_bytes - retained_delta)) + self._realtime_vad_metrics.observe_inference(vad_batch.inference_ms) + prefix_samples = round(server_vad_config.prefix_padding_ms * pipeline.sample_rate_hz / 1000) + speech_started = False + speech_stopped = False + for frame_index, vad_frame in enumerate(vad_batch.frames): + decision = vad_frame.decision + if not native_input: + released_bytes = session.append_server_vad_frame( + vad_frame.samples, + speech_started=decision.speech_started, + speech_stopped=decision.speech_stopped, + prefix_samples=prefix_samples, + ) + session.release_input_bytes(released_bytes) + if decision.speech_started: + speech_started = True + item_id = f"item_{uuid.uuid4().hex}" + session.begin_server_vad_speech(item_id) + await emit_event( + { + "type": "input_audio_buffer.speech_started", + "session_id": session.session_id, + "item_id": item_id, + "audio_start_ms": decision.audio_start_ms or 0, + } + ) + if decision.speech_stopped: + speech_stopped = True + item_id = session.finish_server_vad_speech() or f"item_{uuid.uuid4().hex}" + await emit_event( + { + "type": "input_audio_buffer.speech_stopped", + "session_id": session.session_id, + "item_id": item_id, + "audio_end_ms": decision.audio_end_ms or 0, + } + ) + if decision.endpoint_delay_ms is not None: + self._realtime_vad_metrics.observe_endpoint_delay(decision.endpoint_delay_ms) + if not native_input: + remaining_frame_bytes = sum( + frame.samples.nbytes for frame in vad_batch.frames[frame_index + 1 :] + ) + await commit_server_vad_turn( + item_id=item_id, + create_response=server_vad_config.create_response, + retain_input_bytes=remaining_frame_bytes + pipeline.scratch_bytes, + event_id=realtime_event_id, + ) + if accepted_audio: + session.lock_turn_detection_config() + if not native_input: + continue + speech_active = pipeline.speech_active + contained_speech = speech_was_active or speech_started or speech_stopped or speech_active + payload["is_speech"] = contained_speech + payload["vad"] = { + "backend": "silero", + "is_speech": contained_speech, + "speech_active": speech_active, + "speech_started": speech_started, + "speech_stopped": speech_stopped, + } + if speech_active: + payload["force_listen"] = True # Speech/silence tag for the Stage0 turn-ended latch. payload["is_speech"] = self._input_looks_like_speech(event, payload, session=session) defer_native_append = False buffer_overlap_audio = True - if self._uses_native_input_append(session): + if native_input: mark_pending_silence_superseded() overlap_active = native_response_in_progress() and ( not self._session_auto_responds(session) @@ -1570,7 +1962,7 @@ class DuplexSessionRunnerMixin: payload = pcm_reservation.payload else: session.append_audio(audio, fmt=fmt, sample_rate_hz=sample_rate_hz) - if self._uses_native_input_append(session): + if native_input: await start_native_append( payload, final=False, @@ -1582,11 +1974,23 @@ class DuplexSessionRunnerMixin: continue if event_type in {"input.commit", "input_audio_buffer.commit", "response.create"}: - realtime_item_id = event.get("realtime_item_id") + if ( + event_type == "response.create" + and session.config.server_vad is not None + and actor.active_response_task is not None + ): + await emit_event( + { + "type": "error", + "session_id": session.session_id, + "code": "response_already_active", + "error": "response.create must wait for the active response to complete", + } + ) + continue + item_id = event.get("item_id") realtime_validated_audio_commit = ( - event_type == "input_audio_buffer.commit" - and isinstance(realtime_item_id, str) - and bool(realtime_item_id) + event_type == "input_audio_buffer.commit" and isinstance(item_id, str) and bool(item_id) ) if ( self._uses_native_input_append(session) @@ -1604,6 +2008,7 @@ class DuplexSessionRunnerMixin: { "type": "input.committed", "session_id": session.session_id, + "item_id": item_id, "turn_id": session.turn_id, "epoch": session.epoch, "empty": True, @@ -1699,6 +2104,7 @@ class DuplexSessionRunnerMixin: { "type": "input.committed", "session_id": session.session_id, + "item_id": item_id, "turn_id": session.turn_id, "epoch": session.epoch, "empty": True, @@ -1736,13 +2142,13 @@ class DuplexSessionRunnerMixin: native.speech_since_commit = False committed = self._commit_native_audio_input( session, - realtime_item_id=event.get("realtime_item_id"), + item_id=event.get("item_id"), transcript=event.get("transcript"), ) committed_payload = self._native_audio_committed_payload( session, committed=committed, - realtime_item_id=event.get("realtime_item_id"), + item_id=event.get("item_id"), transcript=event.get("transcript"), ) committed_payload["overlap_deferred"] = True @@ -1777,7 +2183,7 @@ class DuplexSessionRunnerMixin: data_plane_turn_id = session.turn_id committed = self._commit_native_audio_input( session, - realtime_item_id=event.get("realtime_item_id"), + item_id=event.get("item_id"), transcript=event.get("transcript"), turn_id=data_plane_turn_id, ) @@ -1785,7 +2191,7 @@ class DuplexSessionRunnerMixin: self._native_audio_committed_payload( session, committed=committed, - realtime_item_id=event.get("realtime_item_id"), + item_id=event.get("item_id"), transcript=event.get("transcript"), ) ) @@ -1883,14 +2289,14 @@ class DuplexSessionRunnerMixin: native.input_since_commit = False committed = self._commit_native_audio_input( session, - realtime_item_id=event.get("realtime_item_id"), + item_id=event.get("item_id"), transcript=event.get("transcript"), ) await emit_event( self._native_audio_committed_payload( session, committed=committed, - realtime_item_id=event.get("realtime_item_id"), + item_id=event.get("item_id"), transcript=event.get("transcript"), ) ) @@ -1949,7 +2355,7 @@ class DuplexSessionRunnerMixin: committed = ( self._commit_native_audio_input( session, - realtime_item_id=event.get("realtime_item_id"), + item_id=event.get("item_id"), transcript=event.get("transcript"), ) if native_had_uncommitted_audio @@ -1959,7 +2365,7 @@ class DuplexSessionRunnerMixin: self._native_audio_committed_payload( session, committed=committed, - realtime_item_id=event.get("realtime_item_id"), + item_id=event.get("item_id"), transcript=event.get("transcript"), ) ) @@ -1968,19 +2374,20 @@ class DuplexSessionRunnerMixin: { "type": "input.committed", "session_id": session.session_id, + "item_id": item_id, "empty": True, } ) continue if committed is not None: - realtime_item_id = event.get("realtime_item_id") - if isinstance(realtime_item_id, str): - session.register_history_item(realtime_item_id, committed.message) + item_id = event.get("item_id") + if isinstance(item_id, str): + session.register_history_item(item_id, committed.message) await emit_event( self._input_committed_payload( session, committed, - realtime_item_id=realtime_item_id, + item_id=item_id, ) ) if not should_create_response: @@ -1992,7 +2399,10 @@ class DuplexSessionRunnerMixin: emit_event, reason="new_response", ) - actor.active_response_task = asyncio.create_task(self._run_response(session, emit_event)) + if session.config.server_vad is not None and not self._uses_native_input_append(session): + start_chat_response() + else: + actor.active_response_task = asyncio.create_task(self._run_response(session, emit_event)) continue await emit_event( diff --git a/vllm_omni/entrypoints/duplex/vad.py b/vllm_omni/entrypoints/duplex/vad.py deleted file mode 100644 index 362e500ee..000000000 --- a/vllm_omni/entrypoints/duplex/vad.py +++ /dev/null @@ -1,181 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import binascii -from collections.abc import Callable -from dataclasses import dataclass - -import numpy as np -import pybase64 as base64 - -SILERO_VAD_MIN_THRESHOLD = 0.15 - - -class ServerVADUnavailableError(RuntimeError): - pass - - -@dataclass(frozen=True, slots=True) -class SileroVADConfig: - threshold: float = 0.5 - prefix_padding_ms: int = 300 - silence_duration_ms: int = 500 - min_speech_duration_ms: int = 96 - - -@dataclass(frozen=True, slots=True) -class StreamingVADResult: - is_speech: bool - speech_active: bool - speech_started: bool = False - speech_stopped: bool = False - speech_probability: float = 0.0 - speech_start_ms: int | None = None - speech_end_ms: int | None = None - - -class SileroStreamingVAD: - _SAMPLE_RATE_HZ = 16_000 - _WINDOW_SAMPLES = 512 - - def __init__( - self, - config: SileroVADConfig, - *, - frame_scorer: Callable[[np.ndarray], float] | None = None, - ) -> None: - self.config = config - self._frame_scorer = frame_scorer - self._model: object | None = None - self._clear_stream_state() - - def _clear_stream_state(self) -> None: - self._pending = np.empty(0, dtype=np.float32) - self._speech_active = False - self._candidate_samples = 0 - self._candidate_start_sample: int | None = None - self._silence_samples = 0 - self._processed_samples = 0 - self._resample_rate_hz: int | None = None - self._resample_remainder = 0 - - def reset(self) -> None: - self._clear_stream_state() - reset_states = getattr(self._model, "reset_states", None) - if callable(reset_states): - reset_states() - - def process_base64( - self, - audio: object, - *, - fmt: object, - sample_rate_hz: object, - ) -> StreamingVADResult: - if fmt != "pcm_f32le" or not isinstance(audio, str): - raise ValueError("Silero server VAD requires decoded pcm_f32le audio") - try: - raw = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError) as exc: - raise ValueError("Silero server VAD received invalid base64 audio") from exc - if len(raw) % 4: - raise ValueError("Silero server VAD received an incomplete pcm_f32le frame") - samples = np.frombuffer(raw, dtype=" StreamingVADResult: - samples = np.asarray(samples, dtype=np.float32).reshape(-1) - if samples.size: - self._pending = np.concatenate((self._pending, samples)) - - started = False - stopped = False - start_ms: int | None = None - end_ms: int | None = None - max_probability = 0.0 - contained_speech = self._speech_active - min_speech_samples = self.config.min_speech_duration_ms * 16 - min_silence_samples = self.config.silence_duration_ms * 16 - negative_threshold = max(0.0, self.config.threshold - SILERO_VAD_MIN_THRESHOLD) - - while self._pending.size >= self._WINDOW_SAMPLES: - frame = self._pending[: self._WINDOW_SAMPLES] - self._pending = self._pending[self._WINDOW_SAMPLES :] - frame_start = self._processed_samples - self._processed_samples += self._WINDOW_SAMPLES - probability = min(1.0, max(0.0, float(self._score_frame(frame)))) - max_probability = max(max_probability, probability) - - if self._speech_active: - contained_speech = True - if probability < negative_threshold: - self._silence_samples += self._WINDOW_SAMPLES - if self._silence_samples >= min_silence_samples: - speech_end_sample = self._processed_samples - self._silence_samples - self._speech_active = False - self._silence_samples = 0 - stopped = True - end_ms = max(0, int(round(speech_end_sample / 16))) - else: - self._silence_samples = 0 - continue - - if probability >= self.config.threshold: - if self._candidate_samples == 0: - self._candidate_start_sample = frame_start - self._candidate_samples += self._WINDOW_SAMPLES - if self._candidate_samples >= min_speech_samples: - candidate_start = self._candidate_start_sample or 0 - prefix_samples = self.config.prefix_padding_ms * 16 - self._speech_active = True - self._candidate_samples = self._silence_samples = 0 - self._candidate_start_sample = None - contained_speech = True - started = True - start_ms = max(0, int(round((candidate_start - prefix_samples) / 16))) - else: - self._candidate_samples = 0 - self._candidate_start_sample = None - - return StreamingVADResult( - contained_speech, self._speech_active, started, stopped, max_probability, start_ms, end_ms - ) - - def _score_frame(self, frame: np.ndarray) -> float: - if self._frame_scorer is not None: - return float(self._frame_scorer(frame)) - if self._model is None: - try: - import torch - from silero_vad import load_silero_vad - except ImportError as exc: - raise ServerVADUnavailableError( - "server_vad requires the optional 'silero-vad' package; install vllm-omni[server-vad]" - ) from exc - self._model = load_silero_vad(onnx=False) - self._frame_scorer = lambda value: float( - self._model(torch.from_numpy(np.ascontiguousarray(value)), self._SAMPLE_RATE_HZ).item() - ) - return float(self._frame_scorer(frame)) diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 5b13fc246..264b37fb5 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -339,6 +339,7 @@ class OmniBase(PDDisaggregationMixin): allow_delta_coercion: bool = False, ) -> Sequence[Any]: """Resolve request parameters; pipeline sampling constraints override caller values.""" + normalized: Sequence[Any] if sampling_params_list is None: normalized = self.default_sampling_params_list # Set the output kind to delta since no params were specified @@ -366,7 +367,7 @@ class OmniBase(PDDisaggregationMixin): @staticmethod def _get_sampling_constraints_list(stage_configs: Sequence[Any]) -> list[dict[str, Any]]: """Extract each stage's required sampling settings from runtime configs.""" - constraints_list = [] + constraints_list: list[dict[str, Any]] = [] for stage_config in stage_configs: constraints = getattr(stage_config, "sampling_constraints", {}) if not isinstance(constraints, Mapping): @@ -400,7 +401,7 @@ class OmniBase(PDDisaggregationMixin): req_state = self.request_states.get(request_id) prom = getattr(self, "prom_metrics", None) metrics = getattr(req_state, "metrics", None) - if metrics is None or prom is None: + if req_state is None or metrics is None or prom is None: return if str(request_id) in metrics.e2e_done or getattr(req_state, "failure_recorded", False): return diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py old mode 100755 new mode 100644 index 3d9403bf8..d976eb181 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -15,8 +15,8 @@ import tempfile import time import uuid from argparse import Namespace -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager +from collections.abc import AsyncIterator, Mapping, Sequence +from contextlib import asynccontextmanager, suppress from http import HTTPStatus from numbers import Integral from pathlib import Path @@ -131,6 +131,7 @@ from vllm_omni.entrypoints.openai.protocol.images import ( from vllm_omni.entrypoints.openai.protocol.videos import ( SecondStr, SizeStr, + VideoAction, VideoDeleteResponse, VideoError, VideoGenerationRequest, @@ -160,6 +161,8 @@ 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 ( VideoFrames, + _decode_image_bytes, + _validate_image_pixel_limit, decode_audio_url, decode_input_reference, ) @@ -181,6 +184,9 @@ MINIMAX_H3_MAX_REFERENCE_COUNT = 12 MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"}) MINIMAX_H3_REFERENCE_VIDEO_SUFFIXES = frozenset({".mp4", ".mov"}) MINIMAX_H3_REFERENCE_AUDIO_SUFFIXES = frozenset({".wav", ".mp3"}) +CONTROL_REFERENCE_IMAGE_SUFFIXES = frozenset({".bmp", ".gif", ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp"}) +CONTROL_REFERENCE_VIDEO_SUFFIXES = frozenset({".mkv", ".mov", ".mp4", ".webm"}) +CONTROL_REFERENCE_MAX_BYTES = 512 * 1024 * 1024 profiler_router = APIRouter() @@ -197,9 +203,9 @@ def _load_model_chat_template_json(model: str) -> str | None: if template_path is None: try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - template_path = hf_hub_download( + template_path = hf_api().hf_hub_download( repo_id=model, filename="chat_template.json", local_files_only=True, @@ -1286,10 +1292,12 @@ async def omni_init_app_state( state.openai_serving_duplex = None if state.openai_serving_chat is not None and should_enable_duplex_endpoint( state.stage_configs, - config_path=getattr(args, "deploy_config", None), + config_path=getattr(engine_client, "config_path", None) or getattr(args, "deploy_config", None), ): state.openai_serving_duplex = OmniDuplexSessionHandler( chat_service=state.openai_serving_chat, + served_model_name=model_name, + log_stats=state.log_stats, duplex_session_config=getattr(engine_client, "duplex_session_config", None), serving_runtime_adapter_path=getattr(engine_client, "duplex_serving_adapter_path", None), ) @@ -1846,8 +1854,8 @@ async def realtime_websocket(websocket: WebSocket): logger.warning("Duplex warmup still running after 120 s; admitting the client anyway.") duplex_handler = getattr(websocket.app.state, "openai_serving_duplex", None) duplex_query = websocket.query_params.get("duplex") - use_duplex_realtime = ( - duplex_handler is not None and isinstance(duplex_query, str) and duplex_query.lower() in {"1", "true", "on"} + use_duplex_realtime = duplex_handler is not None and ( + duplex_query is None or (isinstance(duplex_query, str) and duplex_query.lower() in {"1", "true", "on"}) ) if use_duplex_realtime and duplex_handler is not None: await duplex_handler.handle_realtime_session(websocket) @@ -1943,12 +1951,26 @@ async def show_available_models(raw_request: Request) -> JSONResponse: # Image generation API endpoints +def _build_image_response_metrics( + *, + response_metrics: Any, + stage_durations: Any, + peak_memory_mb: Any, +) -> dict[str, Any]: + """Merge detailed stage metrics with the legacy image timing fields.""" + metrics = dict(response_metrics) if isinstance(response_metrics, dict) else {} + metrics["stage_durations"] = stage_durations or None + metrics["peak_memory_mb"] = float(peak_memory_mb) if peak_memory_mb else None + return metrics + + def _build_image_generation_response( *, images: list[Image.Image], request: ImageGenerationRequest, stage_durations: Any, peak_memory_mb: Any, + response_metrics: Any = None, ) -> ImageGenerationResponse | StreamingResponse: """Encode generated images and apply the requested response format.""" output_format = _choose_output_format(request.output_format or "png", None) @@ -1963,10 +1985,11 @@ def _build_image_generation_response( "created": int(time.time()), "data": image_data, "output_format": output_format, - "metrics": { - "stage_durations": stage_durations or None, - "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, - }, + "metrics": _build_image_response_metrics( + response_metrics=response_metrics, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ), } if request.size is not None: response_kwargs["size"] = request.size @@ -2063,6 +2086,8 @@ async def generate_images( extra_body["use_system_prompt"] = request.use_system_prompt if request.system_prompt is not None: extra_body["system_prompt"] = request.system_prompt + if request.return_stage_metrics is not None: + extra_body["return_stage_metrics"] = request.return_stage_metrics generation_result = await chat_handler.generate_diffusion_images( prompt=request.prompt, @@ -2076,12 +2101,13 @@ async def generate_images( status_code=generation_result.error.code if generation_result.error else 400, content=generation_result.model_dump(), ) - flat_images, stage_durations, peak_memory_mb, _ = generation_result + flat_images, stage_durations, peak_memory_mb, _, response_metrics = generation_result return _build_image_generation_response( images=flat_images, request=request, stage_durations=stage_durations, peak_memory_mb=peak_memory_mb, + response_metrics=response_metrics, ) # Build params - pass through user values directly @@ -2171,11 +2197,13 @@ async def generate_images( stage_durations = getattr(result, "stage_durations", None) peak_memory_mb = getattr(result, "peak_memory_mb", None) + response_metrics = getattr(result, "metrics", None) if request.return_stage_metrics else None return _build_image_generation_response( images=images, request=request, stage_durations=stage_durations, peak_memory_mb=peak_memory_mb, + response_metrics=response_metrics, ) except (EngineGenerateError, EngineDeadError) as exc: @@ -2507,7 +2535,7 @@ async def edit_images( status_code=generation_result.error.code if generation_result.error else 400, detail=generation_result.message, ) - images, stage_durations, peak_memory_mb, cot_output = generation_result + images, stage_durations, peak_memory_mb, cot_output, response_metrics = generation_result else: # Single-stage diffusion: use the direct path. result = await _generate_with_async_omni( @@ -2520,6 +2548,7 @@ async def edit_images( images = _extract_images_from_result(result) stage_durations = getattr(result, "stage_durations", None) peak_memory_mb = getattr(result, "peak_memory_mb", None) + response_metrics = getattr(result, "metrics", None) if return_stage_metrics else None logger.debug(f"Successfully generated {len(images)} image(s)") @@ -2540,10 +2569,11 @@ async def edit_images( output_format=output_format, size=size_str, cot_output=cot_output, - metrics={ - "stage_durations": stage_durations or None, - "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, - }, + metrics=_build_image_response_metrics( + response_metrics=response_metrics, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ), ) except (EngineGenerateError, EngineDeadError) as exc: @@ -3007,12 +3037,22 @@ def _reference_video_decode_spec( def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> VideoResponse: + video_params = req.resolve_video_params() + duration_s = None + if req.seconds is not None: + duration_s = float(req.seconds) + elif video_params.num_frames is not None and video_params.fps is not None: + duration_s = video_params.num_frames / video_params.fps + resp = VideoResponse( model=model_name, status=VideoGenerationStatus.QUEUED, size=req.size, prompt=req.prompt, quality=req.quality or "default", + fps=video_params.fps, + num_frames=video_params.num_frames, + duration_s=duration_s, ) resp.seconds = str(req.seconds or resp.seconds) return resp @@ -3073,6 +3113,7 @@ async def _cleanup_video(video_id: str): def _cleanup_video_references( reference_video: ReferenceVideo | None, reference_audio: ReferenceAudio | None, + control_path: str | None = None, ) -> None: if reference_video is not None: for path in reference_video.cleanup_paths: @@ -3083,6 +3124,27 @@ def _cleanup_video_references( for path in cleanup_paths: if os.path.exists(path): os.unlink(path) + if control_path is not None and os.path.exists(control_path): + os.unlink(control_path) + + +def _unpack_video_generation_result( + result: Sequence[object], +) -> tuple[bytes, dict[str, float], float, VideoAction | None, dict[str, object]]: + video_metadata: dict[str, object] = {} + if len(result) == 5: + video_bytes, stage_durations, peak_memory_mb, action, raw_metadata = result + if isinstance(raw_metadata, dict): + video_metadata = {str(key): value for key, value in raw_metadata.items()} + else: + video_bytes, stage_durations, peak_memory_mb, action = result + return ( + cast(bytes, video_bytes), + cast(dict[str, float], stage_durations), + float(cast(float, peak_memory_mb)), + cast(VideoAction | None, action), + video_metadata, + ) async def _run_video_generation_job( @@ -3092,22 +3154,26 @@ async def _run_video_generation_job( reference_image: ReferenceImage | None = None, reference_video: ReferenceVideo | None = None, reference_audio: ReferenceAudio | None = None, + control_path: str | None = None, app_state: Any | None = None, ) -> None: job = await VIDEO_STORE.get(video_id) if job is None: logger.warning("Video job %s missing before generation task started; skipping", video_id) + _cleanup_video_references(reference_video, reference_audio, control_path) return await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS}) started_at = time.perf_counter() try: - video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes( - request, - video_id, - reference_image=reference_image, - reference_video=reference_video, - reference_audio=reference_audio, + video_bytes, stage_durations, peak_memory_mb, action, video_metadata = _unpack_video_generation_result( + await handler.generate_video_bytes( + request, + video_id, + reference_image=reference_image, + reference_video=reference_video, + reference_audio=reference_audio, + ) ) save_context = await STORAGE_MANAGER.save(video_bytes, video_id) @@ -3123,6 +3189,7 @@ async def _run_video_generation_job( "peak_memory_mb": peak_memory_mb, "action": action, } + updated_fields.update({key: value for key, value in video_metadata.items() if value is not None}) if save_context.expires_at is not None: updated_fields["expires_at"] = save_context.expires_at @@ -3165,7 +3232,7 @@ async def _run_video_generation_job( await VIDEO_STORE.pop(video_id) raise finally: - _cleanup_video_references(reference_video, reference_audio) + _cleanup_video_references(reference_video, reference_audio, control_path) VIDEO_SYNC_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_VIDEO_SYNC_TIMEOUT", 600.0)) @@ -3191,6 +3258,121 @@ async def _persist_uploaded_video_references(uploads: list[UploadFile]) -> list[ return paths +async def _persist_uploaded_control_reference( + upload: UploadFile, + *, + max_bytes: int = CONTROL_REFERENCE_MAX_BYTES, +) -> str: + """Stream one model control upload to request-scoped local storage.""" + kind = _uploaded_media_kind(upload) + suffix = Path(upload.filename or "").suffix.lower() + supported_suffixes = CONTROL_REFERENCE_IMAGE_SUFFIXES | CONTROL_REFERENCE_VIDEO_SUFFIXES + if kind == "audio" or (suffix and suffix not in supported_suffixes): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="control_reference must be an image or video file.", + ) + if not suffix: + suffix = ".png" if kind == "image" else ".mp4" + + declared_size = getattr(upload, "size", None) + if isinstance(declared_size, Integral) and int(declared_size) > max_bytes: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"control_reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", + ) + + fd, path = tempfile.mkstemp(prefix="vllm_omni_control_reference_", suffix=suffix) + size = 0 + persisted = False + try: + with os.fdopen(fd, "wb") as output: + while chunk := await upload.read(min(1024 * 1024, max_bytes - size + 1)): + size += len(chunk) + if size > max_bytes: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"control_reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", + ) + output.write(chunk) + if size == 0: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="control_reference must not be empty.", + ) + persisted = True + return path + finally: + if not persisted: + with suppress(OSError): + os.unlink(path) + + +def _validate_control_upload( + handler: OmniOpenAIServingVideo, + request: VideoGenerationRequest, + control_reference: UploadFile | None, + control_type: str | None, +) -> str | None: + """Validate a generic uploaded control against model capabilities.""" + if control_reference is None: + if control_type is not None: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="control_type requires a control_reference upload.", + ) + return None + if control_type is None or not control_type.strip(): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="control_reference requires control_type.", + ) + + normalized_type = control_type.strip().lower() + supported_types = frozenset(getattr(handler, "supported_control_upload_types", ())) + if normalized_type not in supported_types: + supported = ", ".join(sorted(supported_types)) or "none" + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=( + f"Control upload type '{normalized_type}' is not supported by this model. " + f"Supported control upload types: {supported}." + ), + ) + + existing = (request.extra_params or {}).get(normalized_type) + if existing is not None: + if not isinstance(existing, Mapping) and existing is not True: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"extra_params.{normalized_type} must be an object when control_reference is uploaded.", + ) + if isinstance(existing, Mapping) and any(existing.get(key) is not None for key in ("control", "control_path")): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=( + "Provide either control_reference or " + f"extra_params.{normalized_type}.control/control_path, not both." + ), + ) + return normalized_type + + +def _attach_control_upload( + request: VideoGenerationRequest, + control_type: str, + control_path: str | None = None, +) -> None: + """Declare an uploaded control and attach its persisted path when available.""" + extra_params = dict(request.extra_params or {}) + existing = extra_params.get(control_type) + control_params = dict(existing) if isinstance(existing, Mapping) else {} + if control_path is not None: + control_params["control_path"] = control_path + extra_params[control_type] = control_params + request.extra_params = extra_params + + def _reference_list(value: Any) -> list[Any]: if value is None: return [] @@ -3258,6 +3440,15 @@ def _validate_minimax_h3_image_payload( try: with Image.open(io.BytesIO(payload)) as image: image_format = str(image.format or "").lower() + if image_format in MINIMAX_H3_REFERENCE_IMAGE_FORMATS: + _validate_image_pixel_limit(image) + except InvalidInputReferenceError as exc: + raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc + except Image.DecompressionBombError as exc: + raise HTTPException( + HTTPStatus.BAD_REQUEST.value, + detail=f"Invalid uploaded image reference: {filename}; image exceeds the decoder pixel limit.", + ) from exc except (OSError, ValueError) as exc: if allow_non_image: return @@ -3298,10 +3489,14 @@ async def _persist_uploaded_media_references( if kind == "image": try: _validate_minimax_h3_image_payload(payload, filename=upload.filename) - with Image.open(io.BytesIO(payload)) as image: - images.append(image.convert("RGB")) - except (OSError, ValueError) as exc: - raise HTTPException(400, detail=f"Invalid uploaded image reference: {upload.filename}") from exc + images.append( + _decode_image_bytes( + payload, + source=f"uploaded image reference: {upload.filename}", + ) + ) + except InvalidInputReferenceError as exc: + raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc continue suffix = Path(upload.filename or "").suffix.lower() if kind == "video" and suffix and suffix not in MINIMAX_H3_REFERENCE_VIDEO_SUFFIXES: @@ -3337,6 +3532,8 @@ async def _parse_video_form( prompt: str = Form(...), input_reference: UploadFile | None = File(default=None), input_references: list[UploadFile] | None = File(default=None), + control_reference: UploadFile | None = File(default=None), + control_type: str | None = Form(default=None), image_reference: str | None = Form(default=None), video_reference: str | None = Form(default=None), audio_reference: str | None = Form(default=None), @@ -3369,6 +3566,7 @@ async def _parse_video_form( frame_interpolation_model_path: str | None = Form(default=None), lora: str | None = Form(default=None), extra_params: str | None = Form(default=None), + return_stage_metrics: bool | None = Form(default=None), ) -> tuple[ VideoGenerationRequest, "OmniOpenAIServingVideo", @@ -3376,6 +3574,7 @@ async def _parse_video_form( ReferenceImage | None, ReferenceVideo | None, ReferenceAudio | None, + str | None, ]: """FastAPI dependency that parses video form data, validates inputs, resolves the handler, and decodes any reference image. @@ -3438,6 +3637,7 @@ async def _parse_video_form( "frame_interpolation_model_path": frame_interpolation_model_path, "lora": _parse_form_json(lora, expected_type=dict), "extra_params": _parse_form_json(extra_params, expected_type=dict), + "return_stage_metrics": return_stage_metrics, } request_data = {k: v for k, v in request_data.items() if v is not None} request = VideoGenerationRequest(**request_data) @@ -3470,6 +3670,13 @@ async def _parse_video_form( detail=f"Video generation setup failed: {str(e)}", ) + normalized_control_type = _validate_control_upload(handler, request, control_reference, control_type) + if normalized_control_type is not None: + # Make the selected transfer mode visible while choosing the model's + # reference-video decode policy. The upload itself stays unpersisted + # until reference parsing succeeds, preserving failure cleanup. + _attach_control_upload(request, normalized_control_type) + supports_mixed_reference_inputs = bool(getattr(handler, "supports_mixed_reference_inputs", False)) if input_reference is not None: input_reference_bytes = await _read_upload_limited( @@ -3592,7 +3799,27 @@ async def _parse_video_form( cleanup_paths=cleanup_paths, ) - return request, handler, effective_model_name, reference_image, reference_video, reference_audio + control_path: str | None = None + if control_reference is not None and normalized_control_type is not None: + try: + control_path = await _persist_uploaded_control_reference( + control_reference, + max_bytes=CONTROL_REFERENCE_MAX_BYTES, + ) + _attach_control_upload(request, normalized_control_type, control_path) + except (asyncio.CancelledError, HTTPException, OSError, TypeError, ValueError): + _cleanup_video_references(reference_video, reference_audio, control_path) + raise + + return ( + request, + handler, + effective_model_name, + reference_image, + reference_video, + reference_audio, + control_path, + ) @router.post( @@ -3613,6 +3840,7 @@ async def create_video( ReferenceImage | None, ReferenceVideo | None, ReferenceAudio | None, + str | None, ] = Depends(_parse_video_form), ) -> VideoResponse: """Create an asynchronous video generation job. @@ -3620,7 +3848,15 @@ async def create_video( Accepts multipart form-data (see ``_parse_video_form`` for parameters), persists a queued job record, and starts generation in the background. """ - request, handler, effective_model_name, reference_image, reference_video, reference_audio = ctx + ( + request, + handler, + effective_model_name, + reference_image, + reference_video, + reference_audio, + control_path, + ) = ctx ref = video_response_from_request(effective_model_name, request) await VIDEO_STORE.upsert(ref.id, ref) task = asyncio.create_task( @@ -3631,6 +3867,7 @@ async def create_video( reference_image, reference_video, reference_audio, + control_path, app_state=raw_request.app.state, ) ) @@ -3656,6 +3893,7 @@ async def create_video_sync( ReferenceImage | None, ReferenceVideo | None, ReferenceAudio | None, + str | None, ] = Depends(_parse_video_form), ) -> Response: """Synchronous video generation endpoint. @@ -3667,20 +3905,30 @@ async def create_video_sync( Metadata is returned via response headers ``X-Request-Id``, ``X-Model``, and ``X-Inference-Time-S``. """ - request, handler, effective_model_name, reference_image, reference_video, reference_audio = ctx + ( + request, + handler, + effective_model_name, + reference_image, + reference_video, + reference_audio, + control_path, + ) = ctx request_id = f"video_sync-{random_uuid()}" raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id) started_at = time.perf_counter() try: - video_bytes, stage_durations, peak_memory_mb, _action = await asyncio.wait_for( - handler.generate_video_bytes( - request, - request_id, - reference_image=reference_image, - reference_video=reference_video, - reference_audio=reference_audio, + video_bytes, stage_durations, peak_memory_mb, _action, _video_metadata = _unpack_video_generation_result( + await asyncio.wait_for( + handler.generate_video_bytes( + request, + request_id, + reference_image=reference_image, + reference_video=reference_video, + reference_audio=reference_audio, + ), + timeout=VIDEO_SYNC_TIMEOUT_S, ), - timeout=VIDEO_SYNC_TIMEOUT_S, ) except asyncio.TimeoutError: raise HTTPException( @@ -3701,7 +3949,7 @@ async def create_video_sync( detail=f"Video generation failed: {str(exc)}", ) from exc finally: - _cleanup_video_references(reference_video, reference_audio) + _cleanup_video_references(reference_video, reference_audio, control_path) inference_time_s = time.perf_counter() - started_at return Response( diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py index 3cfcfc31d..e3bce1a67 100644 --- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import math from io import BytesIO import numpy as np @@ -19,85 +20,152 @@ logger = init_logger(__name__) class StreamingAudioResampler: - """Stateful resampler for streaming mono audio. + """Stateful polyphase resampler for streaming mono float audio. - Only integer downsampling ratios are supported so every input chunk can - produce output without buffering the complete stream. + Retains filter state so output is invariant to input chunk boundaries. """ + _half_filter_width = 10 + _kaiser_beta = 5.0 + _max_polyphase_factor = 2_048 + def __init__(self, source_rate: int, target_rate: int): if source_rate <= 0 or target_rate <= 0: raise ValueError("Audio sample rates must be positive") self.source_rate = source_rate self.target_rate = target_rate - if source_rate < target_rate or source_rate % target_rate != 0: - raise ValueError( - "Streaming audio resampling requires an integer downsampling ratio, " - f"got {source_rate} Hz to {target_rate} Hz" + rate_gcd = math.gcd(source_rate, target_rate) + self._up = target_rate // rate_gcd + self._down = source_rate // rate_gcd + if max(self._up, self._down) > self._max_polyphase_factor: + raise ValueError("source and target sample rates have an unsupported resampling ratio") + self._half_len, self._phase_kernels = self._design_polyphase_filter() + self._history_samples = self._phase_kernels.shape[1] - 1 + self.reset() + + def _design_polyphase_filter(self) -> tuple[int, np.ndarray]: + if self._up == self._down: + return 0, np.ones((1, 1), dtype=np.float32) + + max_rate = max(self._up, self._down) + half_len = self._half_filter_width * max_rate + offsets = np.arange(-half_len, half_len + 1, dtype=np.float64) + cutoff = 1.0 / max_rate + taps = cutoff * np.sinc(cutoff * offsets) + taps *= np.kaiser(taps.size, self._kaiser_beta) + taps /= np.sum(taps) + taps *= self._up + taps = np.ascontiguousarray(taps, dtype=np.float32) + + phases = [taps[phase :: self._up] for phase in range(self._up)] + phase_width = max(phase.size for phase in phases) + kernels = np.zeros((self._up, phase_width), dtype=np.float32) + for phase_index, phase in enumerate(phases): + kernels[phase_index, -phase.size :] = phase[::-1] + return half_len, kernels + + @property + def scratch_bytes(self) -> int: + pending_samples = self._ceil_div(self._input_samples * self._up, self._down) - self._output_samples + return max(0, pending_samples) * np.dtype(np.float32).itemsize + + @staticmethod + def _ceil_div(numerator: int, denominator: int) -> int: + return -(-numerator // denominator) + + def _render_outputs( + self, + combined: np.ndarray, + *, + first_output: int, + output_count: int, + input_start: int, + ) -> np.ndarray: + if output_count <= 0: + return np.empty(0, dtype=np.float32) + + output_indexes = np.arange(first_output, first_output + output_count, dtype=np.int64) + filter_indexes = output_indexes * self._down + self._half_len + source_indexes = filter_indexes // self._up + phases = filter_indexes % self._up + windows = np.lib.stride_tricks.sliding_window_view( + combined, + self._phase_kernels.shape[1], + ) + selected_windows = windows[source_indexes - input_start] + return np.sum( + selected_windows * self._phase_kernels[phases], + axis=1, + dtype=np.float32, + ) + + def _push(self, chunk: np.ndarray) -> np.ndarray: + if chunk.size == 0: + return np.empty(0, dtype=np.float32) + if self._flushed: + raise RuntimeError("cannot process audio after the resampler has been finalized; call reset first") + + combined = np.concatenate((self._history, chunk)) + next_input_samples = self._input_samples + chunk.size + stable_output_samples = max( + 0, + self._ceil_div(next_input_samples * self._up - self._half_len, self._down), + ) + output = self._render_outputs( + combined, + first_output=self._output_samples, + output_count=stable_output_samples - self._output_samples, + input_start=self._input_samples, + ) + if self._history_samples: + self._history = combined[-self._history_samples :].copy() + self._input_samples = next_input_samples + self._output_samples = stable_output_samples + return output + + def _flush(self) -> np.ndarray: + if self._flushed: + return np.empty(0, dtype=np.float32) + + total_output_samples = self._ceil_div(self._input_samples * self._up, self._down) + output_count = total_output_samples - self._output_samples + if output_count: + last_output = total_output_samples - 1 + last_source_index = (last_output * self._down + self._half_len) // self._up + right_padding = max(0, last_source_index - self._input_samples + 1) + combined = np.concatenate((self._history, np.zeros(right_padding, dtype=np.float32))) + output = self._render_outputs( + combined, + first_output=self._output_samples, + output_count=output_count, + input_start=self._input_samples, ) - - self._ratio = source_rate // target_rate - self._buffer = np.empty((0,), dtype=np.float32) - self._buffer_start = 0 - self._total_samples = 0 - self._next_center = 0 - - if self._ratio == 1: - self._taps = np.ones((1,), dtype=np.float32) else: - # Use a Kaiser-windowed sinc and leave a small transition band - # below the target Nyquist frequency. - num_taps = 20 * self._ratio + 1 - half = num_taps // 2 - offsets = np.arange(num_taps, dtype=np.float64) - half - cutoff = 0.95 / self._ratio - taps = cutoff * np.sinc(cutoff * offsets) * np.kaiser(num_taps, 5.0) - self._taps = (taps / taps.sum()).astype(np.float32) + output = np.empty(0, dtype=np.float32) + self._output_samples = total_output_samples + self._flushed = True + return output def process(self, audio: np.ndarray, *, final: bool = False) -> np.ndarray: chunk = np.asarray(audio, dtype=np.float32) if chunk.ndim != 1: raise ValueError(f"Streaming audio resampling only supports mono audio, got shape {chunk.shape}") + chunk = np.ascontiguousarray(chunk, dtype=np.float32) + output = self._push(chunk) + if not final: + return output + tail = self._flush() + if not output.size: + return tail + if not tail.size: + return output + return np.concatenate((output, tail)) - if self._ratio == 1: - return chunk - - if chunk.size: - self._buffer = np.concatenate((self._buffer, chunk)) - self._total_samples += int(chunk.size) - - half = self._taps.size // 2 - max_center = self._total_samples - 1 if final else self._total_samples - 1 - half - if self._next_center > max_center: - return np.empty((0,), dtype=np.float32) - - centers = np.arange(self._next_center, max_center + 1, self._ratio, dtype=np.int64) - first_center = int(centers[0]) - last_center = int(centers[-1]) - window_start = first_center - half - window_end = last_center + half - segment = np.zeros((window_end - window_start + 1,), dtype=np.float32) - - copy_start = max(window_start, self._buffer_start, 0) - copy_end = min(window_end + 1, self._buffer_start + self._buffer.size) - if copy_end > copy_start: - dst_start = copy_start - window_start - src_start = copy_start - self._buffer_start - segment[dst_start : dst_start + copy_end - copy_start] = self._buffer[ - src_start : src_start + copy_end - copy_start - ] - - windows = np.lib.stride_tricks.sliding_window_view(segment, self._taps.size)[:: self._ratio] - output = windows @ self._taps[::-1] - self._next_center = last_center + self._ratio - - keep_from = max(0, self._next_center - half) - drop = min(max(keep_from - self._buffer_start, 0), self._buffer.size) - if drop: - self._buffer = self._buffer[drop:] - self._buffer_start += int(drop) - - return np.asarray(output, dtype=np.float32) + def reset(self) -> None: + self._history = np.zeros(self._history_samples, dtype=np.float32) + self._input_samples = 0 + self._output_samples = 0 + self._flushed = False class AudioMixin: diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index f204ae209..2446114df 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -578,6 +578,22 @@ class StreamingSpeechSessionConfig(BaseModel): "frames (existing behavior)." ), ) + seed: int | None = Field( + default=None, + ge=_INT64_MIN, + le=_INT64_MAX, + description="Random seed forwarded to /v1/audio/speech for this session.", + ) + split_granularity: Literal["none", "sentence", "clause"] = Field( + default="none", + description=( + "How incoming input.text is segmented before TTS. 'none' (default) " + "buffers until input.done and runs one request, matching the " + "long-form timbre-continuity path. 'sentence' emits a request at " + "each sentence boundary (Latin .!? plus CJK/Indic/Arabic marks). " + "'clause' also splits on commas/semicolons for lower TTFA." + ), + ) @model_validator(mode="after") def validate_streaming_constraints(self) -> "StreamingSpeechSessionConfig": diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 2f28b6199..0215e1ce0 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -175,6 +175,10 @@ class ImageGenerationRequest(BaseModel): default=None, description="Output image format: 'png', 'jpeg', or 'webp'. Defaults to 'png'.", ) + return_stage_metrics: bool | None = Field( + default=None, + description="Return stage metrics for benchmark clients.", + ) class ImageData(BaseModel): diff --git a/vllm_omni/entrypoints/openai/protocol/videos.py b/vllm_omni/entrypoints/openai/protocol/videos.py index fb4c395d7..cb8cc198d 100644 --- a/vllm_omni/entrypoints/openai/protocol/videos.py +++ b/vllm_omni/entrypoints/openai/protocol/videos.py @@ -290,6 +290,10 @@ class VideoGenerationRequest(BaseModel): default=None, description=("Optional model-specific parameters passed directly to the model's extra_args. "), ) + return_stage_metrics: bool | None = Field( + default=None, + description="Whether to include server-side stage metrics in async video metadata.", + ) def resolve_video_params( self, @@ -407,6 +411,13 @@ class VideoResponse(BaseModel): description="Filename of the saved output video files for this job.", ) inference_time_s: float | None = Field(default=None, description="End-to-end inference time in seconds.") + fps: float | None = Field(default=None, description="Resolved output video frames per second, if known.") + num_frames: int | None = Field(default=None, description="Resolved number of output video frames, if known.") + duration_s: float | None = Field(default=None, description="Resolved output video duration in seconds, if known.") + metrics: dict[str, Any] | None = Field( + default=None, + description="Optional profiler and stage metrics for benchmark clients.", + ) stage_durations: dict[str, float] = Field( default_factory=dict, description="Profiler stage durations reported by the diffusion pipeline.", diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py old mode 100755 new mode 100644 index e799fce81..bfa241a40 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -3288,7 +3288,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): output_compression: int = 100, size: str = "auto", raw_request: Request | None = None, - ) -> tuple[list[Image.Image], dict[str, Any], float, str | None] | ErrorResponse | AsyncIterator[str]: + ) -> ( + tuple[list[Image.Image], dict[str, Any], float, str | None, dict[str, Any] | None] + | ErrorResponse + | AsyncIterator[str] + ): """Generate diffusion images and return raw images plus generation stats.""" if request_id is None: request_id = f"chatcmpl-{uuid.uuid4().hex[:16]}" @@ -3375,6 +3379,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): images = getattr(result, "images", []) stage_durations = result.stage_durations peak_memory_mb = result.peak_memory_mb + response_metrics = getattr(result, "metrics", None) if return_stage_metrics else None cot_output = None req_out = result @@ -3407,7 +3412,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if isinstance(ar_text, str) and ar_text.strip(): cot_output = ar_text - return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output + return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output, response_metrics async def _stream_diffusion_image_chunks( self, diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 813d916fe..88d9f090d 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -30,11 +30,9 @@ from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, RequestResponseMetadata, ) -from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm.multimodal.media import MediaConnector from vllm.utils import random_uuid -from vllm.utils.async_utils import make_async from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler @@ -63,17 +61,6 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( tts_entry_stage_archs, ) from vllm_omni.entrypoints.utils import coerce_param_message_types -from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( - build_fish_text_only_prompt_ids, - estimate_fish_voice_clone_prompt_len_from_normalized, - normalize_fish_voice_clone_texts, -) -from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( - DEFAULT_PROMPT as MING_DEFAULT_PROMPT, -) -from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( - create_instruction as ming_create_instruction, -) from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -93,12 +80,6 @@ _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS" # ``tts_adapters.detect_tts_model_type``. Adding a TTS model must not require an # edit to this module. # -# 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}" - # Audex contract: zero-codec / invalid generations arrive as empty terminal # payloads and must fail the request, never serialize as a successful empty # WAV. Covers both the TTS ("audex") and TTA ("audex_tta") pipelines. @@ -109,12 +90,7 @@ _REF_AUDIO_METADATA_FETCH_ATTEMPTS = 3 _REMOTE_REF_AUDIO_SCHEMES = frozenset({"http", "https", "data"}) _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256 _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 256 * 1024 * 1024 -_HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES = 256 -_HIGGS_V3_REF_CODE_CACHE_MAX_BYTES = 64 * 1024 * 1024 -_QWEN3_TTS_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key" _TTS_MAX_INSTRUCTIONS_LENGTH = 500 -_TTS_MAX_NEW_TOKENS_MAX = 4096 -_MING_DEFAULT_PROMPT = MING_DEFAULT_PROMPT _DEFAULT_VOICE_NAME = "default" @@ -220,64 +196,6 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool: return False -def _conditioning_cache_salt(request, tts_params: dict | None = None) -> str: - """Stable hash of the real Stage 0 conditioning for the prefix cache. - - The talker's vLLM prompt is placeholder token ids; the real inputs are - rebuilt from text / ref_audio / ref_text into inputs_embeds. vLLM hashes - token ids (folded with cache_salt) for prefix caching, so without a salt - every request collides and a hit could reuse KV from a semantically - different input. Tying the salt to the conditioning keeps a hit safe: - identical conditioning may share the prefix, any difference never does. - - Raw request fields alone are not enough for uploaded voices: the request - only carries the voice *name* (ref_audio/ref_text/task_type are resolved - from stored voice data into ``tts_params``). Delete + re-upload under the - same name leaves every raw field identical, so the resolved conditioning - must also be folded in. ``voice_created_at`` bumps on every (re-)upload, - which uniquely identifies the resolved reference artifact together with - the voice name; the decoded ref_audio array itself need not be hashed. - - Local ``ref_audio`` locators are the same class of gap: the raw request - hashes the path/URI string, so an on-disk rewrite is invisible to the - salt. Callers stash the content-aware key from ``_resolve_ref_audio`` as - ``ref_audio_cache_key`` (and ``ref_audio_2_cache_key`` / emotion keys - when those clips exist) so a file edit cannot reuse KV from the previous - reference. - """ - h = hashlib.sha256() - for part in ( - request.input, - request.task_type, - request.language, - request.voice, - request.ref_text, - request.ref_audio, - request.instructions, - request.x_vector_only_mode, - request.speaker_embedding, - ): - h.update(b"\x00") - if part is not None: - h.update(repr(part).encode("utf-8")) - # Fold resolved conditioning that is auto-derived for uploaded voices and - # absent from the raw request (or that a locator string cannot see). - for key in ( - "voice_created_at", - "task_type", - "speaker", - "ref_text", - "x_vector_only_mode", - "ref_audio_cache_key", - "ref_audio_2_cache_key", - ): - h.update(b"\x00") - value = tts_params.get(key) if tts_params is not None else None - if value is not None: - h.update(repr(value).encode("utf-8")) - return h.hexdigest()[:32] - - class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): _diffusion_mode: bool = False _media_connector: MediaConnector | None = None @@ -306,9 +224,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # 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]] = {} self._speaker_cache = get_speaker_cache() self._last_upload_ts = 0 self._upload_lock = asyncio.Lock() @@ -438,16 +353,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Find and cache the TTS stage (if any) during initialization self._tts_stage = self._find_tts_stage() self._is_tts = self._tts_stage is not None - self._fish_speech_tokenizer = None - self._covo_audio_tokenizer = None - # Cached per process: the CosyVoice3 Qwen tokenizer + resolved model - # path used for dynamic-token sizing. Without this, every request - # re-ran snapshot_download + reloaded the tokenizer (~100 ms on the - # TTFP critical path) in the CosyVoice3 adapter's sampling override. - self._cosyvoice3_tokenizer = None # Determine TTS model type or None self._tts_model_type = self._detect_tts_model_type() + + # Shared executor for blocking adapter preprocessing. It must exist + # before adapter construction so adapters can create their make_async + # wrappers during their own lifecycle initialization. + self._tts_executor = ThreadPoolExecutor(max_workers=1) # Resolve the per-model serving adapter (RFC #4327), keyed on the # detected model-type. Every dedicated TTS model has an adapter; the # adapter owns request validation, prompt/param building, capability @@ -467,37 +380,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): available_speakers = self._get_available_speakers() logger.info("Loaded %d supported speakers: %s", len(available_speakers), sorted(available_speakers)) - # Sub-variant inside the full MOSS-TTS family. We collapse all five - # variants onto the same _tts_model_type="moss_tts" because they share - # the same stage layout, but request validation + param building - # differ per HF repo (voice-clone vs dialogue vs ambient-sound vs - # instruction vs streaming voice-clone). - self._moss_variant = self._detect_moss_variant() if self._tts_model_type == "moss_tts" else None - - # GLM-TTS lazy-cached resources (populated on first GLM-TTS request) - self._glm_tts_text_tokenizer: object | None = None - self._glm_tts_text_frontend: object | None = None - # Cache TTS configuration values (computed once, reused per request) self._max_instructions_length = self._compute_max_instructions_length() self._tts_tokenizer = None - self._voxcpm2_tokenizer = None - self._audex_tokenizer = None - self._audex_tta_rvq = None - self._voxcpm2_split_map: dict[int, list[int]] = {} # Batch configuration self._batch_max_items: int = getattr(self.engine_client, "tts_batch_max_items", 32) - # Shared thread pool executor for blocking TTS preprocessing - # operations. max_workers=1 serializes tokenizer access to avoid - # Rust RefCell "Already borrowed" errors from concurrent use. - self._tts_executor = ThreadPoolExecutor(max_workers=1) - self._build_voxtral_prompt_async = make_async(self._build_voxtral_prompt, executor=self._tts_executor) - self._build_fish_speech_prompt_async = make_async(self._build_fish_speech_prompt, executor=self._tts_executor) - self._estimate_prompt_len_async = make_async(self._estimate_prompt_len, executor=self._tts_executor) - def _get_tts_adapter(self): """Return the per-model serving adapter for the current ``_tts_model_type``. @@ -622,73 +512,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """Get all voice names accepted by the API, including the placeholder default.""" return self._get_available_speakers() | {_DEFAULT_VOICE_NAME} - def _estimate_ref_code_len(self, ref_audio: object) -> int | None: - """Estimate ref_code length from ref_audio waveform without running the codec. - - The codec produces one frame per (output_sample_rate / encode_downsample_rate) - audio samples, so ref_code_len = ceil(duration_seconds * codec_frame_rate). - """ - codec_frame_rate = self._adapter.capabilities.codec_frame_rate - if codec_frame_rate is None: - return None - try: - # ref_audio comes from tts_params as [[wav_array, sr]] or similar nested structure - item = ref_audio - while isinstance(item, list) and item: - if len(item) == 2 and isinstance(item[1], (int, float)): - break - item = item[0] - if isinstance(item, list) and len(item) == 2: - wav, sr = item - elif isinstance(item, tuple) and len(item) == 2: - wav, sr = item - else: - return None - sr = int(sr) - if hasattr(wav, "__len__"): - n_samples = len(wav) - elif hasattr(wav, "shape"): - n_samples = wav.shape[-1] if wav.ndim > 1 else wav.shape[0] - else: - return None - if sr <= 0 or n_samples <= 0: - return None - duration = n_samples / sr - return math.ceil(duration * codec_frame_rate) - except Exception: - return None - - def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int: - """Estimate prompt length so the placeholder matches model-side embeddings.""" - try: - from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( - Qwen3TTSPromptEmbedsBuilder, - ) - - if self._tts_tokenizer is None: - from transformers import AutoTokenizer - - model_name = self.engine_client.model_config.model - self._tts_tokenizer = AutoTokenizer.from_pretrained( - model_name, - trust_remote_code=True, - padding_side="left", - ) - hf_config = self.engine_client.model_config.hf_config - talker_config = hf_config.talker_config - task_type = (tts_params.get("task_type") or ["CustomVoice"])[0] - return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( - additional_information=tts_params, - task_type=task_type, - tokenize_prompt=lambda t: self._tts_tokenizer(t, padding=False)["input_ids"], - codec_language_id=getattr(talker_config, "codec_language_id", None), - spk_is_dialect=getattr(talker_config, "spk_is_dialect", None), - estimate_ref_code_len=self._estimate_ref_code_len, - ) - except Exception as e: - 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. @@ -784,41 +607,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER: str(usage.input_token_details.audio_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 ( - DAC_HOP_LENGTH, - DAC_SAMPLE_RATE, - ) - - if not isinstance(ref_audio, (list, tuple)) or len(ref_audio) != 2: - return None - wav, sr = ref_audio - sr = int(sr) - n_samples = len(wav) - if sr <= 0 or n_samples <= 0: - return None - resampled_len = max(1, math.ceil(n_samples * DAC_SAMPLE_RATE / sr)) - return max(1, math.ceil(resampled_len / DAC_HOP_LENGTH)) - - def _estimate_fish_prompt_len(self, text: str, ref_text: str, ref_audio: object) -> int: - """Estimate Fish Speech clone prompt length without encoding reference audio.""" - try: - from transformers import AutoTokenizer - - if self._fish_speech_tokenizer is None: - model_name = self.engine_client.model_config.model - self._fish_speech_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) - - tokenizer = self._fish_speech_tokenizer - semantic_len = self._estimate_fish_ref_code_len(ref_audio) - if semantic_len is None: - raise ValueError("Failed to estimate Fish Speech semantic token length") - return estimate_fish_voice_clone_prompt_len_from_normalized(tokenizer, text, ref_text, semantic_len) - except Exception as e: - logger.warning("Failed to estimate Fish Speech prompt length, using fallback 2048: %s", e) - return 2048 - def _voice_created_at(self, voice_lower: str) -> int: """Return the upload timestamp of an uploaded voice, or 0 for built-ins. @@ -828,41 +616,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): info = self.uploaded_speakers.get(voice_lower) return int(info.get("created_at", 0)) if info else 0 - async def _build_voxcpm2_prompt( - self, - request: OpenAICreateSpeechRequest, - *, - uploaded_ref: tuple[np.ndarray, int] | None = None, - ) -> dict[str, Any]: - """Build prefill prompt for VoxCPM2 TTS (`prompt_token_ids` padded to full prefill length). - - ``uploaded_ref`` supplies the audio for uploaded voices (no explicit - ``ref_audio`` in the request) so prefill length includes it. - """ - from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import build_voxcpm2_prompt - - self._voxcpm2_encode("") # lazy-init tokenizer + split_map - ref_audio = None - ref_sr = None - voice_profile = None - if request.ref_audio is not None: - ref_audio, ref_sr, _ = await self._resolve_ref_audio(request.ref_audio) - elif uploaded_ref is not None: - wav_np, ref_sr = uploaded_ref - ref_audio = wav_np.tolist() - elif request.voice is not None: - voice_profile = self._adapter.capabilities.precomputed_speakers.get(request.voice.lower()) - return build_voxcpm2_prompt( - hf_config=self.engine_client.model_config.hf_config, - tokenizer=self._voxcpm2_tokenizer, - split_map=self._voxcpm2_split_map, - text=request.input, - ref_audio=ref_audio, - ref_sr=ref_sr, - ref_text=request.ref_text, - voice_profile=voice_profile, - ) - def _load_uploaded_audio(self, voice_name: str) -> tuple[np.ndarray, int] | None: """Load decoded audio samples + sample rate from an uploaded voice's safetensors.""" voice_name_lower = voice_name.lower() @@ -959,6 +712,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if self._tts_model_type in ( "cosyvoice3", "fish_tts", + "audio8_tts", "omnivoice", "moss_tts_nano", "glm_tts", @@ -968,6 +722,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): label = { "cosyvoice3": "CosyVoice3", "fish_tts": "Fish Speech", + "audio8_tts": "Audio8 TTS", "omnivoice": "OmniVoice", "moss_tts_nano": "MOSS-TTS-Nano", "higgs_audio_v2": "Higgs-Audio V2", @@ -1330,25 +1085,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "sample_rate is not supported by the current TTS model" return None - def _voxcpm2_encode(self, text: str) -> list[int]: - """Tokenize text for VoxCPM2, splitting multichar Chinese tokens.""" - from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( - build_cjk_split_map, - split_multichar_chinese, - ) - - if self._voxcpm2_tokenizer is None: - from transformers import AutoTokenizer - - model_name = self.engine_client.model_config.model - tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) - self._voxcpm2_split_map = build_cjk_split_map(tokenizer) - self._voxcpm2_tokenizer = tokenizer - logger.info("VoxCPM2 serving: built multichar split map (%d entries)", len(self._voxcpm2_split_map)) - - ids = self._voxcpm2_tokenizer.encode(text, add_special_tokens=True) - return split_multichar_chinese(ids, self._voxcpm2_split_map) - def _validate_ref_audio_format(self, ref_audio: str) -> str | None: """Validate ref_audio is a supported URI format. Returns error or None.""" if not isinstance(ref_audio, str): @@ -1358,445 +1094,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" 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). - - Detection key is the HF repo path / model_name; matches - ``_try_resolve_omni_model_type`` in entrypoints/utils.py so users get - consistent behaviour no matter how they launched the server (--model - OpenMOSS-Team/MOSS-TTSD-v1.0 vs --deploy-config moss_ttsd.yaml). - """ - try: - name = (self.engine_client.model_config.model or "").lower().replace("-", "").replace("_", "") - except Exception: - name = "" - if "realtime" in name: - return "realtime" - if "local" in name: - return "local" - if "ttsd" in name: - return "ttsd" - if "soundeffect" in name: - return "sound_effect" - if "voicegenerator" in name: - return "voice_generator" - return "tts" - - def _get_moss_processor(self): - """Lazily load the upstream MOSS-TTS processor once per server. - - Cached on ``self._moss_processor_cache``. The processor owns its own - audio_tokenizer (~1.6 B params); we keep it on CPU so it doesn't - compete with the talker (~8 GiB) and codec (~7 GiB) for our 96 GiB - GPU — per-request ref-audio encoding is fast enough on CPU. - """ - cached = getattr(self, "_moss_processor_cache", None) - if cached is not None: - return cached - from transformers import AutoProcessor - - model_id = self.engine_client.model_config.model - proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) - if hasattr(proc, "audio_tokenizer"): - proc.audio_tokenizer = proc.audio_tokenizer.to("cpu").eval() - self._moss_processor_cache = proc - return proc - - def _get_moss_ref_encoder(self): - """Lazily build the per-server MOSS reference-audio encoder once.""" - cached = getattr(self, "_moss_ref_encoder", None) - if cached is not None: - return cached - from vllm_omni.model_executor.models.moss_tts.reference_encoder import build_reference_encoder - - # The variant's encode geometry (n_vq, working sample rate) is derived - # inside the model package; this layer only supplies the processor and - # the process-wide speaker cache. - encoder = build_reference_encoder( - self._get_moss_processor(), - variant=self._moss_variant, - speaker_cache=self._speaker_cache, - ) - self._moss_ref_encoder = encoder - return encoder - - async def _build_moss_tts_params( - self, - request: OpenAICreateSpeechRequest, - *, - has_inline_ref_audio: bool = False, - ) -> dict[str, Any]: - """Build the talker prompt + ``additional_information`` payload for any - MOSS-TTS-family request (nano + 5 full variants). - - For the legacy ``moss_tts_nano`` model_type, keeps the original nano - contract (``{text, mode=voice_clone, prompt_audio_array}``); the - caller still uses a ``[1]`` placeholder prompt. - - For the full MOSS-TTS family (``MossTTSDelay*`` / ``MossTTSRealtime*``) - we **call the upstream processor server-side** to produce the unified - ``(text_ids, audio_codes)`` shape the talker actually consumes — same - flow as ``examples/.../moss_tts/end2end.py:_build_unified_codes``. - Returns ``{prompt_token_ids: list[int], codes.ref: torch.LongTensor, - max_new_frames, ...}``. The caller treats ``prompt_token_ids`` as the - prompt and forwards the rest as ``additional_information``. - """ - import torch # local to avoid pulling torch at module import time - - v = self._moss_variant - - # ---- Legacy nano path (unchanged) ---- - if v is None: # moss_tts_nano - params: dict[str, Any] = { - "text": [request.input or ""], - "mode": ["voice_clone"], - } - if request.max_new_tokens is not None: - params["max_new_frames"] = [request.max_new_tokens] - wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio) - params["prompt_audio_array"] = [[wav_list, sr]] - params["ref_audio_cache_key"] = cache_key - return params - - # ---- MOSS-TTS-Realtime: keep the old prompt_audio_array path ---- - # ``AutoProcessor.from_pretrained`` doesn't auto-discover - # ``MossTTSRealtimeProcessor`` (no ``processor_config.json`` in the - # snapshot), and Realtime's prompt format diverges from MossTTSDelay - # (16-channel grid, separate per-step text feed). The - # ``prompt_audio_array`` shape lines up well enough with what the - # talker reads for short prompts; full Realtime support needs a - # separate processor.from_module path which we don't wire here. - if v == "realtime": - params: dict[str, Any] = { - "text": [request.input or ""], - "mode": ["voice_clone"], - } - if request.max_new_tokens is not None: - params["max_new_frames"] = [request.max_new_tokens] - wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio) - params["prompt_audio_array"] = [[wav_list, sr]] - params["ref_audio_cache_key"] = cache_key - return params - - # ---- 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() - - user_kwargs: dict[str, Any] = {"text": request.input or ""} - resolve_keys: dict[int, str] = {} - if v in ("tts", "local", "ttsd"): - user_kwargs["reference"], resolve_keys = await self._encode_moss_references( - request, - has_inline_ref_audio=has_inline_ref_audio, - two_speaker=(v == "ttsd"), - ) - elif v == "sound_effect": - user_kwargs["text"] = request.input or "" # may be empty - user_kwargs["ambient_sound"] = request.ambient_sound or "" - if request.duration_seconds is not None: - user_kwargs["tokens"] = max(1, int(float(request.duration_seconds) * 12.5)) - elif request.max_new_tokens is not None: - user_kwargs["tokens"] = int(request.max_new_tokens) - elif v == "voice_generator": - user_kwargs["instruction"] = request.instructions or "" - - # Optional language tag for the spoken-text variants. MOSS-TTS-v1.5's - # headline improvement is multilingual synthesis when the language is - # given (build_user_message(..., language=...)); 1.0 ignores it - # gracefully. Sound-effect output is non-verbal, so skip it there. - if v in ("tts", "local", "ttsd", "voice_generator") and getattr(request, "language", None): - user_kwargs["language"] = request.language - - # Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special - # tokens and cols 1..n_vq are the delay-pattern audio code grid (mostly - # audio_pad_code outside the reference block). - user_msg = proc.build_user_message(**user_kwargs) - batch = proc(conversations=[[user_msg]], mode="generation") - unified = batch["input_ids"][0] # torch.LongTensor (L, 1+n_vq) - text_ids: list[int] = unified[:, 0].tolist() - audio_codes: torch.Tensor = unified[:, 1:].contiguous().to(torch.int64) - - params: dict[str, Any] = { - "prompt_token_ids": text_ids, - "codes": {"ref": audio_codes}, - } - if request.max_new_tokens is not None: - params["max_new_frames"] = [request.max_new_tokens] - if 0 in resolve_keys: - params["ref_audio_cache_key"] = resolve_keys[0] - if 1 in resolve_keys: - params["ref_audio_2_cache_key"] = resolve_keys[1] - return params - - async def _encode_moss_references( - self, - request: OpenAICreateSpeechRequest, - *, - has_inline_ref_audio: bool, - two_speaker: bool, - ) -> tuple[list, dict[int, str]]: - """Encode the request's reference clip(s) into MOSS RVQ code tensors. - - Reference encoding + speaker caching lives in the model package - (moss_tts.reference_encoder), mirroring Fish Speech / CosyVoice3 / - Qwen3-TTS which keep reference handling with the model rather than in - this shared serving file. This method only resolves the serving-side - inputs the model helper needs: the generic audio resolver/artifact-key - lookups and whether the request may use the named-voice cache. - - Returns ``(codes_per_speaker, resolve_keys)`` where ``resolve_keys`` - maps the reference slot (0 = ref_audio, 1 = ref_audio_2) to its - content-aware resolve key, for salting the KV prefix cache. - """ - from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_request_references - - # Named-voice caching is only valid for uploaded speakers without an - # inline ref_audio: ``request.voice`` plus a file/URL would otherwise - # key on (name, created_at=0) and skip the content-aware resolve. - raw_voice = getattr(request, "voice", None) - raw_voice = raw_voice.strip() if isinstance(raw_voice, str) else "" - voice_lower = raw_voice.lower() - use_named_voice = bool(voice_lower) and voice_lower in self.uploaded_speakers and not has_inline_ref_audio - voice = voice_lower if use_named_voice else "" - voice_created = self._voice_created_at(voice) if voice else 0 - - return await encode_request_references( - self._get_moss_ref_encoder(), - request.ref_audio, - request.ref_audio_2 if two_speaker else None, - resolve_ref_audio=self._resolve_ref_audio, - get_artifact_key=self._get_resolved_ref_audio_artifact_key, - voice_name=voice or None, - voice_created_at=voice_created, - ) - - async def _build_higgs_audio_v2_params(self, request: OpenAICreateSpeechRequest): - """Build prompt_token_ids for higgs_audio_v2 via the upstream processor. - - Plain-text path: runs ``build_plain_text_prompt`` and returns the - token-only prompt. Voice-clone path (``ref_audio`` + ``ref_text``): - resolves the reference clip via ``_resolve_ref_audio``, runs - ``build_voice_clone_prompt`` (which encodes the clip through HF's - ``HiggsAudioV2TokenizerModel`` loaded from the k2-fsa/OmniVoice - ``audio_tokenizer/`` subdirectory), and attaches the encoded - ``audio_input_ids`` + ``audio_input_ids_mask`` tensors via - ``additional_information`` so the talker substitutes them at the - prompt-side audio placeholders. - """ - from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( - build_plain_text_prompt, - build_voice_clone_prompt, - input_ids_to_python_list, - ) - - processor = await self._resolve_higgs_audio_v2_processor() - - if request.ref_audio is None: - inputs = await asyncio.to_thread(build_plain_text_prompt, processor, request.input) - prompt_token_ids = input_ids_to_python_list(inputs) - return tokens_input(prompt_token_ids=prompt_token_ids) - - wav_list, sr, _ = await self._resolve_ref_audio(request.ref_audio) - wav = np.asarray(wav_list, dtype=np.float32) - out = await asyncio.to_thread( - build_voice_clone_prompt, - processor, - request.input, - wav, - int(sr), - request.ref_text or "", - ) - prompt = tokens_input(prompt_token_ids=out["prompt_token_ids"]) - # Pass tensors at the top level of additional_information (NOT list- - # wrapped). ``vllm_omni.data_entry_keys.serialize_payload`` routes - # bare ``torch.Tensor`` values through ``_serialize_tensor``; a list - # containing tensors would fall into the ``list_data`` field which - # msgspec cannot serialize and the tensors would be dropped over the - # process boundary (silent voice-clone failure). - prompt["additional_information"] = { - "audio_input_ids": out["audio_input_ids"], - "audio_input_ids_mask": out["audio_input_ids_mask"], - } - return prompt - - async def _resolve_higgs_audio_v2_processor(self): - """Lazy-load the AutoProcessor for higgs_audio_v2 (once per serving instance).""" - cached = getattr(self, "_higgs_audio_v2_processor", None) - if cached is not None: - return cached - - from transformers import AutoProcessor - - model_path = None - for stage in self.engine_client.stage_configs: - model_path = getattr(getattr(stage, "engine_args", None), "model", None) - if model_path: - break - if model_path is None: - # Fallback: the orchestrator stores the served model id on the engine - # itself (set by AsyncOmniEngine.__init__). Stage-level engine_args - # may not surface ``model`` when the deploy yaml doesn't set it per - # stage (the CLI-passed model id is the single source of truth). - model_path = getattr(self.engine_client, "model", None) - if model_path is None: - raise RuntimeError("higgs_audio_v2 serving could not resolve the model path from the engine stage configs") - processor = AutoProcessor.from_pretrained(model_path) - self._higgs_audio_v2_processor = processor - return processor - - # ---- higgs-audio v3 ---- - - async def _build_higgs_audio_v3_params(self, request: OpenAICreateSpeechRequest): - """Build prompt_token_ids for higgs_audio_v3. - - Plain-text path: builds ``[tts, text, tokens, audio]``. - Voice-clone path: encodes reference audio, applies delay pattern, and - records reference-audio prompt positions while submitting valid token IDs. - """ - adapter = await self._resolve_higgs_audio_v3_adapter() - - if request.ref_audio is None: - prompt_ids = adapter.build_prompt(request.input) - return tokens_input(prompt_token_ids=prompt_ids) - - # Voice clone - from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( - apply_delay_pattern, - encode_reference_audio, - ) - - wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio) - artifact_key = self._get_resolved_ref_audio_artifact_key(cache_key) - wav = np.asarray(wav_list, dtype=np.float32) - ref_codes_delayed, cache_hit, inflight_wait = await self._resolve_higgs_audio_v3_ref_codes( - artifact_key, - wav, - int(sr), - encode_reference_audio, - apply_delay_pattern, - ) - del cache_hit, inflight_wait - - prompt_ids = adapter.build_prompt( - request.input, - num_ref_tokens=int(ref_codes_delayed.shape[0]), - reference_text=request.ref_text or None, - ) - prompt_ids, audio_placeholder_positions = adapter.prepare_prompt_for_engine(prompt_ids) - prompt = tokens_input(prompt_token_ids=prompt_ids) - import torch - - prompt["additional_information"] = { - "audio_input_ids": ref_codes_delayed.to(torch.long), - "audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool), - "audio_placeholder_positions": audio_placeholder_positions, - "ref_audio_cache_key": cache_key, - } - # Placeholder prompt ids do not change when a same-path file is - # rewritten at the same duration; fold the content-aware resolve key - # into cache_salt so prefix caching cannot reuse the previous KV. - prompt["cache_salt"] = _conditioning_cache_salt(request, prompt["additional_information"]) - return prompt - - async def _resolve_higgs_audio_v3_ref_codes( - self, - artifact_key: str | None, - wav: np.ndarray, - sr: int, - encode_reference_audio, - apply_delay_pattern, - ) -> tuple[torch.Tensor, bool, bool]: - ref_codes_delayed = self._get_higgs_audio_v3_ref_codes(artifact_key) - if ref_codes_delayed is not None: - return ref_codes_delayed, True, False - if not artifact_key: - ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr) - return apply_delay_pattern(ref_codes_raw), False, False - - task = self._higgs_audio_v3_ref_code_inflight.get(artifact_key) - if task is not None: - return (await task).clone(), False, True - - async def _encode_and_cache() -> torch.Tensor: - ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr) - delayed = apply_delay_pattern(ref_codes_raw) - self._put_higgs_audio_v3_ref_codes(artifact_key, delayed) - cached = self._get_higgs_audio_v3_ref_codes(artifact_key) - return cached if cached is not None else delayed.detach().to("cpu", dtype=torch.long).contiguous() - - task = asyncio.create_task(_encode_and_cache()) - self._higgs_audio_v3_ref_code_inflight[artifact_key] = task - try: - return (await task).clone(), False, False - finally: - if self._higgs_audio_v3_ref_code_inflight.get(artifact_key) is task: - self._higgs_audio_v3_ref_code_inflight.pop(artifact_key, None) - - def _get_higgs_audio_v3_ref_codes(self, artifact_key: str | None) -> torch.Tensor | None: - if not artifact_key: - return None - cached = self._higgs_audio_v3_ref_code_cache.get(artifact_key) - if cached is None: - return None - self._higgs_audio_v3_ref_code_cache.move_to_end(artifact_key) - return cached[0].clone() - - def _put_higgs_audio_v3_ref_codes(self, artifact_key: str, codes: torch.Tensor) -> None: - if _HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES <= 0 or _HIGGS_V3_REF_CODE_CACHE_MAX_BYTES <= 0 or not artifact_key: - return - cached_codes = codes.detach().to("cpu", dtype=torch.long).contiguous() - size = int(cached_codes.numel() * cached_codes.element_size()) - if size > _HIGGS_V3_REF_CODE_CACHE_MAX_BYTES: - return - previous = self._higgs_audio_v3_ref_code_cache.pop(artifact_key, None) - if previous is not None: - self._higgs_audio_v3_ref_code_cache_bytes -= previous[1] - self._higgs_audio_v3_ref_code_cache[artifact_key] = (cached_codes, size) - self._higgs_audio_v3_ref_code_cache_bytes += size - while len(self._higgs_audio_v3_ref_code_cache) > _HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES: - _, (_, old_size) = self._higgs_audio_v3_ref_code_cache.popitem(last=False) - self._higgs_audio_v3_ref_code_cache_bytes -= old_size - while self._higgs_audio_v3_ref_code_cache_bytes > _HIGGS_V3_REF_CODE_CACHE_MAX_BYTES: - _, (_, old_size) = self._higgs_audio_v3_ref_code_cache.popitem(last=False) - self._higgs_audio_v3_ref_code_cache_bytes -= old_size - - async def _resolve_higgs_audio_v3_adapter(self): - """Lazy-load the tokenizer adapter for higgs_audio_v3.""" - cached = getattr(self, "_higgs_audio_v3_adapter", None) - if cached is not None: - return cached - - from transformers import AutoTokenizer - - from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( - HiggsAudioV3TokenizerAdapter, - ) - - model_path = None - for stage in self.engine_client.stage_configs: - model_path = getattr(getattr(stage, "engine_args", None), "model", None) - if model_path: - break - if model_path is None: - model_path = getattr(self.engine_client, "model", None) - if model_path is None: - raise RuntimeError("higgs_audio_v3 serving could not resolve model path") - tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) - adapter = HiggsAudioV3TokenizerAdapter(tokenizer) - self._higgs_audio_v3_adapter = adapter - return adapter - @staticmethod def _local_ref_audio_stat_path(ref_audio_str: str) -> str | None: """Filesystem path to stat for a local locator, else ``None``. @@ -2043,14 +1340,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): 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 - 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, x_vector_only: bool = False ) -> None: @@ -2075,145 +1364,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): resolved.append((wav_list, sr)) return resolved - # ---- Ming TTS helpers ---- - - def _coerce_ming_prompt_waveform(self, wav_samples, sample_rate): - from torchaudio.functional import resample as resample_audio - - from vllm_omni.model_executor.models.ming_tts.config_ming_tts import SAMPLE_RATE - - waveform = torch.as_tensor(wav_samples, dtype=torch.float32).reshape(1, -1) - if int(sample_rate) != SAMPLE_RATE: - waveform = resample_audio(waveform, int(sample_rate), SAMPLE_RATE) - return waveform - - def _build_ming_reference_waveforms( - self, - ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None, - ) -> tuple[list[torch.Tensor], list[int]]: - if ref_audio_data is None: - return [], [] - items = ref_audio_data if isinstance(ref_audio_data, list) else [ref_audio_data] - waveforms = [torch.as_tensor(samples, dtype=torch.float32).reshape(1, -1) for samples, _ in items] - sample_rates = [int(sample_rate) for _, sample_rate in items] - if any(sample_rate <= 0 for sample_rate in sample_rates): - raise ValueError(f"Ming reference audio sample rates must be positive, got {sample_rates}") - return waveforms, sample_rates - - def _parse_ming_instruction_fields( - self, - request, - *, - include_voice=False, - plain_text_passthrough=False, - ): - instruction_text = request.instructions.strip() if isinstance(request.instructions, str) else None - instruction_dict: dict[str, Any] = {} - - voice_lower = request.voice.lower() if isinstance(request.voice, str) else None - if include_voice and request.voice and not (voice_lower and voice_lower in self.uploaded_speakers): - instruction_dict["IP"] = request.voice - - if instruction_text: - try: - parsed = json.loads(instruction_text) - except json.JSONDecodeError: - parsed = None - if isinstance(parsed, dict): - instruction_dict.update(parsed) - elif instruction_dict or not plain_text_passthrough: - instruction_dict["风格"] = instruction_text - else: - return instruction_text - - return instruction_dict or None - - def _parse_ming_instruction(self, request: OpenAICreateSpeechRequest) -> Any: - """Build a Ming instruction payload from OpenAI speech fields.""" - return self._parse_ming_instruction_fields( - request, - include_voice=True, - plain_text_passthrough=True, - ) - - def _build_ming_dense_prompt( - self, - request: OpenAICreateSpeechRequest, - *, - ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None = None, - voice_name: str | None = None, - voice_created_at: int = 0, - ) -> dict[str, Any]: - """Build a Ming dense prompt directly from the OpenAI speech request.""" - from transformers import AutoTokenizer - - from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( - KEY_MAX_DECODE_STEPS, - KEY_SPEAKER_SAMPLE_RATES, - KEY_SPEAKER_WAVEFORM, - KEY_SPEAKER_WAVEFORM_LENGTHS, - ) - from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt - - if self._tts_tokenizer is None: - model_name = self.engine_client.model_config.model - 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 - reference_waveforms, reference_sample_rates = self._build_ming_reference_waveforms(ref_audio_data) - reference_waveform = torch.cat(reference_waveforms, dim=-1) if reference_waveforms else None - prompt_waveforms = ( - [ - self._coerce_ming_prompt_waveform(waveform, sample_rate) - for waveform, sample_rate in zip(reference_waveforms, reference_sample_rates, strict=True) - ] - if ref_text is not None - else [] - ) - prompt_waveform = torch.cat(prompt_waveforms, dim=-1) if prompt_waveforms else None - speaker_embedding = request.speaker_embedding - pending_speaker_extraction = speaker_embedding is None and reference_waveform is not None - use_zero_spk_emb = not pending_speaker_extraction and prompt_waveform is None and speaker_embedding is None - - runtime_controls = {} - if request.max_new_tokens is not None: - runtime_controls[KEY_MAX_DECODE_STEPS] = request.max_new_tokens - - prompt_dict = build_ming_dense_prompt( - self._tts_tokenizer, - # bgm / music-prompt mode not supported online; - # requires prompt_mode API extension (deferred). - prompt=_MING_DEFAULT_PROMPT, - text=request.input, - runtime_controls=runtime_controls or None, - instruction=self._parse_ming_instruction(request), - prompt_text=ref_text, - prompt_waveform=prompt_waveform, - speaker_embedding=speaker_embedding, - speaker_count=len(reference_waveforms) if pending_speaker_extraction else None, - use_zero_spk_emb=use_zero_spk_emb, - ) - additional_information = prompt_dict["additional_information"] - if pending_speaker_extraction: - additional_information[KEY_SPEAKER_WAVEFORM] = reference_waveform - additional_information[KEY_SPEAKER_WAVEFORM_LENGTHS] = torch.tensor( - [int(waveform.shape[-1]) for waveform in reference_waveforms], - dtype=torch.int32, - ) - additional_information[KEY_SPEAKER_SAMPLE_RATES] = torch.tensor( - reference_sample_rates, - dtype=torch.int32, - ) - if voice_name: - additional_information["voice_name"] = voice_name - additional_information["voice_created_at"] = int(voice_created_at) - prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"]) - prompt["prompt"] = prompt_dict["prompt"] - prompt["text"] = prompt_dict["text"] - prompt["additional_information"] = prompt_dict["additional_information"] - return prompt - async def _generate_audio_chunks( self, generator, @@ -2542,534 +1692,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None) return mm, key - def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: - """Build TTS parameters from request. - - Processes each parameter if present, skips if not. - Values are wrapped in lists as required by the model. - """ - params: dict[str, Any] = {} - - # Text content (always required) - params["text"] = [request.input] - - # Task type - if request.task_type is not None: - params["task_type"] = [request.task_type] - else: - params["task_type"] = ["CustomVoice"] - - # Language - if request.language is not None: - params["language"] = [request.language] - else: - params["language"] = ["Auto"] - - # Speaker (voice) - if request.voice is not None: - voice_lower = request.voice.lower() - precomputed_speakers = self._adapter.capabilities.precomputed_speakers if self._adapter is not None else {} - params["speaker"] = [request.voice] - params["voice_created_at"] = [self._voice_created_at(voice_lower)] - - # Uploaded voices use task_type="Base" (CustomVoice requires built-in spk_id). - # If ref_text was provided at upload time, use in-context cloning; otherwise x_vector only. - if voice_lower in self.uploaded_speakers and request.ref_audio is None: - speaker_info = self.uploaded_speakers[voice_lower] - - # Check if this voice was uploaded with a pre-computed embedding. - # Populate request.speaker_embedding so the existing code path - # (below) handles voice_clone_prompt and x_vector_only_mode. - embedding = self._get_uploaded_speaker_embedding(request.voice) - if embedding is not None: - request.speaker_embedding = embedding - params["speaker"] = [voice_lower] - params["task_type"] = ["Base"] - logger.info("Auto-set speaker_embedding for uploaded voice: %s", request.voice) - else: - audio_data = self._get_uploaded_audio_data(request.voice) - if not audio_data: - raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing or corrupted") - stored_ref_text = speaker_info.get("ref_text") - params["speaker"] = [voice_lower] - params["ref_audio"] = [audio_data] - params["task_type"] = ["Base"] - if stored_ref_text: - params["ref_text"] = [stored_ref_text] - params["x_vector_only_mode"] = [False] - else: - params["x_vector_only_mode"] = [True] - logger.info( - "Auto-set ref_audio for uploaded voice: %s (icl=%s)", request.voice, bool(stored_ref_text) - ) - elif voice_lower in precomputed_speakers and request.ref_audio is None: - profile = precomputed_speakers[voice_lower] - mode = str(profile.get("mode") or "xvec").lower() - params["speaker"] = [voice_lower] - params["task_type"] = ["Base"] - params["x_vector_only_mode"] = [mode != "icl"] - ref_text = request.ref_text or profile.get("ref_text") - if isinstance(ref_text, str) and ref_text.strip(): - params["ref_text"] = [ref_text] - ref_code_length = profile.get("ref_code_length") - if mode == "icl" and ref_code_length: - params["ref_code_length"] = [int(ref_code_length)] - logger.info("Using precomputed Qwen3-TTS custom voice profile: %s (mode=%s)", voice_lower, mode) - - elif params["task_type"][0] == "CustomVoice": - params["speaker"] = ["Vivian"] # Default for CustomVoice - - # Instructions for style/emotion control - if request.instructions is not None: - params["instruct"] = [request.instructions] - else: - params["instruct"] = [""] - - # Voice clone: ref_audio resolved in create_speech(), not here. - if request.ref_text is not None: - params["ref_text"] = [request.ref_text] - if request.speaker_embedding is not None: - # Store as plain float list (not tensor) so it survives msgspec - # serialization through the EngineCore IPC boundary. The talker's - # _build_prompt_embeds converts it back to a tensor on the GPU. - params["voice_clone_prompt"] = [ - { - "ref_spk_embedding": list(request.speaker_embedding), - } - ] - # speaker_embedding implies x_vector_only_mode - params["x_vector_only_mode"] = [True] - elif request.x_vector_only_mode is not None: - params["x_vector_only_mode"] = [request.x_vector_only_mode] - - # Generation parameters - if request.max_new_tokens is not None: - params["max_new_tokens"] = [request.max_new_tokens] - else: - params["max_new_tokens"] = [2048] - - if request.initial_codec_chunk_frames is not None: - params["initial_codec_chunk_frames"] = [request.initial_codec_chunk_frames] - - if request.non_streaming_mode is not None: - params["non_streaming_mode"] = [request.non_streaming_mode] - # Preserve the legacy VoiceDesign fallback when the request omits an - # explicit override. CustomVoice and Base rely on model defaults - # (True and False respectively). - elif params["task_type"][0] == "VoiceDesign": - params["non_streaming_mode"] = [True] - - return params - - # ---- Voxtral TTS helpers ---- - - def _build_voxtral_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: - """Build Voxtral TTS engine prompt, supporting both preset voices and inline - ``ref_audio`` (base64 or data URI).""" - from mistral_common.protocol.speech.request import SpeechRequest - - text = request.input - voice = request.voice - ref_audio = request.ref_audio - if not voice and not ref_audio: - raise ValueError("Voxtral requires either a voice name or ref_audio.") - # mistral_common expects raw base64 (no data: prefix) - if ref_audio is not None and isinstance(ref_audio, str) and ref_audio.startswith("data:"): - _, _, ref_audio = ref_audio.partition(",") - if self._tts_tokenizer is None: - from vllm.tokenizers import cached_tokenizer_from_config - - mistral_tokenizer = cached_tokenizer_from_config(self.engine_client.model_config) - self._tts_tokenizer = mistral_tokenizer.instruct - if voice is not None: - tokens = self._tts_tokenizer.encode_speech_request(SpeechRequest(input=text, voice=voice)).tokens - prompt = tokens_input(prompt_token_ids=tokens) - prompt["additional_information"] = {"voice": [voice]} - return prompt - else: - tokenized = self._tts_tokenizer.encode_speech_request(SpeechRequest(input=text, ref_audio=ref_audio)) - audio = tokenized.audios[0] - return { - "prompt_token_ids": tokenized.tokens, - "multi_modal_data": {"audio": [(audio.audio_array, audio.sampling_rate)]}, - } - - # ---- Step-Audio2 helpers ---- - - def _build_step_audio2_prompt( - self, - request: OpenAICreateSpeechRequest, - ) -> dict[str, Any]: - """Build prompt for Step-Audio2 TTS. - - Constructs the chat prompt with ```` as the last token - of the assistant turn (without ``<|im_end|>``), so the thinker - continues generating audio tokens. - - Prompt format:: - <|im_start|>system\\n{system_prompt}<|im_end|>\\n - <|im_start|>user\\n{input_text}<|im_end|>\\n - <|im_start|>assistant\\n - """ - system_prompt = getattr(request, "instructions", None) or "You are a voice assistant. Read the text aloud." - text = request.input - - raw_prompt = ( - f"<|im_start|>system\n{system_prompt}<|im_end|>\n" - f"<|im_start|>user\n{text}<|im_end|>\n" - f"<|im_start|>assistant\n" - ) - return {"prompt": raw_prompt} - - # ---- Fish Speech helpers ---- - - def _build_fish_speech_prompt( - self, - request: OpenAICreateSpeechRequest, - ref_audio_data: tuple[list[float], int] | None = None, - *, - has_inline_ref_audio: bool = False, - ) -> dict[str, Any]: - """Build prompt for Fish Speech S2 Pro. - - Without voice cloning: - <|im_start|>system\\nconvert the provided text to speech<|im_end|> - <|im_start|>user\\n{text}<|im_end|>\\n<|im_start|>assistant\\n<|voice|> - - With voice cloning (ref_audio + ref_text): - <|im_start|>system\\nconvert the provided text to speech reference to the following... - <|im_end|>\\n<|im_start|>user\\n{text}<|im_end|>\\n<|im_start|>assistant\\n<|voice|> - """ - from transformers import AutoTokenizer - - if self._fish_speech_tokenizer is None: - model_name = self.engine_client.model_config.model - self._fish_speech_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) - - tokenizer = self._fish_speech_tokenizer - - if ref_audio_data is None or not request.ref_text: - prompt_ids, normalized_text = build_fish_text_only_prompt_ids(tokenizer, request.input) - - # Keep the prompt-dict metadata shape aligned with the existing text-only - # TTS entrypoints: scalar values are wrapped in single-item lists before - # EngineCore serialization. Structured clone below is different because - # model-side preprocess consumes concrete per-request scalar fields. - additional_information: dict[str, Any] = { - "text": [normalized_text], - } - if request.max_new_tokens is not None: - additional_information["max_new_tokens"] = [request.max_new_tokens] - prompt = tokens_input(prompt_token_ids=prompt_ids) - prompt["additional_information"] = additional_information - return prompt - - wav_samples, sr = ref_audio_data - normalized_text, normalized_ref_text = normalize_fish_voice_clone_texts(request.input, request.ref_text) - ph_len = self._estimate_fish_prompt_len(normalized_text, normalized_ref_text, ref_audio_data) - - # Structured clone: scalars (not list-wrapped) because model-side - # preprocess() consumes per-request fields directly. - additional_information: dict[str, Any] = { - "text": normalized_text, - "ref_text": normalized_ref_text, - "ref_audio_wav": torch.from_numpy(np.asarray(wav_samples, dtype=np.float32)), - "ref_audio_sr": int(sr), - "fish_structured_voice_clone": True, - } - # Pass voice identity for model-side DAC code caching. - if request.voice is not None: - voice_lower = request.voice.lower() - if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: - additional_information["voice_name"] = voice_lower - additional_information["voice_created_at"] = self._voice_created_at(voice_lower) - if request.max_new_tokens is not None: - additional_information["max_new_tokens"] = request.max_new_tokens - prompt = tokens_input(prompt_token_ids=[1] * ph_len) - prompt["additional_information"] = additional_information - return prompt - - # ---- CosyVoice3 helpers ---- - - async def _build_cosyvoice3_prompt( - self, - request: OpenAICreateSpeechRequest, - *, - has_inline_ref_audio: bool = False, - ) -> dict[str, Any]: - """Build prompt for CosyVoice3. - - CosyVoice3 uses multimodal input with reference audio for voice cloning. - The prompt format matches the offline example: text prompt + audio data - + mm_processor_kwargs with prompt_text. - """ - # Resolve reference audio - 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": ref_text, - "sample_rate": sr, - } - # Pass voice metadata for caching in the processor - if request.voice: - voice_lower = request.voice.lower() - if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: - mm_kwargs["voice_name"] = voice_lower - mm_kwargs["voice_created_at"] = self._voice_created_at(voice_lower) - - return { - "prompt": request.input, - "multi_modal_data": { - "audio": audio_data, - }, - "mm_processor_kwargs": mm_kwargs, - } - - # ---- Covo-Audio helpers ---- - - def _build_covo_audio_prompt( - self, - request: OpenAICreateSpeechRequest, - ) -> dict[str, Any]: - """Build a chat-style prompt for Covo-Audio-Chat. - - Covo-Audio requires a specific system prompt that instructs the model - to interleave text and audio tokens in its output. We render the - messages through the chat template and pass prompt_token_ids so that - the engine does not need to re-tokenize. - """ - from transformers import AutoTokenizer - - from vllm_omni.model_executor.models.covo_audio.prompt_utils import ( - build_covo_audio_prompt_token_ids, - ) - - if self._covo_audio_tokenizer is None: - model_name = self.engine_client.model_config.model - try: - self._covo_audio_tokenizer = AutoTokenizer.from_pretrained( - model_name, - trust_remote_code=True, - ) - except Exception as exc: - raise RuntimeError(f"Failed to load Covo-Audio tokenizer from '{model_name}': {exc}") from exc - - prompt_ids = build_covo_audio_prompt_token_ids( - self._covo_audio_tokenizer, - request.input, - ) - return {"prompt_token_ids": prompt_ids} - - def _get_audex_tokenizer(self): - """Load (once per process) the Audex thinker tokenizer for CFG/TTA prompts. - - The folder follows the serving stage: the full-pipeline thinker - (``audex_omni``) tokenizes with ``checkpoint_folder_full``; the TTS - and TTA thinkers use ``checkpoint_folder_audiogen``. - """ - if self._audex_tokenizer is None: - from transformers import AutoTokenizer - - from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot - - stage = getattr(getattr(self._tts_stage, "engine_args", None), "model_stage", None) - if stage == "audex_omni": - profile, folder = "full", "checkpoint_folder_full" - elif stage == "audex_tta_thinker": - profile, folder = "tta", "checkpoint_folder_audiogen" - else: - profile, folder = "tts", "checkpoint_folder_audiogen" - # The serving model_config may already hold the RESOLVED per-stage - # checkpoint folder (stage init joins model_subdir onto the - # snapshot root); joining again would produce a non-existent path - # that transformers then mistakes for a repo id. - model_path = os.path.normpath(self.engine_client.model_config.model) - if os.path.basename(model_path).startswith("checkpoint_folder"): - root = os.path.dirname(model_path) - else: - root = ensure_audex_snapshot(model_path, profile=profile) - self._audex_tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, folder)) - return self._audex_tokenizer - - # ---- GLM-TTS helpers ---- - - async def _build_glm_tts_prompt( - self, - request: OpenAICreateSpeechRequest, - *, - has_inline_ref_audio: bool = False, - ) -> dict[str, Any]: - """Build prompt for GLM-TTS. - - Uses the multimodal processor path (same as CosyVoice3): - - prompt: synthesis text - - multi_modal_data["audio"]: (wav_samples, sr) reference audio - - mm_processor_kwargs["prompt_text"]: reference text transcript - - AR preprocess() builds [PromptText | Text | BOA | PromptSpeechTokens + ATS]. - DiT receives prompt_token, prompt_feat, embedding for conditioning. - """ - # Voice cloning requires ref_audio + ref_text - if request.ref_audio is not None and request.ref_text: - wav_samples, sr, _ = await self._resolve_ref_audio(request.ref_audio) - audio_data = (np.asarray(wav_samples, dtype=np.float32), int(sr)) - - mm_kwargs: dict[str, Any] = { - "prompt_text": request.ref_text, - } - if request.voice: - voice_lower = request.voice.lower() - if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: - mm_kwargs["voice_name"] = voice_lower - mm_kwargs["voice_created_at"] = self._voice_created_at(voice_lower) - - return { - "prompt": request.input, - "multi_modal_data": { - "audio": audio_data, - }, - "mm_processor_kwargs": mm_kwargs, - "additional_information": self._build_glm_tts_prefill_metadata( - request.input, - request.ref_text, - ), - } - - raise ValueError("GLM-TTS requires ref_audio and ref_text for voice cloning.") - - def _glm_tts_text_tokenizer_and_frontend(self): - from vllm_omni.model_executor.models.glm_tts.glm_tts import ( - load_glm_tts_tokenizer, - resolve_glm_tts_tokenizer_path, - ) - from vllm_omni.model_executor.models.glm_tts.text_frontend import GLMTTSTextFrontend - - cached = self._glm_tts_text_tokenizer - if cached is None: - model_name_or_path = self.engine_client.model_config.model - tokenizer_path = getattr(self.engine_client.model_config, "tokenizer", None) - if tokenizer_path is None: - tokenizer_path = resolve_glm_tts_tokenizer_path(model_name_or_path) - cached = load_glm_tts_tokenizer( - tokenizer_path, - model_name_or_path=model_name_or_path, - trust_remote_code=bool(getattr(self.engine_client.model_config, "trust_remote_code", False)), - ) - self._glm_tts_text_tokenizer = cached - - frontend = self._glm_tts_text_frontend - if frontend is None: - frontend = GLMTTSTextFrontend() - self._glm_tts_text_frontend = frontend - return cached, frontend - - def _estimate_glm_tts_text_token_len(self, text: str | None, *, add_trailing_space: bool = False) -> int: - """Estimate GLM-TTS normalized text length with the model tokenizer.""" - cached, frontend = self._glm_tts_text_tokenizer_and_frontend() - text = text or "" - normalized = frontend.text_normalize(text) or text - normalized = normalized.strip() - if add_trailing_space and normalized: - normalized = f"{normalized} " - return max(1, len(cached.encode(normalized))) - - def _build_glm_tts_prefill_metadata(self, text: str, prompt_text: str | None) -> dict[str, Any]: - """Build GLM-TTS processor length metadata for additional_information. - - The model preprocess hook runs before postprocess can mirror MM kwargs - into additional_information, so these scalar fields must originate from - the request payload rather than runner-level mm_features. - """ - text_len = self._estimate_glm_tts_text_token_len(text) - prompt_text_len = ( - self._estimate_glm_tts_text_token_len(prompt_text, add_trailing_space=True) if prompt_text else 0 - ) - return { - "glm_tts_text_token_len": [text_len], - "glm_tts_prompt_text_token_len": [prompt_text_len], - "input_len": [prompt_text_len + text_len + 1], - } - - # ---- Ming-flash-omni standalone-talker (TTS) helpers ---- - - def _build_ming_flash_omni_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: - # request.instructions accepts two forms: - # 1. Plain text: mapped to the caption's 风格 (style) field - # 2. JSON object: parsed and splatted into the caption. Unlocks - # Unknown keys are dropped by `ming_create_instruction`. - caption_fields = self._parse_ming_instruction_fields(request) or {} - - has_spk_emb = request.speaker_embedding is not None - - # TTS path applies ming task type `instruct`. - # voice_name enables talker-side voice preset resolution (e.g. "DB30"). - additional_information: dict[str, Any] = { - "ming_task": "instruct", - "prompt": MING_DEFAULT_PROMPT, - "text": request.input, - "instruction": ming_create_instruction(caption_fields), - "voice_name": request.voice or None, - "use_zero_spk_emb": not has_spk_emb, - "max_decode_steps": request.max_new_tokens or _TTS_MAX_NEW_TOKENS_MAX, - "cfg": 2.0, - "sigma": 0.25, - "temperature": 0.0, - } - if has_spk_emb: - # Passed as plain float list - additional_information["spk_emb"] = list(request.speaker_embedding) - prompt = tokens_input(prompt_token_ids=[0]) - prompt["additional_information"] = additional_information - return prompt - - # ---- Common speech generation helpers ---- - - async def _build_qwen3_tts_request( - self, - request: OpenAICreateSpeechRequest, - ) -> tuple[dict[str, Any], dict[str, Any], str | None]: - """Build prompt + tts_params for Qwen3-TTS. - - Called from ``Qwen3TTSAdapter.build``. Returns - ``(prompt, tts_params, warmup_artifact_key)`` where the warmup key is the - Qwen3-TTS ref-audio artifact tracked after ``generate()``. - """ - qwen3_ref_audio_warmup_artifact_key: str | None = None - tts_params = self._build_tts_params(request) - # Resolve ref_audio (explicit or auto-set for uploaded voices) - # to [[wav_list, sr]] so the model doesn't re-decode base64. - ref_audio_source = request.ref_audio - if ref_audio_source is None and isinstance(tts_params.get("ref_audio"), list): - # Uploaded voice: ref_audio was auto-set as [base64_data_url] - ref_audio_source = tts_params["ref_audio"][0] - if ref_audio_source is not None and isinstance(ref_audio_source, str): - wav_list, sr, cache_key = await self._resolve_ref_audio(ref_audio_source) - tts_params["ref_audio_cache_key"] = cache_key - artifact_key = self._get_resolved_ref_audio_artifact_key(cache_key) - if artifact_key: - tts_params[_QWEN3_TTS_REF_AUDIO_CACHE_KEY] = [artifact_key] - ref_code_length = self._estimate_ref_code_len([wav_list, sr]) - if ref_code_length is not None: - tts_params["ref_code_length"] = [int(ref_code_length)] - if self._qwen3_tts_can_use_ref_audio_artifact_only(tts_params, artifact_key): - logger.debug("Using Qwen3-TTS ref_audio artifact-only path: %s", artifact_key) - else: - tts_params["ref_audio"] = [[wav_list, sr]] - qwen3_ref_audio_warmup_artifact_key = artifact_key - - ph_len = await self._estimate_prompt_len_async(tts_params) - prompt = tokens_input(prompt_token_ids=[1] * ph_len) - prompt["additional_information"] = tts_params - prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params) - return prompt, tts_params, qwen3_ref_audio_warmup_artifact_key - async def _prepare_speech_generation( self, request: OpenAICreateSpeechRequest, @@ -3189,18 +1811,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): 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) - 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)) - # When word_timestamps is requested, also ask for the aligner stage's # output so the orchestrator drives the request through the forced-aligner # stage (final_stage_id extends to it). Harmless if no aligner stage exists. @@ -3505,6 +2115,25 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if sampling_params_list[0].extra_args is None: sampling_params_list[0].extra_args = {} sampling_params_list[0].extra_args.update(extra) + + sampling = sampling_params_list[0] + + # This change allows StepScheduler read total_steps from upper + # sampling.num_inference_steps, check diffusion/sched/step_scheduler:_get_total_steps + if "num_inference_steps" in extra: + value = extra["num_inference_steps"] + try: + sampling.num_inference_steps = int(value) + except (TypeError, ValueError) as exc: + raise ValueError("num_inference_steps must be an integer") from exc + + if "guidance_scale" in extra: + value = extra["guidance_scale"] + try: + sampling.guidance_scale = float(value) + except (TypeError, ValueError) as exc: + raise ValueError("guidance_scale must be a number") from exc + logger.info("Applied extra_params to diffusion: %s", extra) generator = self._diffusion_engine.generate( diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 0f883f179..32e9fa40d 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -3,8 +3,10 @@ """WebSocket handler for streaming text input TTS. -Accepts text incrementally via WebSocket, buffers it until input.done, and -generates audio once for the buffered input using the existing TTS pipeline. +Accepts text incrementally via WebSocket. By default (split_granularity=none) +it buffers until input.done and generates audio once for the buffered input. +Opting into split_granularity=sentence or clause emits a TTS request at each +detected boundary so incremental STT/LLM clients can hear audio before flush. input.done is a flush, not a close: it ends the current utterance and the connection stays open, so the next utterance reuses the same connection @@ -13,9 +15,10 @@ session.close, on the idle timeout, or when the client closes the socket. The session config is sticky across flushes and can be replaced by sending another session.config between utterances. -"Utterance" here names the flush unit rather than any linguistic unit: it is -whatever text the client had buffered when it sent input.done, from a word to -several paragraphs, synthesized as a single request. +"Utterance" here names the flush unit. With split_granularity=none it is +whatever text was buffered when input.done arrived, synthesized as one +request. With sentence/clause splitting it is still one input.done cycle, +but sentence_index counts linguistic units inside that flush. Protocol: Client -> Server: @@ -32,11 +35,10 @@ Protocol: {"type": "audio.done", "utterance_index": 0, "sentence_index": 0} {"type": "session.done", "utterance_index": 0, "total_sentences": N} {"type": "error", "message": "..."} - # session.done ends the flushed utterance, not the connection. An - # utterance is just the flush unit: whatever text was buffered when - # input.done arrived, of any length. utterance_index counts those - # flushes across the connection, while sentence_index counts within - # one of them and so pairs with total_sentences. + # session.done ends the flushed utterance, not the connection. + # utterance_index counts input.done flushes. sentence_index counts + # TTS requests inside one flush (always 0 of 1 when split_granularity + # is none; 0..N-1 when sentence/clause splitting is enabled). Server -> Client (when word_timestamps=true): {"type": "audio.start", "utterance_index": 0, "sentence_index": 0, @@ -64,6 +66,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( StreamingSpeechSessionConfig, ) from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.entrypoints.openai.speech_text_splitter import SpeechTextSplitter from vllm_omni.utils.forced_aligner import extract_word_timestamps logger = init_logger(__name__) @@ -79,11 +82,10 @@ _MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024 class OmniStreamingSpeechHandler: """Handles WebSocket sessions for streaming text-input TTS. - A connection carries one or more utterances. Text arrives incrementally, - is buffered until input.done, and audio is generated once for the - buffered input using the existing OmniOpenAIServingSpeech pipeline. The - connection outlives each utterance so a client can keep synthesizing on - it without reconnecting. + A connection carries one or more utterances. Text arrives incrementally. + With split_granularity=none it is buffered until input.done and synthesized + as one request. sentence/clause modes emit a request at each boundary, + including before input.done. The connection outlives each utterance. Args: speech_service: The existing TTS serving instance (reused for @@ -105,17 +107,20 @@ class OmniStreamingSpeechHandler: async def handle_session(self, websocket: WebSocket) -> None: """Main loop for a single WebSocket connection. - Serves any number of utterances: text is buffered until input.done, - which flushes it as one TTS request and then leaves the connection - open for the next one. Rejecting a message is only fatal before the - first valid session.config; afterwards the error is reported and the - connection survives. + Serves any number of utterances. input.done always ends a flush; + split_granularity controls whether that flush is one TTS request or + several sentence/clause requests. """ await websocket.accept() config: StreamingSpeechSessionConfig | None = None - text_parts: list[str] = [] + splitter = SpeechTextSplitter("none") utterance_index = 0 + sentence_index = 0 + # An utterance is in progress from the first input.text until the next + # input.done. Buffered text is not enough to detect it: sentence/clause + # splitting can leave the buffer empty after emitting a unit. + utterance_open = False try: while True: @@ -140,10 +145,10 @@ class OmniStreamingSpeechHandler: msg_type = msg.get("type") if msg_type == "session.config": - if text_parts: + if utterance_open: await self._send_error( websocket, - "session.config cannot be applied while input is buffered; send input.done first", + "session.config cannot be applied while an utterance is in progress; send input.done first", ) continue new_config = await self._build_config(websocket, msg) @@ -152,6 +157,7 @@ class OmniStreamingSpeechHandler: return # Error already sent, connection closing continue # Keep serving with the previous config config = new_config + splitter = SpeechTextSplitter(config.split_granularity) elif config is None: await self._send_error( @@ -165,32 +171,39 @@ class OmniStreamingSpeechHandler: if not isinstance(text, str): await self._send_error(websocket, "input.text requires a string value") continue - text_parts.append(text) - - elif msg_type == "input.done": - full_text = "".join(text_parts).strip() - text_parts.clear() - total_sentences = 0 - if full_text: - # However long the buffered text is, the pipeline takes - # it as one request, so every flush is sentence 0 of 1. + if text: + utterance_open = True + for unit in splitter.feed(text): await self._generate_and_send( websocket, config, - full_text, + unit, utterance_index=utterance_index, - sentence_index=0, + sentence_index=sentence_index, ) - total_sentences = 1 + sentence_index += 1 + + elif msg_type == "input.done": + for unit in splitter.flush(): + await self._generate_and_send( + websocket, + config, + unit, + utterance_index=utterance_index, + sentence_index=sentence_index, + ) + sentence_index += 1 await websocket.send_json( { "type": "session.done", "utterance_index": utterance_index, - "total_sentences": total_sentences, + "total_sentences": sentence_index, } ) utterance_index += 1 + sentence_index = 0 + utterance_open = False elif msg_type == "session.close": await websocket.close() @@ -310,6 +323,7 @@ class OmniStreamingSpeechHandler: speaker_embedding=config.speaker_embedding, stream=config.stream_audio, word_timestamps=config.word_timestamps, + seed=config.seed, ) start_payload = { diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index f3f4b7570..b3afe8b40 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -37,6 +37,7 @@ from vllm_omni.entrypoints.openai.video_api_utils import ( encode_video_base64, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.metrics import count_video_frames from vllm_omni.model_extras import get_video_generation_defaults, should_preserve_reference_image_size from vllm_omni.model_extras.video_generation import VideoGenerationDefaults from vllm_omni.outputs.output_metadata import ( @@ -87,6 +88,27 @@ class VideoGenerationArtifacts: output_fps: float stage_durations: dict[str, float] peak_memory_mb: float + metrics: dict[str, object] | None = None + + +def _video_metadata_from_artifacts(artifacts: VideoGenerationArtifacts) -> dict[str, object]: + metadata: dict[str, object] = {} + if artifacts.output_fps > 0: + metadata["fps"] = artifacts.output_fps + + if not artifacts.videos: + return metadata + + num_frames = count_video_frames(artifacts.videos[0]) + if num_frames is not None and num_frames > 0: + metadata["num_frames"] = num_frames + if artifacts.output_fps > 0: + metadata["duration_s"] = num_frames / artifacts.output_fps + + if artifacts.metrics: + metadata["metrics"] = artifacts.metrics + + return metadata class OmniOpenAIServingVideo: @@ -182,6 +204,35 @@ class OmniOpenAIServingVideo: ) return capability is True or metadata_capability + @property + def supported_control_upload_types(self) -> frozenset[str]: + """Return multipart control types accepted by the active pipeline. + + Unknown pipelines deliberately return an empty set. This keeps the + generic video API isolated from model-specific controls unless a model + explicitly opts into the ``control_path`` contract in metadata. + """ + od_config = self._resolve_diffusion_od_config() + model_archs = [None if od_config is None else getattr(od_config, "model_class_name", None)] + for stage_config in self.stage_configs or (): + stage_get = ( + stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None) + ) + engine_args = stage_get("engine_args") or {} + model_archs.extend( + ( + stage_get("model_arch"), + engine_args.get("model_class_name") + if isinstance(engine_args, Mapping) + else getattr(engine_args, "model_class_name", None), + ) + ) + + supported: set[str] = set() + for model_arch in model_archs: + supported.update(get_diffusion_model_metadata(model_arch).supported_control_upload_types) + return frozenset(supported) + @classmethod def for_diffusion( cls, @@ -393,6 +444,8 @@ class OmniOpenAIServingVideo: 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) + raw_metrics = getattr(result, "metrics", None) if request.return_stage_metrics else None + metrics = {str(key): value for key, value in raw_metrics.items()} if isinstance(raw_metrics, Mapping) else None return VideoGenerationArtifacts( videos=videos, audios=audios, @@ -401,6 +454,7 @@ class OmniOpenAIServingVideo: output_fps=output_fps, stage_durations=self._extract_stage_durations(result), peak_memory_mb=self._extract_peak_memory_mb(result), + metrics=metrics, ) async def generate_videos( @@ -466,7 +520,7 @@ class OmniOpenAIServingVideo: reference_image: ReferenceImage | None = None, reference_video: ReferenceVideo | None = None, reference_audio: ReferenceAudio | None = None, - ) -> tuple[bytes, dict[str, float], float, VideoAction | None]: + ) -> tuple[bytes, dict[str, float], float, VideoAction | None, dict[str, object]]: """Generate a video and return raw MP4 bytes, bypassing base64 encoding.""" artifacts = await self._run_and_extract( request, @@ -489,9 +543,10 @@ class OmniOpenAIServingVideo: video_codec_options = request.extra_params["video_codec_options"] action = artifacts.actions[0] + video_metadata = _video_metadata_from_artifacts(artifacts) if action is not None and isinstance(artifacts.videos[0], dict): logger.info("Action-only video request %s completed; skipping MP4 encoding.", reference_id) - return b"", artifacts.stage_durations, artifacts.peak_memory_mb, action + return b"", artifacts.stage_durations, artifacts.peak_memory_mb, action, video_metadata _t_encode_start = time.perf_counter() video_bytes = _encode_video_bytes( @@ -503,7 +558,7 @@ class OmniOpenAIServingVideo: ) _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000 logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms) - return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb, artifacts.actions[0] + return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb, artifacts.actions[0], video_metadata @staticmethod def _resolve_video_fps_multiplier(result: object) -> int: diff --git a/vllm_omni/entrypoints/openai/speech_text_splitter.py b/vllm_omni/entrypoints/openai/speech_text_splitter.py new file mode 100644 index 000000000..1f25a8dba --- /dev/null +++ b/vllm_omni/entrypoints/openai/speech_text_splitter.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Incremental text splitting for WebSocket TTS input. + +Default serving still synthesizes one request per ``input.done`` (see +``split_granularity='none'``). Opting into ``sentence`` or ``clause`` restores +per-boundary requests so incremental STT/LLM clients can start audio before +the full utterance arrives. + +Terminators cover Latin, CJK, Indic danda, and Arabic question/semicolon marks +so Hindi and similar scripts are not stuck waiting for ``.`` / ``?``. +""" + +from __future__ import annotations + +# ASCII punctuation is ambiguous (abbreviations, decimals, thousands +# separators), so it only closes a unit when whitespace or a flush follows. +# Script-specific marks below are unambiguous and close a unit as soon as the +# character after them is known. +_ASCII_SENTENCE_TERMINATORS = frozenset(".!?") +_ASCII_CLAUSE_TERMINATORS = _ASCII_SENTENCE_TERMINATORS | frozenset(",;") +_SCRIPT_SENTENCE_TERMINATORS = frozenset( + { + "。", + "！", + "？", + "…", + "।", # Devanagari danda + "॥", # Devanagari double danda + "؟", # Arabic question mark + } +) +_SCRIPT_CLAUSE_TERMINATORS = _SCRIPT_SENTENCE_TERMINATORS | { + "，", + "；", + "،", # Arabic comma + "؛", # Arabic semicolon +} + +_SENTENCE_TERMINATORS = _ASCII_SENTENCE_TERMINATORS | _SCRIPT_SENTENCE_TERMINATORS +_CLAUSE_TERMINATORS = _ASCII_CLAUSE_TERMINATORS | _SCRIPT_CLAUSE_TERMINATORS + +# Trailing characters that belong to the unit they close, so `Wait...` and +# `He said "Hello."` stay one TTS request instead of several. +_CLOSING_DELIMITERS = frozenset("\"'”’»)]}』」）】") +_RUN_CHARS = _CLAUSE_TERMINATORS | _CLOSING_DELIMITERS + +# Abbreviations that end in `.` mid-sentence. Single letters are covered +# separately so initials such as `J. R. R.` do not split either. +_ABBREVIATIONS = frozenset( + { + "dr", + "mr", + "mrs", + "ms", + "prof", + "sr", + "jr", + "st", + "mt", + "fig", + "no", + "vs", + "etc", + "approx", + "dept", + "inc", + "ltd", + "co", + "al", + "e.g", + "i.e", + "u.s", + "u.k", + "a.m", + "p.m", + } +) +_MAX_ABBREVIATION_LEN = max(len(abbrev) for abbrev in _ABBREVIATIONS) + + +def _is_token_char(ch: str) -> bool: + return ch.isalnum() or ch == "." + + +def _is_numeric_separator(buffer: str, index: int) -> bool: + """True for the `.`/`,` inside `3.14` or `1,000`.""" + if buffer[index] not in ".,": + return False + if index == 0 or not buffer[index - 1].isdigit(): + return False + return index + 1 < len(buffer) and buffer[index + 1].isdigit() + + +def _is_abbreviation(buffer: str, index: int) -> bool: + """True for the `.` of `Dr.`, `e.g.`, or an initial like `J.`.""" + if buffer[index] != ".": + return False + # Bounded so a frame of `a.a.a...` cannot walk the whole prefix at every + # `.`. Every abbreviation fits inside the limit, so hitting it can only + # rule a token out, never produce a false positive. + limit = max(0, index - _MAX_ABBREVIATION_LEN) + start = index + while start > limit and _is_token_char(buffer[start - 1]): + start -= 1 + if start > 0 and _is_token_char(buffer[start - 1]): + return False + token = buffer[start:index] + if not token: + return False + if len(token) == 1 and token.isalpha(): + return True + return token.lower().strip(".") in _ABBREVIATIONS + + +def extract_complete_units( + buffer: str, + terminators: frozenset[str], + *, + flush: bool, + scan_from: int = 0, +) -> tuple[list[str], str, int]: + """Split ``buffer`` into complete units. + + Returns ``(units, remainder, next_scan)``. ``next_scan`` is an offset into + ``remainder``: everything before it has been examined already, so a caller + appending more text can pass it back as ``scan_from`` instead of rescanning + the whole buffer on every ``input.text`` message. + """ + units: list[str] = [] + last_split = 0 + i = max(scan_from, 0) + length = len(buffer) + while i < length: + ch = buffer[i] + if ch not in terminators or _is_numeric_separator(buffer, i) or _is_abbreviation(buffer, i): + i += 1 + continue + + # Absorb the punctuation/closing-delimiter run so `Wait...` emits once. + run_end = i + 1 + while run_end < length and buffer[run_end] in _RUN_CHARS: + run_end += 1 + + if run_end >= length and not flush: + # The run may still grow; re-examine this terminator next feed. + return units, buffer[last_split:], i - last_split + + ascii_terminator = ch in _ASCII_CLAUSE_TERMINATORS + if run_end < length and ascii_terminator and not buffer[run_end].isspace(): + i = run_end + continue + + piece = buffer[last_split:run_end].strip() + if piece: + units.append(piece) + i = run_end + while i < length and buffer[i].isspace(): + i += 1 + last_split = i + + remainder = buffer[last_split:] + if flush: + tail = remainder.strip() + if tail: + units.append(tail) + return units, "", 0 + return units, remainder, len(remainder) + + +class SpeechTextSplitter: + """Stateful splitter used by one WebSocket utterance.""" + + def __init__(self, granularity: str = "none") -> None: + if granularity not in ("none", "sentence", "clause"): + raise ValueError(f"Unsupported split_granularity: {granularity!r}") + self.granularity = granularity + self._buf = "" + self._scan = 0 + + def _terminators(self) -> frozenset[str] | None: + if self.granularity == "none": + return None + if self.granularity == "clause": + return _CLAUSE_TERMINATORS + return _SENTENCE_TERMINATORS + + def feed(self, text: str) -> list[str]: + self._buf += text + terminators = self._terminators() + if terminators is None: + return [] + units, self._buf, self._scan = extract_complete_units(self._buf, terminators, flush=False, scan_from=self._scan) + return units + + def flush(self) -> list[str]: + terminators = self._terminators() + if terminators is None: + piece = self._buf.strip() + self._buf = "" + self._scan = 0 + return [piece] if piece else [] + units, self._buf, self._scan = extract_complete_units(self._buf, terminators, flush=True, scan_from=self._scan) + return units diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index a09d97b68..39882274d 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -104,6 +104,7 @@ def tts_entry_stage_archs() -> frozenset[str]: from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 audex, audex_tta, + audio8_tts, cosyvoice3, covo_audio, dots_tts, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audex.py b/vllm_omni/entrypoints/openai/tts_adapters/audex.py index ab7c7c25e..0e31b633f 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/audex.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/audex.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Audex (Nemotron-Labs-Audex-2B) TTS serving adapter.""" +import os from typing import TYPE_CHECKING, Any from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter @@ -29,6 +31,30 @@ class AudexAdapter(ARTTSAdapter): stage_keys = frozenset({"audex_thinker", "audex_omni"}) name = "audex" + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._tokenizer = None + + def _get_tokenizer(self): + """Load and cache the tokenizer used only by the Audex TTS adapter.""" + if self._tokenizer is None: + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot + + stage = getattr(getattr(self.ctx.server._tts_stage, "engine_args", None), "model_stage", None) + profile, folder = ( + ("full", "checkpoint_folder_full") if stage == "audex_omni" else ("tts", "checkpoint_folder_audiogen") + ) + model_path = os.path.normpath(self.ctx.engine_client.model_config.model) + root = ( + os.path.dirname(model_path) + if os.path.basename(model_path).startswith("checkpoint_folder") + else ensure_audex_snapshot(model_path, profile=profile) + ) + self._tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, folder)) + return self._tokenizer + @classmethod def stage_serves_speech(cls, model_stage: str | None, all_stage_keys: frozenset[str]) -> bool: """``audex_omni`` covers the audex_s2s deployment, whose TTS pass uses @@ -102,7 +128,6 @@ class AudexAdapter(ARTTSAdapter): element of the engine's shared ``default_sampling_params_list`` — see ``_inject_audex_tta_args``). """ - server = self.ctx.server import copy stage0_params = copy.deepcopy(stage0_params) @@ -121,7 +146,7 @@ class AudexAdapter(ARTTSAdapter): cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt if not isinstance(cond_prompt, str) or not cond_prompt: raise ValueError("Audex CFG requires the adapter-built text prompt") - null_prompt = build_null_prompt(cond_prompt, server._get_audex_tokenizer()) + null_prompt = build_null_prompt(cond_prompt, self._get_tokenizer()) extra_args.update( { "cfg_scale": cfg_scale, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py b/vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py index 33dd2294a..b0af5cdca 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Audex (Nemotron-Labs-Audex-2B) text-to-audio serving adapter. Caption → general audio through the ``audex_tta`` pipeline. @@ -8,6 +9,7 @@ completes the sampling parameters with the RVQ phase-mask contract scale 3.0 — guidance is effectively mandatory for TTA quality). """ +import os from typing import TYPE_CHECKING, Any from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter @@ -32,6 +34,27 @@ class AudexTTAAdapter(ARTTSAdapter): stage_keys = frozenset({"audex_tta_thinker"}) name = "audex_tta" + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._tokenizer = None + self._tta_rvq = None + + def _get_tokenizer(self): + """Load and cache the tokenizer used only by the Audex TTA adapter.""" + if self._tokenizer is None: + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot + + model_path = os.path.normpath(self.ctx.engine_client.model_config.model) + root = ( + os.path.dirname(model_path) + if os.path.basename(model_path).startswith("checkpoint_folder") + else ensure_audex_snapshot(model_path, profile="tta") + ) + self._tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, "checkpoint_folder_audiogen")) + return self._tokenizer + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: if not request.input or not request.input.strip(): return "Audex TTA requires a non-empty caption" @@ -95,7 +118,6 @@ class AudexTTAAdapter(ARTTSAdapter): state (``cfg_pair_id``/``cfg_null_prompt``) into it would race concurrent requests and leak stale CFG metadata into later ones. """ - server = self.ctx.server import copy from vllm_omni.model_executor.models.audex.prompt import build_tta_null_prompt @@ -106,10 +128,10 @@ class AudexTTAAdapter(ARTTSAdapter): raise ValueError("Audex TTA requires the adapter-built caption prompt") stage0_params = copy.deepcopy(stage0_params) - tokenizer = server._get_audex_tokenizer() - if server._audex_tta_rvq is None: + tokenizer = self._get_tokenizer() + if self._tta_rvq is None: phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer) - server._audex_tta_rvq = { + self._tta_rvq = { "phase_token_ids": phase_token_ids, "start_tid": start_tid, "end_tid": end_tid, @@ -123,7 +145,7 @@ class AudexTTAAdapter(ARTTSAdapter): if extra_args is None: extra_args = {} stage0_params.extra_args = extra_args - extra_args["tta_rvq"] = server._audex_tta_rvq + extra_args["tta_rvq"] = self._tta_rvq cfg_scale = extra_args.get("cfg_scale") cfg_scale = 3.0 if cfg_scale is None else float(cfg_scale) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audio8_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/audio8_tts.py new file mode 100644 index 000000000..1706a64d3 --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/audio8_tts.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview serving adapter for ``/v1/audio/speech``.""" + +import copy +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +from vllm.inputs import tokens_input +from vllm.utils.async_utils import make_async + +from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + PreparedRequest, + apply_max_new_tokens, + conditioning_cache_salt, +) +from vllm_omni.model_executor.models.audio8_tts.codec_utils import ( + estimate_reference_code_frames, +) +from vllm_omni.model_executor.models.audio8_tts.prompt_utils import ( + build_text_only_prompt_ids, + estimate_voice_clone_prompt_len, + normalize_text, +) + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + + +@register_tts_adapter +class Audio8TTSAdapter(ARTTSAdapter): + """Audio8 TTS Preview 0.6B: text-only synthesis and zero-shot voice cloning.""" + + stage_keys = frozenset({"audio8_tts_slow_ar"}) + name = "audio8_tts" + + def __init__(self, ctx: Any) -> None: + super().__init__(ctx) + self._cached_tokenizer: Any = None + + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: + server = self.ctx.server + err = server._apply_uploaded_speaker(request) + if err: + return err + 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 + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + sampling_params_list = apply_max_new_tokens(sampling_params_list, request) + # Seed the Fast AR (talker_mtp) residual sampling per request. The Slow AR + # is already seeded via sampling_metadata.generators, but the residual + # codebooks read tts_local_seed; without it they fall back to global RNG, + # so an unseeded request's audio would depend on batch composition. Default + # it from the deploy seed; an explicit request.seed still overrides it + # downstream. Mirrors how serving handles qwen3_tts. + if sampling_params_list: + default_seed = sampling_params_list[0].seed + if default_seed is not None: + sampling_params_list = copy.deepcopy(sampling_params_list) + stage0 = sampling_params_list[0] + if stage0.extra_args is None: + stage0.extra_args = {} + stage0.extra_args.setdefault("tts_local_seed", int(default_seed)) + return sampling_params_list + + async def build( + self, + request: "OpenAICreateSpeechRequest", + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + del sampling_params_list + server = self.ctx.server + ref_audio_data = None + if request.ref_audio is not None: + wav_list, sample_rate = await server._resolve_ref_audio(request.ref_audio) + ref_audio_data = (wav_list, sample_rate) + # Prompt building tokenizes and, for voice clone, allocates tensors; keep + # it off the event loop via the server's single-worker TTS executor. + build_prompt = make_async(self._build_prompt, executor=server._tts_executor) + prompt = await build_prompt( + request, + ref_audio_data=ref_audio_data, + has_inline_ref_audio=has_inline_ref_audio, + ) + tts_params: dict = {} + prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) + return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name) + + def _tokenizer(self) -> Any: + """The Audio8 tokenizer, loaded once per process.""" + if self._cached_tokenizer is None: + from transformers import AutoTokenizer + + model_name = self.ctx.server.engine_client.model_config.model + self._cached_tokenizer = AutoTokenizer.from_pretrained(model_name) + return self._cached_tokenizer + + def _estimate_prompt_len(self, text: str, ref_text: str, ref_audio: object) -> int: + """Exact clone-prompt length, without encoding the reference audio. + + Deliberately not defensive: the placeholder this sizes must match the + length ``preprocess()`` actually builds, so a wrong guess corrupts the + request instead of failing it. Let errors surface to the caller. + """ + if not isinstance(ref_audio, (list, tuple)) or len(ref_audio) != 2: + raise ValueError("Audio8 TTS reference audio must be a (samples, sample_rate) pair") + wav, sample_rate = ref_audio + ref_frames = estimate_reference_code_frames(len(wav), int(sample_rate)) + return estimate_voice_clone_prompt_len(self._tokenizer(), text, ref_text, ref_frames) + + def _build_prompt( + self, + request: "OpenAICreateSpeechRequest", + ref_audio_data: tuple[list[float], int] | None = None, + *, + has_inline_ref_audio: bool = False, + ) -> dict[str, Any]: + """Build the engine prompt for Audio8 TTS Preview. + + Text-only prompts are tokenized here. Voice cloning cannot be: the + prompt embeds the reference audio's own codec codes, so the model-side + ``preprocess`` builds it and this path only reserves a placeholder of the + exact final length. + """ + server = self.ctx.server + if ref_audio_data is None or not request.ref_text: + prompt_ids, normalized_text = build_text_only_prompt_ids(self._tokenizer(), request.input) + # Scalars are list-wrapped for the text-only path, matching the + # other TTS entrypoints' additional_information shape. + additional_information: dict[str, Any] = {"text": [normalized_text]} + if request.max_new_tokens is not None: + additional_information["max_new_tokens"] = [request.max_new_tokens] + prompt = tokens_input(prompt_token_ids=prompt_ids) + prompt["additional_information"] = additional_information + return prompt + + wav_samples, sample_rate = ref_audio_data + normalized_text = normalize_text(request.input) + normalized_ref_text = normalize_text(request.ref_text, add_default_speaker=True) + placeholder_len = self._estimate_prompt_len(normalized_text, normalized_ref_text, ref_audio_data) + + # Structured clone: scalars (not list-wrapped) because model-side + # preprocess() consumes these fields directly. + additional_information = { + "text": normalized_text, + "ref_text": normalized_ref_text, + "ref_audio_wav": torch.from_numpy(np.asarray(wav_samples, dtype=np.float32)), + "ref_audio_sr": int(sample_rate), + "audio8_structured_voice_clone": True, + } + if request.voice is not None: + voice_lower = request.voice.lower() + if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: + # Lets the model cache this speaker's encoded codes. + additional_information["voice_name"] = voice_lower + additional_information["voice_created_at"] = server._voice_created_at(voice_lower) + if request.max_new_tokens is not None: + additional_information["max_new_tokens"] = request.max_new_tokens + prompt = tokens_input(prompt_token_ids=[1] * placeholder_len) + prompt["additional_information"] = additional_information + return prompt + + +__all__ = ["Audio8TTSAdapter"] diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index d4316a2fe..ff43b27a5 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -11,8 +11,9 @@ instead of editing the shared serving module in ~10 scattered places. See the RFC for the full design (issue #4327). """ +import hashlib from abc import ABC, abstractmethod -from collections.abc import Callable, Mapping +from collections.abc import Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar @@ -37,22 +38,52 @@ DEFAULT_TTS_LANGUAGES = frozenset( } ) -_conditioning_cache_salt_fn: "Callable[..., str] | None" = None +def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: dict | None = None) -> str: + """Stable hash of the real Stage 0 conditioning for the prefix cache. -def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: dict) -> str: - """Return the conditioning cache salt for ``request`` + ``tts_params``. + The talker's vLLM prompt is placeholder token ids; the real inputs are + rebuilt from text / ref_audio / ref_text into inputs_embeds. vLLM hashes + token ids (folded with cache_salt) for prefix caching, so without a salt + every request collides and a hit could reuse KV from a semantically + different input. - Lazily imports and caches ``serving_speech._conditioning_cache_salt`` on first - use: the import is deferred to break the adapters<->serving_speech import - cycle, and cached so it resolves once instead of on every ``build()`` call. + Raw request fields alone are not enough for uploaded voices: resolved + conditioning such as ``voice_created_at`` and content-aware ref-audio + cache keys must also be folded in. This distinguishes delete/re-upload of + the same voice and same-path local audio rewrites without hashing decoded + waveform arrays. """ - global _conditioning_cache_salt_fn - if _conditioning_cache_salt_fn is None: - from vllm_omni.entrypoints.openai.serving_speech import _conditioning_cache_salt - - _conditioning_cache_salt_fn = _conditioning_cache_salt - return _conditioning_cache_salt_fn(request, tts_params) + h = hashlib.sha256() + for part in ( + request.input, + request.task_type, + request.language, + request.voice, + request.ref_text, + request.ref_audio, + request.instructions, + request.x_vector_only_mode, + request.speaker_embedding, + ): + h.update(b"\x00") + if part is not None: + h.update(repr(part).encode("utf-8")) + # Fold conditioning derived by adapters and absent from the raw request. + for key in ( + "voice_created_at", + "task_type", + "speaker", + "ref_text", + "x_vector_only_mode", + "ref_audio_cache_key", + "ref_audio_2_cache_key", + ): + h.update(b"\x00") + value = tts_params.get(key) if tts_params is not None else None + if value is not None: + h.update(repr(value).encode("utf-8")) + return h.hexdigest()[:32] def apply_max_new_tokens( diff --git a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py index ec1c145c7..1196c5c70 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """CosyVoice3 TTS serving adapter.""" import os from typing import TYPE_CHECKING, Any +import numpy as np from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter @@ -14,12 +16,46 @@ if TYPE_CHECKING: logger = init_logger(__name__) +# Without this delimiter the talker re-speaks the reference transcript instead +# of emitting target-only speech (issue #4644). Matches the offline example and +# upstream demo. +_PROMPT_DELIMITER = "<|endofprompt|>" +_PROMPT_PREFIX = f"You are a helpful assistant.{_PROMPT_DELIMITER}" + @register_tts_adapter class CosyVoice3Adapter(ARTTSAdapter): stage_keys = frozenset({"cosyvoice3_talker"}) name = "cosyvoice3" + def __init__(self, ctx) -> None: + super().__init__(ctx) + # Cache the resolved Qwen tokenizer used for dynamic-token sizing. + # Rebuilding it would repeat snapshot resolution and add ~100 ms to the + # TTFP-critical sampling override on every request. + self._tokenizer = None + + async def _build_prompt( + self, request: "OpenAICreateSpeechRequest", *, has_inline_ref_audio: bool = False + ) -> dict[str, Any]: + """Build the multimodal CosyVoice3 voice-cloning prompt.""" + server = self.ctx.server + wav_samples, sr, _ = await server._resolve_ref_audio(request.ref_audio) + ref_text = request.ref_text or "" + if _PROMPT_DELIMITER not in ref_text: + ref_text = f"{_PROMPT_PREFIX}{ref_text}" + mm_kwargs: dict[str, Any] = {"prompt_text": ref_text, "sample_rate": sr} + if request.voice: + voice_lower = request.voice.lower() + if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: + mm_kwargs["voice_name"] = voice_lower + mm_kwargs["voice_created_at"] = server._voice_created_at(voice_lower) + return { + "prompt": request.input, + "multi_modal_data": {"audio": (np.asarray(wav_samples, dtype=np.float32), sr)}, + "mm_processor_kwargs": mm_kwargs, + } + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: """Validate CosyVoice3 request parameters. Returns error message or None.""" server = self.ctx.server @@ -51,7 +87,7 @@ class CosyVoice3Adapter(ARTTSAdapter): async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: - prompt = await self.ctx.server._build_cosyvoice3_prompt(request, has_inline_ref_audio=has_inline_ref_audio) + prompt = await self._build_prompt(request, has_inline_ref_audio=has_inline_ref_audio) return PreparedRequest(prompt=prompt, tts_params={}, model_type="cosyvoice3") def apply_sampling_overrides( @@ -76,19 +112,19 @@ class CosyVoice3Adapter(ARTTSAdapter): hf_cfg = server.model_config.hf_config # Build the Qwen tokenizer once per process (resolving the model dir via # snapshot_download at most once) and reuse it across requests. - tokenizer = server._cosyvoice3_tokenizer + tokenizer = self._tokenizer if tokenizer is None: model_path = server.engine_client.model_config.model if not os.path.isdir(model_path): - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_path = snapshot_download(model_path) + model_path = hf_api().snapshot_download(model_path) tokenizer = get_qwen_tokenizer( token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path), skip_special_tokens=hf_cfg.skip_special_tokens, version=hf_cfg.version, ) - server._cosyvoice3_tokenizer = tokenizer + self._tokenizer = tokenizer _, text_token_len = extract_text_token( request.input, tokenizer, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py b/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py index 28e64bc2a..a9840b4a0 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """CoVo-Audio serving adapter.""" from typing import TYPE_CHECKING @@ -15,6 +16,10 @@ class CovoAudioAdapter(ARTTSAdapter): stage_keys = frozenset({"fused_thinker_talker"}) name = "covo_audio" + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._tokenizer = None + @classmethod def matches(cls, model_stage: str | None, model_arch: str | None) -> bool: """``fused_thinker_talker`` is a generic stage key that non-CoVo fused @@ -30,5 +35,21 @@ class CovoAudioAdapter(ARTTSAdapter): async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: - prompt = self.ctx.server._build_covo_audio_prompt(request) - return PreparedRequest(prompt=prompt, tts_params={}, model_type="covo_audio") + """Build the tokenized chat prompt expected by Covo-Audio-Chat. + + The model requires a specific system prompt instructing it to + interleave text and audio tokens. Passing token IDs avoids a second + engine-side tokenization step. + """ + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.covo_audio.prompt_utils import build_covo_audio_prompt_token_ids + + if self._tokenizer is None: + model_name = self.ctx.engine_client.model_config.model + try: + self._tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + except Exception as exc: + raise RuntimeError(f"Failed to load Covo-Audio tokenizer from '{model_name}': {exc}") from exc + prompt_ids = build_covo_audio_prompt_token_ids(self._tokenizer, request.input) + return PreparedRequest(prompt={"prompt_token_ids": prompt_ids}, tts_params={}, model_type="covo_audio") diff --git a/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py b/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py index a6426a13d..f6dc84d52 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py @@ -1,8 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Fish Speech serving adapter (retires the legacy ``_is_fish_speech`` flag).""" +import math from typing import TYPE_CHECKING, Any +import numpy as np +import torch +from vllm.inputs import tokens_input +from vllm.utils.async_utils import make_async + from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ( ARTTSAdapter, @@ -20,6 +27,101 @@ class FishSpeechAdapter(ARTTSAdapter): stage_keys = frozenset({"fish_speech_slow_ar"}) name = "fish_tts" + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._tokenizer = None + self._build_prompt_async = make_async(self._build_prompt, executor=getattr(ctx.server, "_tts_executor", None)) + + @staticmethod + def _estimate_ref_code_len(ref_audio: object) -> int | None: + from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE + + if not isinstance(ref_audio, (list, tuple)) or len(ref_audio) != 2: + return None + wav, sr = ref_audio + sr = int(sr) + if sr <= 0 or len(wav) <= 0: + return None + return max(1, math.ceil(max(1, math.ceil(len(wav) * DAC_SAMPLE_RATE / sr)) / DAC_HOP_LENGTH)) + + def _get_tokenizer(self): + if self._tokenizer is None: + from transformers import AutoTokenizer + + self._tokenizer = AutoTokenizer.from_pretrained( + self.ctx.engine_client.model_config.model, trust_remote_code=True + ) + return self._tokenizer + + def _estimate_prompt_len(self, text: str, ref_text: str, ref_audio: object) -> int: + from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( + estimate_fish_voice_clone_prompt_len_from_normalized, + ) + + try: + semantic_len = self._estimate_ref_code_len(ref_audio) + if semantic_len is None: + raise ValueError("Failed to estimate Fish Speech semantic token length") + return estimate_fish_voice_clone_prompt_len_from_normalized( + self._get_tokenizer(), text, ref_text, semantic_len + ) + except Exception as exc: + import logging + + logging.getLogger(__name__).warning( + "Failed to estimate Fish Speech prompt length, using fallback 2048: %s", exc + ) + return 2048 + + def _build_prompt( + self, + request: "OpenAICreateSpeechRequest", + ref_audio_data: tuple[list[float], int] | None, + has_inline_ref_audio: bool, + ) -> dict[str, Any]: + """Build either the text-only or structured Fish Speech clone prompt. + + Structured clone metadata uses concrete scalar fields because model-side + preprocess consumes one request directly; text-only metadata retains the + legacy single-item list representation used across EngineCore IPC. + """ + from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( + build_fish_text_only_prompt_ids, + normalize_fish_voice_clone_texts, + ) + + tokenizer = self._get_tokenizer() + if ref_audio_data is None or not request.ref_text: + prompt_ids, normalized_text = build_fish_text_only_prompt_ids(tokenizer, request.input) + info: dict[str, Any] = {"text": [normalized_text]} + if request.max_new_tokens is not None: + info["max_new_tokens"] = [request.max_new_tokens] + prompt = tokens_input(prompt_token_ids=prompt_ids) + prompt["additional_information"] = info + return prompt + + wav_samples, sr = ref_audio_data + normalized_text, normalized_ref_text = normalize_fish_voice_clone_texts(request.input, request.ref_text) + ph_len = self._estimate_prompt_len(normalized_text, normalized_ref_text, ref_audio_data) + info = { + "text": normalized_text, + "ref_text": normalized_ref_text, + "ref_audio_wav": torch.from_numpy(np.asarray(wav_samples, dtype=np.float32)), + "ref_audio_sr": int(sr), + "fish_structured_voice_clone": True, + } + server = self.ctx.server + if request.voice is not None: + voice_lower = request.voice.lower() + if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: + info["voice_name"] = voice_lower + info["voice_created_at"] = server._voice_created_at(voice_lower) + if request.max_new_tokens is not None: + info["max_new_tokens"] = request.max_new_tokens + prompt = tokens_input(prompt_token_ids=[1] * ph_len) + prompt["additional_information"] = info + return prompt + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: """Validate Fish Speech request parameters. Returns error message or None.""" server = self.ctx.server @@ -54,9 +156,7 @@ class FishSpeechAdapter(ARTTSAdapter): wav_list, sr, cache_key = await server._resolve_ref_audio(request.ref_audio) ref_audio_data = (wav_list, sr) tts_params["ref_audio_cache_key"] = cache_key - prompt = await server._build_fish_speech_prompt_async( - request, ref_audio_data=ref_audio_data, has_inline_ref_audio=has_inline_ref_audio - ) + prompt = await self._build_prompt_async(request, ref_audio_data, has_inline_ref_audio) prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="fish_speech") diff --git a/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py index e65edb01b..22dbfc6b6 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """GLM-TTS serving adapter.""" from typing import TYPE_CHECKING, Any +import numpy as np from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter @@ -19,6 +21,72 @@ class GlmTTSAdapter(ARTTSAdapter): stage_keys = frozenset({"glm_tts"}) name = "glm_tts" + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._text_tokenizer = None + self._text_frontend = None + + def _text_tokenizer_and_frontend(self): + from vllm_omni.model_executor.models.glm_tts.glm_tts import ( + load_glm_tts_tokenizer, + resolve_glm_tts_tokenizer_path, + ) + from vllm_omni.model_executor.models.glm_tts.text_frontend import GLMTTSTextFrontend + + if self._text_tokenizer is None: + config = self.ctx.engine_client.model_config + tokenizer_path = getattr(config, "tokenizer", None) or resolve_glm_tts_tokenizer_path(config.model) + self._text_tokenizer = load_glm_tts_tokenizer( + tokenizer_path, + model_name_or_path=config.model, + trust_remote_code=bool(getattr(config, "trust_remote_code", False)), + ) + if self._text_frontend is None: + self._text_frontend = GLMTTSTextFrontend() + return self._text_tokenizer, self._text_frontend + + def _estimate_text_token_len(self, text: str | None, *, add_trailing_space: bool = False) -> int: + tokenizer, frontend = self._text_tokenizer_and_frontend() + normalized = (frontend.text_normalize(text or "") or text or "").strip() + if add_trailing_space and normalized: + normalized = f"{normalized} " + return max(1, len(tokenizer.encode(normalized))) + + def _build_prefill_metadata(self, text: str, prompt_text: str | None) -> dict[str, Any]: + text_len = self._estimate_text_token_len(text) + prompt_len = self._estimate_text_token_len(prompt_text, add_trailing_space=True) if prompt_text else 0 + return { + "glm_tts_text_token_len": [text_len], + "glm_tts_prompt_text_token_len": [prompt_len], + "input_len": [prompt_len + text_len + 1], + } + + async def _build_prompt( + self, request: "OpenAICreateSpeechRequest", *, has_inline_ref_audio: bool = False + ) -> dict[str, Any]: + """Build GLM-TTS multimodal voice-cloning input and prefill metadata. + + AR preprocess constructs PromptText, Text, BOA and prompt speech tokens; + the DiT consumes the resulting prompt tokens, features and embedding as + conditioning. + """ + server = self.ctx.server + if request.ref_audio is None or not request.ref_text: + raise ValueError("GLM-TTS requires ref_audio and ref_text for voice cloning.") + wav_samples, sr, _ = await server._resolve_ref_audio(request.ref_audio) + mm_kwargs: dict[str, Any] = {"prompt_text": request.ref_text} + if request.voice: + voice_lower = request.voice.lower() + if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: + mm_kwargs["voice_name"] = voice_lower + mm_kwargs["voice_created_at"] = server._voice_created_at(voice_lower) + return { + "prompt": request.input, + "multi_modal_data": {"audio": (np.asarray(wav_samples, dtype=np.float32), int(sr))}, + "mm_processor_kwargs": mm_kwargs, + "additional_information": self._build_prefill_metadata(request.input, request.ref_text), + } + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: """Validate GLM-TTS request — requires ref_audio for voice cloning.""" server = self.ctx.server @@ -46,7 +114,7 @@ class GlmTTSAdapter(ARTTSAdapter): async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: - prompt = await self.ctx.server._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio) + prompt = await self._build_prompt(request, has_inline_ref_audio=has_inline_ref_audio) return PreparedRequest(prompt=prompt, tts_params={}, model_type="glm_tts") def _load_supported_speakers(self) -> set[str]: @@ -71,9 +139,7 @@ class GlmTTSAdapter(ARTTSAdapter): if isinstance(text_len_value, list) and text_len_value: text_len_value = text_len_value[0] text_token_len = ( - int(text_len_value) - if text_len_value is not None - else server._estimate_glm_tts_text_token_len(request.input) + int(text_len_value) if text_len_value is not None else self._estimate_text_token_len(request.input) ) hf_cfg = server.model_config.hf_config min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2) 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 74ec591cc..207d8b47b 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py @@ -1,8 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Higgs-Audio v2 serving adapter.""" +import asyncio from typing import TYPE_CHECKING, Any +import numpy as np +from vllm.inputs import tokens_input + from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens @@ -12,6 +17,89 @@ if TYPE_CHECKING: @register_tts_adapter class HiggsAudioV2Adapter(ARTTSAdapter): + @property + def engine_client(self): + return self.ctx.engine_client + + async def _resolve_ref_audio(self, ref_audio: str): + return await self.ctx.server._resolve_ref_audio(ref_audio) + + async def _build_higgs_audio_v2_params(self, request: "OpenAICreateSpeechRequest"): + """Build prompt_token_ids for higgs_audio_v2 via the upstream processor. + + Plain-text path: runs ``build_plain_text_prompt`` and returns the + token-only prompt. Voice-clone path (``ref_audio`` + ``ref_text``): + resolves the reference clip via ``_resolve_ref_audio``, runs + ``build_voice_clone_prompt`` (which encodes the clip through HF's + ``HiggsAudioV2TokenizerModel`` loaded from the k2-fsa/OmniVoice + ``audio_tokenizer/`` subdirectory), and attaches the encoded + ``audio_input_ids`` + ``audio_input_ids_mask`` tensors via + ``additional_information`` so the talker substitutes them at the + prompt-side audio placeholders. + """ + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( + build_plain_text_prompt, + build_voice_clone_prompt, + input_ids_to_python_list, + ) + + processor = await self._resolve_higgs_audio_v2_processor() + + if request.ref_audio is None: + inputs = await asyncio.to_thread(build_plain_text_prompt, processor, request.input) + prompt_token_ids = input_ids_to_python_list(inputs) + return tokens_input(prompt_token_ids=prompt_token_ids) + + wav_list, sr, _ = await self._resolve_ref_audio(request.ref_audio) + wav = np.asarray(wav_list, dtype=np.float32) + out = await asyncio.to_thread( + build_voice_clone_prompt, + processor, + request.input, + wav, + int(sr), + request.ref_text or "", + ) + prompt = tokens_input(prompt_token_ids=out["prompt_token_ids"]) + # Pass tensors at the top level of additional_information (NOT list- + # wrapped). ``vllm_omni.data_entry_keys.serialize_payload`` routes + # bare ``torch.Tensor`` values through ``_serialize_tensor``; a list + # containing tensors would fall into the ``list_data`` field which + # msgspec cannot serialize and the tensors would be dropped over the + # process boundary (silent voice-clone failure). + prompt["additional_information"] = { + "audio_input_ids": out["audio_input_ids"], + "audio_input_ids_mask": out["audio_input_ids_mask"], + } + return prompt + + async def _resolve_higgs_audio_v2_processor(self): + """Lazy-load the AutoProcessor for higgs_audio_v2 (once per serving instance).""" + cached = getattr(self, "_higgs_audio_v2_processor", None) + if cached is not None: + return cached + + from transformers import AutoProcessor + + model_path = None + for stage in self.engine_client.stage_configs: + model_path = getattr(getattr(stage, "engine_args", None), "model", None) + if model_path: + break + if model_path is None: + # Fallback: the orchestrator stores the served model id on the engine + # itself (set by AsyncOmniEngine.__init__). Stage-level engine_args + # may not surface ``model`` when the deploy yaml doesn't set it per + # stage (the CLI-passed model id is the single source of truth). + model_path = getattr(self.engine_client, "model", None) + if model_path is None: + raise RuntimeError("higgs_audio_v2 serving could not resolve the model path from the engine stage configs") + processor = AutoProcessor.from_pretrained(model_path) + self._higgs_audio_v2_processor = processor + return processor + + # ---- higgs-audio v3 ---- + stage_keys = frozenset({"higgs_audio_v2"}) name = "higgs_audio_v2" @@ -91,7 +179,7 @@ class HiggsAudioV2Adapter(ARTTSAdapter): self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: server = self.ctx.server - prompt = await server._build_higgs_audio_v2_params(request) + prompt = await self._build_higgs_audio_v2_params(request) if request.voice: voice_lower = request.voice.lower() if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: 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 411d97f31..1e64f7a63 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py @@ -1,10 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Higgs-Audio v3 serving adapter.""" -from typing import TYPE_CHECKING, Any +import asyncio +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, cast + +import numpy as np +import torch +from vllm.inputs import tokens_input from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + PreparedRequest, + apply_max_new_tokens, + conditioning_cache_salt, +) + +_REF_CODE_CACHE_MAX_ENTRIES = 256 +_REF_CODE_CACHE_MAX_BYTES = 64 * 1024 * 1024 if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -12,6 +27,163 @@ if TYPE_CHECKING: @register_tts_adapter class HiggsAudioV3Adapter(ARTTSAdapter): + def __init__(self, ctx) -> None: + super().__init__(ctx) + 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]] = {} + + @property + def engine_client(self): + return self.ctx.engine_client + + async def _resolve_ref_audio(self, ref_audio: str): + return await self.ctx.server._resolve_ref_audio(ref_audio) + + def _get_resolved_ref_audio_artifact_key(self, cache_key: str): + return self.ctx.server._get_resolved_ref_audio_artifact_key(cache_key) + + async def _build_higgs_audio_v3_params(self, request: "OpenAICreateSpeechRequest"): + """Build prompt_token_ids for higgs_audio_v3. + + Plain-text path: builds ``[tts, text, tokens, audio]``. + Voice-clone path: encodes reference audio, applies delay pattern, and + records reference-audio prompt positions while submitting valid token IDs. + """ + adapter = await self._resolve_higgs_audio_v3_adapter() + + if request.ref_audio is None: + prompt_ids = adapter.build_prompt(request.input) + return tokens_input(prompt_token_ids=prompt_ids) + + # Voice clone + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + apply_delay_pattern, + encode_reference_audio, + ) + + wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio)) + artifact_key = self._get_resolved_ref_audio_artifact_key(cache_key) + wav = np.asarray(wav_list, dtype=np.float32) + ref_codes_delayed, cache_hit, inflight_wait = await self._resolve_higgs_audio_v3_ref_codes( + artifact_key, + wav, + int(sr), + encode_reference_audio, + apply_delay_pattern, + ) + del cache_hit, inflight_wait + + prompt_ids = adapter.build_prompt( + request.input, + num_ref_tokens=int(ref_codes_delayed.shape[0]), + reference_text=request.ref_text or None, + ) + prompt_ids, audio_placeholder_positions = adapter.prepare_prompt_for_engine(prompt_ids) + prompt = tokens_input(prompt_token_ids=prompt_ids) + import torch + + prompt["additional_information"] = { + "audio_input_ids": ref_codes_delayed.to(torch.long), + "audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool), + "audio_placeholder_positions": audio_placeholder_positions, + "ref_audio_cache_key": cache_key, + } + prompt["cache_salt"] = conditioning_cache_salt(request, prompt["additional_information"]) + return prompt + + async def _resolve_higgs_audio_v3_ref_codes( + self, + artifact_key: str | None, + wav: np.ndarray, + sr: int, + encode_reference_audio, + apply_delay_pattern, + ) -> tuple[torch.Tensor, bool, bool]: + ref_codes_delayed = self._get_higgs_audio_v3_ref_codes(artifact_key) + if ref_codes_delayed is not None: + return ref_codes_delayed, True, False + if not artifact_key: + ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr) + return apply_delay_pattern(ref_codes_raw), False, False + + task = self._higgs_audio_v3_ref_code_inflight.get(artifact_key) + if task is not None: + return (await asyncio.shield(task)).clone(), False, True + + async def _encode_and_cache() -> torch.Tensor: + ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr) + delayed = apply_delay_pattern(ref_codes_raw) + self._put_higgs_audio_v3_ref_codes(artifact_key, delayed) + cached = self._get_higgs_audio_v3_ref_codes(artifact_key) + return cached if cached is not None else delayed.detach().to("cpu", dtype=torch.long).contiguous() + + task = asyncio.create_task(_encode_and_cache()) + self._higgs_audio_v3_ref_code_inflight[artifact_key] = task + + def _retire(t: asyncio.Task[torch.Tensor]) -> None: + if not t.cancelled(): + t.exception() + if self._higgs_audio_v3_ref_code_inflight.get(artifact_key) is t: + self._higgs_audio_v3_ref_code_inflight.pop(artifact_key, None) + + task.add_done_callback(_retire) + return (await asyncio.shield(task)).clone(), False, False + + def _get_higgs_audio_v3_ref_codes(self, artifact_key: str | None) -> torch.Tensor | None: + if not artifact_key: + return None + cached = self._higgs_audio_v3_ref_code_cache.get(artifact_key) + if cached is None: + return None + self._higgs_audio_v3_ref_code_cache.move_to_end(artifact_key) + return cached[0].clone() + + def _put_higgs_audio_v3_ref_codes(self, artifact_key: str, codes: torch.Tensor) -> None: + if _REF_CODE_CACHE_MAX_ENTRIES <= 0 or _REF_CODE_CACHE_MAX_BYTES <= 0 or not artifact_key: + return + cached_codes = codes.detach().to("cpu", dtype=torch.long).contiguous() + size = int(cached_codes.numel() * cached_codes.element_size()) + if size > _REF_CODE_CACHE_MAX_BYTES: + return + previous = self._higgs_audio_v3_ref_code_cache.pop(artifact_key, None) + if previous is not None: + self._higgs_audio_v3_ref_code_cache_bytes -= previous[1] + self._higgs_audio_v3_ref_code_cache[artifact_key] = (cached_codes, size) + self._higgs_audio_v3_ref_code_cache_bytes += size + while len(self._higgs_audio_v3_ref_code_cache) > _REF_CODE_CACHE_MAX_ENTRIES: + _, (_, old_size) = self._higgs_audio_v3_ref_code_cache.popitem(last=False) + self._higgs_audio_v3_ref_code_cache_bytes -= old_size + while self._higgs_audio_v3_ref_code_cache_bytes > _REF_CODE_CACHE_MAX_BYTES: + _, (_, old_size) = self._higgs_audio_v3_ref_code_cache.popitem(last=False) + self._higgs_audio_v3_ref_code_cache_bytes -= old_size + + async def _resolve_higgs_audio_v3_adapter(self): + """Lazy-load the tokenizer adapter for higgs_audio_v3.""" + cached = getattr(self, "_higgs_audio_v3_adapter", None) + if cached is not None: + return cached + + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import ( + HiggsAudioV3TokenizerAdapter, + ) + + model_path = None + for stage in self.engine_client.stage_configs: + model_path = getattr(getattr(stage, "engine_args", None), "model", None) + if model_path: + break + if model_path is None: + model_path = getattr(self.engine_client, "model", None) + if model_path is None: + raise RuntimeError("higgs_audio_v3 serving could not resolve model path") + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + adapter = HiggsAudioV3TokenizerAdapter(tokenizer) + self._higgs_audio_v3_adapter = adapter + return adapter + stage_keys = frozenset({"higgs_audio_v3"}) name = "higgs_audio_v3" @@ -34,7 +206,7 @@ class HiggsAudioV3Adapter(ARTTSAdapter): async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: - prompt = await self.ctx.server._build_higgs_audio_v3_params(request) + prompt = await self._build_higgs_audio_v3_params(request) return PreparedRequest(prompt=prompt, tts_params={}, model_type="higgs_audio_v3") def apply_sampling_overrides( diff --git a/vllm_omni/entrypoints/openai/tts_adapters/ming_flash_omni_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/ming_flash_omni_tts.py index bc084d6b4..05cd3a64e 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/ming_flash_omni_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/ming_flash_omni_tts.py @@ -1,12 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MingFlashOmniTTS serving adapter.""" from typing import TYPE_CHECKING +from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import parse_ming_instruction_fields if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -61,8 +64,36 @@ class MingFlashOmniTTSAdapter(ARTTSAdapter): async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: + """Build the standalone Ming Flash TTS prompt.""" + from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( + DEFAULT_PROMPT, + create_instruction, + ) + server = self.ctx.server - prompt = server._build_ming_flash_omni_prompt(request) + # ``instructions`` accepts either plain style text or a JSON object. + # Unknown caption keys are filtered by ``create_instruction``. + caption_fields = parse_ming_instruction_fields(server, request) or {} + has_spk_emb = request.speaker_embedding is not None + # Standalone TTS uses the Ming ``instruct`` task. ``voice_name`` lets + # the talker resolve registered presets such as DB30. + info = { + "ming_task": "instruct", + "prompt": DEFAULT_PROMPT, + "text": request.input, + "instruction": create_instruction(caption_fields), + "voice_name": request.voice or None, + "use_zero_spk_emb": not has_spk_emb, + "max_decode_steps": request.max_new_tokens or self.max_new_tokens_max, + "cfg": 2.0, + "sigma": 0.25, + "temperature": 0.0, + } + if has_spk_emb: + # Keep embeddings as plain floats for EngineCore serialization. + info["spk_emb"] = list(request.speaker_embedding) + prompt = tokens_input(prompt_token_ids=[0]) + prompt["additional_information"] = info return PreparedRequest(prompt=prompt, tts_params={}, model_type="ming_flash_omni_tts") def _load_supported_speakers(self) -> set[str]: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py index 613968b8c..8dc5731a0 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py @@ -1,12 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Ming-TTS (dense) serving adapter.""" +import json from typing import TYPE_CHECKING, Any +import torch +from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM if TYPE_CHECKING: @@ -15,8 +20,140 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def parse_ming_instruction_fields(server, request, *, include_voice=False, plain_text_passthrough=False): + instruction_text = request.instructions.strip() if isinstance(request.instructions, str) else None + instruction_dict: dict[str, Any] = {} + 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 server.uploaded_speakers): + instruction_dict["IP"] = request.voice + if instruction_text: + try: + parsed = json.loads(instruction_text) + except json.JSONDecodeError: + parsed = None + if isinstance(parsed, dict): + instruction_dict.update(parsed) + elif instruction_dict or not plain_text_passthrough: + instruction_dict["风格"] = instruction_text + else: + return instruction_text + return instruction_dict or None + + @register_tts_adapter class MingTTSAdapter(ARTTSAdapter): + def _coerce_ming_prompt_waveform(self, wav_samples, sample_rate): + from torchaudio.functional import resample as resample_audio + + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import SAMPLE_RATE + + waveform = torch.as_tensor(wav_samples, dtype=torch.float32).reshape(1, -1) + if int(sample_rate) != SAMPLE_RATE: + waveform = resample_audio(waveform, int(sample_rate), SAMPLE_RATE) + return waveform + + def _build_ming_reference_waveforms( + self, + ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None, + ) -> tuple[list[torch.Tensor], list[int]]: + if ref_audio_data is None: + return [], [] + items = ref_audio_data if isinstance(ref_audio_data, list) else [ref_audio_data] + waveforms = [torch.as_tensor(samples, dtype=torch.float32).reshape(1, -1) for samples, _ in items] + sample_rates = [int(sample_rate) for _, sample_rate in items] + if any(sample_rate <= 0 for sample_rate in sample_rates): + raise ValueError(f"Ming reference audio sample rates must be positive, got {sample_rates}") + return waveforms, sample_rates + + def _parse_ming_instruction(self, request: "OpenAICreateSpeechRequest") -> Any: + """Build a Ming instruction payload from OpenAI speech fields.""" + return parse_ming_instruction_fields( + self.ctx.server, + request, + include_voice=True, + plain_text_passthrough=True, + ) + + def _build_ming_dense_prompt( + self, + request: "OpenAICreateSpeechRequest", + *, + ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None = None, + voice_name: str | None = None, + voice_created_at: int = 0, + ) -> dict[str, Any]: + """Build a Ming dense prompt directly from the OpenAI speech request.""" + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + KEY_MAX_DECODE_STEPS, + KEY_SPEAKER_SAMPLE_RATES, + KEY_SPEAKER_WAVEFORM, + KEY_SPEAKER_WAVEFORM_LENGTHS, + ) + from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt + + if self.ctx.server._tts_tokenizer is None: + model_name = self.ctx.engine_client.model_config.model + trust_remote_code = bool(getattr(self.ctx.engine_client.model_config, "trust_remote_code", False)) + self.ctx.server._tts_tokenizer = AutoTokenizer.from_pretrained( + model_name, trust_remote_code=trust_remote_code + ) + + ref_text = request.ref_text + reference_waveforms, reference_sample_rates = self._build_ming_reference_waveforms(ref_audio_data) + reference_waveform = torch.cat(reference_waveforms, dim=-1) if reference_waveforms else None + prompt_waveforms = ( + [ + self._coerce_ming_prompt_waveform(waveform, sample_rate) + for waveform, sample_rate in zip(reference_waveforms, reference_sample_rates, strict=True) + ] + if ref_text is not None + else [] + ) + prompt_waveform = torch.cat(prompt_waveforms, dim=-1) if prompt_waveforms else None + speaker_embedding = request.speaker_embedding + pending_speaker_extraction = speaker_embedding is None and reference_waveform is not None + use_zero_spk_emb = not pending_speaker_extraction and prompt_waveform is None and speaker_embedding is None + + runtime_controls = {} + if request.max_new_tokens is not None: + runtime_controls[KEY_MAX_DECODE_STEPS] = request.max_new_tokens + + prompt_dict = build_ming_dense_prompt( + self.ctx.server._tts_tokenizer, + # bgm / music-prompt mode not supported online; + # requires prompt_mode API extension (deferred). + prompt=MING_DEFAULT_PROMPT, + text=request.input, + runtime_controls=runtime_controls or None, + instruction=self._parse_ming_instruction(request), + prompt_text=ref_text, + prompt_waveform=prompt_waveform, + speaker_embedding=speaker_embedding, + speaker_count=len(reference_waveforms) if pending_speaker_extraction else None, + use_zero_spk_emb=use_zero_spk_emb, + ) + additional_information = prompt_dict["additional_information"] + if pending_speaker_extraction: + additional_information[KEY_SPEAKER_WAVEFORM] = reference_waveform + additional_information[KEY_SPEAKER_WAVEFORM_LENGTHS] = torch.tensor( + [int(waveform.shape[-1]) for waveform in reference_waveforms], + dtype=torch.int32, + ) + additional_information[KEY_SPEAKER_SAMPLE_RATES] = torch.tensor( + reference_sample_rates, + dtype=torch.int32, + ) + if voice_name: + additional_information["voice_name"] = voice_name + additional_information["voice_created_at"] = int(voice_created_at) + prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"]) + prompt["prompt"] = prompt_dict["prompt"] + prompt["text"] = prompt_dict["text"] + prompt["additional_information"] = prompt_dict["additional_information"] + return prompt + # Ming dense has no dedicated model_stage value, so both stage discovery and # model-type detection go through the architecture. ``ming_flash_omni_tts`` # is the model that owns the ``ming_tts`` *stage key*, and is matched first. @@ -66,7 +203,7 @@ class MingTTSAdapter(ARTTSAdapter): elif ref_audio_source is not None and isinstance(ref_audio_source, str): wav_list, sr, _ = await server._resolve_ref_audio(ref_audio_source) ref_audio_data = (wav_list, sr) - prompt = server._build_ming_dense_prompt( + prompt = self._build_ming_dense_prompt( request, ref_audio_data=ref_audio_data, voice_name=uploaded_audio_voice, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index 8f3ee8c08..07e4bb197 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -2,11 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MOSS-TTS serving adapters (Nano + full family). -Variants share validation and request metadata. Local uses transcript-guided -continuation when the request includes a reference transcript. +Both variants share the same build/validate flow (``_build_moss_tts_params`` +handles each); they are registered under distinct model-type names. """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, cast from vllm.inputs import tokens_input @@ -23,34 +23,261 @@ if TYPE_CHECKING: class _MossTTSAdapterBase(ARTTSAdapter): - async def _build_local_continuation_params( - self, request: "OpenAICreateSpeechRequest", *, has_inline_ref_audio: bool - ) -> dict: - import torch + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._moss_variant = None if self.name == "moss_tts_nano" else self._detect_moss_variant() + self._moss_processor_cache = None - server = self.ctx.server - processor = server._get_moss_processor() - reference, resolve_keys = await server._encode_moss_references( - request, has_inline_ref_audio=has_inline_ref_audio, two_speaker=False + @property + def engine_client(self): + return self.ctx.engine_client + + @property + def uploaded_speakers(self): + return self.ctx.server.uploaded_speakers + + @property + def _speaker_cache(self): + return self.ctx.server._speaker_cache + + async def _resolve_ref_audio(self, ref_audio: str): + return await self.ctx.server._resolve_ref_audio(ref_audio) + + def _voice_created_at(self, voice: str) -> int: + return self.ctx.server._voice_created_at(voice) + + def _get_resolved_ref_audio_artifact_key(self, cache_key: str): + return self.ctx.server._get_resolved_ref_audio_artifact_key(cache_key) + + def _get_moss_ref_encoder(self): + """Lazily build the per-server MOSS reference-audio encoder once.""" + cached = getattr(self, "_moss_ref_encoder", None) + if cached is not None: + return cached + from vllm_omni.model_executor.models.moss_tts.reference_encoder import build_reference_encoder + + # The variant's encode geometry (n_vq, working sample rate) is derived + # inside the model package; this layer only supplies the processor and + # the process-wide speaker cache. + encoder = build_reference_encoder( + self._get_moss_processor(), + variant=cast(str, self._moss_variant), + speaker_cache=self._speaker_cache, ) - assert request.ref_text is not None - user_kwargs = {"text": request.ref_text.strip() + " " + (request.input or "")} - if request.language: + self._moss_ref_encoder = encoder + return encoder + + async def _encode_moss_references( + self, + request: "OpenAICreateSpeechRequest", + *, + has_inline_ref_audio: bool, + two_speaker: bool, + ) -> tuple[list, dict[int, str]]: + """Encode the request's reference clip(s) into MOSS RVQ code tensors. + + Reference encoding + speaker caching lives in the model package + (moss_tts.reference_encoder), mirroring Fish Speech / CosyVoice3 / + Qwen3-TTS which keep reference handling with the model rather than in + this shared serving file. This method only resolves the serving-side + inputs the model helper needs: the generic audio resolver/artifact-key + lookups and whether the request may use the named-voice cache. + + Returns ``(codes_per_speaker, resolve_keys)`` where ``resolve_keys`` + maps the reference slot (0 = ref_audio, 1 = ref_audio_2) to its + content-aware resolve key, for salting the KV prefix cache. + """ + from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_request_references + + # Named-voice caching is only valid for uploaded speakers without an + # inline ref_audio: ``request.voice`` plus a file/URL would otherwise + # key on (name, created_at=0) and skip the content-aware resolve. + raw_voice = getattr(request, "voice", None) + raw_voice = raw_voice.strip() if isinstance(raw_voice, str) else "" + voice_lower = raw_voice.lower() + use_named_voice = bool(voice_lower) and voice_lower in self.uploaded_speakers and not has_inline_ref_audio + voice = voice_lower if use_named_voice else "" + voice_created = self._voice_created_at(voice) if voice else 0 + + return await encode_request_references( + self._get_moss_ref_encoder(), + cast(str, request.ref_audio), + request.ref_audio_2 if two_speaker else None, + resolve_ref_audio=self._resolve_ref_audio, + get_artifact_key=self._get_resolved_ref_audio_artifact_key, + voice_name=voice or None, + voice_created_at=voice_created, + ) + + 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). + + Detection key is the HF repo path / model_name; matches + ``_try_resolve_omni_model_type`` in entrypoints/utils.py so users get + consistent behaviour no matter how they launched the server (--model + OpenMOSS-Team/MOSS-TTSD-v1.0 vs --deploy-config moss_ttsd.yaml). + """ + try: + name = (self.engine_client.model_config.model or "").lower().replace("-", "").replace("_", "") + except Exception: + name = "" + if "realtime" in name: + return "realtime" + if "local" in name: + return "local" + if "ttsd" in name: + return "ttsd" + if "soundeffect" in name: + return "sound_effect" + if "voicegenerator" in name: + return "voice_generator" + return "tts" + + def _get_moss_processor(self): + """Lazily load the upstream MOSS-TTS processor once per server. + + Cached on ``self._moss_processor_cache``. The processor owns its own + audio_tokenizer (~1.6 B params); we keep it on CPU so it doesn't + compete with the talker (~8 GiB) and codec (~7 GiB) for our 96 GiB + GPU — per-request ref-audio encoding is fast enough on CPU. + """ + cached = getattr(self, "_moss_processor_cache", None) + if cached is not None: + return cached + from transformers import AutoProcessor + + model_id = self.engine_client.model_config.model + proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) + if hasattr(proc, "audio_tokenizer"): + proc.audio_tokenizer = proc.audio_tokenizer.to("cpu").eval() + self._moss_processor_cache = proc + return proc + + async def _build_moss_tts_params( + self, + request: "OpenAICreateSpeechRequest", + *, + has_inline_ref_audio: bool = False, + ) -> dict[str, Any]: + """Build the talker prompt + ``additional_information`` payload for any + MOSS-TTS-family request (nano + 5 full variants). + + For the legacy ``moss_tts_nano`` model_type, keeps the original nano + contract (``{text, mode=voice_clone, prompt_audio_array}``); the + caller still uses a ``[1]`` placeholder prompt. + + For the full MOSS-TTS family (``MossTTSDelay*`` / ``MossTTSRealtime*``) + we **call the upstream processor server-side** to produce the unified + ``(text_ids, audio_codes)`` shape the talker actually consumes — same + flow as ``examples/.../moss_tts/end2end.py:_build_unified_codes``. + Returns ``{prompt_token_ids: list[int], codes.ref: torch.LongTensor, + max_new_frames, ...}``. The caller treats ``prompt_token_ids`` as the + prompt and forwards the rest as ``additional_information``. + """ + import torch # local to avoid pulling torch at module import time + + v = self._moss_variant + + params: dict[str, Any] + # ---- Legacy nano path (unchanged) ---- + if v is None: # moss_tts_nano + params = { + "text": [request.input or ""], + "mode": ["voice_clone"], + } + if request.max_new_tokens is not None: + params["max_new_frames"] = [request.max_new_tokens] + wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio)) + params["prompt_audio_array"] = [[wav_list, sr]] + params["ref_audio_cache_key"] = cache_key + return params + + # ---- MOSS-TTS-Realtime: keep the old prompt_audio_array path ---- + # ``AutoProcessor.from_pretrained`` doesn't auto-discover + # ``MossTTSRealtimeProcessor`` (no ``processor_config.json`` in the + # snapshot), and Realtime's prompt format diverges from MossTTSDelay + # (16-channel grid, separate per-step text feed). The + # ``prompt_audio_array`` shape lines up well enough with what the + # talker reads for short prompts; full Realtime support needs a + # separate processor.from_module path which we don't wire here. + if v == "realtime": + params = { + "text": [request.input or ""], + "mode": ["voice_clone"], + } + if request.max_new_tokens is not None: + params["max_new_frames"] = [request.max_new_tokens] + wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio)) + params["prompt_audio_array"] = [[wav_list, sr]] + params["ref_audio_cache_key"] = cache_key + return params + + # ---- 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() + + user_kwargs: dict[str, Any] = {"text": request.input or ""} + resolve_keys: dict[int, str] = {} + if v in ("tts", "local", "ttsd"): + user_kwargs["reference"], resolve_keys = await self._encode_moss_references( + request, + has_inline_ref_audio=has_inline_ref_audio, + two_speaker=(v == "ttsd"), + ) + elif v == "sound_effect": + user_kwargs["text"] = request.input or "" # may be empty + user_kwargs["ambient_sound"] = request.ambient_sound or "" + if request.duration_seconds is not None: + user_kwargs["tokens"] = max(1, int(float(request.duration_seconds) * 12.5)) + elif request.max_new_tokens is not None: + user_kwargs["tokens"] = int(request.max_new_tokens) + elif v == "voice_generator": + user_kwargs["instruction"] = request.instructions or "" + + # Optional language tag for the spoken-text variants. MOSS-TTS-v1.5's + # headline improvement is multilingual synthesis when the language is + # given (build_user_message(..., language=...)); 1.0 ignores it + # gracefully. Sound-effect output is non-verbal, so skip it there. + if v in ("tts", "local", "ttsd", "voice_generator") and getattr(request, "language", None): user_kwargs["language"] = request.language - # Align the reference transcript with an assistant audio prefix. Only - # newly generated frames go to the codec; the prefix is conditioning. - user = processor.build_user_message(**user_kwargs) - assistant = processor.build_assistant_message(audio_codes_list=reference) - batch = processor(conversations=[[user, assistant]], mode="continuation") - unified = batch["input_ids"][0] + + # Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special + # tokens and cols 1..n_vq are the delay-pattern audio code grid (mostly + # audio_pad_code outside the reference block). + if v == "local" and request.ref_text and request.ref_text.strip(): + reference = user_kwargs.pop("reference") + user_kwargs["text"] = request.ref_text.strip() + " " + (request.input or "") + user_msg = proc.build_user_message(**user_kwargs) + assistant_msg = proc.build_assistant_message(audio_codes_list=reference) + batch = proc(conversations=[[user_msg, assistant_msg]], mode="continuation") + else: + user_msg = proc.build_user_message(**user_kwargs) + batch = proc(conversations=[[user_msg]], mode="generation") + unified = batch["input_ids"][0] # torch.LongTensor (L, 1+n_vq) + text_ids: list[int] = unified[:, 0].tolist() + audio_codes: torch.Tensor = unified[:, 1:].contiguous().to(torch.int64) + params = { - "prompt_token_ids": unified[:, 0].tolist(), - "codes": {"ref": unified[:, 1:].contiguous().to(torch.int64)}, + "prompt_token_ids": text_ids, + "codes": {"ref": audio_codes}, } if request.max_new_tokens is not None: params["max_new_frames"] = [request.max_new_tokens] if 0 in resolve_keys: params["ref_audio_cache_key"] = resolve_keys[0] + if 1 in resolve_keys: + params["ref_audio_2_cache_key"] = resolve_keys[1] return params def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: @@ -72,10 +299,10 @@ class _MossTTSAdapterBase(ARTTSAdapter): if not request.input or not request.input.strip(): # SoundEffect can legitimately have empty input (just ambient_sound). - if server._moss_variant != "sound_effect": + if self._moss_variant != "sound_effect": return "Input text cannot be empty" - v = server._moss_variant + v = self._moss_variant if v in (None, "tts", "realtime", "local"): if request.ref_audio is None: label = ( @@ -122,10 +349,7 @@ class _MossTTSAdapterBase(ARTTSAdapter): self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: server = self.ctx.server - if server._moss_variant == "local" and request.ref_text and request.ref_text.strip(): - tts_params = await self._build_local_continuation_params(request, has_inline_ref_audio=has_inline_ref_audio) - else: - tts_params = await server._build_moss_tts_params(request, has_inline_ref_audio=has_inline_ref_audio) + tts_params = await self._build_moss_tts_params(request, has_inline_ref_audio=has_inline_ref_audio) if request.voice: voice_lower = request.voice.lower() if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py index 48d019eeb..3a2bb9860 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py @@ -2,12 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Qwen3-TTS serving adapter.""" +import math from collections.abc import Mapping from pathlib import Path from typing import TYPE_CHECKING, Any import regex as re +from vllm.inputs import tokens_input from vllm.logger import init_logger +from vllm.utils.async_utils import make_async from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ( @@ -15,6 +18,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ( ARTTSAdapter, PreparedRequest, TTSGenerationError, + conditioning_cache_salt, ) from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_precomputed_speakers from vllm_omni.utils.speaker_cache import validate_qwen3_tts_profile @@ -23,6 +27,7 @@ if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest logger = init_logger(__name__) +_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key" QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY = "_qwen3_tts_effective_max_tokens" _MIN_CODEC_FRAMES = 192 @@ -45,6 +50,33 @@ class Qwen3TTSAdapter(ARTTSAdapter): name = "qwen3_tts" supported_output_sample_rates = frozenset({8000, 24000}) + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._estimate_prompt_len_async = make_async( + self._estimate_prompt_len, executor=getattr(ctx.server, "_tts_executor", None) + ) + + def _estimate_ref_code_len(self, ref_audio: object) -> int | None: + codec_frame_rate = self.capabilities.codec_frame_rate + if codec_frame_rate is None: + return None + try: + item = ref_audio + while isinstance(item, list) and item: + if len(item) == 2 and isinstance(item[1], (int, float)): + break + item = item[0] + if not isinstance(item, (list, tuple)) or len(item) != 2: + return None + wav, sr = item + sr = int(sr) + n_samples = len(wav) if hasattr(wav, "__len__") else wav.shape[-1] + if sr <= 0 or n_samples <= 0: + return None + return math.ceil(n_samples / sr * codec_frame_rate) + except Exception: + return None + def _get_model_variant(self) -> str | None: """Return the task supported by the loaded Qwen3-TTS checkpoint. @@ -223,15 +255,208 @@ class Qwen3TTSAdapter(ARTTSAdapter): return None + def _build_tts_params(self, request: "OpenAICreateSpeechRequest") -> dict[str, Any]: + """Build TTS parameters from request. + + Processes each parameter if present, skips if not. + Values are wrapped in lists as required by the model. + """ + server = self.ctx.server + + params: dict[str, Any] = {} + + # Text content (always required) + params["text"] = [request.input] + + # Task type + if request.task_type is not None: + params["task_type"] = [request.task_type] + else: + params["task_type"] = ["CustomVoice"] + + # Language + if request.language is not None: + params["language"] = [request.language] + else: + params["language"] = ["Auto"] + + # Speaker (voice) + if request.voice is not None: + voice_lower = request.voice.lower() + precomputed_speakers = self.capabilities.precomputed_speakers + params["speaker"] = [request.voice] + params["voice_created_at"] = [server._voice_created_at(voice_lower)] + + # Uploaded voices use task_type="Base" (CustomVoice requires built-in spk_id). + # If ref_text was provided at upload time, use in-context cloning; otherwise x_vector only. + if voice_lower in server.uploaded_speakers and request.ref_audio is None: + speaker_info = server.uploaded_speakers[voice_lower] + + # Check if this voice was uploaded with a pre-computed embedding. + # Populate request.speaker_embedding so the existing code path + # (below) handles voice_clone_prompt and x_vector_only_mode. + embedding = server._get_uploaded_speaker_embedding(request.voice) + if embedding is not None: + request.speaker_embedding = embedding + params["speaker"] = [voice_lower] + params["task_type"] = ["Base"] + logger.info("Auto-set speaker_embedding for uploaded voice: %s", request.voice) + else: + audio_data = server._get_uploaded_audio_data(request.voice) + if not audio_data: + raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing or corrupted") + stored_ref_text = speaker_info.get("ref_text") + params["speaker"] = [voice_lower] + params["ref_audio"] = [audio_data] + params["task_type"] = ["Base"] + if stored_ref_text: + params["ref_text"] = [stored_ref_text] + params["x_vector_only_mode"] = [False] + else: + params["x_vector_only_mode"] = [True] + logger.info( + "Auto-set ref_audio for uploaded voice: %s (icl=%s)", request.voice, bool(stored_ref_text) + ) + elif voice_lower in precomputed_speakers and request.ref_audio is None: + profile = precomputed_speakers[voice_lower] + mode = str(profile.get("mode") or "xvec").lower() + params["speaker"] = [voice_lower] + params["task_type"] = ["Base"] + params["x_vector_only_mode"] = [mode != "icl"] + ref_text = request.ref_text or profile.get("ref_text") + if isinstance(ref_text, str) and ref_text.strip(): + params["ref_text"] = [ref_text] + ref_code_length = profile.get("ref_code_length") + if mode == "icl" and ref_code_length: + params["ref_code_length"] = [int(ref_code_length)] + logger.info("Using precomputed Qwen3-TTS custom voice profile: %s (mode=%s)", voice_lower, mode) + + elif params["task_type"][0] == "CustomVoice": + params["speaker"] = ["Vivian"] # Default for CustomVoice + + # Instructions for style/emotion control + if request.instructions is not None: + params["instruct"] = [request.instructions] + else: + params["instruct"] = [""] + + # Voice clone: ref_audio resolved in create_speech(), not here. + if request.ref_text is not None: + params["ref_text"] = [request.ref_text] + if request.speaker_embedding is not None: + # Store as plain float list (not tensor) so it survives msgspec + # serialization through the EngineCore IPC boundary. The talker's + # _build_prompt_embeds converts it back to a tensor on the GPU. + params["voice_clone_prompt"] = [ + { + "ref_spk_embedding": list(request.speaker_embedding), + } + ] + # speaker_embedding implies x_vector_only_mode + params["x_vector_only_mode"] = [True] + elif request.x_vector_only_mode is not None: + params["x_vector_only_mode"] = [request.x_vector_only_mode] + + # Generation parameters + if request.max_new_tokens is not None: + params["max_new_tokens"] = [request.max_new_tokens] + else: + params["max_new_tokens"] = [2048] + + if request.initial_codec_chunk_frames is not None: + params["initial_codec_chunk_frames"] = [request.initial_codec_chunk_frames] + + if request.non_streaming_mode is not None: + params["non_streaming_mode"] = [request.non_streaming_mode] + # Preserve the legacy VoiceDesign fallback when the request omits an + # explicit override. CustomVoice and Base rely on model defaults + # (True and False respectively). + elif params["task_type"][0] == "VoiceDesign": + params["non_streaming_mode"] = [True] + + return params + + def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int: + """Estimate prompt length so the placeholder matches model-side embeddings.""" + try: + from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + Qwen3TTSPromptEmbedsBuilder, + ) + + server = self.ctx.server + if server._tts_tokenizer is None: + from transformers import AutoTokenizer + + model_name = self.ctx.engine_client.model_config.model + server._tts_tokenizer = AutoTokenizer.from_pretrained( + model_name, + trust_remote_code=True, + padding_side="left", + ) + hf_config = self.ctx.engine_client.model_config.hf_config + talker_config = hf_config.talker_config + task_type = (tts_params.get("task_type") or ["CustomVoice"])[0] + return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( + additional_information=tts_params, + task_type=task_type, + tokenize_prompt=lambda t: server._tts_tokenizer(t, padding=False)["input_ids"], + codec_language_id=getattr(talker_config, "codec_language_id", None), + spk_is_dialect=getattr(talker_config, "spk_is_dialect", None), + estimate_ref_code_len=self._estimate_ref_code_len, + ) + except Exception as e: + logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e) + return 2048 + + def _qwen3_tts_can_use_ref_audio_artifact_only(self, tts_params: dict[str, Any], artifact_key: str | None) -> bool: + server = self.ctx.server + x_vector_only = server._tts_x_vector_only(tts_params) + if not artifact_key or (artifact_key, x_vector_only) not in server._ref_audio_model_artifact_ready: + return False + return (tts_params.get("task_type") or ["CustomVoice"])[0] == "Base" + async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: - prompt, tts_params, warmup_key = await self.ctx.server._build_qwen3_tts_request(request) + """Build prompt + tts_params for Qwen3-TTS. + + Called from ``Qwen3TTSAdapter.build``. Returns + ``(prompt, tts_params, warmup_artifact_key)`` where the warmup key is the + Qwen3-TTS ref-audio artifact tracked after ``generate()``. + """ + server = self.ctx.server + qwen3_ref_audio_warmup_artifact_key: str | None = None + tts_params = self._build_tts_params(request) + # Resolve ref_audio (explicit or auto-set for uploaded voices) + # to [[wav_list, sr]] so the model doesn't re-decode base64. + ref_audio_source = request.ref_audio + if ref_audio_source is None and isinstance(tts_params.get("ref_audio"), list): + # Uploaded voice: ref_audio was auto-set as [base64_data_url] + ref_audio_source = tts_params["ref_audio"][0] + if ref_audio_source is not None and isinstance(ref_audio_source, str): + wav_list, sr, cache_key = await server._resolve_ref_audio(ref_audio_source) + tts_params["ref_audio_cache_key"] = cache_key + artifact_key = server._get_resolved_ref_audio_artifact_key(cache_key) + if artifact_key: + tts_params[_REF_AUDIO_CACHE_KEY] = [artifact_key] + ref_code_length = self._estimate_ref_code_len([wav_list, sr]) + if ref_code_length is not None: + tts_params["ref_code_length"] = [int(ref_code_length)] + if self._qwen3_tts_can_use_ref_audio_artifact_only(tts_params, artifact_key): + logger.debug("Using Qwen3-TTS ref_audio artifact-only path: %s", artifact_key) + else: + tts_params["ref_audio"] = [[wav_list, sr]] + qwen3_ref_audio_warmup_artifact_key = artifact_key + + ph_len = await self._estimate_prompt_len_async(tts_params) + prompt = tokens_input(prompt_token_ids=[1] * ph_len) + prompt["additional_information"] = tts_params + prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) return PreparedRequest( prompt=prompt, tts_params=tts_params, model_type=tts_params.get("task_type", ["unknown"])[0], - warmup_artifact_key=warmup_key, + warmup_artifact_key=qwen3_ref_audio_warmup_artifact_key, ) def _get_expected_speaker_embedding_dim(self) -> int: @@ -340,6 +565,16 @@ class Qwen3TTSAdapter(ARTTSAdapter): if isinstance(additional_information, dict) and effective_cap is not None: additional_information[QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] = [effective_cap] + # Propagate the deploy/YAML stage seed to residual MTP sampling. The + # generic serving layer applies an explicit request.seed afterwards, + # so the request value still has higher precedence. + stage0_params = sampling_params_list[0] + default_seed = getattr(stage0_params, "seed", None) + if default_seed is not None: + if stage0_params.extra_args is None: + stage0_params.extra_args = {} + stage0_params.extra_args.setdefault("tts_local_seed", int(default_seed)) + logger.debug( "Qwen3-TTS codec budget: task_type=%s text_tokens=%s dynamic_cap=%s " "configured_cap=%s request_cap=%s effective_cap=%s", diff --git a/vllm_omni/entrypoints/openai/tts_adapters/step_audio2.py b/vllm_omni/entrypoints/openai/tts_adapters/step_audio2.py index 2c720d8e1..7ca5a44ff 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/step_audio2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/step_audio2.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Step-Audio2 serving adapter.""" from typing import TYPE_CHECKING @@ -30,5 +31,17 @@ class StepAudio2Adapter(ARTTSAdapter): async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: - prompt = self.ctx.server._build_step_audio2_prompt(request) + """Build a chat prompt ending at ````. + + The assistant turn deliberately omits ``<|im_end|>`` so the thinker + continues by generating audio tokens. + """ + system_prompt = request.instructions or "You are a voice assistant. Read the text aloud." + prompt = { + "prompt": ( + f"<|im_start|>system\n{system_prompt}<|im_end|>\n" + f"<|im_start|>user\n{request.input}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + } return PreparedRequest(prompt=prompt, tts_params={}, model_type="step_audio2") diff --git a/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py b/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py index 760bb3339..f4574004e 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """VoxCPM2 serving adapter (AR base-LM + diffusion side-computation).""" import time from typing import Any +import numpy as np from vllm.logger import init_logger from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -28,6 +30,63 @@ class VoxCPM2Adapter(ARTTSAdapter): # claims still resolves to VoxCPM2. detect_priority = 10 + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._tokenizer = None + self._split_map: dict[int, list[int]] = {} + + def _encode(self, text: str) -> list[int]: + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + build_cjk_split_map, + split_multichar_chinese, + ) + + if self._tokenizer is None: + self._tokenizer = AutoTokenizer.from_pretrained( + self.ctx.engine_client.model_config.model, trust_remote_code=True + ) + self._split_map = build_cjk_split_map(self._tokenizer) + logger.info("VoxCPM2 serving: built multichar split map (%d entries)", len(self._split_map)) + ids = self._tokenizer.encode(text, add_special_tokens=True) + return split_multichar_chinese(ids, self._split_map) + + async def _build_prompt( + self, + request: "OpenAICreateSpeechRequest", + uploaded_ref: tuple[np.ndarray, int] | None = None, + ) -> dict[str, Any]: + """Build the full-length VoxCPM2 prefill prompt. + + ``uploaded_ref`` supplies an uploaded voice waveform when the request + has no explicit ``ref_audio`` so prefill-length accounting includes it. + """ + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import build_voxcpm2_prompt + + server = self.ctx.server + self._encode("") + ref_audio = None + ref_sr = None + voice_profile = None + if request.ref_audio is not None: + ref_audio, ref_sr, _ = await server._resolve_ref_audio(request.ref_audio) + elif uploaded_ref is not None: + wav_np, ref_sr = uploaded_ref + ref_audio = wav_np.tolist() + elif request.voice is not None: + voice_profile = self.capabilities.precomputed_speakers.get(request.voice.lower()) + return build_voxcpm2_prompt( + hf_config=self.ctx.engine_client.model_config.hf_config, + tokenizer=self._tokenizer, + split_map=self._split_map, + text=request.input, + ref_audio=ref_audio, + ref_sr=ref_sr, + ref_text=request.ref_text, + voice_profile=voice_profile, + ) + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: """Validate VoxCPM2 request parameters. Returns error message or None.""" server = self.ctx.server @@ -66,7 +125,7 @@ class VoxCPM2Adapter(ARTTSAdapter): ) if request.ref_audio is None: uploaded_ref = server._load_uploaded_audio(voice_lower) - prompt = await server._build_voxcpm2_prompt(request, uploaded_ref=uploaded_ref) + prompt = await self._build_prompt(request, uploaded_ref=uploaded_ref) tts_params = {} if request.voice: voice_lower = request.voice.lower() diff --git a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py index c9bea7bb0..f9d7639ca 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py @@ -1,8 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Voxtral TTS serving adapter.""" from typing import TYPE_CHECKING, Any +from vllm.utils.async_utils import make_async + from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers @@ -16,6 +19,39 @@ class VoxtralTTSAdapter(ARTTSAdapter): stage_keys = frozenset({"audio_generation"}) name = "voxtral_tts" + def __init__(self, ctx) -> None: + super().__init__(ctx) + self._build_prompt_async = make_async(self._build_prompt, executor=getattr(ctx.server, "_tts_executor", None)) + + def _build_prompt(self, request: "OpenAICreateSpeechRequest") -> dict[str, Any]: + """Build a Voxtral prompt for either a preset voice or inline reference audio.""" + from mistral_common.protocol.speech.request import SpeechRequest + from vllm.inputs import tokens_input + + ref_audio = request.ref_audio + if not request.voice and not ref_audio: + raise ValueError("Voxtral requires either a voice name or ref_audio.") + if isinstance(ref_audio, str) and ref_audio.startswith("data:"): + _, _, ref_audio = ref_audio.partition(",") + server = self.ctx.server + if server._tts_tokenizer is None: + from vllm.tokenizers import cached_tokenizer_from_config + + server._tts_tokenizer = cached_tokenizer_from_config(self.ctx.engine_client.model_config).instruct + if request.voice is not None: + tokenized = server._tts_tokenizer.encode_speech_request( + SpeechRequest(input=request.input, voice=request.voice) + ) + prompt = tokens_input(prompt_token_ids=tokenized.tokens) + prompt["additional_information"] = {"voice": [request.voice]} + return prompt + tokenized = server._tts_tokenizer.encode_speech_request(SpeechRequest(input=request.input, ref_audio=ref_audio)) + audio = tokenized.audios[0] + return { + "prompt_token_ids": tokenized.tokens, + "multi_modal_data": {"audio": [(audio.audio_array, audio.sampling_rate)]}, + } + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: """Validate Voxtral TTS request parameters. Returns error message or None.""" server = self.ctx.server @@ -48,7 +84,7 @@ class VoxtralTTSAdapter(ARTTSAdapter): async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: - prompt = await self.ctx.server._build_voxtral_prompt_async(request) + prompt = await self._build_prompt_async(request) return PreparedRequest(prompt=prompt, tts_params={}, model_type="voxtral_tts") def _load_supported_speakers(self) -> set[str]: diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 09a3d284d..35a39b2ea 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -71,6 +71,21 @@ class VideoFrames(list[Image.Image]): self.source_path = source_path +class _ImagePixelLimitError(InvalidInputReferenceError): + """An image exceeded a configured or decoder-enforced pixel limit.""" + + +def _validate_image_pixel_limit(image: Image.Image) -> None: + width, height = image.size + max_pixels = envs.VLLM_MAX_IMAGE_PIXELS + if max_pixels > 0 and width * height > max_pixels: + raise _ImagePixelLimitError( + f"Image dimensions {width}x{height} ({width * height} pixels) exceed " + f"the maximum of {max_pixels} pixels. Set " + f"VLLM_MAX_IMAGE_PIXELS to increase this limit." + ) + + def positive_float(value: Any) -> float | None: if value is None: return None @@ -87,7 +102,13 @@ def positive_float(value: Any) -> float | None: def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image: try: - return Image.open(BytesIO(image_bytes)).convert("RGB") + with Image.open(BytesIO(image_bytes)) as image: + _validate_image_pixel_limit(image) + return image.convert("RGB") + except _ImagePixelLimitError: + raise + except Image.DecompressionBombError as exc: + raise _ImagePixelLimitError(f"Invalid {source}: image exceeds the decoder pixel limit.") from exc except (UnidentifiedImageError, OSError, ValueError) as exc: raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc @@ -158,6 +179,8 @@ def _decode_media_bytes( ) -> Image.Image | VideoFrames: try: return _decode_image_bytes(media_bytes, source=source) + except _ImagePixelLimitError: + raise except InvalidInputReferenceError: try: return _decode_video_bytes( diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py index b65264583..9b8bfe0d8 100644 --- a/vllm_omni/entrypoints/openpi/connection.py +++ b/vllm_omni/entrypoints/openpi/connection.py @@ -17,6 +17,13 @@ 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} + +`kind` is required on the `nd` markers because it is what separates them from a +plain user mapping, but its value is dialect-specific: the vLLM-native markers +carry the dtype kind character while the `msgpack-numpy` package leaves it empty +for plain dtypes and sets "V" for structured ones. Both spellings are accepted. +Scalars packed by `msgpack-numpy` omit `kind` entirely and are therefore left as +mappings. """ from __future__ import annotations @@ -83,9 +90,14 @@ def _decode_vllm_numpy_marker(obj: dict[Any, Any]) -> Any: 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: + # Structured markers carry a dtype descriptor list in `type`, so reject them + # before `np.dtype()` sees something it cannot parse. + if kind_text == "V": + raise ValueError("Unsupported dtype: structured arrays") + + dtype_obj = np.dtype(_decode_marker_text(dtype)) + if kind_text not in ("", dtype_obj.kind): 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}") diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index 4f1d58081..f2448dfa1 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Serving layer for robot policy inference via `/v1/realtime/robot/openpi`. @@ -12,7 +12,7 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass from itertools import count -from typing import Any +from typing import Any, TypeAlias import numpy as np from omegaconf import OmegaConf @@ -20,7 +20,7 @@ from vllm.logger import init_logger logger = init_logger(__name__) -ActionOutput = np.ndarray | dict[str, np.ndarray] +ActionOutput: TypeAlias = np.ndarray | dict[str, np.ndarray] def _to_builtin_container(value: Any) -> Any: @@ -150,6 +150,9 @@ class ServingRealtimeRobotOpenPI: from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniDiffusionSamplingParams + # An optional integer ``seed`` in the inference message becomes the engine request's + # ``sampling_params.seed``; omitted, ``OmniDiffusionRequest`` assigns a random one. + seed = obs.pop("seed", None) extra_args = { "reset": reset, "session_id": session_id, @@ -157,7 +160,10 @@ class ServingRealtimeRobotOpenPI: } prompt = obs.get("prompt", "") - sampling_params = OmniDiffusionSamplingParams(extra_args=extra_args) + sampling_params = OmniDiffusionSamplingParams( + seed=int(seed) if seed is not None else None, + extra_args=extra_args, + ) return OmniDiffusionRequest( prompt=prompt, sampling_params=sampling_params, diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 1f4f78c4e..39131fd58 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -4,45 +4,18 @@ import json import os import types -from collections import Counter from dataclasses import fields, is_dataclass -from pathlib import Path from typing import Any, get_args, get_origin from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind, SamplingParams -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.config_factory import ( - StageConfigFactory, - _materialize_object_storage_configs, - _name_match_candidate, - with_trust_remote_code_override, -) -from vllm_omni.config.pipeline_registry import OMNI_PIPELINES -from vllm_omni.config.stage_config import _DEPLOY_DIR -from vllm_omni.config.yaml_util import create_config, load_yaml_config -from vllm_omni.diffusion.utils.hf_utils import ( - _looks_like_dreamzero, - get_diffusion_model_index, - resolve_native_diffusion_model_class, -) from vllm_omni.entrypoints.stage_utils import _to_dict from vllm_omni.inputs.data import OmniSamplingParams -from vllm_omni.platforms import current_omni_platform - -# Get the project root directory (2 levels up from this file) -PROJECT_ROOT = Path(__file__).parent.parent.parent logger = init_logger(__name__) -_DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = { - "GlmImagePipeline": "glm_image", -} - - def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: str, omni_to: str) -> None: """Inject connector configuration into stage engine arguments.""" # Prepare omni_kv_config dict @@ -76,501 +49,9 @@ def inject_omni_kv_config(stage: Any, omni_conn_cfg: dict[str, Any], omni_from: logger.error(f"Failed to inject omni connector config into stage: {e}") -def _try_get_class_name_from_diffusers_config(model: str) -> str | None: - """Try to get class name from diffusers model configuration files. - - Args: - model: Model name or path - - Returns: - Model type string if found, None otherwise - """ - model_index = get_diffusion_model_index(model) - if model_index and isinstance(model_index, dict) and "_class_name" in model_index: - logger.debug(f"Found Diffusers model type '{model_index['_class_name']}'") - return model_index["_class_name"] - - return None - - -def _filter_dict_like_object(obj: dict | Any) -> dict: - """Filter dict-like object by removing callables and recursively converting values. - - Converts dict-like objects to regular dicts while filtering out callable values - that are incompatible with OmegaConf. Recursively processes values through - _convert_dataclasses_to_dict for nested object conversion. - - Args: - obj: Dict or dict-like object to filter - - Returns: - Regular dict with callables filtered out and values recursively converted - - Raises: - TypeError: If obj doesn't support .items() method - ValueError: If dict conversion fails unexpectedly - """ - - def _is_callable_value(value: Any) -> bool: - if callable(value): - return True - return isinstance( - value, - types.FunctionType | types.MethodType | types.BuiltinFunctionType | types.BuiltinMethodType, - ) - - result = {} - filtered_keys = [] - for k, v in obj.items(): - # Preserve class objects by converting to a fully qualified name string - # so callers that resolve via import path (e.g. custom_pipeline_args.pipeline_class) - # still work after OmegaConf round-trip. - if isinstance(v, type): - module = getattr(v, "__module__", None) - qualname = getattr(v, "__qualname__", getattr(v, "__name__", None)) - if module and qualname and module != "builtins": - result[k] = f"{module}.{qualname}" - else: - result[k] = qualname - elif _is_callable_value(v): - filtered_keys.append(str(k)) - else: - result[k] = _convert_dataclasses_to_dict(v) - if filtered_keys: - logger.warning( - f"Filtered out {len(filtered_keys)} callable object(s) from base_engine_args " - f"that are not compatible with OmegaConf: {filtered_keys}. " - ) - return result - - -def _convert_dataclasses_to_dict(obj: Any) -> Any: - """Recursively convert non-serializable objects to OmegaConf-compatible types. - - This is needed because OmegaConf cannot handle: - - Dataclass objects with Literal type annotations (e.g., StructuredOutputsConfig) - - Counter objects (from collections or vllm.utils) - - Set objects - - Callable objects (functions, methods, etc.) - - Other non-primitive types - """ - # IMPORTANT: Check Counter BEFORE dict, since Counter is a subclass of dict - # Handle Counter objects (convert to dict) - # Check by class name first to catch both collections.Counter and vllm.utils.Counter - if hasattr(obj, "__class__") and obj.__class__.__name__ == "Counter": - try: - return dict(obj) - except (TypeError, ValueError): - # If Counter can't be converted to dict, return empty dict - return {} - # Also check isinstance for collections.Counter (must be before dict check) - if isinstance(obj, Counter): - return dict(obj) - # Handle set objects (convert to list) - if isinstance(obj, set): - return list(obj) - # Handle dataclass objects - # Use field iteration instead of asdict() to: - # 1. Only include init fields (non-init fields cause "unexpected kwarg" errors) - # 2. Skip None values matching field defaults (avoids Pydantic validation - # failures when None is explicitly passed for non-Optional typed fields, - # e.g. CompilationConfig.cudagraph_capture_sizes: list[int] = None) - if is_dataclass(obj) and not isinstance(obj, type): - result = {} - for f in fields(obj): - if not f.init: - continue - value = getattr(obj, f.name) - if value is None and f.default is None: - continue - result[f.name] = _convert_dataclasses_to_dict(value) - return result - # Handle dictionaries (recurse into values) and filter out callables(cause error in OmegaConf.create) - # Note: This must come AFTER Counter check since Counter is a dict subclass - if isinstance(obj, dict): - return _filter_dict_like_object(obj) - # Preserve class objects by converting to a fully qualified name string. - if isinstance(obj, type): - module = getattr(obj, "__module__", None) - qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", None)) - if module and qualname and module != "builtins": - return f"{module}.{qualname}" - return qualname - # Handle callable objects (functions, methods, etc.) - skip them - # Note: This comes after dict/list checks to avoid misclassifying dict-like objects - if callable(obj): - return None - # Handle lists and tuples (recurse into items) - if isinstance(obj, list | tuple): - return type(obj)(_convert_dataclasses_to_dict(item) for item in obj if not callable(item)) - # Try to convert any dict-like object (has keys/values methods) to dict - if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, str | bytes): - try: - return _filter_dict_like_object(obj) - except (TypeError, ValueError, AttributeError): - # If conversion fails, return as-is - return obj - # Primitive types and other objects that OmegaConf can handle - return obj - - -def _try_resolve_omni_model_type(model: str) -> str | None: - """Try to resolve model_type for omni models with empty config.json. - - Searches both the legacy ``stage_configs/*.yaml`` directory and the - migrated ``deploy/*.yaml`` directory for a stem that substring-matches - the model path basename (e.g. ``cosyvoice3`` in - ``FunAudioLLM/Fun-CosyVoice3-0.5B-2512``). Only the basename is scanned - so URI segments such as the bucket name cannot select an unrelated - pipeline. The longest match wins so ``cosyvoice3`` beats ``cosyvoice`` - and ``bagel_single_stage`` beats ``bagel``. - """ - model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "") - best_match: str | None = None - best_len = 0 - for subdir in ("model_executor/stage_configs", "deploy"): - config_dir = PROJECT_ROOT / "vllm_omni" / subdir - if not config_dir.exists(): - continue - for config_file in sorted(config_dir.glob("*.yaml")): - candidate = config_file.stem.replace("-", "").replace("_", "") - if candidate and candidate in model_lower and len(candidate) > best_len: - best_match = config_file.stem - best_len = len(candidate) - return best_match - - -def _registry_default_deploy_path(model: str) -> str | None: - """Return the registered pipeline's default deploy YAML path, if any. - - A checkpoint's HF ``model_type`` is not always the registered Omni pipeline - key. MiniCPM-o 4.5, for example, reports ``minicpmo`` while its - predicate-selected pipeline is ``minicpmo_4_5``. Stage construction already - loads that pipeline's default deploy config, so connector discovery must - resolve the same path. - """ - pipeline_config = StageConfigFactory.get_pipeline_config( - model=model, - trust_remote_code=True, - ) - if pipeline_config is None or pipeline_config.default_deploy_config_name is None: - return None - default_deploy_path = _DEPLOY_DIR / pipeline_config.default_deploy_config_name - if default_deploy_path.is_file(): - return str(default_deploy_path) - return None - - -def resolve_model_config_path(model: str) -> str | None: - """Resolve the stage/deploy config file path from the model name. - - Resolves configuration path based on the model type and device type. - Order: - 1. Device-specific stage config under ``stage_configs/{device_type}/`` - 2. If HF ``model_type`` is not an ``OMNI_PIPELINES`` key, the registered - pipeline's ``default_deploy_config_name`` (keeps connectors aligned with - stage construction when HF type and pipeline key differ) - 3. ``deploy/{model_type}.yaml`` - 4. Legacy ``stage_configs/{model_type}.yaml`` - 5. Registered pipeline default deploy config (final fallthrough) - - Args: - model: Model name or path (used to determine model_type) - - Returns: - String path to the stage/deploy configuration file, or ``None`` if no - matching config file exists. - - Raises: - ValueError: If model_type cannot be determined - """ - # Object-storage URIs are streamed by vLLM's Run:AI streamer only after - # each stage builds its ModelConfig, so config resolution here reads the - # local copy that vLLM-Omni materializes for such URIs. Name-based - # fallbacks keep the original string (the materialized path is a hash). - config_source = _materialize_object_storage_configs(model) - # Try to get config from standard transformers format first - try: - hf_config = get_config(config_source, trust_remote_code=True) - model_type = hf_config.model_type - except (ValueError, Exception): - # If standard transformers format fails, try diffusers format - native_model_class = resolve_native_diffusion_model_class(config_source) - if native_model_class is not None: - model_type = native_model_class - elif get_diffusion_model_index(config_source) is not None: - model_type = _try_get_class_name_from_diffusers_config(config_source) - if model_type is None: - raise ValueError( - f"Could not determine model_type for diffusers model: {model}. " - "Please ensure its Diffusers pipeline index contains '_class_name'" - ) - elif file_or_path_exists(config_source, "config.json", revision=None): - # Try to read config.json manually for custom models like Bagel that fail get_config - # but have a valid config.json with model_type - try: - config_dict = get_hf_file_to_dict("config.json", config_source, revision=None) - if config_dict and "model_type" in config_dict: - model_type = config_dict["model_type"] - else: - # For models with empty config.json (e.g. CosyVoice3), - # try matching against registered omni stage configs. - model_type = _try_resolve_omni_model_type(model) - if model_type is None: - raise ValueError(f"config.json found but missing 'model_type' for model: {model}") - except Exception as e: - raise ValueError(f"Failed to read config.json for model: {model}. Error: {e}") from e - else: - # No config.json at repo root (e.g. GLM-TTS stores configs in - # subdirectories only). Try matching against registered deploy - # YAML filenames before giving up. - model_type = _try_resolve_omni_model_type(model) - if model_type is None: - raise ValueError( - f"Could not determine model_type for model: {model}. " - "Model is not in standard transformers or Diffusers format. " - f"Please ensure the model has proper configuration files with 'model_type' field" - ) - - default_config_path = current_omni_platform.get_default_stage_config_path() - if model_type == "vla" and _looks_like_dreamzero(config_source): - model_type = "dreamzero" - - if model_type in _DIFFUSERS_CLASS_TO_CONFIG: - normalized_model_type = _DIFFUSERS_CLASS_TO_CONFIG[model_type] - else: - normalized_model_type = model_type.replace("-", "_") - model_type_str = f"{normalized_model_type}.yaml" - complete_config_path = PROJECT_ROOT / default_config_path / model_type_str - if os.path.exists(complete_config_path): - return str(complete_config_path) - - # Prefer the registry default before deploy/.yaml when the - # HF type is not itself a pipeline key. Otherwise a future - # deploy/minicpmo.yaml would desync connectors from stages for MiniCPM-o 4.5. - if normalized_model_type not in OMNI_PIPELINES: - registry_path = _registry_default_deploy_path(model) - if registry_path is not None: - return registry_path - - deploy_config_path = _DEPLOY_DIR / model_type_str - if os.path.exists(deploy_config_path): - return str(deploy_config_path) - - stage_config_file = f"vllm_omni/model_executor/stage_configs/{normalized_model_type}.yaml" - stage_config_path = PROJECT_ROOT / stage_config_file - if os.path.exists(stage_config_path): - return str(stage_config_path) - - return _registry_default_deploy_path(model) - - -def load_stage_configs_from_model( - model: str, - *, - trust_remote_code: bool | None, - base_engine_args: dict | None = None, - deploy_config_path: str | None = None, - stage_overrides: dict[str, dict[str, Any]] | None = None, - 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, uses - ``StageConfigFactory.create_legacy_stage_configs_from_model()`` which merges - PipelineConfig + DeployConfig + CLI overrides. - - Models that cannot be resolved through the registry return no stages so the - caller can apply its default stage configuration, when available. - - 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. - - Returns: - ``(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)) - # A False inherited from the engine-args dump is the store_true flag's - # default, not an explicit choice — drop it so only the tri-state - # parameter below decides (see with_trust_remote_code_override). - if not cli_overrides.get("trust_remote_code"): - cli_overrides.pop("trust_remote_code", None) - cli_overrides = with_trust_remote_code_override(cli_overrides, 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 - - # 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=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], omni_lb_policy - - strategy_note = "" - if strategy_config_path is not None: - strategy_note = f" Strategy config {strategy_config_path!r} was not applied." - logger.warning( - "No registered PipelineConfig resolved for model %r. Legacy `stage_args` " - "YAMLs are no longer supported; register the pipeline and provide deployment " - "overrides through `deploy_config`.%s", - model, - strategy_note, - ) - return [], None - - -def filter_stages( - config_path: str | None, - stage_configs: list, - kwargs: dict | None, -) -> list: - """Filter stage configs by mode when YAML defines a `modes` section. - - The YAML can define, e.g.: - - modes: - - mode: text-to-image - stages: [1] - - mode: image-to-text - stages: [0] - - When users pass `mode="image-to-text"` into Omni(**kwargs), only the stages - listed for that mode are returned. If no mode is provided, defaults to - "text-to-image". If no modes are defined or filtering fails, returns the - original stage_configs unchanged. - - Args: - config_path: Path to the YAML config (used to read `modes`). - stage_configs: Loaded list of stage configs. - kwargs: Engine/caller kwargs; may contain "mode". - - Returns: - Filtered list of stage configs (or original list if filtering not applied). - """ - if not stage_configs or config_path is None: - return stage_configs - - try: - cfg = load_yaml_config(config_path) - yaml_modes = getattr(cfg, "modes", None) - if yaml_modes is None: - return stage_configs - - mode_to_stage_ids: dict[str, list[int]] = {} - if yaml_modes is not None: - for entry in yaml_modes: - mode_name = None - stages = None - if hasattr(entry, "mode") or hasattr(entry, "stages"): - mode_name = getattr(entry, "mode", None) - stages = getattr(entry, "stages", None) - elif isinstance(entry, dict): - mode_name = entry.get("mode") - stages = entry.get("stages") - - if mode_name is None or stages is None: - continue - - if isinstance(stages, int): - stage_list = [stages] - else: - stage_list = list(stages) - - mode_to_stage_ids[str(mode_name)] = [int(sid) for sid in stage_list] - - # No modes section or empty mapping: use all stages and return early. - active_mode: str | None = None - if isinstance(kwargs, dict): - active_mode = kwargs.get("mode") - - if active_mode is None: - active_mode = "text-to-image" - - if active_mode not in mode_to_stage_ids: - logger.warning( - "Requested mode '%s' not found in config '%s'; available modes: %s. Using all stages.", - active_mode, - config_path, - sorted(mode_to_stage_ids.keys()), - ) - return stage_configs - - allowed_ids = set(mode_to_stage_ids[active_mode]) - filtered_stage_configs = [sc for sc in stage_configs if getattr(sc, "stage_id", None) in allowed_ids] - if not filtered_stage_configs: - logger.warning( - "Mode '%s' in config '%s' resolved to stage ids %s, but none matched loaded stage_args. " - "Falling back to all stages.", - active_mode, - config_path, - sorted(allowed_ids), - ) - return stage_configs - - return filtered_stage_configs - except Exception as e: - logger.warning("Failed to apply mode-based stage filtering: %s", e) - 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. - - Only the **shape** of the override mapping is validated here: - - - top-level must be a dict - - keys must be non-negative ASCII integer strings (stage ids) - - values must be dicts - - Field-name / type / range validation is intentionally not done at this - layer. The downstream resolver forwards every surviving key as - ``stage__`` into ``cli_overrides`` (see - ``load_stage_configs_from_model``); field semantics live in - ``StageConfigFactory`` for registered pipelines, and the - default-diffusion fallback in ``async_omni_engine.py`` reads - ``stage_zero_overrides["extras"]`` for unregistered ones. Unknown - engine-arg keys that don't match ``OmniEngineArgs`` are dropped with - a warning at ``filter_dataclass_kwargs`` rather than failing here, - so a typo surfaces in the log instead of being silently lost. - - Raises: - ValueError: when ``value`` is not a valid per-stage override - mapping (invalid JSON, non-dict top level, non-ASCII-digit - stage id, non-dict value). - """ - if not value: + """Parse and validate the shape of per-stage JSON overrides.""" + if value is None: return None if isinstance(value, str): try: @@ -582,93 +63,22 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None: if not isinstance(parsed, dict): raise ValueError( - f"--stage-overrides must be a JSON object mapping stage_id -> overrides, " + "--stage-overrides must be a JSON object mapping stage_id -> overrides, " f"got {type(parsed).__name__}: {parsed!r}" ) - if not parsed: - return None - - for stage_id_str, overrides in parsed.items(): - # ``str.isdigit()`` accepts Unicode digit classes (e.g. fullwidth - # ``"０"``); downstream's ``_STAGE_OVERRIDE_PATTERN`` is ASCII-only - # (``^stage_(\d+)_(.+)$``), so a non-ASCII-digit key would parse - # here and silently misroute as a global key. ``isascii()`` closes - # that hole. - if not isinstance(stage_id_str, str) or not stage_id_str.isdigit() or not stage_id_str.isascii(): + for stage_id, overrides in parsed.items(): + if not isinstance(stage_id, str) or not stage_id.isascii() or not stage_id.isdigit(): raise ValueError( - f"--stage-overrides keys must be non-negative integer stage ids (as strings), got {stage_id_str!r}" + f"--stage-overrides keys must be non-negative integer stage ids (as strings), got {stage_id!r}" ) if not isinstance(overrides, dict): raise ValueError( - f"--stage-overrides[{stage_id_str!r}] must be an object, got {type(overrides).__name__}: {overrides!r}" + f"--stage-overrides[{stage_id!r}] must be an object, got {type(overrides).__name__}: {overrides!r}" ) return parsed -def load_and_resolve_stage_configs( - model: str, - kwargs: dict | None, - *, - trust_remote_code: bool | None, - default_stage_cfg_factory: Any = None, - deploy_config_path: str | None = None, - stage_overrides: dict[str, dict[str, Any]] | None = None, - strategy_config_path: str | None = None, -) -> tuple[str, list, str | None]: - """Load stage configurations from a deploy YAML or model defaults. - - Args: - model: Model name or path - 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 a deploy YAML. - 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, 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. - """ - config_path = deploy_config_path if deploy_config_path is not None else resolve_model_config_path(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: - default_stage_cfg = default_stage_cfg_factory() - stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg)) - else: - stage_configs = [] - - stage_configs = filter_stages(config_path, stage_configs, kwargs) - - # ``revision`` selects the model source for every stage. Keep the global - # CLI pin on each stage's engine args so diffusion config construction can - # pass it through to all component loaders. The registry-to-legacy-stage - # conversion currently drops this inherited vLLM EngineArgs field. - revision = (kwargs or {}).get("revision") - if revision is not None: - for stage_config in stage_configs: - engine_args = stage_config.get("engine_args") if hasattr(stage_config, "get") else None - if engine_args is not None and engine_args.get("revision") is None: - engine_args["revision"] = revision - - logger.debug(f"stage_configs: {stage_configs}") - - return config_path, stage_configs, omni_lb_policy - - def get_final_stage_id_for_e2e( output_modalities: list[str] | None, default_modalities: list[str], stage_list: list ) -> int: @@ -839,7 +249,7 @@ def detect_pid_host() -> bool: if not ic: return True - return has_pid_host() + return has_pid_host() is True ### Helpers for handling delta messages diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py index 0f01e36da..2b58b77e8 100644 --- a/vllm_omni/host_weight_runtime/filesystem/store.py +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -1247,8 +1247,10 @@ class FilesystemHostWeightStore: retryable=True, ) ) - store_bytes = self._tree_bytes(self.root) - if policy.max_store_bytes is not None and store_bytes + additional_bytes > policy.max_store_bytes: + if ( + policy.max_store_bytes is not None + and self._tree_bytes(self.root) + additional_bytes > policy.max_store_bytes + ): raise HostWeightError( _failure( ResolutionStage.CAPACITY, diff --git a/vllm_omni/metrics/__init__.py b/vllm_omni/metrics/__init__.py index 7b6a2979f..c510c943b 100644 --- a/vllm_omni/metrics/__init__.py +++ b/vllm_omni/metrics/__init__.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from .prometheus import OmniPrometheusMetrics, OmniRequestCounter from .stats import OrchestratorAggregator, StageRequestStats, StageStats from .utils import ( @@ -5,6 +8,7 @@ from .utils import ( count_audio_frames, count_image_pixels, count_tokens_from_outputs, + count_video_frames, ) __all__ = [ @@ -17,4 +21,5 @@ __all__ = [ "count_audio_frames", "count_image_pixels", "count_tokens_from_outputs", + "count_video_frames", ] diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index 68adb44ca..6cc2825c1 100644 --- a/vllm_omni/metrics/definitions.py +++ b/vllm_omni/metrics/definitions.py @@ -76,6 +76,29 @@ MEDIAN_IMAGE_GENERATION_MS = f"median_{IMAGE_GENERATION}_ms" STD_IMAGE_GENERATION_MS = f"std_{IMAGE_GENERATION}_ms" PERCENTILES_IMAGE_GENERATION_MS = f"percentiles_{IMAGE_GENERATION}_ms" +VIDEO_DURATION = "video_duration" +VIDEO_RTF = "video_rtf" +VIDEO_GENERATION = "video_generation" +VIDEO_GENERATION_TIME_MS = f"{VIDEO_GENERATION}_time_ms" +VIDEO_FRAMES = "video_frames" +TOTAL_VIDEO_DURATION_S = f"total_{VIDEO_DURATION}_s" +TOTAL_VIDEO_FRAMES = f"total_{VIDEO_FRAMES}" +VIDEO_THROUGHPUT = "video_throughput" +MEAN_VIDEO_RTF = f"mean_{VIDEO_RTF}" +MEDIAN_VIDEO_RTF = f"median_{VIDEO_RTF}" +STD_VIDEO_RTF = f"std_{VIDEO_RTF}" +PERCENTILES_VIDEO_RTF = f"percentiles_{VIDEO_RTF}" +MEAN_VIDEO_GENERATION_MS = f"mean_{VIDEO_GENERATION}_ms" +MEDIAN_VIDEO_GENERATION_MS = f"median_{VIDEO_GENERATION}_ms" +STD_VIDEO_GENERATION_MS = f"std_{VIDEO_GENERATION}_ms" +PERCENTILES_VIDEO_GENERATION_MS = f"percentiles_{VIDEO_GENERATION}_ms" + +PEAK_MEMORY_MB = "peak_memory_mb" +MEAN_PEAK_MEMORY_MB = f"mean_{PEAK_MEMORY_MB}" +MEDIAN_PEAK_MEMORY_MB = f"median_{PEAK_MEMORY_MB}" +STD_PEAK_MEMORY_MB = f"std_{PEAK_MEMORY_MB}" +PERCENTILES_PEAK_MEMORY_MB = f"percentiles_{PEAK_MEMORY_MB}" + # Stage snapshot / StageBenchmarkMetrics field names. TOTAL_OUTPUT = "total_output" TTFTS = "ttfts" @@ -138,6 +161,12 @@ AUDIO_UNDERRUN_S = METRIC_PREFIX + AUDIO_UNDERRUN + "_s" AUDIO_CONTINUITY_OK_METRIC = METRIC_PREFIX + AUDIO_CONTINUITY_OK AUDIO_SKIPPED_REQUESTS_METRIC = METRIC_PREFIX + AUDIO_SKIPPED_REQUESTS +# Realtime Server VAD serving metrics. +REALTIME_VAD_ACTIVE_SESSIONS = METRIC_PREFIX + "realtime_vad_active_sessions" +REALTIME_VAD_INFERENCE_LATENCY_S = METRIC_PREFIX + "realtime_vad_inference_latency_s" +REALTIME_VAD_ENDPOINT_DELAY_S = METRIC_PREFIX + "realtime_vad_endpoint_delay_s" +REALTIME_VAD_ERRORS = METRIC_PREFIX + "realtime_vad_errors" + # ============================================================================ # Diffusion family (per-stage + per-replica diffusion timing breakdowns) @@ -166,7 +195,7 @@ STAGE_WAITING_REQUESTS = METRIC_PREFIX + "stage_waiting_requests" NUM_INFERENCE_STEPS = METRIC_PREFIX + "num_inference_steps" IMAGE_COUNT_METRIC = METRIC_PREFIX + IMAGE_COUNT IMAGE_PIXELS_METRIC = METRIC_PREFIX + IMAGE_PIXELS -PEAK_MEMORY_MB = METRIC_PREFIX + "peak_memory_mb" +PEAK_MEMORY_MB_METRIC = METRIC_PREFIX + "peak_memory_mb" REQUESTS_FAILED = METRIC_PREFIX + "requests_failed" KV_WAIT_S = METRIC_PREFIX + "kv_wait_s" DIFFUSION_FORWARD_S = METRIC_PREFIX + "diffusion_forward_s" diff --git a/vllm_omni/metrics/prometheus.py b/vllm_omni/metrics/prometheus.py index 911d21e64..2e68045c1 100644 --- a/vllm_omni/metrics/prometheus.py +++ b/vllm_omni/metrics/prometheus.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from prometheus_client import Counter, Gauge, Histogram from vllm_omni.metrics import definitions as defs @@ -72,7 +75,7 @@ _image_pixels_family = Histogram( labelnames=_labelnames, ) _peak_memory_family = Gauge( - defs.PEAK_MEMORY_MB, + defs.PEAK_MEMORY_MB_METRIC, "Peak GPU memory in MB observed during a stage's generation.", labelnames=list(defs.DIFFUSION_LABELS), ) diff --git a/vllm_omni/metrics/realtime.py b/vllm_omni/metrics/realtime.py new file mode 100644 index 000000000..a82bb5b17 --- /dev/null +++ b/vllm_omni/metrics/realtime.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from prometheus_client import Counter, Gauge, Histogram + +from vllm_omni.metrics import definitions as defs + +_active_sessions = Gauge( + defs.REALTIME_VAD_ACTIVE_SESSIONS, + "Number of active Realtime sessions using server-side VAD.", + labelnames=list(defs.PIPELINE_LABELS), +) +_inference_latency = Histogram( + defs.REALTIME_VAD_INFERENCE_LATENCY_S, + "Server VAD inference time for one appended audio batch, in seconds.", + labelnames=list(defs.PIPELINE_LABELS), + buckets=defs.SECONDS_FAST_BUCKETS, +) +_endpoint_delay = Histogram( + defs.REALTIME_VAD_ENDPOINT_DELAY_S, + "Detected silence between the last speech frame and the committed endpoint, in seconds.", + labelnames=list(defs.PIPELINE_LABELS), + buckets=defs.SECONDS_FAST_BUCKETS, +) +_errors = Counter( + defs.REALTIME_VAD_ERRORS, + "Server VAD errors and input-overflow rejections by reason.", + labelnames=list(defs.FAILED_LABELS), +) + + +class RealtimeVADMetrics: + def __init__(self, model_name: str, log_stats: bool = True) -> None: + self._model_name = model_name + self._log_stats = log_stats + self._active_sessions = _active_sessions.labels(model_name=model_name) + self._inference_latency = _inference_latency.labels(model_name=model_name) + self._endpoint_delay = _endpoint_delay.labels(model_name=model_name) + + def session_started(self) -> None: + if self._log_stats: + self._active_sessions.inc() + + def session_finished(self) -> None: + if self._log_stats: + self._active_sessions.dec() + + def observe_inference(self, latency_ms: float) -> None: + if self._log_stats: + self._inference_latency.observe(latency_ms / 1000) + + def observe_endpoint_delay(self, delay_ms: int) -> None: + if self._log_stats: + self._endpoint_delay.observe(delay_ms / 1000) + + def error(self, reason: str) -> None: + if self._log_stats: + _errors.labels(model_name=self._model_name, reason=reason).inc() diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 9e9191eaa..a782cca88 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -48,7 +48,7 @@ class StageRequestStats: finish_reason: str | None = None request_id: str | None = None postprocess_time_ms: float = 0.0 - diffusion_metrics: dict[str, float] = None + diffusion_metrics: dict[str, float] | None = None audio_generated_frames: int = 0 audio_sample_rate: int = 0 audio_duration_s: float = 0.0 @@ -181,11 +181,11 @@ class OrchestratorAggregator: self.e2e_total_ms = 0.0 self.e2e_total_tokens = 0 self.e2e_count = 0 - self.e2e_done = set() + self.e2e_done: set[str] = set() self.wall_start_ts = float(wall_start_ts) self.last_finish_ts = float(wall_start_ts) - self.stage_first_ts = [None for _ in range(self.num_stages)] - self.stage_last_ts = [None for _ in range(self.num_stages)] + self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)] + self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)] self.accumulated_gen_time_ms: defaultdict[str, defaultdict[int, float]] = defaultdict( lambda: defaultdict(float) ) # {request_id: {stage_id:accumulated_gen_time_ms}} @@ -618,6 +618,7 @@ class OrchestratorAggregator: final_output_type: str | None = None, ) -> None: stats = self._as_stage_request_stats(stage_id, req_id, metrics, final_output_type) + assert stats.stage_id is not None self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_out) if stats.stage_id == 0: self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_in) @@ -774,10 +775,8 @@ class OrchestratorAggregator: final_stage_id_map = self.final_stage_id_for_e2e stage_wall_time_ms = [ - ((self.stage_last_ts[i] - self.stage_first_ts[i]) * 1000.0) - if (self.stage_first_ts[i] is not None and self.stage_last_ts[i] is not None) - else 0.0 - for i in range(self.num_stages) + ((last - first) * 1000.0) if first is not None and last is not None else 0.0 + for first, last in zip(self.stage_first_ts, self.stage_last_ts, strict=True) ] overall_summary = { diff --git a/vllm_omni/metrics/utils.py b/vllm_omni/metrics/utils.py index 10cdab94a..ad58e4cec 100644 --- a/vllm_omni/metrics/utils.py +++ b/vllm_omni/metrics/utils.py @@ -163,6 +163,24 @@ def extract_diffusion_denoise_ms(output: Any) -> float | None: return sum_diffusion_stage_durations_ms(output, ".diffuse") +def coerce_bool(value: object, *, default: bool = False) -> bool: + """Coerce JSON / form / env-like values to bool. + + Accepts ``bool``, numeric 0/1, and common truthy strings + (``\"1\"`` / ``\"true\"`` / ``\"yes\"`` / ``\"on\"``, case-insensitive). + ``None`` and unrecognized types return ``default``. + """ + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return default + + def coerce_positive_int_scalar(value: object) -> int | None: """Coerce a value to a positive int without importing tensor libs. @@ -199,6 +217,34 @@ def coerce_positive_int_scalar(value: object) -> int | None: return value if value > 0 else None +def coerce_positive_float_scalar(value: object) -> float | None: + """Coerce a scalar-like value to a positive float. + + Handles plain numeric values, one-element tensor/numpy-like values via + ``.item()``, and list/tuple wrappers. Returns ``None`` when no positive + float can be extracted. + """ + if value is None: + return None + if isinstance(value, (list, tuple)): + for item in value: + coerced = coerce_positive_float_scalar(item) + if coerced is not None: + return coerced + return None + item = getattr(value, "item", None) + if callable(item): + try: + value = item() + except Exception: + return None + try: + value = float(value) + except (TypeError, ValueError): + return None + return value if value > 0 else None + + def resolve_int_by_sequential_keys( source: Mapping[str, object] | object | None, keys: Sequence[str], @@ -370,6 +416,43 @@ def count_image_pixels(value: object) -> int: return dims[-2] * dims[-1] +def count_video_frames(video: object) -> int | None: + """Return the frame count for common nested and tensor video layouts.""" + if isinstance(video, Mapping): + for key in ("video", "frames", "images"): + if key in video: + return count_video_frames(video[key]) + return None + + if isinstance(video, (list, tuple)): + return len(video) + + shape = getattr(video, "shape", None) + if shape is None: + return None + try: + dims = tuple(int(dim) for dim in shape) + except (TypeError, ValueError): + return None + + if len(dims) == 5: + # Common layouts: [B, C, F, H, W] and [B, F, H, W, C]. + if dims[-1] in (1, 3, 4): + return dims[1] + if dims[1] in (1, 3, 4): + return dims[2] + return dims[1] + if len(dims) == 4: + # Common layouts: [F, H, W, C] and [C, F, H, W]. + if dims[-1] in (1, 3, 4): + return dims[0] + if dims[0] in (1, 3, 4): + return dims[1] + return dims[0] + + return None + + def _build_field_defs( cls: type, exclude: set[str], diff --git a/vllm_omni/model_executor/model_loader/weight_utils.py b/vllm_omni/model_executor/model_loader/weight_utils.py index 59c2cd860..559a90b50 100644 --- a/vllm_omni/model_executor/model_loader/weight_utils.py +++ b/vllm_omni/model_executor/model_loader/weight_utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import os import time from pathlib import Path @@ -10,7 +13,11 @@ from vllm.model_executor.model_loader.weight_utils import DisabledTqdm, get_lock if envs.VLLM_USE_MODELSCOPE: from modelscope.hub.snapshot_download import snapshot_download else: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api + + def snapshot_download(*args, **kwargs): + return hf_api().snapshot_download(*args, **kwargs) + logger = init_logger(__name__) diff --git a/vllm_omni/model_executor/models/audex/checkpoint.py b/vllm_omni/model_executor/models/audex/checkpoint.py index 38b0531e6..ae8230e05 100644 --- a/vllm_omni/model_executor/models/audex/checkpoint.py +++ b/vllm_omni/model_executor/models/audex/checkpoint.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Audex checkpoint layout preparation. The nvidia/Nemotron-Labs-Audex-2B repo deduplicates weights: the @@ -57,12 +57,12 @@ def _download_audiogen_index(model: str) -> str: metadata round trip per stage-engine boot; a cache miss falls back to the network. """ - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api try: - return hf_hub_download(model, _AUDIOGEN_INDEX_FILE, local_files_only=True) + return hf_api().hf_hub_download(model, _AUDIOGEN_INDEX_FILE, local_files_only=True) except Exception: - return hf_hub_download(model, _AUDIOGEN_INDEX_FILE) + return hf_api().hf_hub_download(model, _AUDIOGEN_INDEX_FILE) def _dedup_shard_patterns(model: str) -> list[str]: @@ -111,13 +111,13 @@ def ensure_audex_snapshot(model: str, profile: str = "tts") -> str: if profile in ("tts", "tta"): patterns += _dedup_shard_patterns(model) - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api try: - return snapshot_download(model, allow_patterns=patterns) + return hf_api().snapshot_download(model, allow_patterns=patterns) except Exception as download_exc: try: - return snapshot_download(model, allow_patterns=patterns, local_files_only=True) + return hf_api().snapshot_download(model, allow_patterns=patterns, local_files_only=True) except Exception: raise RuntimeError( f"Could not resolve the Audex repo {model!r} (profile {profile!r}): the download " @@ -135,13 +135,13 @@ def ensure_xcodec1_snapshot(model: str | None) -> str: if os.path.isdir(model): return model - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api try: - return snapshot_download(model) + return hf_api().snapshot_download(model) except Exception as download_exc: try: - return snapshot_download(model, local_files_only=True) + return hf_api().snapshot_download(model, local_files_only=True) except Exception: raise RuntimeError( f"Could not resolve the XCodec1 checkpoint {model!r}: the download failed and " diff --git a/vllm_omni/model_executor/models/audio8_tts/__init__.py b/vllm_omni/model_executor/models/audio8_tts/__init__.py new file mode 100644 index 000000000..a6db04780 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview (DualAR, 44.1 kHz codec) model implementation.""" diff --git a/vllm_omni/model_executor/models/audio8_tts/audio8_tts_codec_decoder.py b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_codec_decoder.py new file mode 100644 index 000000000..6f93d31a2 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_codec_decoder.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview -- codec decoder (Stage 1). + +Consumes frames of ``num_codebooks`` codec codes and decodes them to a 44.1 kHz +waveform. Streaming contract is delta: with ``async_chunk`` on, each chunk is +``[left_context | new_frames]``; the whole window is decoded so the causal +codec keeps continuous context, then only the new tail is returned. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.audio8_tts.codec_utils import load_arktts_codec +from vllm_omni.model_executor.models.audio8_tts.configuration_audio8_tts import ( + ARKTTS_CODEC_FRAME_SIZE, + ARKTTS_CODEC_SAMPLE_RATE, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + +_DTYPES = { + "fp16": torch.float16, + "float16": torch.float16, + "half": torch.float16, + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp32": torch.float32, + "float32": torch.float32, + "": torch.float32, +} + + +def _connector_extra_config(vllm_config: VllmConfig) -> dict[str, Any]: + model_config = getattr(vllm_config, "model_config", None) + connector_cfg = getattr(model_config, "stage_connector_config", None) + if isinstance(connector_cfg, dict): + return connector_cfg.get("extra", connector_cfg) + extra = getattr(connector_cfg, "extra", None) + return extra if isinstance(extra, dict) else {} + + +def _resolve_dtype(extra_cfg: dict[str, Any]) -> torch.dtype: + raw = str(extra_cfg.get("audio8_tts_codec_dtype", "float32")).strip().lower() + if raw not in _DTYPES: + raise ValueError(f"Invalid Audio8 TTS codec dtype: {raw!r}") + return _DTYPES[raw] + + +class Audio8TTSCodecDecoder(nn.Module): + """Stage 1: codec codes -> waveform (runs under the generation runner).""" + + input_modalities = "audio" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + del prefix + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + config = vllm_config.model_config.hf_config + + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + self._num_codebooks = int(getattr(config, "num_codebooks", 10)) + self._sample_rate = int(getattr(config, "codec_sample_rate", ARKTTS_CODEC_SAMPLE_RATE)) + self._frame_size = int(getattr(config, "codec_frame_size", ARKTTS_CODEC_FRAME_SIZE)) + self._codec_kwargs = { + "post_n_layer": int(getattr(config, "codec_post_n_layer", 8)), + "post_n_head": int(getattr(config, "codec_post_n_head", 16)), + "post_n_local_heads": int(getattr(config, "codec_post_n_local_heads", 8)), + "post_intermediate_size": int(getattr(config, "codec_post_intermediate_size", 1216)), + } + # float32 by default: the codec's Snake activations and weight-normed + # convs lose audible headroom in bf16. + self._codec_dtype = _resolve_dtype(_connector_extra_config(vllm_config)) + self._codec: nn.Module | None = None + self._codec_device: torch.device | None = None + self._logged_stats = False + + # -------------------- codec -------------------- + + def _ensure_codec_loaded(self) -> None: + if self._codec is not None: + return + device = self.vllm_config.device_config.device + self._codec = load_arktts_codec( + self.model_path, + device=device, + dtype=self._codec_dtype, + role="decode", + **self._codec_kwargs, + ) + self._codec_device = torch.device(device) + + # -------------------- 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) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None: + return None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """No-op: the codec comes from ``codec.pth``, not the main checkpoint.""" + del weights + return set() + + # -------------------- request splitting -------------------- + + def _split_request_ids( + self, + ids: torch.Tensor, + seq_token_counts: list[int] | None = None, + ) -> list[torch.Tensor]: + if seq_token_counts is not None and len(seq_token_counts) > 1: + boundaries = [0] + for count in seq_token_counts: + boundaries.append(boundaries[-1] + count) + total = ids.numel() + return [ids[boundaries[i] : min(boundaries[i + 1], total)] for i in range(len(seq_token_counts))] + if is_forward_context_available(): + slices = get_forward_context().ubatch_slices + if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices): + boundaries = [0] + for size in slices: + boundaries.append(boundaries[-1] + size) + return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)] + return [ids] + + def _codes_from_runtime_info( + self, + info: dict[str, Any] | None, + device: torch.device, + ) -> torch.Tensor | None: + """Prefer the tensor payload the connector carried over input_ids.""" + if not isinstance(info, dict): + return None + codes = info.get("codes", {}).get("audio") + if not isinstance(codes, torch.Tensor) or codes.numel() == 0: + return None + codes = codes.to(device=device, dtype=torch.long, non_blocking=True).contiguous() + if codes.ndim != 2 or int(codes.shape[0]) != self._num_codebooks: + logger.warning( + "Audio8 TTS codec decoder expected codes of shape [%d, frames], got %s", + self._num_codebooks, + tuple(codes.shape), + ) + return None + return codes + + # -------------------- decode -------------------- + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + """Decode one chunk per request. + + ``input_ids`` is codebook-major flat codes + (``[cb0_f0 ... cb0_fN, cb1_f0 ...]``); tensor codes on the connector + win over ``input_ids``. Every return path must set ``model_outputs`` or + the serving layer silently drops a chunk of audio. + """ + del positions, intermediate_tensors, inputs_embeds + sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32) + + self._ensure_codec_loaded() + codec = self._codec + assert codec is not None + + if input_ids is None or input_ids.numel() == 0: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]}, + ) + + ids = input_ids.reshape(-1).to(dtype=torch.long) + per_request_ids = self._split_request_ids(ids, kwargs.get("seq_token_counts")) + num_req = len(per_request_ids) + + context_frames = [0] * num_req + if runtime_additional_information is not None: + for index, info in enumerate(runtime_additional_information[:num_req]): + meta = info.get("meta", {}) if isinstance(info, dict) else {} + context_frames[index] = int(meta.get("left_context_size", 0) or 0) + + audios: list[torch.Tensor] = [empty] * num_req + for index, request_ids in enumerate(per_request_ids): + codes_qf = ( + self._codes_from_runtime_info(runtime_additional_information[index], ids.device) + if runtime_additional_information is not None and index < len(runtime_additional_information) + else None + ) + if codes_qf is None: + if request_ids.numel() < self._num_codebooks: + continue + if int(request_ids.numel()) % self._num_codebooks != 0: + logger.warning( + "Audio8 TTS codec decoder got %d ids, not divisible by num_codebooks=%d; skipping request", + int(request_ids.numel()), + self._num_codebooks, + ) + continue + codes_qf = request_ids.reshape(self._num_codebooks, -1) + + total_frames = int(codes_qf.shape[1]) + if total_frames <= 0: + continue + new_frames = total_frames - context_frames[index] + if new_frames <= 0: + logger.warning( + "Audio8 TTS codec decoder chunk has no new frames (total=%d, context=%d)", + total_frames, + context_frames[index], + ) + continue + + if not self._logged_stats: + self._logged_stats = True + logger.info( + "Audio8 TTS codec decoder: frames=%d codebooks=%d dtype=%s sr=%d", + total_frames, + self._num_codebooks, + self._codec_dtype, + self._sample_rate, + ) + + with torch.amp.autocast("cuda", enabled=False): + # Move codes to the codec's device explicitly: a mismatch here + # silently falls back to CPU or raises deep inside the decoder. + waveform = codec.decode(codes_qf.unsqueeze(0).to(device=self._codec_device)) + waveform = waveform.reshape(-1).to(dtype=torch.float32) + + # Trim the left-context prefix proportionally: the codec's padding / + # rounding means the decoded length is not exactly + # frames * frame_size. + if context_frames[index] > 0: + cut = int(round(context_frames[index] / total_frames * int(waveform.shape[0]))) + cut = max(0, min(cut, int(waveform.shape[0]))) + waveform = waveform[cut:] + if waveform.numel() > 0: + audios[index] = waveform.contiguous() + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audios, "sr": [sr_tensor] * num_req}, + ) + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + del kwargs + if isinstance(model_outputs, OmniOutput): + return model_outputs + if not (isinstance(model_outputs, tuple) and len(model_outputs) == 2): + raise TypeError(f"Audio8TTSCodecDecoder expected (audio_tensor, sr), got {type(model_outputs)}") + audio_tensor, sample_rate = model_outputs + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audio_tensor, "sr": sample_rate}, + ) + + +__all__ = ["Audio8TTSCodecDecoder"] diff --git a/vllm_omni/model_executor/models/audio8_tts/audio8_tts_fast_ar.py b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_fast_ar.py new file mode 100644 index 000000000..5b4ff9696 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_fast_ar.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Fast AR: the residual-codebook predictor. + +Runs once per Slow AR step over ``num_codebooks + 1`` positions: position 0 is +the projected Slow AR hidden state (primes the KV cache, logits discarded), +position 1 is the sampled semantic code (not predicted, only fed in), positions +2.. are the previously sampled residual codes; only codes 1..num_codebooks-1 +are sampled here. + +Not paged: at most 10 tokens per call, SDPA over a per-call pre-allocated KV +cache. Matches the reference ``ArkttsModel._generate_codebooks``. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.config import VllmConfig +from vllm.config.vllm import set_current_vllm_config +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from .configuration_audio8_tts import Audio8TTSFastARConfig, Audio8TTSSlowARConfig +from .sampling import sample_scores + +logger = init_logger(__name__) + + +class _FastARAttention(nn.Module): + """GQA self-attention over the short codebook sequence. + + ``forward_one`` decodes a single position against a caller-owned KV cache; + there is no batched-prefill path because the codebook loop is inherently + sequential. + """ + + def __init__(self, config: Audio8TTSFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self._use_gqa = self.num_kv_heads != self.num_heads + + self.qkv_proj = QKVParallelLinear( + hidden_size=self.hidden_size, + head_size=self.head_dim, + total_num_heads=self.num_heads, + total_num_kv_heads=self.num_kv_heads, + bias=config.attention_qkv_bias, + prefix=f"{prefix}.qkv_proj", + disable_tp=True, + ) + self.o_proj = RowParallelLinear( + input_size=self.num_heads * self.head_dim, + output_size=self.hidden_size, + bias=False, + prefix=f"{prefix}.o_proj", + disable_tp=True, + ) + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + # Audio8 TTS uses interleaved (GPT-J) RoPE, not NeoX. + is_neox_style=False, + rope_parameters={"rope_theta": config.rope_theta, "rope_type": "default"}, + ) + if config.attention_qk_norm: + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + else: + self.q_norm = None + self.k_norm = None + + def forward_one( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_pos: int, + ) -> torch.Tensor: + bsz = int(hidden_states.shape[0]) + + qkv, _ = self.qkv_proj(hidden_states.reshape(bsz, -1)) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + if self.q_norm is not None: + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape) + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape) + + q, k = self.rotary_emb(position_ids.reshape(-1), q, k) + + q = q.view(bsz, self.num_heads, self.head_dim).unsqueeze(2) + k_cache[:bsz, :, cache_pos, :] = k.view(bsz, self.num_kv_heads, self.head_dim) + v_cache[:bsz, :, cache_pos, :] = v.view(bsz, self.num_kv_heads, self.head_dim) + + attn_out = F.scaled_dot_product_attention( + q, + k_cache[:bsz, :, : cache_pos + 1, :], + v_cache[:bsz, :, : cache_pos + 1, :], + scale=self.scaling, + is_causal=False, + enable_gqa=self._use_gqa, + ) + attn_out = attn_out.transpose(1, 2).reshape(bsz, -1) + output, _ = self.o_proj(attn_out) + return output + + +class _FastARMLP(nn.Module): + """SiLU-gated MLP (``w1``/``w3`` gate-up, ``w2`` down).""" + + def __init__(self, config: Audio8TTSFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[config.intermediate_size] * 2, + bias=False, + prefix=f"{prefix}.gate_up_proj", + disable_tp=True, + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + prefix=f"{prefix}.down_proj", + disable_tp=True, + ) + + 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.down_proj(F.silu(gate) * up) + return x + + +class _FastARDecoderLayer(nn.Module): + def __init__(self, config: Audio8TTSFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.self_attn = _FastARAttention(config, prefix=f"{prefix}.self_attn") + self.mlp = _FastARMLP(config, prefix=f"{prefix}.mlp") + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward_one( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_pos: int, + ) -> torch.Tensor: + hidden_states = hidden_states + self.self_attn.forward_one( + self.input_layernorm(hidden_states), position_ids, k_cache, v_cache, cache_pos + ) + return hidden_states + self.mlp(self.post_attention_layernorm(hidden_states)) + + +class Audio8TTSFastAR(nn.Module): + """Residual-codebook predictor for Audio8 TTS. + + ``forward`` returns all ``num_codebooks`` codes for one frame, with code 0 + copied from the already-sampled semantic token. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + config: Audio8TTSFastARConfig, + slow_ar_config: Audio8TTSSlowARConfig, + prefix: str = "fast_ar", + ) -> None: + super().__init__() + self._vllm_config = vllm_config + self.config = config + self.slow_ar_config = slow_ar_config + + self.layers = nn.ModuleList( + [_FastARDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)] + ) + self.fast_embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.fast_output = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.fast_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if slow_ar_config.hidden_size != config.hidden_size: + self.fast_project_in = nn.Linear(slow_ar_config.hidden_size, config.hidden_size, bias=True) + else: + # 0.6b: dim == fast_dim == 896, and the checkpoint has no + # fast_project_in tensor. + self.fast_project_in = nn.Identity() + + self._num_codebooks = int(config.num_codebooks) + self._fast_dim = int(config.hidden_size) + self._k_cache: torch.Tensor | None = None + self._v_cache: torch.Tensor | None = None + self._pos_ids: torch.Tensor | None = None + + def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None: + max_seq = self._num_codebooks + 1 + if ( + self._k_cache is not None + and self._pos_ids is not None + and self._k_cache.shape[1] >= bsz + and self._k_cache.device == device + and self._k_cache.dtype == dtype + and self._pos_ids.shape[0] >= bsz + ): + return + self._k_cache = torch.empty( + self.config.num_hidden_layers, + bsz, + self.config.num_key_value_heads, + max_seq, + self.config.head_dim, + dtype=dtype, + device=device, + ) + self._v_cache = torch.empty_like(self._k_cache) + self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device).unsqueeze(0).expand(bsz, -1).contiguous() + + def _forward_one(self, embed: torch.Tensor, position_ids: torch.Tensor, cache_pos: int) -> torch.Tensor: + assert self._k_cache is not None and self._v_cache is not None + hidden = embed + for layer_idx, layer in enumerate(self.layers): + hidden = layer.forward_one( + hidden, + position_ids, + self._k_cache[layer_idx], + self._v_cache[layer_idx], + cache_pos, + ) + return hidden + + @torch.inference_mode() + def forward( + self, + slow_ar_hidden: torch.Tensor, + semantic_token_id: torch.Tensor, + *, + do_sample: bool = True, + temperature: float = 0.7, + top_k: int = 50, + top_p: float = 0.9, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Predict one frame of codec codes. + + ``slow_ar_hidden`` must already be post-final-norm when + ``norm_fastlayer_input`` is set (it is for the released checkpoint) -- + vLLM's Qwen2Model returns exactly that. Returns ``[B, num_codebooks]`` + int64 codes; code 0 is the semantic code, codes 1.. come from this module. + """ + bsz = int(slow_ar_hidden.shape[0]) + num_cb = self._num_codebooks + device = slow_ar_hidden.device + dtype = slow_ar_hidden.dtype + + semantic_begin = int(self.slow_ar_config.semantic_begin_id) + codebook_size = int(self.slow_ar_config.codebook_size) + # Non-semantic tokens (EOS) clamp to code 0; those frames are dropped + # downstream because the request finishes on that token. + semantic_code = (semantic_token_id.reshape(bsz) - semantic_begin).clamp(min=0, max=codebook_size - 1) + + codes = torch.empty(bsz, num_cb, dtype=torch.long, device=device) + codes[:, 0] = semantic_code + + self._ensure_buffers(bsz, device, dtype) + pos_ids = self._pos_ids + assert pos_ids is not None + + # Position 0 primes the cache; its logits are unused. + projected = self.fast_project_in(slow_ar_hidden.reshape(bsz, -1)) + self._forward_one(projected, pos_ids[:bsz, 0], 0) + + embed = self.fast_embeddings(semantic_code) + for step in range(1, num_cb): + hidden = self._forward_one(embed, pos_ids[:bsz, step], step) + logits = self.fast_output(self.fast_norm(hidden)).float() + next_ids = sample_scores( + logits, + temperature=temperature, + top_k=top_k, + top_p=top_p, + do_sample=do_sample, + generator=generator, + ) + codes[:, step] = next_ids + if step < num_cb - 1: + embed = self.fast_embeddings(next_ids) + + return codes + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load already-remapped Fast AR weights (see the Slow AR remapper).""" + with set_current_vllm_config(self._vllm_config): + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded: set[str] = set() + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + for name, loaded_weight in weights: + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped not in params_dict: + continue + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded.add(mapped) + handled = True + break + if handled or 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.add(name) + return loaded + + +__all__ = ["Audio8TTSFastAR"] diff --git a/vllm_omni/model_executor/models/audio8_tts/audio8_tts_slow_ar.py b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_slow_ar.py new file mode 100644 index 000000000..9d2cd884e --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_slow_ar.py @@ -0,0 +1,755 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview -- Slow AR model (Stage 0). + +Text (+ optional reference voice) -> one semantic token per codec frame; the +frame's remaining 9 codebooks are produced by the nested Fast AR. + +Backbone is vLLM's ``Qwen2Model`` (qkv bias, no q/k norm), with three +adjustments: interleaved (GPT-J) RoPE, multi-codebook input embedding at +semantic positions (summed, unlike Fish Speech's ``1/sqrt(Q+1)`` rescale), and +logits masked to the semantic range + ``<|im_end|>`` then Repetition-Aware +sampled. + +Streaming contract is delta: each decode step appends one frame of +``num_codebooks`` codes to ``codes.audio``. +""" + +from __future__ import annotations + +import copy +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from transformers import AutoTokenizer +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.qwen2 import Qwen2Model +from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix +from vllm.sequence import IntermediateTensors +from vllm.v1.outputs import SamplerOutput + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.utils.speaker_cache import get_speaker_cache + +from .audio8_tts_fast_ar import Audio8TTSFastAR +from .codec_utils import encode_reference_audio_codes +from .configuration_audio8_tts import ( + Audio8TTSConfig, + Audio8TTSFastARConfig, + Audio8TTSSlowARConfig, +) +from .prompt_utils import build_voice_clone_prompt_ids +from .sampling import SAMPLING_EPS, ras_sample_batch + +logger = init_logger(__name__) + +#: Defaults from ``generation_config.json`` of the released checkpoint. +DEFAULT_TEMPERATURE = 0.7 +DEFAULT_TOP_K = 50 +DEFAULT_TOP_P = 0.9 + + +def _remap_audio8_tts_weights( + weights: Iterable[tuple[str, torch.Tensor]], + *, + q_size: int, + kv_size: int, + fast_q_size: int, + fast_kv_size: int, +) -> Iterable[tuple[str, torch.Tensor]]: + """Rename Audio8 checkpoint tensors to vLLM / Qwen2 names. + + The checkpoint is flat (no ``text_model.`` prefix). Slow AR tensors move + under ``model.``, Fast AR tensors under ``fast_ar.``, and ``wqkv`` splits + into ``{q,k,v}_proj``. See the body for the exact mapping table. + """ + + def rewrite_block(suffix: str) -> str: + for old, new in ( + (".attention.wo.", ".self_attn.o_proj."), + (".attention.q_norm.", ".self_attn.q_norm."), + (".attention.k_norm.", ".self_attn.k_norm."), + (".attention_norm.", ".input_layernorm."), + (".feed_forward.w1.", ".mlp.gate_proj."), + (".feed_forward.w3.", ".mlp.up_proj."), + (".feed_forward.w2.", ".mlp.down_proj."), + (".ffn_norm.", ".post_attention_layernorm."), + ): + suffix = suffix.replace(old, new) + return suffix + + for name, tensor in weights: + # RoPE tables are recomputed by vLLM. + if name.endswith(("freqs_cis", "fast_freqs_cis")) or "rotary_emb.inv_freq" in name: + continue + + if name.startswith("fast_layers."): + suffix = name[len("fast_layers.") :] + if ".attention.wqkv." in suffix: + layer_prefix, _, param = suffix.partition(".attention.wqkv.") + yield f"fast_ar.layers.{layer_prefix}.self_attn.q_proj.{param}", tensor[:fast_q_size] + yield ( + f"fast_ar.layers.{layer_prefix}.self_attn.k_proj.{param}", + tensor[fast_q_size : fast_q_size + fast_kv_size], + ) + yield ( + f"fast_ar.layers.{layer_prefix}.self_attn.v_proj.{param}", + tensor[fast_q_size + fast_kv_size :], + ) + continue + yield f"fast_ar.layers.{rewrite_block(suffix)}", tensor + continue + + if name.startswith(("fast_embeddings.", "fast_output.", "fast_norm.", "fast_project_in.")): + yield f"fast_ar.{name}", tensor + continue + + if name.startswith("layers."): + suffix = name[len("layers.") :] + if ".attention.wqkv." in suffix: + layer_prefix, _, param = suffix.partition(".attention.wqkv.") + yield f"model.layers.{layer_prefix}.self_attn.q_proj.{param}", tensor[:q_size] + yield ( + f"model.layers.{layer_prefix}.self_attn.k_proj.{param}", + tensor[q_size : q_size + kv_size], + ) + yield f"model.layers.{layer_prefix}.self_attn.v_proj.{param}", tensor[q_size + kv_size :] + continue + yield f"model.layers.{rewrite_block(suffix)}", tensor + continue + + if name == "embeddings.weight": + yield "model.embed_tokens.weight", tensor + continue + if name == "norm.weight": + yield "model.norm.weight", tensor + continue + + yield name, tensor + + +class Audio8TTSSlowARForConditionalGeneration(nn.Module): + """Stage 0: text -> semantic tokens + residual codec codes.""" + + prefer_model_sampler = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + + config: Audio8TTSConfig = vllm_config.model_config.hf_config # type: ignore[assignment] + self.config = config + self.text_config: Audio8TTSSlowARConfig = config.get_text_config() + self.fast_ar_config: Audio8TTSFastARConfig = config.fast_ar_config + + self._semantic_begin_id = int(self.text_config.semantic_begin_id) + self._semantic_end_id = int(self.text_config.semantic_end_id) + self._num_semantic_ids = self._semantic_end_id - self._semantic_begin_id + 1 + self._eos_token_id = int(self.text_config.eos_token_id) + self._pad_token_id = int(self.text_config.pad_token_id) + self._codebook_size = int(self.text_config.codebook_size) + self._num_codebooks = int(self.text_config.num_codebooks) + self._ras_window_size = int(config.ras_window_size) + self._ras_temperature = float(config.ras_temperature) + self._ras_top_p = float(config.ras_top_p) + + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + self.mtp_hidden_size = int(self.text_config.hidden_size) + self.talker_mtp_output_key = ("codes", "audio") + self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("hidden_states", "last")} + self.talker_mtp_graph_safe = True + + self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + self._fix_rope_style() + + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.text_config.vocab_size, + self.text_config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(self.text_config.vocab_size) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + # Summed multi-codebook input embedding table. + self.codebook_embeddings = nn.Embedding( + self._codebook_size * self._num_codebooks, + self.text_config.hidden_size, + ) + + # Fast AR shares the parent's VllmConfig but must not share the + # compilation context: see the Fish Speech note on why this is + # copy.copy and not dataclasses.replace (pydantic validators re-run and + # reject an already-rebound compilation backend). + fast_ar_compilation = copy.copy(vllm_config.compilation_config) + fast_ar_compilation.static_forward_context = {} + self._fast_ar_vllm_config = copy.copy(vllm_config) + self._fast_ar_vllm_config.compilation_config = fast_ar_compilation + from vllm.config.vllm import set_current_vllm_config + + with set_current_vllm_config(self._fast_ar_vllm_config): + self.fast_ar = Audio8TTSFastAR( + vllm_config=self._fast_ar_vllm_config, + config=self.fast_ar_config, + slow_ar_config=self.text_config, + prefix="fast_ar", + ) + + # Constant logits mask: semantic codes plus <|im_end|>. Safe as a + # non-persistent buffer under vLLM (module built on the target device), + # unlike under ``PreTrainedModel.from_pretrained`` on transformers >= 5, + # which meta-inits and then empty_like's every non-persistent buffer. + vocab = int(self.text_config.vocab_size) + allowed = torch.zeros((vocab,), dtype=torch.bool) + allowed[self._semantic_begin_id : min(self._semantic_end_id + 1, vocab)] = True + if self._eos_token_id < vocab: + allowed[self._eos_token_id] = True + self.register_buffer("_semantic_allowed_mask", allowed, persistent=False) + + self._speaker_cache = get_speaker_cache() + self._tokenizer = None + + def _fix_rope_style(self) -> None: + """Rebuild RoPE as interleaved (GPT-J); vLLM's Qwen2 defaults to NeoX. + + This also replaces the reference implementation's ``freqs_cis`` table + entirely: ``get_rope`` computes its cos/sin cache in the constructor, so + there is no checkpoint-absent buffer for a loader to leave uninitialised. + """ + from vllm.model_executor.layers.rotary_embedding import get_rope + + rope_parameters = dict(self.text_config.rope_parameters or {}) + rope_parameters.setdefault("rope_type", "default") + for layer in self.model.layers: + attn = layer.self_attn + attn.rotary_emb = get_rope( + head_size=attn.head_dim, + max_position=self.text_config.max_position_embeddings, + is_neox_style=False, + rope_parameters=rope_parameters, + ) + logger.info("Audio8 TTS: switched %d layers to interleaved RoPE", len(self.model.layers)) + + # -------------------- vLLM hooks -------------------- + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_input_ids(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, positions, intermediate_tensors, inputs_embeds) + + 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: + return None + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is None: + return None + return logits.masked_fill(~self._semantic_allowed_mask, float("-inf")) + + # -------------------- Repetition-Aware Sampling -------------------- + + def _compact_semantic_logits(self, logits: torch.Tensor) -> torch.Tensor: + """Slice the allowed vocabulary into ``[B, num_semantic_ids + 1]``. + + Local ids ``0..num_semantic_ids-1`` are semantic codes; the last column + is ``<|im_end|>``. Sorting 4097 columns instead of the full 155776 + vocabulary is the difference between a cheap and an expensive AR step, + and is exactly equivalent because everything else is already ``-inf``. + """ + semantic = logits[:, self._semantic_begin_id : self._semantic_end_id + 1] + eos = logits[:, self._eos_token_id : self._eos_token_id + 1] + return torch.cat((semantic, eos), dim=-1) + + def _local_to_token_id(self, local_ids: torch.Tensor) -> torch.Tensor: + is_eos = local_ids >= self._num_semantic_ids + return torch.where(is_eos, torch.full_like(local_ids, self._eos_token_id), local_ids + self._semantic_begin_id) + + def _recent_local_ids(self, output_token_ids: list[list[int]], num_reqs: int, device: torch.device): + """Build the ``[B, W]`` RAS window in compact id space. + + Reads the host-side decoded-token history the AR runner already + materialised, so this costs one small H2D copy and never syncs the GPU. + Returns ``None`` when no request has history yet. + """ + window = self._ras_window_size + if window <= 0: + return None + rows: list[list[int]] = [] + any_history = False + for req_idx in range(num_reqs): + history = output_token_ids[req_idx] if req_idx < len(output_token_ids) else [] + recent = [int(token) for token in history[-window:]] + local = [ + token - self._semantic_begin_id if self._semantic_begin_id <= token <= self._semantic_end_id else -1 + for token in recent + ] + any_history = any_history or bool(local) + rows.append([-1] * (window - len(local)) + local) + if not any_history: + return None + return torch.tensor(rows, dtype=torch.long, device=device) + + def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> SamplerOutput | None: + """Sample the semantic token with the reference filter order + RAS. + + Returns ``None`` to fall back to vLLM's sampler when the request mix + needs features RAS does not model (logprobs, penalties, bad words). + """ + if logits is None or logits.numel() == 0: + return None + if sampling_metadata.max_num_logprobs is not None: + return None + if not sampling_metadata.no_penalties: + return None + if sampling_metadata.bad_words_token_ids: + return None + + logits = logits.to(torch.float32) + if sampling_metadata.allowed_token_ids_mask is not None: + num_reqs = int(logits.shape[0]) + logits.masked_fill_(sampling_metadata.allowed_token_ids_mask[:num_reqs], float("-inf")) + for processor in sampling_metadata.logitsprocs.non_argmax_invariant: + logits = processor.apply(logits) + + compact = self._compact_semantic_logits(logits) + num_reqs = int(compact.shape[0]) + device = compact.device + + temperature = sampling_metadata.temperature + if sampling_metadata.all_greedy or temperature is None: + local_ids = compact.argmax(dim=-1) + else: + temperature = temperature[:num_reqs] + top_p = DEFAULT_TOP_P if sampling_metadata.top_p is None else sampling_metadata.top_p[:num_reqs] + top_k = DEFAULT_TOP_K if sampling_metadata.top_k is None else sampling_metadata.top_k[:num_reqs] + recent = self._recent_local_ids(sampling_metadata.output_token_ids, num_reqs, device) + local_ids = ras_sample_batch( + compact, + recent, + temperature=temperature, + top_k=top_k, + top_p=top_p, + ras_temperature=self._ras_temperature, + ras_top_p=self._ras_top_p, + num_semantic_ids=self._num_semantic_ids, + generators=sampling_metadata.generators, + num_reqs=num_reqs, + ) + + token_ids = self._local_to_token_id(local_ids).to(dtype=torch.int32) + return SamplerOutput(sampled_token_ids=token_ids.unsqueeze(-1), logprobs_tensors=None) + + # -------------------- Omni multimodal plumbing -------------------- + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + + hidden = model_outputs + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") or [] + + audio_codes_list: list[torch.Tensor] = [] + for info in info_dicts: + if not isinstance(info, dict): + continue + codes = info.get("codes", {}).get("audio") + if isinstance(codes, torch.Tensor): + audio_codes_list.append(codes) + + if not audio_codes_list: + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + audio_codes = torch.cat(audio_codes_list, dim=0) + hidden = hidden[: int(audio_codes.shape[0])] + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={"audio_codes": audio_codes}) + + # -------------------- 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]]: + 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 key, value in additional_information.items(): + merged.setdefault(key, value) + info_dict = merged + + span_len = int(input_ids.shape[0]) + if span_len <= 0: + embeds = input_embeds if input_embeds is not None else self.embed_input_ids(input_ids) + return input_ids, embeds, {} + + if span_len > 1: + return self._preprocess_prefill(input_ids, info_dict, span_len) + return self._preprocess_decode(input_ids, info_dict) + + def _prefill_chunk( + self, + prompt_embeds_buf: torch.Tensor, + start: int, + span_len: int, + device: torch.device, + ) -> torch.Tensor: + """Slice ``span_len`` rows out of the pinned CPU prompt-embed buffer.""" + total = int(prompt_embeds_buf.shape[0]) + begin = max(0, min(start, total)) + end = max(0, min(start + span_len, total)) + take = prompt_embeds_buf[begin:end] + if int(take.shape[0]) < span_len: + pad_rows = span_len - int(take.shape[0]) + pad_embed = self.embed_input_ids( + torch.tensor([self._pad_token_id], device=device, dtype=torch.long) + ).reshape(1, -1) + take = torch.cat([take.to(device=device, dtype=torch.bfloat16), pad_embed.expand(pad_rows, -1)], dim=0) + return take.to(dtype=torch.bfloat16) + return take.to(device=device, dtype=torch.bfloat16, non_blocking=True) + + def _preprocess_prefill( + self, + input_ids: torch.Tensor, + info_dict: dict[str, Any], + span_len: int, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + device = input_ids.device + embed = info_dict.get("embed", {}) + prompt_embeds_buf = embed.get("prefill") + is_first_chunk = not isinstance(prompt_embeds_buf, torch.Tensor) or prompt_embeds_buf.ndim != 2 + + if is_first_chunk: + if bool(info_dict.get("audio8_structured_voice_clone", False)): + prompt_embeds = self._build_voice_clone_prefill_embeds(info_dict) + else: + prompt_embeds = self.embed_input_ids(input_ids.reshape(1, -1).to(torch.long)).squeeze(0) + prompt_embeds_buf = prompt_embeds.detach().to("cpu", dtype=torch.bfloat16).contiguous() + if not prompt_embeds_buf.is_pinned(): + prompt_embeds_buf = prompt_embeds_buf.pin_memory() + offset = 0 + else: + offset = int(info_dict.get("meta", {}).get("prefill_offset", 0) or 0) + + total_prompt_len = int(prompt_embeds_buf.shape[0]) + chunk = self._prefill_chunk(prompt_embeds_buf, offset, span_len, device) + next_offset = min(offset + span_len, total_prompt_len) + + # No codes are emitted during prefill; a zero frame per position keeps + # the codes buffer aligned with the hidden-state span. + zeros = torch.zeros((chunk.shape[0], self._num_codebooks), device=device, dtype=torch.long) + info_update = { + "embed": {"prefill": prompt_embeds_buf if next_offset < total_prompt_len else None}, + "meta": {"prefill_offset": next_offset}, + "codes": {"audio": zeros}, + } + return torch.full_like(input_ids, self._pad_token_id), chunk, info_update + + def _preprocess_decode( + self, + input_ids: torch.Tensor, + info_dict: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + device = input_ids.device + token_embed = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).to( + device=device, dtype=torch.bfloat16 + ) + last_hidden = info_dict.get("hidden_states", {}).get("last") + if not isinstance(last_hidden, torch.Tensor): + # First decode step right after prefill: no Fast AR input yet. + logger.warning("Audio8 TTS preprocess: hidden_states.last missing; emitting text-only embedding") + return input_ids, token_embed.reshape(1, -1), {} + + # Codebook embeddings are added in talker_mtp, using this step's Fast AR + # output; adding them here would use the previous step's codes. + info_update = { + "mtp_inputs": ( + last_hidden.to(device=device, dtype=torch.bfloat16).reshape(1, -1), + torch.zeros(1, self.text_config.hidden_size, device=device, dtype=torch.bfloat16), + ), + } + return input_ids, token_embed.reshape(1, -1), info_update + + 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().contiguous()}} + + # -------------------- voice cloning -------------------- + + def _get_tokenizer(self): + if self._tokenizer is None: + self._tokenizer = AutoTokenizer.from_pretrained(self.model_path) + return self._tokenizer + + def _codec_kwargs(self) -> dict[str, int]: + return { + "post_n_layer": int(self.config.codec_post_n_layer), + "post_n_head": int(self.config.codec_post_n_head), + "post_n_local_heads": int(self.config.codec_post_n_local_heads), + "post_intermediate_size": int(self.config.codec_post_intermediate_size), + } + + def _build_voice_clone_prefill_embeds(self, info_dict: dict[str, Any]) -> torch.Tensor: + """Encode the reference audio, build the clone prompt, embed the codes.""" + text = info_dict.get("text") + ref_text = info_dict.get("ref_text") + if not isinstance(text, str) or not isinstance(ref_text, str): + raise ValueError("Audio8 TTS voice cloning requires string 'text' and 'ref_text'") + + device = self.codebook_embeddings.weight.device + cache_key = None + voice_name = info_dict.get("voice_name") + if isinstance(voice_name, str) and voice_name: + cache_key = self._speaker_cache.make_cache_key( + voice_name, + model_type="audio8_tts", + created_at=int(info_dict.get("voice_created_at") or 0), + ) + cached = self._speaker_cache.get(cache_key) + if cached is not None: + logger.debug("Speaker cache HIT for Audio8 TTS speaker '%s'", voice_name) + return self._embed_voice_clone_prompt( + text, ref_text, cached["ref_codes_fq"].to(device=device, dtype=torch.long) + ) + + ref_audio_sr = info_dict.get("ref_audio_sr") + if not isinstance(ref_audio_sr, int): + raise ValueError("Audio8 TTS voice cloning requires integer 'ref_audio_sr'") + ref_audio_wav = info_dict.get("ref_audio_wav") + if ref_audio_wav is None: + raise ValueError("Audio8 TTS voice cloning requires 'ref_audio_wav'") + if not isinstance(ref_audio_wav, torch.Tensor): + ref_audio_wav = torch.from_numpy(np.asarray(ref_audio_wav, dtype=np.float32)) + + ref_codes_fq = encode_reference_audio_codes( + self.model_path, + ref_audio_wav, + ref_audio_sr, + device=device, + **self._codec_kwargs(), + ) + if cache_key is not None: + self._speaker_cache.put(cache_key, {"ref_codes_fq": ref_codes_fq.detach().cpu()}) + logger.debug("Speaker cache STORE for Audio8 TTS speaker '%s'", voice_name) + return self._embed_voice_clone_prompt(text, ref_text, ref_codes_fq) + + def _embed_voice_clone_prompt( + self, + text: str, + ref_text: str, + ref_codes_fq: torch.Tensor, + ) -> torch.Tensor: + """Embed the clone prompt, summing codebooks over the reference frames. + + Args: + ref_codes_fq: ``[frames, num_codebooks]`` reference codec codes. + """ + device = self.codebook_embeddings.weight.device + ref_codes_fq = ref_codes_fq.to(device=device, dtype=torch.long) + semantic_token_ids = (ref_codes_fq[:, 0] + self._semantic_begin_id).tolist() + prompt_ids, ref_start, _, _ = build_voice_clone_prompt_ids( + self._get_tokenizer(), text, ref_text, semantic_token_ids + ) + prompt = torch.tensor(prompt_ids, dtype=torch.long, device=device) + embeds = self.embed_input_ids(prompt.unsqueeze(0)).squeeze(0).to(dtype=torch.bfloat16) + + frames = int(ref_codes_fq.shape[0]) + if frames <= 0: + return embeds + codebooks = min(int(ref_codes_fq.shape[1]), self._num_codebooks) + offsets = (torch.arange(codebooks, device=device, dtype=torch.long) * self._codebook_size).unsqueeze(0) + codes = ref_codes_fq[:, :codebooks].clamp(min=0, max=self._codebook_size - 1) + offsets + codebook_sum = self.codebook_embeddings(codes).sum(dim=1).to(dtype=embeds.dtype) + + result = embeds.clone() + # Audio8 TTS sums the embeddings directly -- no 1/sqrt(Q+1) rescale. + result[ref_start : ref_start + frames] += codebook_sum + return result.to(dtype=torch.bfloat16) + + # -------------------- GPU-side Fast AR fast path -------------------- + + @torch.inference_mode() + def talker_mtp( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor, + last_talker_hidden: torch.Tensor, + text_step: torch.Tensor, + seed: int | None = None, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the Fast AR and fold its codes into this step's input embedding. + + Returns: + ``(inputs_embeds, audio_codes)`` where ``audio_codes`` is + ``[B, num_codebooks]``. + """ + del text_step + bsz = int(input_ids.shape[0]) + device = input_embeds.device + + token_ids = input_ids.reshape(bsz).to(dtype=torch.long, device=device) + past_hidden = last_talker_hidden.reshape(bsz, -1).to(dtype=torch.bfloat16, device=device) + + temperature = kwargs.get("temperature") + temperature = DEFAULT_TEMPERATURE if temperature is None else float(temperature) + do_sample = kwargs.get("do_sample") + generator = kwargs.get("generator") + if generator is None and seed is not None: + generator = torch.Generator(device=device) + generator.manual_seed(int(seed)) + + _top_k = kwargs.get("top_k") + _top_p = kwargs.get("top_p") + audio_codes = self.fast_ar( + slow_ar_hidden=past_hidden, + semantic_token_id=token_ids, + do_sample=(temperature >= SAMPLING_EPS) if do_sample is None else bool(do_sample), + temperature=temperature, + top_k=int(DEFAULT_TOP_K if _top_k is None else _top_k), + top_p=float(DEFAULT_TOP_P if _top_p is None else _top_p), + generator=generator, + ) + + embeds = input_embeds.reshape(bsz, -1) + offsets = (torch.arange(self._num_codebooks, device=device, dtype=torch.long) * self._codebook_size).unsqueeze( + 0 + ) + codebook_sum = self.codebook_embeddings(audio_codes + offsets).sum(dim=1).to(dtype=embeds.dtype) + is_semantic = (token_ids >= self._semantic_begin_id) & (token_ids <= self._semantic_end_id) + embeds = torch.where(is_semantic.unsqueeze(-1), embeds + codebook_sum, embeds) + return embeds, audio_codes.to(dtype=torch.long) + + # -------------------- prompt length estimation -------------------- + + @staticmethod + def estimate_prompt_len_from_additional_information( + additional_information: dict[str, Any] | None, + **kwargs: Any, + ) -> int: + """Upper-bound the text-only prompt length for placeholder allocation.""" + del kwargs + info = additional_information or {} + text = info.get("text", "") + if isinstance(text, list): + text = text[0] if text else "" + # ~1 token per character is a safe bound for CJK; +64 covers the + # chat-template scaffolding. + return max(2, len(str(text)) + 64) + + # -------------------- weight loading -------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + text_config = self.text_config + fast_config = self.fast_ar_config + remapped = _remap_audio8_tts_weights( + weights, + q_size=text_config.num_attention_heads * text_config.head_dim, + kv_size=text_config.num_key_value_heads * text_config.head_dim, + fast_q_size=fast_config.num_attention_heads * fast_config.head_dim, + fast_kv_size=fast_config.num_key_value_heads * fast_config.head_dim, + ) + + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + unexpected: list[str] = [] + + for name, loaded_weight in remapped: + if name == "model.embed_tokens.weight" and self.text_config.tie_word_embeddings: + lm_key = "lm_head.weight" + if lm_key in params_dict: + param = params_dict[lm_key] + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + loaded_params.add(lm_key) + + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped not in params_dict: + continue + param = params_dict[mapped] + loader = getattr(param, "weight_loader", default_weight_loader) + if loader == default_weight_loader: + loader(param, loaded_weight) + else: + loader(param, loaded_weight, shard_id) + loaded_params.add(mapped) + handled = True + break + if handled: + continue + + if name in params_dict: + param = params_dict[name] + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + loaded_params.add(name) + else: + unexpected.append(name) + + if unexpected: + raise ValueError( + f"Audio8 TTS Slow AR received {len(unexpected)} unmapped checkpoint tensors, " + f"e.g. {sorted(unexpected)[:5]}. The weight remapper is out of sync with the checkpoint." + ) + missing = sorted(set(params_dict) - loaded_params) + if missing: + raise ValueError(f"Audio8 TTS Slow AR is missing weights for {missing[:5]} ({len(missing)} total)") + logger.info("Loaded %d weights for Audio8TTSSlowARForConditionalGeneration", len(loaded_params)) + + # The reference precomputes RoPE in bf16; keeping fp32 cos/sin here + # shifts the logits enough to trigger early EOS (same failure mode as + # Fish Speech). + truncated = 0 + for module in self.modules(): + cache = getattr(module, "cos_sin_cache", None) + if isinstance(cache, torch.Tensor): + module.cos_sin_cache = cache.to(torch.bfloat16).to(cache.dtype) + truncated += 1 + if truncated: + logger.info("Audio8 TTS: truncated %d RoPE cos_sin_cache buffers to bf16", truncated) + + return loaded_params + + +__all__ = ["Audio8TTSSlowARForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/audio8_tts/codec.py b/vllm_omni/model_executor/models/audio8_tts/codec.py new file mode 100644 index 000000000..71f43642b --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/codec.py @@ -0,0 +1,600 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS neural audio codec (44.1 kHz, 10 codebooks). + +Vendored from the reference ``modeling_arktts_codec.py`` shipped with +``Audio8/Audio8-TTS-Preview-0.6b`` (Apache-2.0). The module tree and attribute +names are reproduced verbatim so ``codec.pth`` loads under ``strict=True``. + +Inference-only deviations: no ``torch.jit.script`` on Snake (it fights +``torch.compile`` / CUDA graphs), RoPE cached in a plain dict instead of a +non-persistent buffer (the pattern that leaves the reference LM with +uninitialised RoPE under ``from_pretrained`` on transformers >= 5), and +:func:`build_arktts_codec` prunes the unused encoder or decoder half per stage. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.nn.utils import weight_norm as legacy_weight_norm +from torch.nn.utils.parametrizations import weight_norm + +# The codec's own transformer blocks always use base 10000, independent of the +# 1e6 rope_base of the DualAR language model. +_CODEC_ROPE_BASE = 10000.0 + + +def _rope_table(length: int, head_dim: int, base: float, device: torch.device) -> Tensor: + frequencies = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) + phases = torch.outer(torch.arange(length, device=device), frequencies) + values = torch.polar(torch.ones_like(phases), phases) + return torch.stack((values.real, values.imag), dim=-1).to(torch.bfloat16) + + +def _apply_rope(x: Tensor, values: Tensor) -> Tensor: + shaped = x.float().reshape(*x.shape[:-1], -1, 2) + values = values.view(1, shaped.shape[1], 1, shaped.shape[3], 2) + output = torch.stack( + ( + shaped[..., 0] * values[..., 0] - shaped[..., 1] * values[..., 1], + shaped[..., 1] * values[..., 0] + shaped[..., 0] * values[..., 1], + ), + dim=-1, + ) + return output.flatten(3).to(x.dtype) + + +class ArkttsCodecRMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-5) -> None: + super().__init__() + self.eps = float(eps) + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + output = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) + return output.to(x.dtype) * self.weight + + +class ArkttsCodecLayerScale(nn.Module): + def __init__(self, dim: int, init_values: float = 1e-2) -> None: + super().__init__() + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x * self.gamma + + +class ArkttsCodecTransformerConfig: + """Plain option holder for the codec's transformer blocks. + + Deliberately not a dataclass / msgspec Struct: it mirrors a reference + dataclass whose only job is to carry constructor arguments, and it is + instantiated from Python literals inside this module only. + """ + + def __init__( + self, + n_layer: int, + n_head: int, + dim: int, + intermediate_size: int, + n_local_heads: int = -1, + head_dim: int = 64, + rope_base: float = _CODEC_ROPE_BASE, + norm_eps: float = 1e-5, + channels_first: bool = True, + ) -> None: + self.n_layer = int(n_layer) + self.n_head = int(n_head) + self.dim = int(dim) + self.intermediate_size = int(intermediate_size) + self.n_local_heads = int(n_head) if n_local_heads == -1 else int(n_local_heads) + self.head_dim = int(head_dim) + self.rope_base = float(rope_base) + self.norm_eps = float(norm_eps) + self.channels_first = bool(channels_first) + + +class ArkttsCodecAttention(nn.Module): + def __init__(self, config: ArkttsCodecTransformerConfig) -> None: + super().__init__() + total = (config.n_head + 2 * config.n_local_heads) * config.head_dim + self.wqkv = nn.Linear(config.dim, total, bias=False) + self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False) + self.n_head = config.n_head + self.n_local_heads = config.n_local_heads + self.head_dim = config.head_dim + + def forward(self, x: Tensor, rope_values: Tensor, mask: Tensor) -> Tensor: + batch, length, _ = x.shape + query_size = self.n_head * self.head_dim + kv_size = self.n_local_heads * self.head_dim + query, key, value = self.wqkv(x).split((query_size, kv_size, kv_size), dim=-1) + query = query.view(batch, length, self.n_head, self.head_dim) + key = key.view(batch, length, self.n_local_heads, self.head_dim) + value = value.view(batch, length, self.n_local_heads, self.head_dim) + query = _apply_rope(query, rope_values).transpose(1, 2) + key = _apply_rope(key, rope_values).transpose(1, 2) + value = value.transpose(1, 2) + output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=mask, + enable_gqa=self.n_local_heads != self.n_head, + ) + output = output.transpose(1, 2).contiguous().view(batch, length, query_size) + return self.wo(output) + + +class ArkttsCodecFeedForward(nn.Module): + def __init__(self, config: ArkttsCodecTransformerConfig) -> 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)) + + +class ArkttsCodecTransformerBlock(nn.Module): + def __init__(self, config: ArkttsCodecTransformerConfig) -> None: + super().__init__() + self.attention = ArkttsCodecAttention(config) + self.feed_forward = ArkttsCodecFeedForward(config) + self.ffn_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps) + self.attention_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps) + self.attention_layer_scale = ArkttsCodecLayerScale(config.dim) + self.ffn_layer_scale = ArkttsCodecLayerScale(config.dim) + + def forward(self, x: Tensor, rope_values: Tensor, mask: Tensor) -> Tensor: + hidden = x + self.attention_layer_scale(self.attention(self.attention_norm(x), rope_values, mask)) + return hidden + self.ffn_layer_scale(self.feed_forward(self.ffn_norm(hidden))) + + +class ArkttsCodecWindowTransformer(nn.Module): + """Causal transformer with a sliding attention window.""" + + def __init__( + self, + config: ArkttsCodecTransformerConfig, + input_dim: int, + window_size: int | None, + ) -> None: + super().__init__() + self.layers = nn.ModuleList([ArkttsCodecTransformerBlock(config) for _ in range(config.n_layer)]) + self.norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps) + self.window_size = window_size + self.channels_first = config.channels_first + self.input_proj = nn.Linear(input_dim, config.dim) if input_dim != config.dim else nn.Identity() + self.output_proj = nn.Linear(config.dim, input_dim) if input_dim != config.dim else nn.Identity() + self.look_ahead_conv = nn.Identity() + self.head_dim = config.head_dim + self.rope_base = config.rope_base + self._mask_cache: dict[tuple[int, str], Tensor] = {} + self._rope_cache: dict[tuple[int, str], Tensor] = {} + + def _mask(self, length: int, device: torch.device) -> Tensor: + key = (length, str(device)) + cached = self._mask_cache.get(key) + if cached is not None: + return cached + row = torch.arange(length, device=device)[:, None] + column = torch.arange(length, device=device)[None, :] + mask = column <= row + if self.window_size is not None: + mask &= column >= (row - self.window_size + 1).clamp_min(0) + mask = mask[None, None] + self._mask_cache[key] = mask + return mask + + def _rope(self, length: int, device: torch.device) -> Tensor: + key = (length, str(device)) + cached = self._rope_cache.get(key) + if cached is not None: + return cached + values = _rope_table(length, self.head_dim, self.rope_base, device) + self._rope_cache[key] = values + return values + + def forward(self, x: Tensor) -> Tensor: + if self.channels_first: + x = x.transpose(1, 2) + x = self.look_ahead_conv(self.input_proj(x)) + length = int(x.shape[1]) + mask = self._mask(length, x.device) + rope_values = self._rope(length, x.device) + for layer in self.layers: + x = layer(x, rope_values, mask) + x = self.output_proj(self.norm(x)) + return x.transpose(1, 2) if self.channels_first else x + + +def _extra_padding(x: Tensor, kernel_size: int, stride: int, padding_total: int = 0) -> int: + length = x.shape[-1] + frames = (length - kernel_size + padding_total) / stride + 1 + ideal = (math.ceil(frames) - 1) * stride + kernel_size - padding_total + return ideal - length + + +class ArkttsCausalConv1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + dilation: int = 1, + stride: int = 1, + groups: int = 1, + ) -> None: + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, groups=groups) + self.stride = stride + self.kernel_size = (kernel_size - 1) * dilation + 1 + self.padding = self.kernel_size - self.stride + + def forward(self, x: Tensor) -> Tensor: + right = _extra_padding(x, self.kernel_size, self.stride, self.padding) + return self.conv(F.pad(x, (self.padding, right))).contiguous() + + def apply_weight_norm(self) -> ArkttsCausalConv1d: + self.conv = weight_norm(self.conv) + return self + + +class ArkttsCausalConvTranspose1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + ) -> None: + super().__init__() + self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation) + self.stride = stride + self.kernel_size = kernel_size + + def forward(self, x: Tensor) -> Tensor: + x = self.conv(x) + crop = self.kernel_size - self.stride + return x[..., : x.shape[-1] - crop].contiguous() if crop else x.contiguous() + + def apply_weight_norm(self) -> ArkttsCausalConvTranspose1d: + self.conv = weight_norm(self.conv) + return self + + +def _causal_wn_conv(*args: Any, **kwargs: Any) -> ArkttsCausalConv1d: + return ArkttsCausalConv1d(*args, **kwargs).apply_weight_norm() + + +def _causal_wn_transpose(*args: Any, **kwargs: Any) -> ArkttsCausalConvTranspose1d: + return ArkttsCausalConvTranspose1d(*args, **kwargs).apply_weight_norm() + + +class ArkttsSnake1d(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x: Tensor) -> Tensor: + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2) + return x.reshape(shape) + + +class ArkttsResidualUnit(nn.Module): + def __init__(self, dim: int, dilation: int) -> None: + super().__init__() + self.block = nn.Sequential( + ArkttsSnake1d(dim), + _causal_wn_conv(dim, dim, kernel_size=7, dilation=dilation), + ArkttsSnake1d(dim), + _causal_wn_conv(dim, dim, kernel_size=1), + ) + + def forward(self, x: Tensor) -> Tensor: + output = self.block(x) + difference = x.shape[-1] - output.shape[-1] + if difference > 0: + x = x[..., :-difference] + return x + output + + +class ArkttsEncoderBlock(nn.Module): + def __init__(self, dim: int, stride: int, transformer_layers: int) -> None: + super().__init__() + modules: list[nn.Module] = [ + ArkttsResidualUnit(dim // 2, 1), + ArkttsResidualUnit(dim // 2, 3), + ArkttsResidualUnit(dim // 2, 9), + ArkttsSnake1d(dim // 2), + _causal_wn_conv(dim // 2, dim, kernel_size=2 * stride, stride=stride), + ] + if transformer_layers: + config = ArkttsCodecTransformerConfig( + n_layer=transformer_layers, + n_head=dim // 64, + dim=dim, + intermediate_size=dim * 3, + ) + modules.append(ArkttsCodecWindowTransformer(config, dim, window_size=512)) + else: + modules.append(nn.Identity()) + self.block = nn.Sequential(*modules) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class ArkttsEncoder(nn.Module): + def __init__(self) -> None: + super().__init__() + dim = 64 + modules: list[nn.Module] = [_causal_wn_conv(1, dim, kernel_size=7)] + for stride, transformer_layers in zip((2, 4, 8, 8), (0, 0, 0, 4), strict=True): + dim *= 2 + modules.append(ArkttsEncoderBlock(dim, stride, transformer_layers)) + modules.extend((ArkttsSnake1d(dim), _causal_wn_conv(dim, 1024, kernel_size=3))) + self.block = nn.Sequential(*modules) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class ArkttsDecoderBlock(nn.Module): + def __init__(self, input_dim: int, output_dim: int, stride: int) -> None: + super().__init__() + self.block = nn.Sequential( + ArkttsSnake1d(input_dim), + _causal_wn_transpose(input_dim, output_dim, kernel_size=2 * stride, stride=stride), + ArkttsResidualUnit(output_dim, 1), + ArkttsResidualUnit(output_dim, 3), + ArkttsResidualUnit(output_dim, 9), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class ArkttsDecoder(nn.Module): + def __init__(self) -> None: + super().__init__() + channels = 1536 + modules: list[nn.Module] = [_causal_wn_conv(1024, channels, kernel_size=7)] + output_dim = channels + for index, stride in enumerate((8, 8, 4, 2)): + input_dim = channels // (2**index) + output_dim = channels // (2 ** (index + 1)) + modules.append(ArkttsDecoderBlock(input_dim, output_dim, stride)) + modules.extend( + ( + ArkttsSnake1d(output_dim), + _causal_wn_conv(output_dim, 1, kernel_size=7), + nn.Tanh(), + ) + ) + self.model = nn.Sequential(*modules) + + def forward(self, x: Tensor) -> Tensor: + return self.model(x) + + +class ArkttsVectorQuantizer(nn.Module): + def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int) -> None: + super().__init__() + self.codebook_size = int(codebook_size) + self.codebook_dim = int(codebook_dim) + self.in_proj = legacy_weight_norm(nn.Conv1d(input_dim, codebook_dim, kernel_size=1)) + self.out_proj = legacy_weight_norm(nn.Conv1d(codebook_dim, input_dim, kernel_size=1)) + self.codebook = nn.Embedding(codebook_size, codebook_dim) + + def decode_code(self, indices: Tensor) -> Tensor: + return F.embedding(indices, self.codebook.weight).transpose(1, 2) + + def decode_latents(self, latents: Tensor) -> tuple[Tensor, Tensor]: + batch, _, length = latents.shape + flattened = F.normalize(latents.transpose(1, 2).reshape(batch * length, -1)) + codebook = F.normalize(self.codebook.weight) + distances = ( + flattened.pow(2).sum(1, keepdim=True) + - 2 * flattened @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + indices = (-distances).argmax(1).view(batch, length) + return self.decode_code(indices), indices + + def forward(self, z: Tensor) -> tuple[Tensor, Tensor]: + projected = self.in_proj(z) + quantized, indices = self.decode_latents(projected) + return self.out_proj(quantized), indices + + +class ArkttsResidualQuantizer(nn.Module): + def __init__(self, input_dim: int, n_codebooks: int, codebook_size: int, codebook_dim: int) -> None: + super().__init__() + self.n_codebooks = int(n_codebooks) + self.codebook_size = int(codebook_size) + self.quantizers = nn.ModuleList( + [ArkttsVectorQuantizer(input_dim, codebook_size, codebook_dim) for _ in range(n_codebooks)] + ) + + def forward(self, z: Tensor) -> tuple[Tensor, Tensor]: + quantized_sum: Tensor | float = 0.0 + residual = z + codes: list[Tensor] = [] + for quantizer in self.quantizers: + quantized, indices = quantizer(residual) + quantized_sum = quantized_sum + quantized + residual = residual - quantized + codes.append(indices) + return quantized_sum, torch.stack(codes, dim=1) + + def from_codes(self, codes: Tensor) -> Tensor: + output: Tensor | float = 0.0 + for index in range(codes.shape[1]): + projected = self.quantizers[index].decode_code(codes[:, index]) + output = output + self.quantizers[index].out_proj(projected) + return output + + +class ArkttsConvNeXtBlock(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dwconv = ArkttsCausalConv1d(dim, dim, kernel_size=7, groups=dim) + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.pwconv2 = nn.Linear(4 * dim, dim) + self.gamma = nn.Parameter(1e-6 * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + residual = x + x = self.dwconv(x).permute(0, 2, 1) + x = self.pwconv2(self.act(self.pwconv1(self.norm(x)))) + x = (self.gamma * x).permute(0, 2, 1) + return residual + x + + +class ArkttsDownsampleQuantizer(nn.Module): + """1 semantic codebook (4096) + 9 residual codebooks (1024), 4x downsampled.""" + + def __init__( + self, + post_n_layer: int = 8, + post_n_head: int = 16, + post_n_local_heads: int = 8, + post_intermediate_size: int = 1216, + ) -> None: + super().__init__() + self.semantic_quantizer = ArkttsResidualQuantizer(1024, 1, 4096, 8) + self.quantizer = ArkttsResidualQuantizer(1024, 9, 1024, 8) + self.downsample = nn.Sequential( + nn.Sequential(ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)), + nn.Sequential(ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)), + ) + self.upsample = nn.Sequential( + nn.Sequential(ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)), + nn.Sequential(ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)), + ) + self.pre_module = ArkttsCodecWindowTransformer( + ArkttsCodecTransformerConfig(n_layer=8, n_head=16, dim=1024, intermediate_size=3072), + 1024, + window_size=128, + ) + self.post_module = ArkttsCodecWindowTransformer( + ArkttsCodecTransformerConfig( + n_layer=post_n_layer, + n_head=post_n_head, + n_local_heads=post_n_local_heads, + dim=1024, + intermediate_size=post_intermediate_size, + ), + 1024, + window_size=128, + ) + self.semantic_predictor_module = nn.Identity() + + def forward(self, z: Tensor) -> tuple[Tensor, Tensor]: + original_length = z.shape[-1] + z = self.pre_module(self.downsample(z)) + semantic, semantic_codes = self.semantic_quantizer(z) + residual, residual_codes = self.quantizer(z - semantic) + z = self.upsample(self.post_module(semantic + residual)) + difference = original_length - z.shape[-1] + if difference > 0: + z = F.pad(z, (difference, 0)) + elif difference < 0: + z = z[..., -difference:] + return z, torch.cat((semantic_codes, residual_codes), dim=1) + + def decode(self, indices: Tensor) -> Tensor: + # Cloned because the clamps below are in-place and this is the caller's + # codes buffer; once per decode chunk, not per AR step. + indices = indices.clone() + indices[:, 0].clamp_(0, self.semantic_quantizer.codebook_size - 1) + indices[:, 1:].clamp_(0, self.quantizer.codebook_size - 1) + semantic = self.semantic_quantizer.from_codes(indices[:, :1]) + residual = self.quantizer.from_codes(indices[:, 1:]) + return self.upsample(self.post_module(semantic + residual)) + + +class ArkttsCodec(nn.Module): + """Audio8 TTS codec: waveform <-> 10 codebooks at ~21.5 frames/s.""" + + sample_rate = 44100 + frame_length = 2048 + + def __init__( + self, + post_n_layer: int = 8, + post_n_head: int = 16, + post_n_local_heads: int = 8, + post_intermediate_size: int = 1216, + ) -> None: + super().__init__() + self.encoder = ArkttsEncoder() + self.quantizer = ArkttsDownsampleQuantizer( + post_n_layer=post_n_layer, + post_n_head=post_n_head, + post_n_local_heads=post_n_local_heads, + post_intermediate_size=post_intermediate_size, + ) + self.decoder = ArkttsDecoder() + + @torch.inference_mode() + def encode(self, audio: Tensor, audio_lengths: Tensor | None = None) -> tuple[Tensor, Tensor]: + """Encode ``[B, 1, samples]`` (or ``[B, samples]``) to ``[B, 10, frames]``.""" + if audio.ndim == 2: + audio = audio[:, None] + if audio.ndim != 3 or audio.shape[1] != 1: + raise ValueError(f"audio must have shape [B, 1, samples], got {tuple(audio.shape)}") + original_length = int(audio.shape[-1]) + right = math.ceil(original_length / self.frame_length) * self.frame_length - original_length + if right: + audio = F.pad(audio, (0, right)) + if audio_lengths is None: + audio_lengths = torch.full((audio.shape[0],), original_length, device=audio.device, dtype=torch.long) + _, codes = self.quantizer(self.encoder(audio)) + code_lengths = torch.ceil(audio_lengths.float() / self.frame_length).long() + max_codes = int(codes.shape[-1]) + code_lengths = code_lengths.clamp_max(max_codes) + frame_index = torch.arange(max_codes, device=codes.device) + valid = frame_index[None, None, :] < code_lengths[:, None, None] + return torch.where(valid, codes, torch.full_like(codes, -1)), code_lengths + + @torch.inference_mode() + def decode(self, codes: Tensor) -> Tensor: + """Decode ``[B, 10, frames]`` to ``[B, 1, samples]``.""" + return self.decoder(self.quantizer.decode(codes.long())) + + +def build_arktts_codec( + *, + post_n_layer: int = 8, + post_n_head: int = 16, + post_n_local_heads: int = 8, + post_intermediate_size: int = 1216, +) -> ArkttsCodec: + """Construct the codec with uninitialised weights, on CPU, in eval mode.""" + codec = ArkttsCodec( + post_n_layer=post_n_layer, + post_n_head=post_n_head, + post_n_local_heads=post_n_local_heads, + post_intermediate_size=post_intermediate_size, + ) + codec.eval() + return codec + + +__all__ = ["ArkttsCodec", "build_arktts_codec"] diff --git a/vllm_omni/model_executor/models/audio8_tts/codec_utils.py b/vllm_omni/model_executor/models/audio8_tts/codec_utils.py new file mode 100644 index 000000000..64f9c0afb --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/codec_utils.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Loading / caching helpers for the Audio8 TTS ``codec.pth`` checkpoint. + +Both stages need the codec but for opposite halves: Stage 0 encodes reference +audio for zero-shot voice cloning, Stage 1 decodes generated codes to waveform. +The unused half is pruned before the module is moved to the device so a stage +never pays GPU memory for weights it cannot reach. +""" + +from __future__ import annotations + +import os +from functools import lru_cache + +import numpy as np +import torch +import torch.nn as nn +from torch.nn.utils.parametrize import remove_parametrizations +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.audio8_tts.codec import ArkttsCodec, build_arktts_codec +from vllm_omni.model_executor.models.audio8_tts.configuration_audio8_tts import ( + ARKTTS_CODEC_FRAME_SIZE, + ARKTTS_CODEC_SAMPLE_RATE, +) + +logger = init_logger(__name__) + +CODEC_FILENAME = "codec.pth" + +_codec_cache: dict[tuple[str, str, str, str], ArkttsCodec] = {} + + +def resolve_codec_path(model_path: str, filename: str = CODEC_FILENAME) -> str: + """Return a local path to ``filename``, downloading from the Hub if needed.""" + local = os.path.join(model_path, filename) + if os.path.exists(local): + return local + + from transformers.utils.hub import cached_file + + cached = cached_file(model_path, filename) + if cached is None or not os.path.exists(cached): + raise FileNotFoundError( + f"{filename} not found for {model_path}; the Audio8 TTS checkpoint must ship its neural audio codec." + ) + return cached + + +def bake_weight_norm(module: nn.Module) -> int: + """Fold weight-norm parametrisations into plain weights (inference only).""" + baked = 0 + for submodule in module.modules(): + parametrizations = getattr(submodule, "parametrizations", None) + if not parametrizations: + continue + for name in list(parametrizations.keys()): + remove_parametrizations(submodule, name, leave_parametrized=True) + baked += 1 + return baked + + +def load_arktts_codec( + model_path: str, + *, + device: torch.device | str = "cpu", + dtype: torch.dtype = torch.float32, + role: str = "decode", + post_n_layer: int = 8, + post_n_head: int = 16, + post_n_local_heads: int = 8, + post_intermediate_size: int = 1216, +) -> ArkttsCodec: + """Load ``codec.pth`` for ``role`` in {"encode", "decode", "both"}. + + Instances are cached per (path, device, dtype, role) because both the Stage-0 + voice-clone path and the Stage-1 decoder may ask for the same codec many + times per request. + """ + if role not in {"encode", "decode", "both"}: + raise ValueError(f"role must be encode/decode/both, got {role!r}") + device = torch.device(device) + cache_key = (model_path, str(device), str(dtype), role) + cached_codec = _codec_cache.get(cache_key) + if cached_codec is not None: + return cached_codec + + codec_path = resolve_codec_path(model_path) + codec = build_arktts_codec( + post_n_layer=post_n_layer, + post_n_head=post_n_head, + post_n_local_heads=post_n_local_heads, + post_intermediate_size=post_intermediate_size, + ) + + state = torch.load(codec_path, map_location="cpu", weights_only=True) + if "state_dict" in state: + state = state["state_dict"] + if any("generator." in key for key in state): + state = {key.replace("generator.", ""): value for key, value in state.items() if "generator." in key} + state = {key: value for key, value in state.items() if not key.endswith(("freqs_cis", "causal_mask"))} + codec.load_state_dict(state, strict=True) + + baked = bake_weight_norm(codec) + if role == "encode": + codec.decoder = None + elif role == "decode": + codec.encoder = None + codec.quantizer.pre_module = None + codec.quantizer.downsample = None + + codec = codec.to(device=device, dtype=dtype) + codec.eval() + _codec_cache[cache_key] = codec + logger.info( + "Loaded Audio8 TTS codec from %s (role=%s, device=%s, dtype=%s, baked_weight_norms=%d)", + codec_path, + role, + device, + dtype, + baked, + ) + return codec + + +@lru_cache(maxsize=8) +def _resample_transform(source_sr: int, target_sr: int, device: torch.device, dtype: torch.dtype): + import torchaudio + + return torchaudio.transforms.Resample(source_sr, target_sr).to(device=device, dtype=dtype) + + +def prepare_reference_waveform( + wav_samples: list[float] | np.ndarray | torch.Tensor, + sample_rate: int, + *, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Down-mix to mono and resample reference audio to the codec's rate. + + Returns a 1-D waveform tensor. + """ + if isinstance(wav_samples, torch.Tensor): + wav = wav_samples.detach() + else: + wav = torch.as_tensor(np.asarray(wav_samples)) + wav = wav.to(device=device, dtype=dtype) + + if wav.ndim == 2: + # Accept [channels, samples] and [samples, channels]. + if wav.shape[0] <= 8 and wav.shape[1] > wav.shape[0]: + wav = wav.mean(dim=0) + elif wav.shape[-1] <= 8 and wav.shape[0] > wav.shape[-1]: + wav = wav.mean(dim=-1) + else: + wav = wav.mean(dim=0) + elif wav.ndim > 2: + wav = wav.reshape(-1, wav.shape[-1]).mean(dim=0) + wav = wav.reshape(-1) + + if wav.numel() == 0: + raise ValueError("Reference audio must not be empty") + if int(sample_rate) <= 0: + raise ValueError(f"Reference audio sample rate must be positive, got {sample_rate}") + if int(sample_rate) != ARKTTS_CODEC_SAMPLE_RATE: + transform = _resample_transform(int(sample_rate), ARKTTS_CODEC_SAMPLE_RATE, wav.device, wav.dtype) + wav = transform(wav.unsqueeze(0)).squeeze(0) + return wav.contiguous() + + +@torch.no_grad() +def encode_reference_audio_codes( + model_path: str, + wav_samples: list[float] | np.ndarray | torch.Tensor, + sample_rate: int, + *, + device: torch.device | str, + **codec_kwargs: int, +) -> torch.Tensor: + """Encode reference audio into codec codes. + + Returns: + ``[frames, num_codebooks]`` int64 codes on ``device``. + """ + device = torch.device(device) + codec = load_arktts_codec(model_path, device=device, dtype=torch.float32, role="encode", **codec_kwargs) + wav = prepare_reference_waveform(wav_samples, sample_rate, device=device) + lengths = torch.tensor([wav.numel()], device=device, dtype=torch.long) + codes, code_lengths = codec.encode(wav.reshape(1, 1, -1), lengths) + frames = int(code_lengths[0].item()) + codes_fq = codes[0, :, :frames].transpose(0, 1).to(dtype=torch.long).contiguous() + logger.info( + "Encoded Audio8 TTS reference audio: %d samples @ %d Hz -> frames=%d codebooks=%d", + int(wav.numel()), + int(sample_rate), + int(codes_fq.shape[0]), + int(codes_fq.shape[1]), + ) + return codes_fq + + +def estimate_reference_code_frames(num_samples: int, sample_rate: int) -> int: + """Frames the codec will emit for ``num_samples`` at ``sample_rate``. + + Used by the serving layer to size the prompt placeholder without paying for + a codec encode on the API thread. + """ + if num_samples <= 0 or sample_rate <= 0: + raise ValueError("Reference audio must have a positive length and sample rate") + resampled = max(1, -(-num_samples * ARKTTS_CODEC_SAMPLE_RATE // sample_rate)) + return max(1, -(-resampled // ARKTTS_CODEC_FRAME_SIZE)) + + +__all__ = [ + "CODEC_FILENAME", + "bake_weight_norm", + "encode_reference_audio_codes", + "estimate_reference_code_frames", + "load_arktts_codec", + "prepare_reference_waveform", + "resolve_codec_path", +] diff --git a/vllm_omni/model_executor/models/audio8_tts/configuration_audio8_tts.py b/vllm_omni/model_executor/models/audio8_tts/configuration_audio8_tts.py new file mode 100644 index 000000000..de44ab2a2 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/configuration_audio8_tts.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Configuration classes for Audio8 TTS Preview (``model_type = "arktts"``). + +Audio8's HuggingFace config is flat and uses DualAR field names (``dim``, +``n_layer``, ``n_head``, ``fast_dim``, ...); this module re-exports them under +Qwen2 attribute names so vLLM's ``Qwen2Model`` consumes the Slow AR half, and +splits the Fast AR half into its own sub-config. The backbone is Qwen2 rather +than Qwen3 (Fish Speech) because Audio8 uses ``qkv_bias=true`` with no q/k +norm -- exactly the Qwen2 attention shape. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import PretrainedConfig + +#: Codec frame stride in waveform samples (encoder rates 2*4*8*8 = 512 times the +#: quantizer's 2*2 downsample). 44100 / 2048 ~= 21.5 frames per second. +ARKTTS_CODEC_FRAME_SIZE = 2048 +ARKTTS_CODEC_SAMPLE_RATE = 44100 +#: Only the first codebook is a 4096-entry semantic codebook; codebooks 1..N-1 +#: are 1024-entry residual codebooks (see ``ArkttsDownsampleQuantizer``). +ARKTTS_SEMANTIC_CODEBOOK_SIZE = 4096 +ARKTTS_RESIDUAL_CODEBOOK_SIZE = 1024 + + +class Audio8TTSSlowARConfig(PretrainedConfig): + """Slow AR config exposed with Qwen2-compatible attribute names.""" + + model_type = "arktts_slow_ar" + + def __init__( + self, + vocab_size: int = 155776, + dim: int = 896, + n_head: int = 14, + n_local_heads: int = 2, + head_dim: int = 64, + n_layer: int = 24, + intermediate_size: int = 4864, + max_seq_len: int = 2048, + rope_base: float = 1_000_000.0, + norm_eps: float = 1e-6, + attention_qkv_bias: bool = True, + attention_qk_norm: bool = False, + tie_word_embeddings: bool = True, + codebook_size: int = ARKTTS_SEMANTIC_CODEBOOK_SIZE, + num_codebooks: int = 10, + semantic_begin_id: int = 151678, + semantic_end_id: int = 155773, + eos_token_id: int = 151645, + pad_token_id: int = 151643, + **kwargs: Any, + ) -> None: + # Audio8 field names -> standard Transformers / Qwen2 names. + self.hidden_size = int(dim) + self.num_attention_heads = int(n_head) + self.num_key_value_heads = int(n_local_heads) + self.head_dim = int(head_dim) + self.num_hidden_layers = int(n_layer) + self.intermediate_size = int(intermediate_size) + self.max_position_embeddings = int(max_seq_len) + self.rms_norm_eps = float(norm_eps) + self.hidden_act = "silu" + # Qwen2Attention hardcodes qkv bias=True / o_proj bias=False, which is + # what Audio8 TTS ships. Keep the flags so a future checkpoint that + # flips them fails loudly in the model rather than loading silently. + self.attention_qkv_bias = bool(attention_qkv_bias) + self.attention_qk_norm = bool(attention_qk_norm) + # Set explicitly (rather than relying on vLLM's set_default_rope_theta) + # so that a checkpoint with a non-default rope_base is honoured. + self.rope_parameters = {"rope_type": "default", "rope_theta": float(rope_base)} + + # Codec / codebook fields. + self.codebook_size = int(codebook_size) + self.num_codebooks = int(num_codebooks) + self.semantic_begin_id = int(semantic_begin_id) + self.semantic_end_id = int(semantic_end_id) + + super().__init__( + vocab_size=int(vocab_size), + tie_word_embeddings=bool(tie_word_embeddings), + eos_token_id=int(eos_token_id), + pad_token_id=int(pad_token_id), + **kwargs, + ) + + +class Audio8TTSFastARConfig(PretrainedConfig): + """Fast AR config: the ``n_fast_layer`` residual-codebook predictor.""" + + model_type = "arktts_fast_ar" + + def __init__( + self, + codebook_size: int = ARKTTS_SEMANTIC_CODEBOOK_SIZE, + num_codebooks: int = 10, + fast_dim: int = 896, + fast_n_head: int = 14, + fast_n_local_heads: int = 2, + fast_head_dim: int = 64, + n_fast_layer: int = 4, + fast_intermediate_size: int = 4864, + fast_attention_qkv_bias: bool = False, + fast_attention_qk_norm: bool = False, + rope_base: float = 1_000_000.0, + norm_eps: float = 1e-6, + **kwargs: Any, + ) -> None: + self.hidden_size = int(fast_dim) + self.num_attention_heads = int(fast_n_head) + self.num_key_value_heads = int(fast_n_local_heads) + self.head_dim = int(fast_head_dim) + self.num_hidden_layers = int(n_fast_layer) + self.intermediate_size = int(fast_intermediate_size) + # The Fast AR sequence is [slow hidden state, code_0, ..., code_{N-1}]. + self.max_position_embeddings = int(num_codebooks) + self.rms_norm_eps = float(norm_eps) + self.hidden_act = "silu" + self.attention_qkv_bias = bool(fast_attention_qkv_bias) + self.attention_qk_norm = bool(fast_attention_qk_norm) + self.rope_theta = float(rope_base) + self.num_codebooks = int(num_codebooks) + + super().__init__(vocab_size=int(codebook_size), **kwargs) + + +class Audio8TTSConfig(PretrainedConfig): + """Top-level Audio8 TTS config (``model_type = "arktts"``). + + Accepts the flat HF checkpoint fields and derives ``text_config`` (Slow AR) + and ``fast_ar_config`` (Fast AR). ``get_text_config()`` returns the Slow AR + config, which is what ``Qwen2Model`` reads. + """ + + model_type = "arktts" + sub_configs = { + "text_config": Audio8TTSSlowARConfig, + "fast_ar_config": Audio8TTSFastARConfig, + } + + def __init__( + self, + text_config: dict | Audio8TTSSlowARConfig | None = None, + fast_ar_config: dict | Audio8TTSFastARConfig | None = None, + vocab_size: int = 155776, + dim: int = 896, + n_head: int = 14, + n_local_heads: int = 2, + head_dim: int = 64, + n_layer: int = 24, + intermediate_size: int = 4864, + max_seq_len: int = 2048, + rope_base: float = 1_000_000.0, + norm_eps: float = 1e-6, + attention_qkv_bias: bool = True, + attention_qk_norm: bool = False, + tie_word_embeddings: bool = True, + codebook_size: int = ARKTTS_SEMANTIC_CODEBOOK_SIZE, + num_codebooks: int = 10, + semantic_begin_id: int = 151678, + semantic_end_id: int = 155773, + n_fast_layer: int = 4, + fast_dim: int = 896, + fast_n_head: int = 14, + fast_n_local_heads: int = 2, + fast_head_dim: int = 64, + fast_intermediate_size: int = 4864, + fast_attention_qkv_bias: bool = False, + fast_attention_qk_norm: bool = False, + norm_fastlayer_input: bool = True, + codec_filename: str = "codec.pth", + codec_sample_rate: int = ARKTTS_CODEC_SAMPLE_RATE, + codec_frame_size: int = ARKTTS_CODEC_FRAME_SIZE, + codec_post_n_layer: int = 8, + codec_post_n_head: int = 16, + codec_post_n_local_heads: int = 8, + codec_post_intermediate_size: int = 1216, + ras_window_size: int = 10, + ras_temperature: float = 1.0, + ras_top_p: float = 0.9, + eos_token_id: int = 151645, + pad_token_id: int = 151643, + **kwargs: Any, + ) -> None: + if isinstance(text_config, dict): + text_config = Audio8TTSSlowARConfig(**text_config) + if isinstance(fast_ar_config, dict): + fast_ar_config = Audio8TTSFastARConfig(**fast_ar_config) + + self.text_config = text_config or Audio8TTSSlowARConfig( + vocab_size=vocab_size, + dim=dim, + n_head=n_head, + n_local_heads=n_local_heads, + head_dim=head_dim, + n_layer=n_layer, + intermediate_size=intermediate_size, + max_seq_len=max_seq_len, + rope_base=rope_base, + norm_eps=norm_eps, + attention_qkv_bias=attention_qkv_bias, + attention_qk_norm=attention_qk_norm, + tie_word_embeddings=tie_word_embeddings, + codebook_size=codebook_size, + num_codebooks=num_codebooks, + semantic_begin_id=semantic_begin_id, + semantic_end_id=semantic_end_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + ) + self.fast_ar_config = fast_ar_config or Audio8TTSFastARConfig( + codebook_size=codebook_size, + num_codebooks=num_codebooks, + fast_dim=fast_dim, + fast_n_head=fast_n_head, + fast_n_local_heads=fast_n_local_heads, + fast_head_dim=fast_head_dim, + n_fast_layer=n_fast_layer, + fast_intermediate_size=fast_intermediate_size, + fast_attention_qkv_bias=fast_attention_qkv_bias, + fast_attention_qk_norm=fast_attention_qk_norm, + rope_base=rope_base, + norm_eps=norm_eps, + ) + + # Fields the stages read directly off the top-level config. + self.codebook_size = int(codebook_size) + self.num_codebooks = int(num_codebooks) + self.semantic_begin_id = int(semantic_begin_id) + self.semantic_end_id = int(semantic_end_id) + self.norm_fastlayer_input = bool(norm_fastlayer_input) + self.codec_filename = str(codec_filename) + self.codec_sample_rate = int(codec_sample_rate) + self.codec_frame_size = int(codec_frame_size) + self.codec_post_n_layer = int(codec_post_n_layer) + self.codec_post_n_head = int(codec_post_n_head) + self.codec_post_n_local_heads = int(codec_post_n_local_heads) + self.codec_post_intermediate_size = int(codec_post_intermediate_size) + # Repetition-Aware Sampling (RAS): resample a repeated semantic token + # from a flatter distribution instead of masking it. + self.ras_window_size = int(ras_window_size) + self.ras_temperature = float(ras_temperature) + self.ras_top_p = float(ras_top_p) + + super().__init__( + eos_token_id=int(eos_token_id), + pad_token_id=int(pad_token_id), + tie_word_embeddings=bool(tie_word_embeddings), + **kwargs, + ) + + def get_text_config(self, *args: Any, **kwargs: Any) -> Audio8TTSSlowARConfig: + return self.text_config + + +__all__ = [ + "ARKTTS_CODEC_FRAME_SIZE", + "ARKTTS_CODEC_SAMPLE_RATE", + "ARKTTS_RESIDUAL_CODEBOOK_SIZE", + "ARKTTS_SEMANTIC_CODEBOOK_SIZE", + "Audio8TTSConfig", + "Audio8TTSFastARConfig", + "Audio8TTSSlowARConfig", +] diff --git a/vllm_omni/model_executor/models/audio8_tts/pipeline.py b/vllm_omni/model_executor/models/audio8_tts/pipeline.py new file mode 100644 index 000000000..836ab71d9 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/pipeline.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview pipeline topology. + +Stage 0: ``audio8_tts_slow_ar`` -- text -> semantic tokens + codec codes. +Stage 1: ``audio8_tts_codec_decoder`` -- codec codes -> 44.1 kHz waveform. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.audio8_tts" + +#: ``<|im_end|>``: the Slow AR ends the utterance with an end-of-turn token. +AUDIO8_TTS_EOS_TOKEN_ID = 151645 + +AUDIO8_TTS_PIPELINE = PipelineConfig( + model_type="arktts", + default_deploy_config_name="audio8_tts.yaml", + model_arch="Audio8TTSSlowARForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="audio8_tts_slow_ar", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + async_chunk_process_next_stage_input_func=(f"{_PROC}.slow_ar_to_codec_decoder_async_chunk"), + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [AUDIO8_TTS_EOS_TOKEN_ID], + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="audio8_tts_codec_decoder", + model_arch="Audio8TTSCodecDecoder", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + sampling_constraints={"detokenize": True}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/audio8_tts/prompt_utils.py b/vllm_omni/model_executor/models/audio8_tts/prompt_utils.py new file mode 100644 index 000000000..f0660ac5c --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/prompt_utils.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS prompt construction. + +Reproduces ``ArkttsProcessor._prompt_segments`` from the reference checkpoint, +so tokenisation is byte-identical to HF inference. +""" + +from __future__ import annotations + +from typing import Any + +import regex as re + +SYSTEM_PROMPT_TEXT_ONLY = "convert the provided text to speech" +SYSTEM_PROMPT_CLONE_PREFIX = "convert the provided text to speech reference to the following:\n\nText:\n" +SYSTEM_PROMPT_CLONE_SUFFIX = "\n\nSpeech:\n" + +_SPEAKER_TAG_PATTERN = re.compile(r"<\|speaker:\d+\|>") +_LEGACY_SPEAKER_TAG_PATTERN = re.compile(r"") +_CONTROL_TOKEN_PATTERN = re.compile(r"<\|[^>]+\|>") + + +def clean_text(text: str) -> str: + """Collapse all whitespace runs to single spaces (reference ``_clean_text``).""" + return " ".join(str(text).strip().split()) + + +def normalize_text(text: str, *, add_default_speaker: bool = False) -> str: + """Clean text, normalise speaker tags, reject other control tokens. + + ``<|speaker:N|>`` is allowed; any other ``<|...|>`` token raises so callers + cannot inject ``<|semantic:...|>`` / ``<|im_end|>`` into a prompt. + """ + normalized = clean_text(_LEGACY_SPEAKER_TAG_PATTERN.sub(r"<|speaker:\1|>", text)) + + disallowed = [ + token for token in _CONTROL_TOKEN_PATTERN.findall(normalized) if not _SPEAKER_TAG_PATTERN.fullmatch(token) + ] + if disallowed: + tokens = ", ".join(sorted(set(disallowed))) + raise ValueError(f"Audio8 TTS input contains unsupported control token(s): {tokens}") + + if add_default_speaker and not _SPEAKER_TAG_PATTERN.search(normalized): + normalized = f"<|speaker:0|>{normalized}" + return normalized + + +def _encode(tokenizer: Any, text: str) -> list[int]: + return list(tokenizer.encode(text, add_special_tokens=False)) + + +def _encode_segments(tokenizer: Any, segments: list[str]) -> list[int]: + ids: list[int] = [] + for segment in segments: + ids.extend(_encode(tokenizer, segment)) + return ids + + +def build_text_only_prompt_ids(tokenizer: Any, text: str) -> tuple[list[int], str]: + """Build the no-reference prompt. + + Returns: + ``(prompt_token_ids, normalized_text)``. + """ + normalized = normalize_text(text) + if not normalized: + raise ValueError("Audio8 TTS input text must not be empty") + ids = _encode_segments( + tokenizer, + [ + "<|im_start|>system\n", + SYSTEM_PROMPT_TEXT_ONLY, + "<|im_end|>\n", + "<|im_start|>user\n", + normalized, + "<|im_end|>\n", + "<|im_start|>assistant\n<|voice|>", + ], + ) + return ids, normalized + + +def build_voice_clone_prompt_parts( + tokenizer: Any, + text: str, + ref_text: str, +) -> tuple[list[int], list[int], str, str]: + """Build the clone prompt around the reference-code slot. + + Caller splices ``[code + semantic_begin_id for code in ref_codes[:, 0]]`` + between the returned prefix and suffix. + """ + normalized = normalize_text(text) + if not normalized: + raise ValueError("Audio8 TTS input text must not be empty") + normalized_ref = normalize_text(ref_text, add_default_speaker=True) + if not normalized_ref: + raise ValueError("Audio8 TTS voice cloning requires a non-empty reference transcript") + + prefix = _encode_segments( + tokenizer, + [ + "<|im_start|>system\n", + SYSTEM_PROMPT_CLONE_PREFIX, + normalized_ref, + SYSTEM_PROMPT_CLONE_SUFFIX, + ], + ) + suffix = _encode_segments( + tokenizer, + [ + "<|im_end|>\n", + "<|im_start|>user\n", + normalized, + "<|im_end|>\n", + "<|im_start|>assistant\n<|voice|>", + ], + ) + return prefix, suffix, normalized, normalized_ref + + +def build_voice_clone_prompt_ids( + tokenizer: Any, + text: str, + ref_text: str, + semantic_token_ids: list[int], +) -> tuple[list[int], int, str, str]: + """Build the full clone prompt. + + ``ref_start_index`` is the position of the first reference frame, i.e. + where the residual codebook embeddings must be added. + """ + prefix, suffix, normalized, normalized_ref = build_voice_clone_prompt_parts(tokenizer, text, ref_text) + return prefix + list(semantic_token_ids) + suffix, len(prefix), normalized, normalized_ref + + +def estimate_voice_clone_prompt_len( + tokenizer: Any, + normalized_text: str, + normalized_ref_text: str, + ref_frames: int, +) -> int: + """Exact clone prompt length, without encoding the reference audio.""" + prefix, suffix, _, _ = build_voice_clone_prompt_parts(tokenizer, normalized_text, normalized_ref_text) + return len(prefix) + max(0, int(ref_frames)) + len(suffix) + + +__all__ = [ + "SYSTEM_PROMPT_CLONE_PREFIX", + "SYSTEM_PROMPT_CLONE_SUFFIX", + "SYSTEM_PROMPT_TEXT_ONLY", + "build_text_only_prompt_ids", + "build_voice_clone_prompt_ids", + "build_voice_clone_prompt_parts", + "clean_text", + "estimate_voice_clone_prompt_len", + "normalize_text", +] diff --git a/vllm_omni/model_executor/models/audio8_tts/sampling.py b/vllm_omni/model_executor/models/audio8_tts/sampling.py new file mode 100644 index 000000000..638408411 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/sampling.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS sampling primitives, ported 1:1 from the reference checkpoint. + +Filter (top-k/top-p on **raw** logits) -> divide surviving logits by +temperature -> Gumbel-max draw. This is the reverse of the usual filter order +and is audibly different at ``temperature != 1``, which is why this cannot +reuse ``models/common/nucleus_ras_sampling.py``. Per-row scalars are accepted +so mixed batches run in one shot without a GPU->CPU sync. +""" + +from __future__ import annotations + +import torch + +#: Below this temperature a request is treated as greedy, matching vLLM. +SAMPLING_EPS = 1e-5 + + +def _as_column( + value: float | int | torch.Tensor, + batch_size: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Broadcast a scalar or ``[B]`` tensor to a ``[B, 1]`` tensor.""" + if isinstance(value, torch.Tensor): + return value.to(device=device, dtype=dtype).reshape(-1, 1).expand(batch_size, 1) + return torch.full((batch_size, 1), float(value), device=device, dtype=dtype) + + +def filter_top_k_top_p( + logits: torch.Tensor, + top_k: int | torch.Tensor, + top_p: float | torch.Tensor, +) -> torch.Tensor: + """Mask candidates outside the top-k / top-p set of the raw logits. + + ``top_k <= 0`` and ``top_p >= 1`` disable each filter. The highest-scoring + candidate is always kept, so no row can become all ``-inf``. + """ + if logits.ndim != 2: + raise ValueError(f"logits must be 2-D [B, V], got {tuple(logits.shape)}") + batch_size, vocab = logits.shape + scalar_k = not isinstance(top_k, torch.Tensor) + scalar_p = not isinstance(top_p, torch.Tensor) + if scalar_k and scalar_p and (int(top_k) <= 0 or int(top_k) >= vocab) and float(top_p) >= 1.0: + return logits + + sorted_scores, sorted_indices = torch.sort(logits, descending=True, dim=-1) + cumulative = torch.cumsum(torch.softmax(sorted_scores, dim=-1), dim=-1) + positions = torch.arange(vocab, device=logits.device).unsqueeze(0) + + keep_k = _as_column(top_k, batch_size, logits.device, torch.long) + # top_k <= 0 means "no cut"; represent it as the full vocabulary. + keep_k = torch.where(keep_k <= 0, torch.full_like(keep_k, vocab), keep_k) + threshold_p = _as_column(top_p, batch_size, logits.device, cumulative.dtype) + + remove_sorted = (cumulative > threshold_p) | (positions >= keep_k) + # Always keep the argmax so at least one candidate survives. + remove_sorted[:, 0] = False + remove = torch.zeros_like(remove_sorted).scatter(-1, sorted_indices, remove_sorted) + return logits.masked_fill(remove, float("-inf")) + + +def gumbel_argmax_sample( + scores: torch.Tensor, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Draw one index per row via Gumbel-max over ``softmax(scores)``.""" + probabilities = torch.softmax(scores, dim=-1) + uniform = torch.rand( + probabilities.shape, + dtype=probabilities.dtype, + device=probabilities.device, + generator=generator, + ) + # -log(u) is Exp(1); argmax(p / Exp(1)) samples from p. + noise = -torch.log(uniform.clamp_min(torch.finfo(probabilities.dtype).tiny)) + return torch.argmax(probabilities / noise, dim=-1) + + +def sample_scores( + logits: torch.Tensor, + *, + temperature: float | torch.Tensor, + top_k: int | torch.Tensor, + top_p: float | torch.Tensor, + do_sample: bool = True, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Filter -> temper -> draw, in the reference order. + + Rows whose temperature is below :data:`SAMPLING_EPS` fall back to argmax. + """ + filtered = filter_top_k_top_p(logits, top_k, top_p) + greedy = filtered.argmax(dim=-1) + if not do_sample: + return greedy + if not isinstance(temperature, torch.Tensor): + if float(temperature) < SAMPLING_EPS: + return greedy + return gumbel_argmax_sample(filtered / float(temperature), generator=generator) + + temperature_col = _as_column(temperature, filtered.shape[0], filtered.device, filtered.dtype) + safe_temperature = torch.where( + temperature_col < SAMPLING_EPS, + torch.ones_like(temperature_col), + temperature_col, + ) + sampled = gumbel_argmax_sample(filtered / safe_temperature, generator=generator) + return torch.where(temperature_col.reshape(-1) < SAMPLING_EPS, greedy, sampled) + + +def ras_sample_semantic( + logits: torch.Tensor, + recent_ids: torch.Tensor | None, + *, + temperature: float | torch.Tensor, + top_k: int | torch.Tensor, + top_p: float | torch.Tensor, + ras_temperature: float, + ras_top_p: float, + do_sample: bool = True, + num_semantic_ids: int | None = None, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Repetition-Aware Sampling for the semantic (Slow AR) token. + + Draws with the request's ``temperature`` / ``top_p``; if the drawn id is + already in ``recent_ids`` (pad with a value that cannot be sampled, e.g. + ``-1``), redraws from the flatter RAS distribution instead. The + ``num_semantic_ids`` slot and beyond (EOS) are never substituted, so a + repeat EOS still ends the utterance. + + DualAR RAS *resamples* repeats rather than masking them — masking the + dominant token inflates EOS probability and truncates the utterance. + """ + normal = sample_scores( + logits, + temperature=temperature, + top_k=top_k, + top_p=top_p, + do_sample=do_sample, + generator=generator, + ) + if not do_sample or recent_ids is None or recent_ids.numel() == 0: + return normal + + high = sample_scores( + logits, + temperature=ras_temperature, + top_k=top_k, + top_p=ras_top_p, + do_sample=True, + generator=generator, + ) + repeated = (recent_ids == normal[:, None]).any(dim=1) + if num_semantic_ids is not None: + repeated &= normal < int(num_semantic_ids) + return torch.where(repeated, high, normal) + + +def _row_value(value: object, index: int) -> object: + """Slice a per-row value for request ``index``; scalars pass through.""" + return value[index : index + 1] if isinstance(value, torch.Tensor) else value + + +def ras_sample_batch( + logits: torch.Tensor, + recent_ids: torch.Tensor | None, + *, + temperature: float | torch.Tensor, + top_k: int | torch.Tensor, + top_p: float | torch.Tensor, + ras_temperature: float, + ras_top_p: float, + num_semantic_ids: int | None, + generators: dict, + num_reqs: int, +) -> torch.Tensor: + """Batched RAS that honors each request's own RNG generator. + + vLLM keys ``generators`` by batch slot, so a single shared generator would + leak slot 0's RNG stream to every row and make a prompt's audio depend on + its batch neighbours (undercutting per-request ``seed``). The single-request + case keeps the batched fast path; larger batches loop so each row draws from + its own (possibly unseeded) generator. + """ + common = dict( + ras_temperature=ras_temperature, + ras_top_p=ras_top_p, + num_semantic_ids=num_semantic_ids, + ) + if num_reqs == 1: + return ras_sample_semantic( + logits, + recent_ids, + temperature=temperature, + top_k=top_k, + top_p=top_p, + generator=generators.get(0), + **common, + ) + rows = [ + ras_sample_semantic( + logits[i : i + 1], + None if recent_ids is None else recent_ids[i : i + 1], + temperature=_row_value(temperature, i), + top_k=_row_value(top_k, i), + top_p=_row_value(top_p, i), + generator=generators.get(i), + **common, + ) + for i in range(num_reqs) + ] + return torch.cat(rows, dim=0) + + +__all__ = [ + "SAMPLING_EPS", + "filter_top_k_top_p", + "gumbel_argmax_sample", + "ras_sample_batch", + "ras_sample_semantic", + "sample_scores", +] diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index a3dfd7bf1..495e0037f 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable, Mapping, Sequence from math import isqrt from typing import Any @@ -122,9 +125,9 @@ class OmniBagelProcessingInfo(BaseProcessingInfo): if p.is_dir(): index_path = p / "model.safetensors.index.json" else: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - index_path = Path(hf_hub_download(model_name, "model.safetensors.index.json")) + index_path = Path(hf_api().hf_hub_download(model_name, "model.safetensors.index.json")) if not index_path.exists(): return 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 0308cda7b..60b4d4525 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Qwen3 Code Predictor -- optimized re-prefill, no KV cache. Shared by Qwen3-Omni and Qwen3-TTS talker models. @@ -97,7 +100,11 @@ class _RotaryEmbedding(CustomOp): "head_dim", config.hidden_size // config.num_attention_heads, ) - rope_theta = getattr(config, "rope_theta", 10000.0) + rope_parameters = getattr(config, "rope_parameters", None) or {} + rope_theta = rope_parameters.get( + "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) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 81bae6d6c..84652cada 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -10,7 +10,6 @@ import numpy as np import onnxruntime import torch import torch.nn as nn -from huggingface_hub import snapshot_download from scipy.signal import resample_poly from transformers import Qwen2Config from transformers.feature_extraction_utils import BatchFeature @@ -52,6 +51,7 @@ from vllm_omni.model_executor.models.cosyvoice3.utils 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.cosyvoice3 import CosyVoice3Config +from vllm_omni.transformers_utils.repo_utils import hf_api from vllm_omni.utils.speaker_cache import get_speaker_cache logger = init_logger(__name__) @@ -137,7 +137,7 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal """Build the per-model runtime components once (cached process-wide).""" # If model_dir is an HF repo ID (not a local path), resolve to cache. if not os.path.isdir(model_dir): - model_dir = snapshot_download(model_dir) + model_dir = hf_api().snapshot_download(model_dir) tokenizer = get_qwen_tokenizer( token_path=os.path.join(model_dir, config.qwen_pretrain_path), @@ -463,7 +463,7 @@ class CosyVoice3Model( self.model_stage = vllm_config.model_config.model_stage model_dir = vllm_config.model_config.model if not os.path.isdir(model_dir): - model_dir = snapshot_download(model_dir) + model_dir = hf_api().snapshot_download(model_dir) self.model_dir = model_dir self.model = None if self.model_stage == "cosyvoice3_talker": @@ -918,7 +918,7 @@ class CosyVoice3Model( repo = getattr(self.config, "flow_estimator_onnx_repo", None) if repo: try: - fetched_dir = snapshot_download(repo, allow_patterns=[fp16_name]) + fetched_dir = hf_api().snapshot_download(repo, allow_patterns=[fp16_name]) fetched = os.path.join(fetched_dir, fp16_name) if os.path.exists(fetched): return fetched diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py b/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py index 4b0b2fe4b..15cd0875e 100644 --- a/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py +++ b/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 Tencent. import os from collections.abc import Iterable @@ -25,9 +28,9 @@ class CovoAudioCode2WavForConditionalGeneration(nn.Module, SupportsPP): if os.path.isdir(model_name): model_path = model_name else: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_path = snapshot_download(model_name) + model_path = hf_api().snapshot_download(model_name) token2wav_path = os.path.join(model_path, "token2wav") code2wav_config = CovoAudioCode2WavConfig() diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py index ff2f460d6..862af8e17 100644 --- a/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """dots.tts talker — vLLM-native AR base LM + audio side path. Mirrors upstream rednote-hilab/dots.tts (pinned @ a393d2e): @@ -516,9 +516,9 @@ class DotsTTSForConditionalGeneration(nn.Module): if os.path.exists(local): return local try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return hf_hub_download( + return hf_api().hf_hub_download( repo_id=model_arg, filename="latent_stats.pt", local_files_only=True, diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py index 5c67decaa..125822c1e 100644 --- a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py +++ b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import hashlib @@ -20,9 +23,9 @@ from vllm.logger import init_logger logger = init_logger(__name__) try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except Exception: # pragma: no cover - snapshot_download = None + hf_api = None class DetokTarget(IntEnum): @@ -397,13 +400,13 @@ def _find_dynin_config_under_root(root: Path) -> Path | None: @lru_cache(maxsize=16) def _resolve_dynin_config_from_hf_repo(repo_id: str) -> str | None: - if not _looks_like_hf_repo_id(repo_id) or snapshot_download is None: + if not _looks_like_hf_repo_id(repo_id) or hf_api is None: return None try: snapshot_dir = ( Path( - snapshot_download( + hf_api().snapshot_download( repo_id=repo_id, repo_type="model", allow_patterns=list(_DYNIN_CONFIG_CANDIDATE_RELPATHS), @@ -719,7 +722,7 @@ def _resolve_remote_snapshot_dir( if source_path.is_dir(): return str(source_path.resolve()) - if snapshot_download is None: + if hf_api is None: raise RuntimeError("huggingface_hub is required to load remote code.") kwargs: dict[str, Any] = { @@ -732,10 +735,10 @@ def _resolve_remote_snapshot_dir( kwargs["revision"] = revision try: - return str(snapshot_download(**kwargs)) + return str(hf_api().snapshot_download(**kwargs)) except TypeError: kwargs.pop("local_files_only", None) - return str(snapshot_download(**kwargs)) + return str(hf_api().snapshot_download(**kwargs)) def _ensure_remote_package(snapshot_dir: str) -> str: diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py index 80bc0a78b..dbbb63524 100644 --- a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py +++ b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import os @@ -80,7 +83,7 @@ def _ensure_remote_s2u_vendor_root( return str(existing_path) try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except Exception as e: logger.warning("huggingface_hub unavailable; cannot fetch s2u_vendor from %s: %s", repo_id, e) return None @@ -91,7 +94,7 @@ def _ensure_remote_s2u_vendor_root( for revision in revisions: try: - snapshot_dir = snapshot_download( + snapshot_dir = hf_api().snapshot_download( repo_id=repo_id, revision=revision, allow_patterns=["s2u_vendor/**"], @@ -99,7 +102,7 @@ def _ensure_remote_s2u_vendor_root( ) except TypeError: try: - snapshot_dir = snapshot_download( + snapshot_dir = hf_api().snapshot_download( repo_id=repo_id, revision=revision, allow_patterns=["s2u_vendor/**"], diff --git a/vllm_omni/model_executor/models/fish_speech/dac_modules/__init__.py b/vllm_omni/model_executor/models/fish_speech/dac_modules/__init__.py new file mode 100644 index 000000000..654916598 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_modules/__init__.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Vendored DAC codec modules for Fish Speech S2 Pro. + +Adopted from the ``fish-speech`` 0.1.0 PyPI release +(https://pypi.org/project/fish-speech/0.1.0/, Apache-2.0) and +descript-audio-codec (https://github.com/descriptinc/descript-audio-codec, +MIT), trimmed to inference-only so the codec loads without external model +packages. Module attribute paths and weight-norm styles are kept identical +to upstream so ``codec.pth`` checkpoints load unchanged. +""" + +from vllm_omni.model_executor.models.fish_speech.dac_modules.codec import DAC +from vllm_omni.model_executor.models.fish_speech.dac_modules.rvq import ( + DownsampleResidualVectorQuantize, +) +from vllm_omni.model_executor.models.fish_speech.dac_modules.transformer import ( + ModelArgs, + WindowLimitedTransformer, +) + +__all__ = [ + "DAC", + "DownsampleResidualVectorQuantize", + "ModelArgs", + "WindowLimitedTransformer", +] diff --git a/vllm_omni/model_executor/models/fish_speech/dac_modules/codec.py b/vllm_omni/model_executor/models/fish_speech/dac_modules/codec.py new file mode 100644 index 000000000..4bc8483d0 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_modules/codec.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from the fish-speech 0.1.0 PyPI release (Apache-2.0) +# https://pypi.org/project/fish-speech/0.1.0/ +# Copyright (c) Fish Audio +"""DAC codec model (encoder / quantizer / decoder). + +Differences from upstream ``fish_speech.models.dac.modded_dac.DAC``, none of +which affect the parameter tree or the inference math: + +- Base class is plain ``nn.Module`` instead of ``(audiotools BaseModel, + descript CodecMixin)`` -- both bases only add training/CLI conveniences. +- ``self.delay = self.get_delay()`` (a CodecMixin helper) is dropped; nothing + in the inference paths reads it. +- The training-only ``forward``/``preprocess`` methods are dropped; only + ``encode`` and ``decode`` are used at inference time. +""" + +import math + +import torch +from torch import nn + +from vllm_omni.model_executor.models.fish_speech.dac_modules.conv import ( + CausalWNConv1d, + CausalWNConvTranspose1d, +) +from vllm_omni.model_executor.models.fish_speech.dac_modules.layers import ( + Snake1d, + WNConv1d, + WNConvTranspose1d, +) +from vllm_omni.model_executor.models.fish_speech.dac_modules.transformer import ( + ModelArgs, + WindowLimitedTransformer, +) + + +def init_weights(m): + if isinstance(m, nn.Conv1d): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.constant_(m.bias, 0) + + +class ResidualUnit(nn.Module): + def __init__(self, dim: int = 16, dilation: int = 1, causal: bool = False): + super().__init__() + conv_class = CausalWNConv1d if causal else WNConv1d + pad = ((7 - 1) * dilation) // 2 + self.block = nn.Sequential( + Snake1d(dim), + conv_class(dim, dim, kernel_size=7, dilation=dilation, padding=pad), + Snake1d(dim), + conv_class(dim, dim, kernel_size=1), + ) + self.causal = causal + + def forward(self, x): + y = self.block(x) + pad = x.shape[-1] - y.shape[-1] + if pad > 0: + if self.causal: + x = x[..., :-pad] + else: + x = x[..., pad // 2 : -pad // 2] + return x + y + + +class EncoderBlock(nn.Module): + def __init__( + self, + dim: int = 16, + stride: int = 1, + causal: bool = False, + n_t_layer: int = 0, + transformer_general_config=None, + ): + super().__init__() + conv_class = CausalWNConv1d if causal else WNConv1d + transformer_module = ( + nn.Identity() + if n_t_layer == 0 + else ( + WindowLimitedTransformer( + causal=causal, + input_dim=dim, + window_size=512, + config=transformer_general_config( + n_layer=n_t_layer, + n_head=dim // 64, + dim=dim, + intermediate_size=dim * 3, + ), + ) + ) + ) + self.block = nn.Sequential( + ResidualUnit(dim // 2, dilation=1, causal=causal), + ResidualUnit(dim // 2, dilation=3, causal=causal), + ResidualUnit(dim // 2, dilation=9, causal=causal), + Snake1d(dim // 2), + conv_class( + dim // 2, + dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ), + transformer_module, + ) + + def forward(self, x): + return self.block(x) + + +class Encoder(nn.Module): + def __init__( + self, + d_model: int = 64, + strides: list = [2, 4, 8, 8], + d_latent: int = 64, + n_transformer_layers: list = [0, 0, 4, 4], + transformer_general_config: ModelArgs = None, + causal: bool = False, + ): + super().__init__() + conv_class = CausalWNConv1d if causal else WNConv1d + # Create first convolution + self.block = [conv_class(1, d_model, kernel_size=7, padding=3)] + + # Create EncoderBlocks that double channels as they downsample by `stride` + for stride, n_t_layer in zip(strides, n_transformer_layers): + d_model *= 2 + self.block += [ + EncoderBlock( + d_model, + stride=stride, + causal=causal, + n_t_layer=n_t_layer, + transformer_general_config=transformer_general_config, + ) + ] + + # Create last convolution + self.block += [ + Snake1d(d_model), + conv_class(d_model, d_latent, kernel_size=3, padding=1), + ] + + # Wrap black into nn.Sequential + self.block = nn.Sequential(*self.block) + self.enc_dim = d_model + + def forward(self, x): + return self.block(x) + + +class DecoderBlock(nn.Module): + def __init__( + self, + input_dim: int = 16, + output_dim: int = 8, + stride: int = 1, + causal: bool = False, + n_t_layer: int = 0, + transformer_general_config=None, + ): + super().__init__() + conv_trans_class = CausalWNConvTranspose1d if causal else WNConvTranspose1d + # NOTE: upstream constructs a WindowLimitedTransformer from + # ``n_t_layer`` here but never registers it (the line placing it in + # ``self.block`` is commented out upstream), so decoder blocks carry no + # transformer parameters. It is not built here at all -- same module + # tree, without allocating throwaway weights. + del n_t_layer, transformer_general_config + self.block = nn.Sequential( + Snake1d(input_dim), + conv_trans_class( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ), + ResidualUnit(output_dim, dilation=1, causal=causal), + ResidualUnit(output_dim, dilation=3, causal=causal), + ResidualUnit(output_dim, dilation=9, causal=causal), + ) + + def forward(self, x): + return self.block(x) + + +class Decoder(nn.Module): + def __init__( + self, + input_channel, + channels, + rates, + d_out: int = 1, + causal: bool = False, + n_transformer_layers: list = [0, 0, 0, 0], + transformer_general_config=None, + ): + super().__init__() + conv_class = CausalWNConv1d if causal else WNConv1d + # Add first conv layer + layers = [conv_class(input_channel, channels, kernel_size=7, padding=3)] + + # Add upsampling + MRF blocks + for i, (stride, n_t_layer) in enumerate(zip(rates, n_transformer_layers)): + input_dim = channels // 2**i + output_dim = channels // 2 ** (i + 1) + layers += [ + DecoderBlock( + input_dim, + output_dim, + stride, + causal=causal, + n_t_layer=n_t_layer, + transformer_general_config=transformer_general_config, + ) + ] + + # Add final conv layer + layers += [ + Snake1d(output_dim), + conv_class(output_dim, d_out, kernel_size=7, padding=3), + nn.Tanh(), + ] + + self.model = nn.Sequential(*layers) + + def forward(self, x): + return self.model(x) + + +class DAC(nn.Module): + def __init__( + self, + encoder_dim: int = 64, + encoder_rates: list[int] = [2, 4, 8, 8], + latent_dim: int = None, + decoder_dim: int = 1536, + decoder_rates: list[int] = [8, 8, 4, 2], + quantizer: torch.nn.Module = None, + sample_rate: int = 44100, + causal: bool = True, + encoder_transformer_layers: list[int] = [0, 0, 0, 0], + decoder_transformer_layers: list[int] = [0, 0, 0, 0], + transformer_general_config=None, + ): + super().__init__() + + self.encoder_dim = encoder_dim + self.encoder_rates = encoder_rates + self.decoder_dim = decoder_dim + self.decoder_rates = decoder_rates + self.sample_rate = sample_rate + + if latent_dim is None: + latent_dim = encoder_dim * (2 ** len(encoder_rates)) + + self.latent_dim = latent_dim + + self.hop_length = int(math.prod(encoder_rates)) + self.encoder = Encoder( + encoder_dim, + encoder_rates, + latent_dim, + causal=causal, + n_transformer_layers=encoder_transformer_layers, + transformer_general_config=transformer_general_config, + ) + + self.quantizer = quantizer + + self.decoder = Decoder( + latent_dim, + decoder_dim, + decoder_rates, + causal=causal, + n_transformer_layers=decoder_transformer_layers, + transformer_general_config=transformer_general_config, + ) + self.sample_rate = sample_rate + self.apply(init_weights) + + self.frame_length = self.hop_length * 4 + + def encode( + self, + audio_data: torch.Tensor, + audio_lengths: torch.Tensor = None, + n_quantizers: int = None, + **kwargs, + ): + """Encode given audio data and return quantized latent codes. + + Parameters + ---------- + audio_data : Tensor[B x 1 x T] + Audio data to encode + audio_lengths : Tensor[B], optional + Valid sample counts per batch entry + n_quantizers : int, optional + Number of quantizers to use, by default None (all quantizers) + + Returns + ------- + Tuple of (codes Tensor[B x N x F], code lengths Tensor[B]). + """ + # pad to multiple of self.frame_length + if audio_data.ndim == 2: + audio_data = audio_data.unsqueeze(1) + length = audio_data.shape[-1] + right_pad = math.ceil(length / self.frame_length) * self.frame_length - length + audio_data = nn.functional.pad(audio_data, (0, right_pad)) + if audio_lengths is None: + audio_lengths = torch.LongTensor([length + right_pad]).to(audio_data.device) + + z = self.encoder(audio_data) + vq_results = self.quantizer(z, n_quantizers, **kwargs) + indices = vq_results.codes + indices_lens = torch.ceil(audio_lengths / self.frame_length).long() + return indices, indices_lens + + def decode(self, indices: torch.Tensor, feature_lengths): + """Decode codebook indices [B x N x F] to waveform [B x 1 x T].""" + if indices.ndim == 2: + indices = indices[None] + + z = self.quantizer.decode(indices) + audio_lengths = feature_lengths * self.frame_length + return self.decoder(z), audio_lengths diff --git a/vllm_omni/model_executor/models/fish_speech/dac_modules/conv.py b/vllm_omni/model_executor/models/fish_speech/dac_modules/conv.py new file mode 100644 index 000000000..2b75197a3 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_modules/conv.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from the fish-speech 0.1.0 PyPI release (Apache-2.0) +# https://pypi.org/project/fish-speech/0.1.0/ +# Copyright (c) Fish Audio +"""Causal convolution primitives shared by the DAC encoder/decoder and the +downsample RVQ. Upstream duplicates these between ``modded_dac.py`` and +``rvq.py``; they are consolidated here (the copies are identical).""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.parametrizations import weight_norm +from torch.nn.utils.parametrize import remove_parametrizations + + +def unpad1d(x: torch.Tensor, paddings: tuple[int, int]): + """Remove padding from x, handling properly zero padding. Only for 1d!""" + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + assert (padding_left + padding_right) <= x.shape[-1] + end = x.shape[-1] - padding_right + return x[..., padding_left:end] + + +def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0) -> int: + """See `pad_for_conv1d`.""" + 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 = "zeros", + value: float = 0.0, +): + """Tiny wrapper around F.pad, just to allow for reflect padding on small input. + If this is the case, we insert extra 0 padding to the right + before the reflection happen. + """ + 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 CausalConvNet(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + dilation=1, + stride=1, + groups=1, + padding=None, + ): + super().__init__() + self.conv = nn.Conv1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + dilation=dilation, + groups=groups, + ) + self.stride = stride + self.kernel_size = (kernel_size - 1) * dilation + 1 + self.dilation = dilation + self.padding = self.kernel_size - self.stride + + def forward(self, x): + pad = self.padding + extra_padding = get_extra_padding_for_conv1d(x, self.kernel_size, self.stride, pad) + x = pad1d(x, (pad, extra_padding), mode="constant", value=0) + return self.conv(x).contiguous() + + def weight_norm(self, name="weight", dim=0): + self.conv = weight_norm(self.conv, name=name, dim=dim) + return self + + def remove_weight_norm(self): + self.conv = remove_parametrizations(self.conv) + return self + + +class CausalTransConvNet(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, dilation=1, stride=1, padding=None): + super().__init__() + self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation) + self.stride = stride + self.kernel_size = kernel_size + + def forward(self, x): + x = self.conv(x) + pad = self.kernel_size - self.stride + padding_right = math.ceil(pad) + padding_left = pad - padding_right + x = unpad1d(x, (padding_left, padding_right)) + return x.contiguous() + + def weight_norm(self, name="weight", dim=0): + self.conv = weight_norm(self.conv, name=name, dim=dim) + return self + + def remove_weight_norm(self): + self.conv = remove_parametrizations(self.conv) + return self + + +def CausalWNConv1d(*args, **kwargs): + return CausalConvNet(*args, **kwargs).weight_norm() + + +def CausalWNConvTranspose1d(*args, **kwargs): + return CausalTransConvNet(*args, **kwargs).weight_norm() + + +class ConvNeXtBlock(nn.Module): + r"""ConvNeXt Block. There are two equivalent implementations: + (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) + (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back + We use (2) as we find it slightly faster in PyTorch + + Args: + dim (int): Number of input channels. + layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0. + kernel_size (int): Kernel size for depthwise conv. Default: 7. + dilation (int): Dilation for depthwise conv. Default: 1. + """ # noqa: E501 + + def __init__( + self, + dim: int, + layer_scale_init_value: float = 1e-6, + mlp_ratio: float = 4.0, + kernel_size: int = 7, + dilation: int = 1, + ): + super().__init__() + convnet_type = CausalConvNet + self.dwconv = convnet_type( + dim, + dim, + kernel_size=kernel_size, + groups=dim, + dilation=dilation, + ) # depthwise conv + self.norm = nn.LayerNorm(dim, eps=1e-6) + # pointwise/1x1 convs, implemented with linear layers + self.pwconv1 = nn.Linear(dim, int(mlp_ratio * dim)) + self.act = nn.GELU() + self.pwconv2 = nn.Linear(int(mlp_ratio * 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, apply_residual: bool = True): + input = x + + x = self.dwconv(x) + x = x.permute(0, 2, 1) # (N, C, L) -> (N, L, C) + 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.permute(0, 2, 1) # (N, L, C) -> (N, C, L) + + if apply_residual: + x = input + x + + return x diff --git a/vllm_omni/model_executor/models/fish_speech/dac_modules/layers.py b/vllm_omni/model_executor/models/fish_speech/dac_modules/layers.py new file mode 100644 index 000000000..9486f6f7a --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_modules/layers.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/descriptinc/descript-audio-codec (MIT License) +# Copyright (c) 2023-present, Descript +"""Descript audio codec building blocks (Snake activation, weight-normed convs, +vector quantizers). + +Uses the legacy ``torch.nn.utils.weight_norm`` (``weight_g``/``weight_v`` keys) +exactly like upstream so quantizer projection weights in ``codec.pth`` load +unchanged. The ``einops.rearrange`` calls are replaced with equivalent +reshape/permute ops; ``from_latents`` (training-only) is dropped. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils import weight_norm + + +def WNConv1d(*args, **kwargs): + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def WNConvTranspose1d(*args, **kwargs): + return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) + + +# Scripting this brings model speed up 1.4x +@torch.jit.script +def snake(x, alpha): + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + x = x.reshape(shape) + return x + + +class Snake1d(nn.Module): + def __init__(self, channels): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x): + return snake(x, self.alpha) + + +class VectorQuantize(nn.Module): + """ + Implementation of VQ similar to Karpathy's repo: + https://github.com/karpathy/deep-vector-quantization + Additionally uses following tricks from Improved VQGAN + (https://arxiv.org/pdf/2110.04627.pdf): + 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space + for improved codebook usage + 2. l2-normalized codes: Converts euclidean distance to cosine similarity which + improves training stability + """ + + def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int): + super().__init__() + self.codebook_size = codebook_size + self.codebook_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) + self.codebook = nn.Embedding(codebook_size, codebook_dim) + + def forward(self, z): + """Quantized the input tensor using a fixed codebook and returns + the corresponding codebook vectors + + Parameters + ---------- + z : Tensor[B x D x T] + + Returns + ------- + Tensor[B x D x T] + Quantized continuous representation of input + Tensor[1] + Commitment loss to train encoder to predict vectors closer to codebook + entries + Tensor[1] + Codebook loss to update the codebook + Tensor[B x T] + Codebook indices (quantized discrete representation of input) + Tensor[B x D x T] + Projected latents (continuous representation of input before quantization) + """ + + # Factorized codes (ViT-VQGAN) Project input into low-dimensional space + z_e = self.in_proj(z) # z_e : (B x D x T) + z_q, indices = self.decode_latents(z_e) + + commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2]) + codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2]) + + z_q = z_e + (z_q - z_e).detach() # noop in forward pass, straight-through gradient estimator in backward pass + + z_q = self.out_proj(z_q) + + return z_q, commitment_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): + batch, dim, time = latents.shape + # rearrange "b d t -> (b t) d" + encodings = latents.permute(0, 2, 1).reshape(batch * time, dim) + codebook = self.codebook.weight # codebook: (N x D) + + # L2 normalize encodings and codebook (ViT-VQGAN) + encodings = F.normalize(encodings) + codebook = F.normalize(codebook) + + # Compute euclidean distance with codebook + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + # rearrange "(b t) -> b t" + indices = (-dist).max(1)[1].reshape(batch, time) + z_q = self.decode_code(indices) + return z_q, indices + + +class ResidualVectorQuantize(nn.Module): + """ + Introduced in SoundStream: An end2end neural audio codec + https://arxiv.org/abs/2107.03312 + """ + + def __init__( + self, + input_dim: int = 512, + n_codebooks: int = 9, + codebook_size: int = 1024, + codebook_dim: int | list = 8, + quantizer_dropout: float = 0.0, + ): + super().__init__() + if isinstance(codebook_dim, int): + codebook_dim = [codebook_dim for _ in range(n_codebooks)] + + self.n_codebooks = n_codebooks + self.codebook_dim = codebook_dim + self.codebook_size = codebook_size + + self.quantizers = nn.ModuleList( + [VectorQuantize(input_dim, codebook_size, codebook_dim[i]) for i in range(n_codebooks)] + ) + self.quantizer_dropout = quantizer_dropout + + def forward(self, z, n_quantizers: int = None): + """Quantized the input tensor using a fixed set of `n` codebooks and returns + the corresponding codebook vectors + + Parameters + ---------- + z : Tensor[B x D x T] + n_quantizers : int, optional + No. of quantizers to use + (n_quantizers < self.n_codebooks ex: for quantizer dropout) + Note: if `self.quantizer_dropout` is True, this argument is ignored + when in training mode, and a random number of quantizers is used. + + Returns + ------- + Tuple of (z_q, codes, latents, commitment_loss, codebook_loss). + """ + z_q = 0 + residual = z + commitment_loss = 0 + codebook_loss = 0 + + codebook_indices = [] + latents = [] + + if n_quantizers is None: + n_quantizers = self.n_codebooks + if self.training: + n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1 + dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],)) + n_dropout = int(z.shape[0] * self.quantizer_dropout) + n_quantizers[:n_dropout] = dropout[:n_dropout] + n_quantizers = n_quantizers.to(z.device) + + for i, quantizer in enumerate(self.quantizers): + if self.training is False and i >= n_quantizers: + break + + z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(residual) + + # Create mask to apply quantizer dropout + mask = torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers + z_q = z_q + z_q_i * mask[:, None, None] + residual = residual - z_q_i + + # Sum losses + commitment_loss += (commitment_loss_i * mask).mean() + codebook_loss += (codebook_loss_i * mask).mean() + + codebook_indices.append(indices_i) + latents.append(z_e_i) + + codes = torch.stack(codebook_indices, dim=1) + latents = torch.cat(latents, dim=1) + + return z_q, codes, latents, commitment_loss, codebook_loss + + def from_codes(self, codes: torch.Tensor): + """Given the quantized codes, reconstruct the continuous representation + + Parameters + ---------- + codes : Tensor[B x N x T] + Quantized discrete representation of input + + Returns + ------- + Tensor[B x D x T] + Quantized continuous representation of input + """ + z_q = 0.0 + z_p = [] + n_codebooks = codes.shape[1] + for i in range(n_codebooks): + z_p_i = self.quantizers[i].decode_code(codes[:, i, :]) + z_p.append(z_p_i) + + z_q_i = self.quantizers[i].out_proj(z_p_i) + z_q = z_q + z_q_i + return z_q, torch.cat(z_p, dim=1), codes diff --git a/vllm_omni/model_executor/models/fish_speech/dac_modules/rvq.py b/vllm_omni/model_executor/models/fish_speech/dac_modules/rvq.py new file mode 100644 index 000000000..25b0a89ee --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_modules/rvq.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from the fish-speech 0.1.0 PyPI release (Apache-2.0) +# https://pypi.org/project/fish-speech/0.1.0/ +# Copyright (c) Fish Audio +"""Downsampled residual vector quantizer (1 semantic + N residual codebooks).""" + +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm_omni.model_executor.models.fish_speech.dac_modules.conv import ( + CausalConvNet, + CausalTransConvNet, + ConvNeXtBlock, +) +from vllm_omni.model_executor.models.fish_speech.dac_modules.layers import ( + ResidualVectorQuantize, +) + + +@dataclass +class VQResult: + z: torch.Tensor + codes: torch.Tensor + latents: torch.Tensor + codebook_loss: torch.Tensor + commitment_loss: torch.Tensor + semantic_distill_z: torch.Tensor | None = None + + +class DownsampleResidualVectorQuantize(nn.Module): + def __init__( + self, + input_dim: int = 1024, + n_codebooks: int = 9, + codebook_dim: int = 8, + quantizer_dropout: float = 0.5, + codebook_size: int = 1024, + semantic_codebook_size: int = 4096, + downsample_factor: tuple[int] = (2, 2), + downsample_dims: tuple[int] | None = None, + pre_module: nn.Module | None = None, + post_module: nn.Module | None = None, + semantic_predictor_module: nn.Module | None = None, + ): + super().__init__() + + if downsample_dims is None: + downsample_dims = [input_dim for _ in range(len(downsample_factor))] + + all_dims = (input_dim,) + tuple(downsample_dims) + + self.semantic_quantizer = ResidualVectorQuantize( + input_dim=input_dim, + n_codebooks=1, + codebook_size=semantic_codebook_size, + codebook_dim=codebook_dim, + quantizer_dropout=0.0, + ) + + self.quantizer = ResidualVectorQuantize( + input_dim=input_dim, + n_codebooks=n_codebooks, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + quantizer_dropout=quantizer_dropout, + ) + + self.downsample_factor = downsample_factor + self.downsample_dims = downsample_dims + + convnet_type = CausalConvNet + transconvnet_type = CausalTransConvNet + + self.downsample = nn.Sequential( + *[ + nn.Sequential( + convnet_type( + all_dims[idx], + all_dims[idx + 1], + kernel_size=factor, + stride=factor, + ), + ConvNeXtBlock(dim=all_dims[idx + 1]), + ) + for idx, factor in enumerate(downsample_factor) + ] + ) + + self.upsample = nn.Sequential( + *[ + nn.Sequential( + transconvnet_type( + all_dims[idx + 1], + all_dims[idx], + kernel_size=factor, + stride=factor, + ), + ConvNeXtBlock(dim=all_dims[idx]), + ) + for idx, factor in reversed(list(enumerate(downsample_factor))) + ] + ) + self.apply(self._init_weights) + # leave for transformer, LSTM or Mamba or something else + self.pre_module = pre_module if pre_module is not None else nn.Identity() + self.post_module = post_module if post_module is not None else nn.Identity() + self.semantic_predictor_module = ( + semantic_predictor_module if semantic_predictor_module is not None else nn.Identity() + ) + + 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, z, n_quantizers: int = None, semantic_len: torch.Tensor = None, **kwargs): + # z: (B, D, T) + original_shape = z.shape + if semantic_len is None: + semantic_len = torch.LongTensor([z.shape[-1]]) + z = self.downsample(z) + z = self.pre_module(z) # B, T, D + ( + semantic_z, + semantic_codes, + semantic_latents, + semantic_commitment_loss, + semantic_codebook_loss, + ) = self.semantic_quantizer(z) + residual_z = z - semantic_z + residual_z, codes, latents, commitment_loss, codebook_loss = self.quantizer( + residual_z, n_quantizers=n_quantizers + ) + z = semantic_z + residual_z + commitment_loss = commitment_loss + semantic_commitment_loss + codebook_loss = codebook_loss + semantic_codebook_loss + codes = torch.cat([semantic_codes, codes], dim=1) + latents = torch.cat([semantic_latents, latents], dim=1) + z = self.post_module(z) + z = self.upsample(z) + # z: (B, D, T) + + # Pad or crop z to match original shape + diff = original_shape[-1] - z.shape[-1] + right = 0 + left = abs(diff) - right + + if diff > 0: + z = F.pad(z, (left, right)) + elif diff < 0: + z = z[..., left:] + + results = VQResult( + z=z, + codes=codes, + latents=latents, + commitment_loss=commitment_loss, + codebook_loss=codebook_loss, + ) + + return results + + def decode(self, indices: torch.Tensor): + new_indices = torch.zeros_like(indices) + new_indices[:, 0] = torch.clamp(indices[:, 0], max=self.semantic_quantizer.codebook_size - 1) + new_indices[:, 1:] = torch.clamp(indices[:, 1:], max=self.quantizer.codebook_size - 1) + + z_q_semantic = self.semantic_quantizer.from_codes(new_indices[:, :1])[0] + z_q_residual = self.quantizer.from_codes(new_indices[:, 1:])[0] + z_q = z_q_semantic + z_q_residual + z_q = self.post_module(z_q) + z_q = self.upsample(z_q) + return z_q diff --git a/vllm_omni/model_executor/models/fish_speech/dac_modules/transformer.py b/vllm_omni/model_executor/models/fish_speech/dac_modules/transformer.py new file mode 100644 index 000000000..cb9521287 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_modules/transformer.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from the fish-speech 0.1.0 PyPI release (Apache-2.0) +# https://pypi.org/project/fish-speech/0.1.0/ +# Copyright (c) Fish Audio +"""Window-limited transformer used inside the DAC codec (encoder blocks, +decoder blocks, and RVQ pre/post modules).""" + +import math +from dataclasses import dataclass + +import torch +from torch import Tensor, nn +from torch.nn import functional as F + + +def find_multiple(n: int, k: int) -> int: + if n % k == 0: + return n + return n + k - (n % k) + + +@dataclass +class ModelArgs: + block_size: int = 2048 + n_layer: int = 8 + n_head: int = 8 + dim: int = 512 + intermediate_size: int = 1536 + n_local_heads: int = -1 + head_dim: int = 64 + rope_base: float = 10000 + norm_eps: float = 1e-5 + dropout_rate: float = 0.1 + attn_dropout_rate: float = 0.1 + channels_first: bool = True # to be compatible with conv1d input/output + pos_embed_type: str = "rope" # can be "rope" or "conformer" + max_relative_position: int = 128 # for conformer-style relative position embedding + + 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) + assert self.pos_embed_type in [ + "rope", + "conformer", + ], "pos_embed_type must be either 'rope' or 'conformer'" + + +class KVCache(nn.Module): + def __init__(self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16): + super().__init__() + cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim) + self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype)) + self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype)) + + def update(self, input_pos, k_val, v_val): + # input_pos: [S], k_val: [B, H, S, D] + assert input_pos.shape[0] == k_val.shape[2] + + k_out = self.k_cache + v_out = self.v_cache + k_out[:, :, input_pos] = k_val + v_out[:, :, input_pos] = v_val + + return ( + k_out[:, :, : input_pos.max() + 1, :], + v_out[:, :, : input_pos.max() + 1, :], + ) + + def clear_cache(self, prompt_len): + self.k_cache[:, :, prompt_len:, :].fill_(0) + self.v_cache[:, :, prompt_len:, :].fill_(0) + + +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 = RMSNorm(config.dim, eps=config.norm_eps) + + # Only compute RoPE frequencies if using RoPE + if config.pos_embed_type == "rope": + freqs_cis = precompute_freqs_cis(self.config.block_size, self.config.head_dim, self.config.rope_base) + self.register_buffer("freqs_cis", freqs_cis) + else: + self.register_buffer("freqs_cis", None) + + causal_mask = torch.tril(torch.ones(self.config.block_size, self.config.block_size, dtype=torch.bool)) + self.register_buffer("causal_mask", causal_mask) + + self.max_batch_size = -1 + self.max_seq_length = -1 + self.use_kv_cache = False + + def setup_caches(self, max_batch_size, max_seq_length): + """ + This method will only be called during inference when using KV cache. + """ + head_dim = self.config.dim // self.config.n_head + 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.weight.dtype + device = self.norm.weight.device + + for b in self.layers: + b.attention.kv_cache = KVCache( + max_batch_size, + max_seq_length, + self.config.n_local_heads, + head_dim, + dtype, + ).to(device) + + self.use_kv_cache = True + + def forward( + self, + x: Tensor, + input_pos: Tensor | None = None, + mask: Tensor | None = None, + ) -> Tensor: + if self.config.pos_embed_type == "rope": + assert self.freqs_cis is not None, "RoPE frequencies must be initialized for RoPE positional embedding" + freqs_cis = self.freqs_cis[input_pos] + else: + freqs_cis = None + + 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] + mask = mask[..., : input_pos.max() + 1] + else: + mask = self.causal_mask[None, None, input_pos] + mask = mask[..., input_pos] + + for i, layer in enumerate(self.layers): + x = layer(x, input_pos, freqs_cis, mask) + x = self.norm(x) + 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 = RMSNorm(config.dim, eps=config.norm_eps) + self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps) + self.attention_layer_scale = LayerScale(config.dim, inplace=True) + self.ffn_layer_scale = LayerScale(config.dim, inplace=True) + + def forward( + self, + x: Tensor, + input_pos: Tensor, + freqs_cis: Tensor, + mask: Tensor, + ) -> Tensor: + h = x + self.attention_layer_scale(self.attention(self.attention_norm(x), freqs_cis, mask, input_pos)) + out = h + self.ffn_layer_scale(self.feed_forward(self.ffn_norm(h))) + return out + + +class Attention(nn.Module): + def __init__(self, config: ModelArgs): + 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 + 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 + self.attn_dropout_rate = config.attn_dropout_rate + self.pos_embed_type = config.pos_embed_type + + # Add relative position embedding for conformer-style + if self.pos_embed_type == "conformer": + self.max_relative_position = config.max_relative_position + num_pos_embeddings = 2 * config.max_relative_position + 1 + self.rel_pos_embeddings = nn.Parameter(torch.zeros(num_pos_embeddings, self.head_dim)) + nn.init.normal_(self.rel_pos_embeddings, mean=0.0, std=0.02) + + def _compute_conformer_pos_scores(self, q: Tensor, seqlen: int) -> Tensor: + # q: [B, H, S, D] + # Returns: [B, H, S, S] + positions = torch.arange(seqlen, device=q.device) + relative_positions = positions.unsqueeze(1) - positions.unsqueeze(0) # [S, S] + relative_positions = torch.clamp( + relative_positions + self.max_relative_position, + 0, + 2 * self.max_relative_position, + ) + rel_embeddings = self.rel_pos_embeddings[relative_positions] # [S, S, D] + + # Compute attention scores with relative position embeddings + q = q.transpose(1, 2) # [B, S, H, D] + rel_logits = torch.matmul(q, rel_embeddings.transpose(-2, -1)) # [B, S, H, S] + rel_logits = rel_logits.transpose(1, 2) # [B, H, S, S] + return rel_logits + + def forward( + self, + x: Tensor, + freqs_cis: Tensor, + mask: Tensor, + input_pos: Tensor | None = None, + ) -> Tensor: + bsz, seqlen, _ = x.shape + + kv_size = self.n_local_heads * self.head_dim + q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1) + context_seqlen = seqlen + + 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) + + if self.pos_embed_type == "rope": + q = apply_rotary_emb(q, freqs_cis) + k = apply_rotary_emb(k, freqs_cis) + + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + + if self.kv_cache is not None: + k, v = self.kv_cache.update(input_pos, k, v) + + k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1) + v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1) + + if self.pos_embed_type == "conformer": + # Compute attention scores + scale = 1.0 / math.sqrt(self.head_dim) + scores = torch.matmul(q, k.transpose(-2, -1)) * scale + + # Add relative position embeddings for conformer-style + rel_scores = self._compute_conformer_pos_scores(q, seqlen) + scores = scores + rel_scores + + # Apply attention + if mask is not None: + scores = scores.masked_fill(~mask, float("-inf")) + + attn = F.softmax(scores, dim=-1) + if self.attn_dropout_rate > 0 and self.training: + attn = F.dropout(attn, p=self.attn_dropout_rate) + + y = torch.matmul(attn, v) + else: + y = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.attn_dropout_rate if self.training else 0.0, + attn_mask=mask, + ) + y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.head_dim * self.n_head) + 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) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward(self, x: Tensor) -> Tensor: + return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x))) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def _norm(self, x): + return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps) + + def forward(self, x: Tensor) -> Tensor: + output = self._norm(x.float()).type_as(x) + return output * self.weight + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: float | Tensor = 1e-2, + inplace: bool = False, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma + + +class WindowLimitedTransformer(Transformer): + """ + Transformer with window limited attention, causal. + """ + + def __init__( + self, + config: ModelArgs, + input_dim: int = 512, + window_size: int | None = None, + causal: bool = True, + look_ahead_conv: nn.Module = None, + ): + super().__init__(config) + self.window_size = window_size + self.causal = causal + self.channels_first = config.channels_first + self.look_ahead_conv = look_ahead_conv if look_ahead_conv is not None else nn.Identity() + self.input_proj = nn.Linear(input_dim, config.dim) if input_dim != config.dim else nn.Identity() + self.output_proj = nn.Linear(config.dim, input_dim) if input_dim != config.dim else nn.Identity() + + def make_window_limited_mask( + self, + max_length: int, + x_lens: Tensor | None = None, + ) -> Tensor: + """ + Make mask to form window limited attention. + """ + if self.causal: + mask = torch.tril(torch.ones(max_length, max_length)) + row_indices = torch.arange(max_length).view(-1, 1) + window_size = self.window_size or max_length + valid_range = (row_indices - window_size + 1).clamp(min=0) + column_indices = torch.arange(max_length) + mask = (column_indices >= valid_range) & mask.bool() + else: + raise NotImplementedError + mask = mask.bool()[None, None] + return mask + + def make_mask( + self, + max_length: int, + x_lens: Tensor | None = None, + ) -> Tensor: + """ + Make ordinary mask if window size is not specified. + """ + if self.causal: + mask = torch.tril(torch.ones(max_length, max_length)) + else: + mask = torch.ones(max_length, max_length) + mask = mask.bool()[None, None] + for i, x_len in enumerate(x_lens): + mask[:x_len, i] = 0 + mask = mask.bool()[None, None] + return mask + + def forward( + self, + x: Tensor, + x_lens: Tensor | None = None, + ) -> Tensor: + if self.channels_first: + x = x.transpose(1, 2) + x = self.input_proj(x) # (B, T, D) + x = self.look_ahead_conv(x) + input_pos = torch.arange(x.shape[1], device=x.device) + # construct mask to form window limited attention + max_length = x.shape[1] + if self.window_size is not None: + mask = self.make_window_limited_mask(max_length, x_lens) + else: + mask = self.make_mask(max_length, x_lens) + mask = mask.to(x.device) + x = super().forward(x, input_pos, mask) + x = self.output_proj(x) # (B, T, D) + if self.channels_first: + x = x.transpose(1, 2) + return 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/fish_speech/dac_utils.py b/vllm_omni/model_executor/models/fish_speech/dac_utils.py index f1f2590a3..b0d04e086 100644 --- a/vllm_omni/model_executor/models/fish_speech/dac_utils.py +++ b/vllm_omni/model_executor/models/fish_speech/dac_utils.py @@ -19,12 +19,12 @@ def build_dac_codec() -> nn.Module: Returns the model on CPU in eval mode -- caller is responsible for loading weights and moving to the target device. """ - from fish_speech.models.dac.modded_dac import ( + from vllm_omni.model_executor.models.fish_speech.dac_modules import ( DAC, + DownsampleResidualVectorQuantize, ModelArgs, WindowLimitedTransformer, ) - from fish_speech.models.dac.rvq import DownsampleResidualVectorQuantize base_transformer_kwargs = dict( block_size=16384, 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 9627dc7ce..984d3f638 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 @@ -5,8 +5,8 @@ Loads the DAC codec from ``codec.pth`` and decodes codebook indices Analogous to ``Qwen3TTSCode2Wav`` in qwen3_tts. -Requires the ``fish-speech`` package for the DAC model architecture. -Install with: ``pip install fish-speech`` +The DAC model architecture is vendored in ``dac_modules`` -- no external +``fish-speech`` package is required. """ from __future__ import annotations 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 664a794c4..fb3517036 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """GLM-TTS AR Model (Stage 0): Text → Speech Tokens. Based on Llama architecture, generates speech token sequences from input text. @@ -113,9 +113,9 @@ def resolve_glm_tts_tokenizer_path(model_name_or_path: Any) -> str: return candidate return model_path - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - local_dir = snapshot_download( + local_dir = hf_api().snapshot_download( model_path, allow_patterns=[ f"{_GLM_TTS_TOKENIZER_SUBDIR}/tokenizer*", @@ -169,9 +169,9 @@ def resolve_glm_tts_model_dir( except Exception: pass - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download(model_name_or_path) + return hf_api().snapshot_download(model_name_or_path) def _first_glm_tts_value(value: Any) -> Any: 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 91609e1a4..c6a54a393 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 @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage 1 (codec decoder) for higgs-audio v2. Two surfaces exposed: @@ -231,9 +231,9 @@ class HiggsAudioV2Code2Wav(nn.Module): if os.path.isfile(os.path.join(candidate, "config.json")): return candidate - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id=repo_id) + return hf_api().snapshot_download(repo_id=repo_id) # ------------------------------------------------------ direct decode API @torch.inference_mode() diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py index c3f5024b3..587bcda67 100755 --- 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 @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """TTS prompt builder + scope validators for higgs-audio v2. vllm-omni's higgs path supports two request shapes: @@ -248,9 +248,9 @@ def _load_audio_tokenizer(): audio_tokenizer_dir = _resolve_audio_tokenizer_dir() if audio_tokenizer_dir is None: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - repo_path = snapshot_download( + repo_path = hf_api().snapshot_download( _K2_OMNIVOICE_REPO, allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"], ) diff --git a/vllm_omni/model_executor/models/indextts2/preprocess_utils.py b/vllm_omni/model_executor/models/indextts2/preprocess_utils.py index 88e9a0639..827291602 100644 --- a/vllm_omni/model_executor/models/indextts2/preprocess_utils.py +++ b/vllm_omni/model_executor/models/indextts2/preprocess_utils.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """External model loading, audio I/O, and emotion conditioning for IndexTTS2.""" from __future__ import annotations @@ -171,9 +171,10 @@ def load_semantic_codec( 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") + from vllm_omni.transformers_utils.repo_utils import hf_api + + ckpt_path = hf_api().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 @@ -198,9 +199,11 @@ def load_campplus(model_path: str, device: torch.device): if ckpt_path is not None: break if ckpt_path is None: - from huggingface_hub import hf_hub_download + if os.path.isdir(model_path): + raise FileNotFoundError(f"IndexTTS CAMPPlus checkpoint is missing from local bundle {model_path!r}") + from vllm_omni.transformers_utils.repo_utils import hf_api - ckpt_path = hf_hub_download("funasr/campplus", filename="campplus_cn_common.bin") + ckpt_path = hf_api().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)) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py index 5876bb147..048d1de2d 100644 --- a/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright (c) 2024 NVIDIA CORPORATION. # Licensed under the MIT license. @@ -9,13 +12,14 @@ import os import torch import torch.nn as nn -from huggingface_hub import PyTorchModelHubMixin, hf_hub_download +from huggingface_hub import PyTorchModelHubMixin 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 +from vllm_omni.transformers_utils.repo_utils import hf_api # --------------------------------------------------------------------------- # Helpers (inlined from env.py / utils.py) @@ -328,7 +332,7 @@ class BigVGAN( def _resolve(filename): if os.path.isdir(model_id): return os.path.join(model_id, filename) - return hf_hub_download( + return hf_api().hf_hub_download( repo_id=model_id, filename=filename, **{ diff --git a/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py b/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py index 6eae97703..ba27cad62 100644 --- a/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py +++ b/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """IndexTTS 2.5 multilingual tiktoken vocabulary. The vocabulary layout mirrors the official IndexTTS 2.5 tokenizer: 58,836 @@ -234,9 +234,9 @@ def resolve_indextts25_tokenizer_file( return str(model_path) try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return hf_hub_download(repo_id=model_dir, filename=tokenizer_file) + return hf_api().hf_hub_download(repo_id=model_dir, filename=tokenizer_file) except Exception as exc: raise FileNotFoundError( f"Could not resolve IndexTTS 2.5 tokenizer {tokenizer_file!r} from {model_dir!r}" diff --git a/vllm_omni/model_executor/models/joyai_vl_interaction/__init__.py b/vllm_omni/model_executor/models/joyai_vl_interaction/__init__.py new file mode 100644 index 000000000..b38884c62 --- /dev/null +++ b/vllm_omni/model_executor/models/joyai_vl_interaction/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""JoyAI-VL-Interaction native pipeline.""" diff --git a/vllm_omni/model_executor/models/joyai_vl_interaction/pipeline.py b/vllm_omni/model_executor/models/joyai_vl_interaction/pipeline.py new file mode 100644 index 000000000..478ce4a11 --- /dev/null +++ b/vllm_omni/model_executor/models/joyai_vl_interaction/pipeline.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""JoyAI-VL-Interaction -> Qwen3-TTS native all-sync pipeline.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_JOYAI_INPUT_PROCESSOR = "vllm_omni.model_executor.stage_input_processors.joyai_vl_interaction" +_QWEN3_TTS_INPUT_PROCESSOR = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" + + +JOYAI_VL_INTERACTION_PIPELINE = PipelineConfig( + model_type="joyai_vl_interaction", + model_arch="Qwen3VLForConditionalGeneration", + default_deploy_config_name="joyai_vl_interaction.yaml", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="joyvl", + 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", + sampling_constraints={ + "detokenize": True, + "skip_special_tokens": False, + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="qwen3_tts", + execution_type=StageExecutionType.LLM_AR, + input_sources=(0,), + owns_tokenizer=True, + engine_output_type="latent", + model_arch="Qwen3TTSTalkerForConditionalGeneration", + custom_process_input_func=f"{_JOYAI_INPUT_PROCESSOR}.joyai_action_to_tts", + custom_process_next_stage_input_func=f"{_QWEN3_TTS_INPUT_PROCESSOR}.talker2code2wav_full_payload", + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [2150], + }, + ), + StagePipelineConfig( + stage_id=2, + model_stage="code2wav", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(1,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="Qwen3TTSCode2Wav", + sync_process_input_func=f"{_QWEN3_TTS_INPUT_PROCESSOR}.talker2code2wav_token_only", + requires_full_payload_input=True, + sampling_constraints={"detokenize": True}, + extras={"tts_args": {"max_instructions_length": 500}}, + ), + ), +) 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 019670cf9..5423788a0 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 @@ -99,8 +99,10 @@ class MiMoAudioTokenizerWorker: self.audio_tokenizer.config.nfft, ) mel_start = time.monotonic() - self.mel_transform = ( - MelSpectrogram( + # Build on CPU first: torchaudio MelSpectrogram init can mix CPU/CUDA + # tensors on newer PyTorch when default device is CUDA. + with torch.device("cpu"): + self.mel_transform = MelSpectrogram( sample_rate=self.audio_tokenizer.config.sampling_rate, n_fft=self.audio_tokenizer.config.nfft, hop_length=self.audio_tokenizer.config.hop_length, @@ -110,10 +112,9 @@ class MiMoAudioTokenizerWorker: n_mels=self.audio_tokenizer.config.n_mels, power=1.0, center=True, - ) - .to(self.device) - .to(torch.float32) - ) + ).to(torch.float32) + if self.device != "cpu": + self.mel_transform = self.mel_transform.to(self.device) logger.info( "[tokenizer worker] MelSpectrogram ready in %.2fs", time.monotonic() - mel_start, @@ -382,12 +383,19 @@ def extract_audio_code_tensor( def _normalize_tokenizer_worker_cache_key( - device: torch.device, + device: torch.device | str, config_path: str | None, audio_tokenizer_path: str, ) -> tuple[str, str, str]: """Normalize cache key so that same tokenizer always hits the same cache entry.""" - device_type = device.type if isinstance(device, torch.device) else str(device).split(":")[0] + dev = device if isinstance(device, torch.device) else torch.device(str(device)) + # Keep the index: an explicit "cuda:1" must not collapse onto the process's + # current CUDA device. Resolve a bare "cuda" so it shares the cache entry + # with the equivalent indexed spelling. + if dev.type == "cuda" and dev.index is None and torch.cuda.is_available(): + # Repo bans torch.cuda.current_device; match mimo_audio.py. + dev = torch.device("cuda", torch.accelerator.current_device_index()) + device_key = str(dev) # Use realpath so symlinks / trailing slash don't create duplicate entries ap = audio_tokenizer_path or "" if ap and os.path.exists(ap): @@ -398,22 +406,21 @@ def _normalize_tokenizer_worker_cache_key( if not cp and ap: cp = os.path.dirname(ap) - return (device_type, cp, ap) + return (device_key, cp, ap) _TOKENIZER_WORKER_CACHE: dict[tuple[str, str, str], MiMoAudioTokenizerWorker] = {} def get_tokenizer_worker( - device: torch.device, - config_path: str, + device: torch.device | str, + config_path: str | None, audio_tokenizer_path: str, ) -> MiMoAudioTokenizerWorker: key = _normalize_tokenizer_worker_cache_key(device, config_path, audio_tokenizer_path) if key not in _TOKENIZER_WORKER_CACHE: - device_type = key[0] _TOKENIZER_WORKER_CACHE[key] = MiMoAudioTokenizerWorker( - device_str=device_type, + device_str=key[0], config_path=config_path, audio_tokenizer_path=audio_tokenizer_path, ) @@ -469,8 +476,14 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): or self.config.name_or_path ) + _tok_dev = os.environ.get("MIMO_AUDIO_TOKENIZER_DEVICE") + if _tok_dev: + tokenizer_device = torch.device("cpu" if _tok_dev.lower() == "cpu" else _tok_dev) + else: + tokenizer_device = self.device + self._tokenizer_service: MiMoAudioTokenizerWorker | None = get_tokenizer_worker( - device=self.device, + device=tokenizer_device, config_path=self.tokenizer_config_path, audio_tokenizer_path=self.audio_tokenizer_path, ) @@ -723,7 +736,7 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): Instead of calling _decode_waveform_from_codes per request (which incurs 4 GPU↔CPU round-trips each), this method: 1. Extracts audio_codes for all valid requests on CPU (cheap). - 2. Runs quantizer.decode_vq (embedding lookup) for each on GPU. + 2. Runs quantizer.decode_vq (embedding lookup) on the tokenizer device. 3. Packs all hidden-states into one tensor and calls decoder(packed_hs, input_lengths) once. 4. Splits the output waveforms back to per-request tensors. @@ -734,6 +747,9 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): return [empty] tokenizer = self._tokenizer_service.audio_tokenizer + # Codes/weights must share the tokenizer device (may be CPU via + # MIMO_AUDIO_TOKENIZER_DEVICE), not the stage CUDA device. + tok_device = torch.device(self._tokenizer_service.device) group_size = self.streamer_config.group_size audio_channels = self.streamer_config.audio_channels @@ -768,7 +784,7 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): use_cuda_graph = cg_ready and len(extracted) == 1 if use_cuda_graph: - wav_out = cg_wrapper.decode(extracted[0][1].to(self.device)) + wav_out = cg_wrapper.decode(extracted[0][1].to(tok_device)) wav = wav_out.squeeze(0).squeeze(0) cfg = tokenizer.config frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length @@ -782,7 +798,7 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): hidden_list: list[torch.Tensor] = [] lengths: list[int] = [] for _, audio_codes in extracted: - hs = tokenizer.encoder.decode_vq(audio_codes.to(self.device)) + hs = tokenizer.encoder.decode_vq(audio_codes.to(tok_device)) hidden_list.append(hs) lengths.append(hs.size(0)) @@ -790,7 +806,7 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): packed_hs = hidden_list[0] else: packed_hs = torch.cat(hidden_list, dim=0) - input_lengths = torch.tensor(lengths, device=self.device) + input_lengths = torch.tensor(lengths, device=tok_device) recon_wav = tokenizer.decoder(packed_hs, input_lengths) @@ -831,6 +847,9 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): return [empty] tokenizer = self._tokenizer_service.audio_tokenizer + # Codes/weights must share the tokenizer device (may be CPU via + # MIMO_AUDIO_TOKENIZER_DEVICE), not the stage CUDA device. + tok_device = torch.device(self._tokenizer_service.device) group_size = self.streamer_config.group_size audio_channels = self.streamer_config.audio_channels @@ -886,7 +905,7 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length if use_cuda_graph: - wav_out = cg_wrapper.decode(extracted[0][1].to(self.device)) + wav_out = cg_wrapper.decode(extracted[0][1].to(tok_device)) wav = wav_out.squeeze(0).squeeze(0) valid_len = extracted[0][1].shape[-1] * frames_per_token if wav.numel() > valid_len: @@ -903,7 +922,7 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): hidden_list: list[torch.Tensor] = [] lengths: list[int] = [] for _, audio_codes in extracted: - hs = tokenizer.encoder.decode_vq(audio_codes.to(self.device)) + hs = tokenizer.encoder.decode_vq(audio_codes.to(tok_device)) hidden_list.append(hs) lengths.append(hs.size(0)) @@ -911,7 +930,7 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): packed_hs = hidden_list[0] else: packed_hs = torch.cat(hidden_list, dim=0) - input_lengths = torch.tensor(lengths, device=self.device) + input_lengths = torch.tensor(lengths, device=tok_device) recon_wav = tokenizer.decoder(packed_hs, input_lengths) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 8b3926cc0..f701e8b26 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -39,9 +39,9 @@ def _resolve_model_dir(model_ref: str, revision: str | None = None) -> str: """ if Path(model_ref).is_dir(): return model_ref - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download(model_ref, revision=revision, allow_patterns=["assets/*"]) + return hf_api().snapshot_download(model_ref, revision=revision, allow_patterns=["assets/*"]) def _batch_error(reason: str, **details: Any) -> RuntimeError: 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 b0fa9d203..9dfff5589 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 @@ -3310,7 +3310,7 @@ class MiniCPMOAudioEmbeddingItems(DictEmbeddingItems): ) -> None: super().__init__( data, - modality="image", + modality="audio", required_fields={"audio_embeds"}, fields_factory=fields_factory, ) @@ -3733,7 +3733,8 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn if isinstance(audios, MiniCPMOAudioEmbeddingItems): single_audio_embeds = audios.get(item_idx)["audio_embeds"] - audio_len = self.info.get_audio_len_by_num_chunks(sum(map(len, single_audio_embeds))) + # One item is ``(s, h)``, so its leading dim is the audio embedding count. + audio_len = self.info.get_audio_len_by_num_chunks(len(single_audio_embeds)) else: audio_len = audios.get_audio_length(item_idx) @@ -3796,11 +3797,17 @@ class MiniCPMOAudioFeatureInputs(TensorSchema): audio_feature_lens: Annotated[ torch.Tensor | list[torch.Tensor], - TensorShape("bn", "s"), + TensorShape("bn", "s", dynamic_dims={"s"}), ] """ This should be feature length of each audio slice, which equals to `audio_features.shape[-1]` + + Each audio in the batch may be split into a different number of slices + (e.g. audio >30s splits into multiple slices while shorter audio doesn't), + so `s` must be dynamic: batching audios with different slice counts is a + normal, valid input and is already handled below via `hstack` + per-audio + iteration, not a schema violation. """ diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index c167aa1c0..26a6f6b59 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ OmniVoice model for vLLM-Omni two-stage TTS pipeline. @@ -487,9 +487,9 @@ class OmniVoiceModel( if os.path.isdir(model_dir): return model_dir # HF hub model ID — resolve to local cache - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download(model_dir) + return hf_api().snapshot_download(model_dir) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: try: diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py index a75c4e6f7..965d021b8 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py @@ -6,9 +6,8 @@ OmniVoice Generator (Stage 0) - Iterative unmasking with Qwen3 backbone. Generates 8-codebook audio tokens from text via 32-step non-autoregressive iterative masked prediction with classifier-free guidance. -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. +Uses backend-dispatched variable-length full bidirectional attention over +packed conditional and unconditional sequences, with an SDPA mask fallback. """ from __future__ import annotations @@ -22,7 +21,11 @@ import torch import torch.nn as nn import torch.nn.functional as F from vllm.logger import init_logger +from vllm.utils.math_utils import round_up +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.model_executor.models.omnivoice.fused_qkv_rope import fused_qkv_norm_rope from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig @@ -241,7 +244,7 @@ def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Ge # --------------------------------------------------------------------------- -# Qwen3-style transformer blocks using PyTorch SDPA +# Qwen3-style transformer blocks using PyTorch variable-length attention # --------------------------------------------------------------------------- @@ -261,9 +264,9 @@ class OmniVoiceRMSNorm(nn.Module): class OmniVoiceAttention(nn.Module): - """Qwen3-style GQA attention using PyTorch SDPA (full bidirectional).""" + """Qwen3-style GQA using packed, full-bidirectional varlen attention.""" - def __init__(self, config: OmniVoiceConfig): + def __init__(self, config: OmniVoiceConfig, layer_idx: int): super().__init__() self.hidden_size = config.llm_hidden_size self.num_heads = config.llm_num_attention_heads @@ -282,16 +285,27 @@ class OmniVoiceAttention(nn.Module): self.k_norm = OmniVoiceRMSNorm(self.head_dim) self.scale = 1.0 / math.sqrt(self.head_dim) + self.attention_op = Attention( + num_heads=self.num_heads, + num_kv_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=self.scale, + prefix=f"layers.{layer_idx}.self_attn.attention_op", + qkv_layout="BSND", + skip_sequence_parallel=True, + ) + self.use_packed_varlen = self.attention_op.attn_backend.supports_multi_doc_packed_varlen() def forward( self, hidden_states: torch.Tensor, rope_table: torch.Tensor, - attention_mask: torch.Tensor | None = None, + attn_metadata: AttentionMetadata, ) -> torch.Tensor: - batch_size, seq_len, _ = hidden_states.shape + seq_len, _ = hidden_states.shape - qkv = self.qkv_proj(hidden_states).view(batch_size, seq_len, self.num_qkv_heads, self.head_dim) + qkv = self.qkv_proj(hidden_states).view(1, seq_len, self.num_qkv_heads, self.head_dim) # One kernel for the whole prologue: split the packed projection, RMSNorm # Q and K per head, rotate both, broadcast K and V across their query @@ -306,22 +320,16 @@ class OmniVoiceAttention(nn.Module): self.num_kv_heads, ) - # Caller passes a float mask; materialize float form if a bool slips through. - sdpa_mask = attention_mask - if sdpa_mask is not None and sdpa_mask.dtype == torch.bool: - sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf")) + # The fused prologue emits [B, N, S, D]; Omni attention backends use + # [B, S, N, D]. Keep the conversion outside the backend. + q, k, v = (tensor.permute(0, 2, 1, 3).contiguous().to(torch.bfloat16) for tensor in (q, k, v)) + if self.use_packed_varlen: + out = self.attention_op(q, k, v, attn_metadata) + else: + out = self.attention_op.sdpa_fallback.forward(q, k, v, attn_metadata) - out = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=sdpa_mask, - scale=self.scale, - ) - - # Back to (batch, seq, heads * head_dim) - out = out.permute(0, 2, 1, 3).contiguous() - out = out.view(batch_size, seq_len, self.num_heads * self.head_dim) + out = out.squeeze(0).to(hidden_states.dtype) + out = out.reshape(seq_len, self.num_heads * self.head_dim) return self.o_proj(out) @@ -351,19 +359,19 @@ _FUSED_PROJECTIONS: dict[str, tuple[str, ...]] = { class OmniVoiceTransformerBlock(nn.Module): - """Single Qwen3 transformer block with PyTorch SDPA attention.""" + """Single Qwen3 transformer block with variable-length attention.""" - def __init__(self, config: OmniVoiceConfig): + def __init__(self, config: OmniVoiceConfig, layer_idx: int): super().__init__() self.input_layernorm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps) - self.self_attn = OmniVoiceAttention(config) + self.self_attn = OmniVoiceAttention(config, layer_idx) self.post_attention_layernorm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps) self.mlp = OmniVoiceMLP(config) def forward( self, hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, + attn_metadata: AttentionMetadata, rope_table: torch.Tensor | None = None, residual: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -391,7 +399,7 @@ class OmniVoiceTransformerBlock(nn.Module): residual = hidden_states hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, rope_table, attention_mask=attention_mask) + hidden_states = self.self_attn(hidden_states, rope_table, attn_metadata) if _TRITON_AVAILABLE: # Fused: (attn_out + residual) + RMSNorm in one kernel @@ -432,6 +440,63 @@ def _precompute_rope_table( return torch.cat([freqs.cos(), freqs.sin()], dim=-1) +def _position_ids_from_cu_seqs(cu_seqs: torch.Tensor, seq_len: int) -> torch.Tensor: + """Return each packed token's zero-based position within its sequence.""" + token_ids = torch.arange(seq_len, device=cu_seqs.device, dtype=cu_seqs.dtype) + sequence_ids = torch.searchsorted(cu_seqs[1:], token_ids, right=True) + sequence_starts = cu_seqs.index_select(0, sequence_ids.to(torch.long)) + return (token_ids - sequence_starts).to(torch.long) + + +def _attention_metadata_from_cu_seqs( + cu_seqs: torch.Tensor, + seq_len: int, + *, + needs_sdpa_mask: bool, + max_seqlen: int | None = None, +) -> AttentionMetadata: + """Build packed-varlen metadata and, when needed, an SDPA block mask.""" + # Local benchmarks found no measurable performance difference between the + # exact longest segment and the packed total. Eager uses the exact bound; + # CUDA Graph uses the fixed token-bucket maximum. + kernel_max_seqlen = seq_len if max_seqlen is None else max_seqlen + extra = { + "cu_seqlens_q": cu_seqs, + "cu_seqlens_k": cu_seqs, + "max_seqlen_q": kernel_max_seqlen, + "max_seqlen_k": kernel_max_seqlen, + } + if not needs_sdpa_mask: + return AttentionMetadata(extra=extra) + + token_ids = torch.arange(seq_len, device=cu_seqs.device, dtype=cu_seqs.dtype) + sequence_ids = torch.searchsorted(cu_seqs[1:], token_ids, right=True) + attn_mask = sequence_ids.view(1, 1, seq_len, 1) == sequence_ids.view(1, 1, 1, seq_len) + return AttentionMetadata(attn_mask=attn_mask, extra=extra) + + +def _build_cu_seqs( + cond_lens: list[int], + uncond_lens: list[int], + device: torch.device, + *, + tail_end: int | None = None, +) -> torch.Tensor: + """Build request-major [cond0, uncond0, ...] cumulative offsets.""" + if len(cond_lens) != len(uncond_lens): + raise ValueError(f"Mismatched cond/uncond lengths: {len(cond_lens)} != {len(uncond_lens)}.") + offsets = [0] + for cond_len, uncond_len in zip(cond_lens, uncond_lens): + offsets.append(offsets[-1] + cond_len) + offsets.append(offsets[-1] + uncond_len) + if tail_end is None: + tail_end = offsets[-1] + if offsets[-1] > tail_end: + raise ValueError(f"Packed length {offsets[-1]} exceeds tail end {tail_end}.") + offsets.append(tail_end) + return torch.tensor(offsets, device=device, dtype=torch.int32) + + # --------------------------------------------------------------------------- # TF32 opt-in (process-wide; default off) # --------------------------------------------------------------------------- @@ -461,23 +526,6 @@ def _maybe_enable_tf32() -> None: # --------------------------------------------------------------------------- -def _additive_float_mask(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - """Convert a boolean attention mask to its additive float form. - - ``True`` (attend) maps to ``0.0`` and ``False`` (masked) to ``-inf``. A bool - mask must never be copied straight into a float buffer: the implicit cast - maps True/False to 1.0/0.0, which leaves masked positions at 0.0 and so - silently *unmasks* them. - - ``dtype`` is required rather than defaulting to float32: SDPA rejects an - additive mask whose dtype differs from the query, so the mask has to follow - the model dtype and a default would just hide that coupling. - """ - if mask.dtype != torch.bool: - return mask - return torch.zeros_like(mask, dtype=dtype).masked_fill_(~mask, float("-inf")) - - class _OmniVoiceCUDAGraphForward: """Pre-captures CUDA graphs for predefined sequence-length buckets. @@ -486,24 +534,47 @@ class _OmniVoiceCUDAGraphForward: (one step at a time) means pool sharing is safe. """ - # Default bucket count is 10; 16 gives modest headroom for edge cases - # (seq_len > max bucket or non-CFG batch) without unbounded GPU growth. - _MAX_LAZY_GRAPHS: int = 16 + _MAX_LAZY_GRAPHS = 16 + _STATIC_CAPTURE_TOKEN_LIMIT = 1024 + _LAZY_CAPTURE_ALIGNMENT = 128 def __init__(self, generator: OmniVoiceGenerator, capture_sizes: list[int]) -> None: self._gen = generator - self._capture_sizes = sorted(capture_sizes) - # Pre-warmed graphs keyed by (two_b, bucket); fixed set, never evicted. + self.capture_batch_sizes = self._derive_capture_batch_size() + self.capture_bucket_sizes_by_batch = self._derive_capture_bucket_sizes(capture_sizes) self._graphs: dict[tuple[int, int], dict] = {} - # Lazy-captured graphs for oversized / non-CFG shapes; capped via LRU. self._lazy_graphs: OrderedDict[tuple[int, int], dict] = OrderedDict() self._lock = threading.Lock() # Per-instance pool handle: isolates OmniVoice CUDA memory from other # vllm modules while still allowing safe re-use across sequential replays. self._pool_handle: int | None = None - def _find_bucket(self, seq_len: int) -> int | None: - for bucket in self._capture_sizes: + def _derive_capture_batch_size(self) -> list[int]: + return list(range(1, self._gen.od_config.max_num_seqs + 1)) + + def _derive_capture_bucket_sizes(self, capture_sizes: list[int]) -> dict[int, list[int]]: + """Build a triangular, batch-aware token-bucket capture plan.""" + base_sizes = sorted(set(capture_sizes)) + if not base_sizes: + return {batch_size: [] for batch_size in self.capture_batch_sizes} + alignment = base_sizes[0] + single_request_cap = min(base_sizes[-1], 512) + max_graph_tokens = min(base_sizes[-1], self._STATIC_CAPTURE_TOKEN_LIMIT) + growth_per_batch = max(alignment, single_request_cap // 2) + candidates = sorted(set(base_sizes) | set(range(alignment, max_graph_tokens + alignment, alignment))) + return { + batch_size: [ + bucket + for bucket in candidates + if alignment * batch_size + <= bucket + <= min(single_request_cap + (batch_size - 1) * growth_per_batch, max_graph_tokens) + ] + for batch_size in self.capture_batch_sizes + } + + def _find_bucket(self, batch_size: int, seq_len: int) -> int | None: + for bucket in self.capture_bucket_sizes_by_batch.get(batch_size, ()): if bucket >= seq_len: return bucket return None @@ -512,57 +583,36 @@ class _OmniVoiceCUDAGraphForward: self, input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, bucket: int, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - S = input_ids.shape[-1] - if S == bucket: - return input_ids, audio_mask, attention_mask - - two_b = input_ids.shape[0] - num_cb = input_ids.shape[1] - - ids_padded = torch.zeros(two_b, num_cb, bucket, dtype=input_ids.dtype, device=input_ids.device) - ids_padded[:, :, :S] = input_ids - - mask_padded = torch.zeros(two_b, bucket, dtype=torch.bool, device=audio_mask.device) - mask_padded[:, :S] = audio_mask - - if attention_mask is not None: - # Callers normalize to the additive float form first, so pad with -inf. - attn_padded = torch.full( - (two_b, 1, bucket, bucket), - float("-inf"), - dtype=attention_mask.dtype, - device=attention_mask.device, - ) - attn_padded[:, :, :S, :S] = attention_mask - else: - attn_padded = None - - return ids_padded, mask_padded, attn_padded + ) -> tuple[torch.Tensor, torch.Tensor]: + seq_len = input_ids.shape[0] + if seq_len == bucket: + return input_ids, audio_mask + return ( + F.pad(input_ids, (0, 0, 0, bucket - seq_len), value=0), + F.pad(audio_mask, (0, bucket - seq_len), value=False), + ) def _capture_for_key( self, key: tuple[int, int], input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, + cu_seqs: torch.Tensor, ) -> dict: _, bucket = key device = input_ids.device - static_rope_table = self._gen._rope_table_for(bucket, device, self._gen.model_dtype) - static_input_ids = input_ids.clone() static_audio_mask = audio_mask.clone() - static_attn_mask = attention_mask.clone() if attention_mask is not None else None + static_cu_seqs = cu_seqs.clone() + static_rope_table = self._gen._rope_table_for(bucket, device, self._gen.model_dtype) with torch.no_grad(): _ = self._gen._step_forward( static_input_ids, static_audio_mask, - static_attn_mask, + static_cu_seqs, static_rope_table, ) torch.accelerator.synchronize(device) @@ -580,7 +630,7 @@ class _OmniVoiceCUDAGraphForward: static_output = self._gen._step_forward( static_input_ids, static_audio_mask, - static_attn_mask, + static_cu_seqs, static_rope_table, ) @@ -588,87 +638,83 @@ class _OmniVoiceCUDAGraphForward: "graph": graph, "static_input_ids": static_input_ids, "static_audio_mask": static_audio_mask, - "static_attn_mask": static_attn_mask, + "static_cu_seqs": static_cu_seqs, "static_rope_table": static_rope_table, "static_output": static_output, } logger.info("OmniVoice CUDA Graph captured for key %s", key) return entry + def make_capture_cu_seq(self, batch_size: int, bucket_size: int, device: torch.device) -> torch.Tensor: + num_real_sequences = 2 * batch_size + base, rem = divmod(bucket_size, num_real_sequences) + lengths = torch.full((num_real_sequences,), base, dtype=torch.int32, device=device) + lengths[:rem] += 1 + cu_seq = torch.empty(num_real_sequences + 2, dtype=torch.int32, device=device) + cu_seq[0] = 0 + cu_seq[1:-1] = lengths.cumsum(0) + cu_seq[-1] = bucket_size + return cu_seq + def warmup(self, device: torch.device) -> None: - """Pre-capture graphs for all bucket sizes with B=1 (two_b=2 for CFG).""" + """Pre-capture common request batch sizes for every token bucket.""" if not torch.cuda.is_available(): return logger.info( - "OmniVoice CUDA Graph warmup: capturing %d bucket sizes %s", - len(self._capture_sizes), - self._capture_sizes, + "OmniVoice CUDA Graph warmup: batch-aware capture plan %s", + self.capture_bucket_sizes_by_batch, ) - two_b = 2 num_cb = self._gen.config.num_audio_codebook - for bucket in self._capture_sizes: - key = (two_b, bucket) - dummy_ids = torch.zeros(two_b, num_cb, bucket, dtype=torch.long, device=device) - dummy_mask = torch.zeros(two_b, bucket, dtype=torch.bool, device=device) - # Capture with a float mask to match what forward() feeds at replay time, - # in the model dtype so replay can copy_ into it without a cast. - dummy_attn = torch.zeros(two_b, 1, bucket, bucket, dtype=self._gen.model_dtype, device=device) - self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_attn) + for batch_size in self.capture_batch_sizes: + for bucket in self.capture_bucket_sizes_by_batch[batch_size]: + key = (batch_size, bucket) + dummy_ids = torch.zeros(bucket, num_cb, dtype=torch.long, device=device) + dummy_mask = torch.zeros(bucket, dtype=torch.bool, device=device) + dummy_cu_seqs = self.make_capture_cu_seq(batch_size, bucket, device) + self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_cu_seqs) logger.info("OmniVoice CUDA Graph warmup complete (%d graphs)", len(self._graphs)) def __call__( self, input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, + cu_seqs: torch.Tensor, + batch_size: int, ) -> torch.Tensor: if torch.cuda.is_current_stream_capturing(): - rope_table = self._gen._rope_table_for(input_ids.shape[-1], input_ids.device, self._gen.model_dtype) - return self._gen._step_forward(input_ids, audio_mask, attention_mask, rope_table) - - seq_len = input_ids.shape[-1] - two_b = input_ids.shape[0] - bucket = self._find_bucket(seq_len) if two_b == 2 else None - - # Graphs are captured with (and their static buffers hold) the additive - # float mask, so normalize here, before padding or any copy_ into them. - if attention_mask is not None: - attention_mask = _additive_float_mask(attention_mask, self._gen.model_dtype) + rope_table = self._gen._rope_table_for(input_ids.shape[0], input_ids.device, self._gen.model_dtype) + return self._gen._step_forward(input_ids, audio_mask, cu_seqs, rope_table) + seq_len = input_ids.shape[0] + bucket = self._find_bucket(batch_size, seq_len) + is_lazy = bucket is None if bucket is None: - # Lazy capture: oversized sequence or non-unit batch (no pre-warmed bucket). - # Lock prevents concurrent threads from double-capturing the same key. - # _lazy_graphs is capped at _MAX_LAZY_GRAPHS with LRU eviction to - # prevent unbounded GPU memory growth when seq_len varies widely. - key = (two_b, seq_len) - ids_in, mask_in, attn_in = input_ids, audio_mask, attention_mask - with self._lock: - entry = self._lazy_graphs.get(key) - if entry is None: - entry = self._capture_for_key(key, ids_in, mask_in, attn_in) - if len(self._lazy_graphs) >= self._MAX_LAZY_GRAPHS: - evicted_key, _ = self._lazy_graphs.popitem(last=False) - logger.warning("OmniVoice CUDA Graph lazy cache full; evicted key %s", evicted_key) - self._lazy_graphs[key] = entry - else: - key = (two_b, bucket) - ids_in, mask_in, attn_in = self._pad_inputs(input_ids, audio_mask, attention_mask, bucket) - with self._lock: - entry = self._graphs.get(key) - if entry is None: - entry = self._capture_for_key(key, ids_in, mask_in, attn_in) - self._graphs[key] = entry + bucket = round_up(seq_len, self._LAZY_CAPTURE_ALIGNMENT) + ids_in, mask_in = self._pad_inputs(input_ids, audio_mask, bucket) + runtime_cu_seqs = cu_seqs.clone() + runtime_cu_seqs[-1] = bucket + key = (batch_size, bucket) + cache = self._lazy_graphs if is_lazy else self._graphs + with self._lock: + entry = cache.get(key) + if entry is None: + if is_lazy and len(self._lazy_graphs) >= self._MAX_LAZY_GRAPHS: + evicted_key, _ = self._lazy_graphs.popitem(last=False) + logger.info("Evicted OmniVoice lazy CUDA Graph key %s", evicted_key) + entry = self._capture_for_key(key, ids_in, mask_in, runtime_cu_seqs) + cache[key] = entry + elif is_lazy: + self._lazy_graphs.move_to_end(key) entry["static_input_ids"].copy_(ids_in) entry["static_audio_mask"].copy_(mask_in) - if attn_in is not None and entry["static_attn_mask"] is not None: - entry["static_attn_mask"].copy_(attn_in) + entry["static_cu_seqs"].copy_(runtime_cu_seqs) entry["graph"].replay() output = entry["static_output"] - if bucket is not None and bucket != seq_len: - output = output[:, :, :seq_len, :] + if bucket != seq_len: + output = output[:, :seq_len, :] return output def clear(self) -> None: @@ -692,17 +738,17 @@ class OmniVoiceGenerator(nn.Module): - 32-step iterative unmasking with classifier-free guidance Optimizations: - - Full bidirectional attention via PyTorch SDPA (no auto-selected - FlashAttn/SageAttn/DiffusionAttention backend) + - Packed full-bidirectional varlen attention with SDPA fallback - regionally_compile() compatible for torch.compile on repeated blocks """ # For regionally_compile() support _repeated_blocks = ["layers"] - def __init__(self, config: OmniVoiceConfig): + def __init__(self, config: OmniVoiceConfig, od_config: OmniDiffusionConfig): super().__init__() self.config = config + self.od_config = od_config # Opt-in TF32; must run before any CUDA-graph capture so captured kernels honour it. if getattr(config, "enable_tf32", False): @@ -722,7 +768,10 @@ class OmniVoiceGenerator(nn.Module): ) # Transformer layers - self.layers = nn.ModuleList([OmniVoiceTransformerBlock(config) for _ in range(config.llm_num_hidden_layers)]) + self.layers = nn.ModuleList( + [OmniVoiceTransformerBlock(config, layer_idx) for layer_idx in range(config.llm_num_hidden_layers)] + ) + self._needs_sdpa_mask = any(not layer.self_attn.use_packed_varlen for layer in self.layers) self.norm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps) # Prediction head: hidden → 8 * 1025 @@ -775,23 +824,23 @@ class OmniVoiceGenerator(nn.Module): """Prepare mixed text+audio embeddings. Args: - input_ids: [B, 8, S] - text tokens replicated across codebooks, + input_ids: [T, 8] - text tokens replicated across codebooks, audio positions have per-codebook token IDs - audio_mask: [B, S] - True for audio positions, False for text - text_embeds: optional cached [B, S, H] text-position embeddings - audio_mask_3d: optional cached [B, S, 1] audio_mask.unsqueeze(-1) + audio_mask: [T] - True for audio positions, False for text + text_embeds: optional cached [T, H] text-position embeddings + audio_mask_3d: optional cached [T, 1] audio_mask.unsqueeze(-1) Returns: - embeddings: [B, S, hidden_size] + embeddings: [T, hidden_size] """ # Cached across the denoising loop since text ids don't change. if text_embeds is None: - text_embeds = self.text_embedding(input_ids[:, 0, :]) + text_embeds = self.text_embedding(input_ids[:, 0]) if audio_mask_3d is None: audio_mask_3d = audio_mask.unsqueeze(-1) # Audio embeddings: offset per codebook, then sum across codebooks - shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1, 1) + shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1) audio_embeds = self.audio_embeddings(shifted_ids).sum(dim=1) # Merge: audio where audio_mask=True, text elsewhere @@ -800,35 +849,39 @@ class OmniVoiceGenerator(nn.Module): def _transformer_forward( self, inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor | None = None, + cu_seqs: torch.Tensor, + max_seqlen: int | None = None, rope_table: torch.Tensor | None = None, ) -> torch.Tensor: """Run through transformer layers. Args: - inputs_embeds: [B, S, hidden_size] - attention_mask: [B, 1, S, S] or None - rope_table: optional precomputed [B * S, head_dim] RoPE table + inputs_embeds: [T, hidden_size] + cu_seqs: cumulative boundaries for packed sequences + rope_table: optional base [T, head_dim] RoPE table Returns: - hidden_states: [B, S, hidden_size] + hidden_states: [T, hidden_size] """ hidden_states = inputs_embeds if rope_table is None: - rope_table = self._rope_table_for(inputs_embeds.shape[1], inputs_embeds.device, hidden_states.dtype) - - # Safety: convert bool mask if caller hasn't (e.g. external paths beyond forward()). - if attention_mask is not None and attention_mask.dtype == torch.bool: - attention_mask = torch.zeros_like(attention_mask, dtype=hidden_states.dtype).masked_fill_( - ~attention_mask, float("-inf") - ) + rope_table = self._rope_table_for(inputs_embeds.shape[0], inputs_embeds.device, hidden_states.dtype) + seq_len = inputs_embeds.shape[0] + position_ids = _position_ids_from_cu_seqs(cu_seqs, seq_len) + packed_rope_table = rope_table.index_select(0, position_ids).contiguous() + attn_metadata = _attention_metadata_from_cu_seqs( + cu_seqs, + seq_len, + needs_sdpa_mask=self._needs_sdpa_mask, + max_seqlen=max_seqlen, + ) residual = None for layer in self.layers: hidden_states, residual = layer( hidden_states, - attention_mask=attention_mask, - rope_table=rope_table, + attn_metadata=attn_metadata, + rope_table=packed_rope_table, residual=residual, ) @@ -838,44 +891,77 @@ class OmniVoiceGenerator(nn.Module): """Project hidden states to per-codebook logits. Args: - hidden_states: [B, S, hidden_size] + hidden_states: [T, hidden_size] Returns: - logits: [B, 8, S, 1025] + logits: [8, T, 1025] """ - batch_size, seq_len, _ = hidden_states.shape - logits_flat = self.audio_heads(hidden_states) # [B, S, 8*1025] + seq_len, _ = hidden_states.shape + logits_flat = self.audio_heads(hidden_states) # [T, 8*1025] return logits_flat.view( - batch_size, seq_len, self.config.num_audio_codebook, self.config.audio_vocab_size, - ).permute(0, 2, 1, 3) # [B, 8, S, 1025] + ).permute(1, 0, 2) # [8, T, 1025] def _step_forward( self, input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, + cu_seqs: torch.Tensor, rope_table: torch.Tensor, ) -> torch.Tensor: """Single unmasking-step forward using a pre-cast RoPE table (CUDA graph safe).""" hidden_states = self._prepare_embeddings(input_ids, audio_mask) - residual = None - for layer in self.layers: - hidden_states, residual = layer( - hidden_states, attention_mask=attention_mask, rope_table=rope_table, residual=residual + hidden_states = self._transformer_forward(hidden_states, cu_seqs, rope_table=rope_table) + return self._get_logits(hidden_states) + + def _unmask_one_request( + self, + c_logits: torch.Tensor, + u_logits: torch.Tensor, + sample_tokens: torch.Tensor, + *, + num_to_unmask: int | torch.Tensor, + guidance_scale: float, + generator: torch.Generator, + class_temperature: float, + position_temperature: float, + layer_penalty_factor: float, + layer_ids: torch.Tensor, + ) -> None: + """Sample and unmask one request in place from FP32 target logits.""" + mask_id = self.config.audio_mask_id + if guidance_scale != 0: + log_probs = F.log_softmax( + (1.0 + guidance_scale) * c_logits - guidance_scale * u_logits, + dim=-1, ) - return self._get_logits(self.norm(hidden_states + residual)) + else: + log_probs = F.log_softmax(c_logits, dim=-1) + log_probs[..., mask_id] = -float("inf") + if class_temperature > 0.0: + pred_tokens = _gumbel_sample(log_probs, class_temperature, generator).argmax(dim=-1) + else: + pred_tokens = log_probs.argmax(dim=-1) + scores = log_probs.max(dim=-1)[0] + scores = scores - (layer_ids * layer_penalty_factor) + if position_temperature > 0.0: + scores = _gumbel_sample(scores, position_temperature, generator) + scores.masked_fill_(sample_tokens != mask_id, -float("inf")) + _, topk_idx = torch.topk(scores.flatten(), num_to_unmask) + flat_tokens = sample_tokens.flatten() + flat_tokens[topk_idx] = pred_tokens.flatten()[topk_idx] + sample_tokens.copy_(flat_tokens.view_as(sample_tokens)) @torch.inference_mode() def forward( self, input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor, + cond_lens: list[int], target_lens: list[int], - seed: int | None = None, + seed: int | list[int | None] | None = None, num_step: int = 32, guidance_scale: float = 2.0, t_shift: float = 0.1, @@ -886,36 +972,51 @@ class OmniVoiceGenerator(nn.Module): """Run the full 32-step iterative unmasking generation. Args: - input_ids: [2*B, 8, S] - conditional (0:B) + unconditional (B:2B) - audio_mask: [2*B, S] - True for audio positions - attention_mask: [2*B, 1, S, S] - attention mask - target_lens: List of target audio lengths per batch item - num_step: Number of unmasking steps - guidance_scale: CFG scale - t_shift: Time shift for schedule - layer_penalty_factor: Penalty for later codebooks - position_temperature: Gumbel temperature for position selection - class_temperature: Temperature for token prediction (0=greedy) + input_ids: Packed token IDs with shape ``[total_seq_len, 8]`` in + request-major ``[cond0, uncond0, ...]`` order. + audio_mask: Boolean audio-position mask with shape + ``[total_seq_len]``. + cond_lens: Conditional sequence length for each request. + target_lens: Target length for each request; also the corresponding + unconditional sequence length. + seed: One seed per request, a shared scalar seed, or ``None``. + num_step: Number of iterative unmasking steps. + guidance_scale: Classifier-free guidance scale. + t_shift: Time shift used to construct the unmasking schedule. + layer_penalty_factor: Penalty applied to later codebooks. + position_temperature: Gumbel temperature for position selection. + class_temperature: Token sampling temperature; zero selects greedy + decoding. Returns: - tokens: [B, 8, max_target_len] - generated audio tokens + Packed generated audio tokens with shape + ``[1, 8, sum(target_lens)]``. """ B = len(target_lens) device = input_ids.device - max_target_len = max(target_lens) + total_target_lens = sum(target_lens) mask_id = self.config.audio_mask_id num_codebooks = self.config.num_audio_codebook - if seed is None: - seed = random.randint(0, 2**63 - 1) - generator = torch.Generator(device=device).manual_seed(seed) + seeds = seed if isinstance(seed, list) else [seed] * B + generators = [ + torch.Generator(device=device).manual_seed( + request_seed if request_seed is not None else random.randint(0, 2**63 - 1) + ) + for request_seed in seeds + ] # Initialize all target tokens as [MASK] - tokens = torch.full( - (B, num_codebooks, max_target_len), - mask_id, - dtype=torch.long, - device=device, - ) + tokens = torch.full((1, num_codebooks, total_target_lens), mask_id, dtype=torch.long, device=device) + target_offsets: list[int] = [] + target_offset = 0 + sequence_offsets: list[int] = [] + sequence_offset = 0 + for cond_len, target_len in zip(cond_lens, target_lens): + target_offsets.append(target_offset) + target_offset += target_len + sequence_offsets.append(sequence_offset) + sequence_offset += cond_len + target_len + cu_seqs = _build_cu_seqs(cond_lens, target_lens, device) # Compute unmasking schedule timesteps = _get_time_steps(0.0, 1.0, num_step + 1, t_shift).tolist() @@ -939,45 +1040,46 @@ class OmniVoiceGenerator(nn.Module): layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1) - # Single D2H pull for all conditional lengths instead of B per-item .item() syncs. - c_lens = attention_mask[:B, 0, 0].sum(dim=-1).tolist() - - # Materialize the SDPA float mask once so the captured graph (and eager path) skip per-layer conversion. - sdpa_attn_mask = _additive_float_mask(attention_mask, self.model_dtype) - use_cuda_graph = self._cuda_graph_fwd is not None and input_ids.is_cuda if not use_cuda_graph: # Eager-path-only constants (the cuda-graph captures its own). - text_embeds_cached = self.text_embedding(input_ids[:, 0, :]) + text_embeds_cached = self.text_embedding(input_ids[:, 0]) audio_mask_3d = audio_mask.unsqueeze(-1) - rope_table = self._rope_table_for(input_ids.shape[-1], device, text_embeds_cached.dtype) + rope_table = self._rope_table_for(input_ids.shape[0], device, text_embeds_cached.dtype) # Main iterative loop for step in range(num_step): if use_cuda_graph: # Float mask skips per-layer conversion; fp32 cast deferred to the per-item slices below. - batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, sdpa_attn_mask) + batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, cu_seqs, B) else: # Eager fallback reuses hoisted constants (text embeds, sdpa mask, rope table). inputs_embeds = self._prepare_embeddings( input_ids, audio_mask, text_embeds=text_embeds_cached, audio_mask_3d=audio_mask_3d ) - hidden_states = self._transformer_forward(inputs_embeds, sdpa_attn_mask, rope_table=rope_table) + hidden_states = self._transformer_forward( + inputs_embeds, + cu_seqs, + max_seqlen=max(cond_lens), + rope_table=rope_table, + ) # fp32 cast deferred to the per-item slices below. batch_logits = self._get_logits(hidden_states) - # batch_logits: [2*B, 8, S, 1025] + # batch_logits: [8, T, 1025] for i in range(B): k = schedules[i][step] if k <= 0: continue - c_len = c_lens[i] + c_len = cond_lens[i] t_len = target_lens[i] + request_start = sequence_offsets[i] + cond_end = request_start + c_len # Extract logits for target region; upcast only the slices we actually consume. - c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :].to(torch.float32) - u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :].to(torch.float32) + c_logits = batch_logits[:, cond_end - t_len : cond_end, :].unsqueeze(0).to(torch.float32) + u_logits = batch_logits[:, cond_end : cond_end + t_len, :].unsqueeze(0).to(torch.float32) # Classifier-free guidance. Fuse the chain: the two inner # log_softmax normalizers are per-position scalars that the final @@ -996,7 +1098,7 @@ class OmniVoiceGenerator(nn.Module): # Token prediction if class_temperature > 0.0: - pred_tokens = _gumbel_sample(log_probs, class_temperature, generator).argmax(dim=-1) + pred_tokens = _gumbel_sample(log_probs, class_temperature, generators[i]).argmax(dim=-1) else: pred_tokens = log_probs.argmax(dim=-1) # [1, 8, T] @@ -1008,10 +1110,11 @@ class OmniVoiceGenerator(nn.Module): # Gumbel noise for position selection if position_temperature > 0.0: - scores = _gumbel_sample(scores, position_temperature, generator) + scores = _gumbel_sample(scores, position_temperature, generators[i]) # Mask out already unmasked positions - sample_tokens = tokens[i : i + 1, :, :t_len] + target_start = target_offsets[i] + sample_tokens = tokens[:, :, target_start : target_start + t_len] scores.masked_fill_(sample_tokens != mask_id, -float("inf")) # Select top-k positions to unmask. .flatten() on this non-contiguous view already copies. @@ -1021,8 +1124,9 @@ class OmniVoiceGenerator(nn.Module): sample_tokens.copy_(flat_tokens.view_as(sample_tokens)) # Mirror update into both cond and uncond input_ids halves for the next step. - input_ids[i, :, c_len - t_len : c_len] = sample_tokens.squeeze(0) - input_ids[B + i, :, :t_len] = sample_tokens.squeeze(0) + packed_sample_tokens = sample_tokens.squeeze(0).transpose(0, 1) + input_ids[cond_end - t_len : cond_end] = packed_sample_tokens + input_ids[cond_end : cond_end + t_len] = packed_sample_tokens return tokens diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py index 48df20fab..3b8813faf 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Streaming Mimi codec for PersonaPlex (moshi-free). Frame-clocked duplex needs a codec that encodes/decodes exactly one 80 ms frame @@ -288,10 +288,11 @@ class PersonaPlexMimiCodec(nn.Module): def __init__(self, hf_repo: str = DEFAULT_HF_REPO, checkpoint: str | None = None, device: str = "cuda") -> None: super().__init__() - from huggingface_hub import hf_hub_download from safetensors.torch import load_file from transformers import MimiConfig, MimiModel + from vllm_omni.transformers_utils.repo_utils import hf_api + self.device = torch.device(device) # The PersonaPlex repo ships the reference mimi checkpoint in the moshi @@ -300,7 +301,7 @@ class PersonaPlexMimiCodec(nn.Module): # transformer stacks remain absent because the streaming implementations # below replace them with the checkpoint's fused QKV layout. if checkpoint is None: - checkpoint = hf_hub_download( + checkpoint = hf_api().hf_hub_download( "nvidia/personaplex-7b-v1", "tokenizer-e351c8d8-checkpoint125.safetensors", ) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index cb5167597..b9373e5f2 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -304,6 +304,17 @@ _OMNI_MODELS = { "gepard_talker", "GepardTalkerForConditionalGeneration", ), + ## audio8_tts (Audio8 TTS Preview 0.6B) + "Audio8TTSSlowARForConditionalGeneration": ( + "audio8_tts", + "audio8_tts_slow_ar", + "Audio8TTSSlowARForConditionalGeneration", + ), + "Audio8TTSCodecDecoder": ( + "audio8_tts", + "audio8_tts_codec_decoder", + "Audio8TTSCodecDecoder", + ), ## VoxCPM2 "VoxCPM2TalkerForConditionalGeneration": ( "voxcpm2", 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 e9ca78b21..8e3c56b32 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 @@ -461,9 +461,9 @@ class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): if model_path is None: # Resolve HF repo names to local cache path if not os.path.isdir(model_name_or_path): - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_name_or_path = snapshot_download(model_name_or_path) + model_name_or_path = hf_api().snapshot_download(model_name_or_path) model_path = f"{model_name_or_path}/token2wav" float16 = getattr(self.config, "token2wav_float16", False) diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index f8dc962f3..5febb8912 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """VoxCPM2 AR talker — PagedAttention pipeline with per-request state. Architecture: @@ -158,7 +158,7 @@ def build_voxcpm2_prompt( """Build a VoxCPM2 prefill prompt whose ``prompt_token_ids`` length matches the talker-side prefill length. - Used by both online serving (``serving_speech._build_voxcpm2_prompt``) and + Used by both online serving (``VoxCPM2Adapter._build_prompt``) and the offline example, so the talker-side length assertion never fires. """ ids = split_multichar_chinese(tokenizer.encode(text, add_special_tokens=True), split_map) @@ -2208,7 +2208,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): assert scaffold_len == tts_len, ( f"voxcpm2 prefill length mismatch: scaffold_len={scaffold_len} tts_len={tts_len}; " "caller must pad prompt_token_ids to the full prefill length " - "(see serving_speech._build_voxcpm2_prompt or the offline example)." + "(see VoxCPM2Adapter._build_prompt or the offline example)." ) enc_out = base_lm_out.unsqueeze(0) 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 69cfa6d26..c61aeaaad 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable, Mapping from dataclasses import replace from functools import cached_property @@ -6,7 +9,6 @@ from pathlib import Path import regex as re import torch import torch.nn as nn -from huggingface_hub import hf_hub_download from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsMultiModal @@ -28,6 +30,7 @@ from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation im VoxtralTTSMultiModalProcessor, VoxtralTTSProcessingInfo, ) +from vllm_omni.transformers_utils.repo_utils import hf_api logger = init_logger(__name__) @@ -133,7 +136,7 @@ class VoxtralTTSForConditionalGeneration( self.voice_to_embedding = {} for sid in speaker_id: if self.is_hf_model: - path = hf_hub_download(repo_id=self.repo_id, filename=f"voice_embedding/{sid}.pt") + path = hf_api().hf_hub_download(repo_id=self.repo_id, filename=f"voice_embedding/{sid}.pt") else: path = Path(self.repo_id) / "voice_embedding" / f"{sid}.pt" if Path(path).exists(): diff --git a/vllm_omni/model_executor/stage_input_processors/audio8_tts.py b/vllm_omni/model_executor/stage_input_processors/audio8_tts.py new file mode 100644 index 000000000..4b453c13a --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/audio8_tts.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview: Slow AR -> codec decoder stage input processor.""" + +from collections.abc import Mapping +from typing import Any + +import torch + +from vllm_omni.data_entry_keys import ( + CodesStruct, + MetaStruct, + OmniPayloadStruct, +) + +#: Chunking defaults; ~21.5 codec frames per second, so 25 frames ~= 1.16 s. +DEFAULT_CHUNK_FRAMES = 25 +DEFAULT_LEFT_CONTEXT_FRAMES = 25 + + +def _connector_extra(transfer_manager: Any) -> dict[str, Any]: + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + return raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + + +def _cfg_int(cfg: dict[str, Any], key: str, default: int) -> int: + value = cfg.get(key, default) + try: + return int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid Audio8 TTS integer config {key}={value!r}") from exc + + +def _request_initial_chunk_frames(request: Any, default: int) -> int: + """Per-request override of ``initial_codec_chunk_frames``, if present.""" + additional_information = getattr(request, "additional_information", None) + entries = getattr(additional_information, "entries", None) + if not entries or "initial_codec_chunk_frames" not in entries: + return default + entry = entries["initial_codec_chunk_frames"] + if entry.list_data is not None and len(entry.list_data) == 1: + return int(entry.list_data[0]) + return default + + +def extract_last_frame(multimodal_output: Mapping[str, Any]) -> torch.Tensor | None: + """Return the newest frame of codec codes, or ``None`` if this step had none. + + Prefill steps emit an all-zero placeholder frame, which must not be + forwarded to the decoder. + """ + audio_codes = multimodal_output.get("audio_codes") + if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: + return None + if audio_codes.ndim == 1: + return audio_codes.to(device="cpu", dtype=torch.long).reshape(-1) + if audio_codes.ndim != 2: + raise ValueError(f"Invalid audio_codes shape for Audio8 TTS async_chunk: {tuple(audio_codes.shape)}") + + frame = audio_codes[-1] + if frame.numel() == 0: + return None + valid = multimodal_output.get("audio_code_valid") + if isinstance(valid, torch.Tensor) and valid.numel() > 0: + is_valid = bool(valid.reshape(-1)[-1].item()) + elif valid is not None: + is_valid = bool(valid) + else: + is_valid = bool(frame.any().item()) + if not is_valid: + return None + return frame.to(device="cpu", dtype=torch.long).reshape(-1) + + +def slow_ar_to_codec_decoder_async_chunk( + transfer_manager: Any, + multimodal_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Emit ``[left_context | new_frames]`` chunks as the Slow AR produces codes. + + Codes accumulate across AR steps and are shipped as a + ``[num_codebooks, frames]`` tensor. ``meta.left_context_size`` tells the + decoder how many leading frames to drop, keeping chunk boundaries + artefact-free. Returns ``None`` when the current chunk is not full yet. + """ + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + cfg = _connector_extra(transfer_manager) + + appended_this_call = False + if isinstance(multimodal_output, Mapping): + frame = extract_last_frame(multimodal_output) + if frame is not None: + transfer_manager.code_prompt_token_ids[request_id].append(frame.detach()) + appended_this_call = True + elif not finished: + return None + + chunk_frames = _cfg_int(cfg, "codec_chunk_frames", DEFAULT_CHUNK_FRAMES) + left_context_frames = _cfg_int(cfg, "codec_left_context_frames", DEFAULT_LEFT_CONTEXT_FRAMES) + initial_chunk_frames = _request_initial_chunk_frames(request, _cfg_int(cfg, "initial_codec_chunk_frames", 0)) + if chunk_frames <= 0 or left_context_frames < 0 or initial_chunk_frames < 0: + raise ValueError( + "Invalid Audio8 TTS codec chunk config: " + f"codec_chunk_frames={chunk_frames}, codec_left_context_frames={left_context_frames}, " + f"initial_codec_chunk_frames={initial_chunk_frames}" + ) + initial_chunk_frames = min(initial_chunk_frames, chunk_frames) + + frames = transfer_manager.code_prompt_token_ids[request_id] + length = len(frames) + if length <= 0: + if finished: + # Still emit a terminal payload so the decoder can close the stream. + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) + return None + + # A smaller first chunk trades a little quality for time-to-first-audio. + if initial_chunk_frames > 0 and length <= chunk_frames: + already_sent = transfer_manager.put_req_chunk[request_id] * initial_chunk_frames + pending = length - already_sent + if pending <= 0: + return None + if pending < initial_chunk_frames and not finished: + return None + new_frames = min(pending, initial_chunk_frames) + left_context_size = max(0, length - new_frames) + window = frames[:length] + else: + initial_coverage = ( + (chunk_frames // initial_chunk_frames) * initial_chunk_frames if initial_chunk_frames > 0 else 0 + ) + pending = (length - initial_coverage) % chunk_frames + if pending != 0 and not finished: + return None + if pending == 0 and finished and not appended_this_call: + # This boundary's chunk was already shipped when the stream first + # reached it; a terminal call carrying no new frame must close the + # stream, not re-send that window (otherwise the last chunk_frames + # frames ship twice). + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) + new_frames = pending if pending != 0 else chunk_frames + end_index = min(length, left_context_frames + new_frames) + left_context_size = max(0, end_index - new_frames) + window = frames[-end_index:] + + # Codebook-major [num_codebooks, frames]; the tensor payload avoids + # expanding codec indices into Python ints across the connector boundary. + codes_qf = torch.stack(window, dim=0).transpose(0, 1).contiguous() + return OmniPayloadStruct( + codes=CodesStruct(audio=codes_qf), + meta=MetaStruct( + left_context_size=left_context_size, + finished=torch.tensor(finished, dtype=torch.bool), + ), + ) + + +__all__ = [ + "DEFAULT_CHUNK_FRAMES", + "DEFAULT_LEFT_CONTEXT_FRAMES", + "extract_last_frame", + "slow_ar_to_codec_decoder_async_chunk", +] diff --git a/vllm_omni/model_executor/stage_input_processors/joyai_vl_interaction.py b/vllm_omni/model_executor/stage_input_processors/joyai_vl_interaction.py new file mode 100644 index 000000000..77b86691c --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/joyai_vl_interaction.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Build Qwen3-TTS Talker inputs from completed JoyAI actions.""" + +from typing import Any + +from vllm.outputs import RequestOutput +from vllm.tokenizers import cached_tokenizer_from_config + +from vllm_omni.errors import OmniClientError +from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import ( + parse_action, +) +from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + Qwen3TTSPromptEmbedsBuilder, + first_value, +) + +# Match the standalone Qwen3-TTS default attributes +_DEFAULT_TTS_TASK_TYPE = "CustomVoice" +_DEFAULT_TTS_LANGUAGE = "Auto" +_DEFAULT_TTS_SPEAKER = "Vivian" + + +def _extract_completed_action_text(joyai_output: RequestOutput) -> str: + """Return the final JoyAI action text, using cumulative text when available.""" + completion = joyai_output.outputs[0] + action_text = getattr(completion, "cumulative_text", None) or completion.text + return action_text if isinstance(action_text, str) else "" + + +def _build_tts_metadata(request_prompt: object, spoken_text: str) -> dict[str, list[str]]: + """Build Qwen3-TTS metadata for the text JoyAI chose to speak.""" + prompt_dict = request_prompt if isinstance(request_prompt, dict) else {} + raw_additional_info = prompt_dict.get("additional_information") + additional_info = raw_additional_info if isinstance(raw_additional_info, dict) else {} + + # The native pipeline loads CustomVoice checkpoints for both TTS stages. + # Keep client metadata from selecting an incompatible Qwen3-TTS task. + task_type = _DEFAULT_TTS_TASK_TYPE + request_language = first_value(additional_info.get("language"), _DEFAULT_TTS_LANGUAGE) + language = first_value(additional_info.get("tts_language"), request_language) + request_speaker = first_value(additional_info.get("speaker"), _DEFAULT_TTS_SPEAKER) + speaker = first_value(additional_info.get("tts_speaker"), request_speaker) + instruction = first_value(additional_info.get("tts_instruct"), "") + + return { + "task_type": [str(task_type)], + "language": [str(language)], + "speaker": [str(speaker)], + "instruct": [str(instruction)], + "text": [spoken_text], + } + + +def _validate_talker_speaker(tts_metadata: dict[str, list[str]], talker_model_config: Any) -> None: + """Validate the selected voice using the destination Talker's speaker table.""" + talker_config = getattr(talker_model_config.hf_config, "talker_config", None) + if talker_config is None: + raise ValueError("The target stage is not a Qwen3-TTS Talker.") + + # Match CustomVoice prompt construction, including case and whitespace. + speaker = str(first_value(tts_metadata.get("speaker"), "")).lower().strip() + speakers = {name.lower() for name in (getattr(talker_config, "spk_id", None) or {})} + if not speaker or speaker not in speakers: + raise OmniClientError(f"Unsupported speaker: {speaker!r}") + + +def _compute_talker_prompt_length( + tts_metadata: dict[str, list[str]], + talker_model_config: Any, +) -> int: + """Compute the exact number of prompt positions used by the Talker.""" + task_type = str(first_value(tts_metadata.get("task_type"), _DEFAULT_TTS_TASK_TYPE)) + if task_type != _DEFAULT_TTS_TASK_TYPE: + raise ValueError("JoyAI native TTS supports only CustomVoice.") + + talker_tokenizer = cached_tokenizer_from_config(talker_model_config) + if talker_tokenizer is None: + raise ValueError("The Qwen3-TTS Talker must have an initialized tokenizer.") + + talker_config = getattr(talker_model_config.hf_config, "talker_config", None) + if talker_config is None: + raise ValueError("The target stage is not a Qwen3-TTS Talker.") + + def tokenize_prompt(text: str) -> list[int]: + return talker_tokenizer.encode(text, add_special_tokens=True) + + return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( + additional_information=tts_metadata, + task_type=task_type, + tokenize_prompt=tokenize_prompt, + codec_language_id=getattr(talker_config, "codec_language_id", None), + spk_is_dialect=getattr(talker_config, "spk_is_dialect", None), + ) + + +def joyai_action_to_tts( + source_outputs: list[RequestOutput], + prompt: object | None = None, + requires_multimodal_data: bool = False, + *, + target_model_config: Any, +) -> list[OmniTokensPrompt]: + """Build Talker inputs for JoyAI actions that contain text to speak.""" + del requires_multimodal_data + + request_prompts = prompt if isinstance(prompt, list) else [prompt] * len(source_outputs) + talker_inputs: list[OmniTokensPrompt] = [] + for request_index, joyai_output in enumerate(source_outputs): + parsed_action = parse_action(_extract_completed_action_text(joyai_output)) + if not parsed_action.spoke or not parsed_action.text: + continue + + request_prompt = request_prompts[request_index] if request_index < len(request_prompts) else None + tts_metadata = _build_tts_metadata(request_prompt, parsed_action.text) + _validate_talker_speaker(tts_metadata, target_model_config) + talker_prompt_length = _compute_talker_prompt_length(tts_metadata, target_model_config) + talker_inputs.append( + OmniTokensPrompt( + # OmniTokensPrompt reserves vllm scheduler positions with zero IDs; + # the Talker builds the actual prompt from tts_metadata. + prompt_token_ids=[0] * talker_prompt_length, + additional_information=tts_metadata, + ) + ) + + return talker_inputs diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index 90197f6e6..e74931a20 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from dataclasses import dataclass, field from typing import Any @@ -52,7 +55,7 @@ class OmniModelRunnerOutput(ModelRunnerOutput): """ multimodal_outputs: list[dict[str, object]] | None = None - inter_stage_outputs: list[dict[str, object]] | None = None + inter_stage_outputs: list[dict[str, Any] | None] | 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 diff --git a/vllm_omni/platforms/npu/models/minimax_h3.py b/vllm_omni/platforms/npu/models/minimax_h3.py index c8889f18f..a89768096 100644 --- a/vllm_omni/platforms/npu/models/minimax_h3.py +++ b/vllm_omni/platforms/npu/models/minimax_h3.py @@ -100,7 +100,7 @@ def npu_swiglu_from_packed(gate_up: torch.Tensor) -> torch.Tensor: def _forward_minimax_h3_qwen3vl_text_mlp_npu(self: Any, x: torch.Tensor) -> torch.Tensor: """Run the Qwen3-VL MLP with one packed GEMM and fused SwiGLU.""" - gate_up = F.linear(x, self.gate_up_proj.weight) + gate_up = self.gate_up_proj(x) return self.down_proj(npu_swiglu_from_packed(gate_up)) diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 42d6d8d45..6beb91018 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Custom model configs that cannot be expressed via HuggingFace Transformers alone, following the same pattern as vllm.transformers_utils.configs. @@ -10,6 +10,9 @@ from __future__ import annotations import importlib _CLASS_TO_MODULE: dict[str, str] = { + "Audio8TTSConfig": "vllm_omni.transformers_utils.configs.audio8_tts", + "Audio8TTSSlowARConfig": "vllm_omni.transformers_utils.configs.audio8_tts", + "Audio8TTSFastARConfig": "vllm_omni.transformers_utils.configs.audio8_tts", "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", @@ -35,6 +38,9 @@ _CLASS_TO_MODULE: dict[str, str] = { } __all__ = [ + "Audio8TTSConfig", + "Audio8TTSFastARConfig", + "Audio8TTSSlowARConfig", "HiggsAudioV3Config", "Mammothmoda2Config", "Mammothmoda2Qwen2_5_VLConfig", @@ -75,6 +81,7 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. +from vllm_omni.transformers_utils.configs import audio8_tts as _audio8_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/audio8_tts.py b/vllm_omni/transformers_utils/configs/audio8_tts.py new file mode 100644 index 000000000..42e920419 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/audio8_tts.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS config registration with transformers AutoConfig. + +Registers ``model_type = "arktts"`` (and the two sub-config types) so +``AutoConfig.from_pretrained("Audio8/Audio8-TTS-Preview-0.6b")`` resolves to +vllm-omni's Qwen2-shaped config instead of the checkpoint's remote code. + +This only wins when ``trust_remote_code`` is **off** -- transformers prefers a +checkpoint's ``auto_map`` over registered classes otherwise -- which is why +``deploy/audio8_tts.yaml`` sets ``trust_remote_code: false``. +""" + +from transformers import AutoConfig + +from vllm_omni.model_executor.models.audio8_tts.configuration_audio8_tts import ( + Audio8TTSConfig, + Audio8TTSFastARConfig, + Audio8TTSSlowARConfig, +) + +AutoConfig.register("arktts", Audio8TTSConfig) +AutoConfig.register("arktts_slow_ar", Audio8TTSSlowARConfig) +AutoConfig.register("arktts_fast_ar", Audio8TTSFastARConfig) + +__all__ = [ + "Audio8TTSConfig", + "Audio8TTSFastARConfig", + "Audio8TTSSlowARConfig", +] diff --git a/vllm_omni/transformers_utils/repo_utils.py b/vllm_omni/transformers_utils/repo_utils.py new file mode 100644 index 000000000..e780c4b53 --- /dev/null +++ b/vllm_omni/transformers_utils/repo_utils.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Utilities for model repo interaction.""" + +from huggingface_hub import HfApi + +from vllm_omni.version import __version__ as VLLM_OMNI_VERSION + +_hf_api: HfApi | None = None + + +def hf_api() -> HfApi: + """Return a shared HfApi instance tagged with vLLM-Omni's library info.""" + global _hf_api + if _hf_api is None: + _hf_api = HfApi( + library_name="vllm-omni", + library_version=VLLM_OMNI_VERSION, + ) + return _hf_api diff --git a/vllm_omni/utils/qwen3_force_align_processor.py b/vllm_omni/utils/qwen3_force_align_processor.py index e70d270db..76c5b9084 100644 --- a/vllm_omni/utils/qwen3_force_align_processor.py +++ b/vllm_omni/utils/qwen3_force_align_processor.py @@ -65,18 +65,18 @@ _LANG_ALIASES = { def build_prompt(words: list[str]) -> str: - """Wrap segmented words in the Qwen3 aligner prompt template. + """Build the Qwen3 aligner prompt exactly as qwen_asr's + ``Qwen3ForceAlignProcessor.encode_timestamp`` does: the audio placeholder + followed by the words, each with two trailing ```` markers + (start + end) that the model classifies into audio time bins. - Each word gets two trailing ```` markers (start + end); the - model classifies each marker into an audio time bin. + No chat template. The official aligner feeds this string straight to the + tokenizer; wrapping it in ``<|im_start|>user ... <|im_start|>assistant`` + puts three tokens in front of the audio and shifts almost every predicted + marker one 80 ms bin later than the official output. """ - 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" + body = f"{TIMESTAMP_TOKEN}{TIMESTAMP_TOKEN}".join(words) + f"{TIMESTAMP_TOKEN}{TIMESTAMP_TOKEN}" + return f"{AUDIO_PLACEHOLDER}{body}" # --- word segmentation (port of Qwen3ForceAlignProcessor) ---------------- diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 6618ba920..811f86f17 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -15,7 +15,7 @@ from collections.abc import Callable, Sequence from contextlib import nullcontext from copy import copy from dataclasses import replace -from typing import Any, NamedTuple +from typing import Any, NamedTuple, cast import numpy as np import torch @@ -212,7 +212,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): if kwargs: raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}") - self._model_runner_output = None + self._model_runner_output: OmniModelRunnerOutput | None = None self._invalid_req_indices = invalid_req_indices self.async_copy_ready_event = torch.Event() @@ -242,7 +242,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): self._num_nans_cpu = self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None else None self.async_copy_ready_event.record() - self._model_runner_output_builder = model_runner_output_builder + self._model_runner_output_builder: Callable[[], OmniModelRunnerOutput] | None = model_runner_output_builder self._background_exception: BaseException | None = None self._background_thread: threading.Thread | None = None self._cuda_device = cuda_device @@ -257,6 +257,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): if self._model_runner_output is not None: return with record_function_or_nullcontext("omni_async_output:get_output/build_model_runner_output"): + assert self._model_runner_output_builder is not None self._model_runner_output = self._model_runner_output_builder() self._model_runner_output_builder = None @@ -352,6 +353,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex outputs per request while keeping Async output semantics. """ + execute_model_state: ExecuteModelState | None + kv_extracted_req_ids: list[str] | None + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) @@ -467,6 +471,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if not isinstance(info, dict): return None val = info.get("omni_final_stage_id") + if val is None: + return None try: return int(val) except (TypeError, ValueError): @@ -1516,7 +1522,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex ) selected_indices.append(idx) selected_req_ids.append(req_id) - selected_req_infos.append(req_info) + selected_req_infos.append(cast(dict[str, Any], req_info)) if not selected_indices: return multimodal_outputs @@ -1583,7 +1589,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex def _build_multimodal_outputs( self, - per_req_payloads: list[dict[str, object] | None] | None, + per_req_payloads: Sequence[dict[str, object] | None] | None, ) -> list[dict[str, torch.Tensor] | None] | None: """Build per-request multimodal output payloads (dedicated channel). @@ -1679,8 +1685,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex def _build_omni_step_outputs( self, - pooler_inter: list[dict[str, object]], - pooler_client: list[dict[str, object]], + pooler_inter: Sequence[dict[str, object] | None] | None, + pooler_client: Sequence[dict[str, object] | None] | None, *, defer_full_payload_d2h: bool, ) -> tuple[ @@ -1966,6 +1972,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex ) pooler_output.append(flatten_payload(payload)) + pooler_inter: Sequence[dict[str, object] | None] | None + pooler_client: Sequence[dict[str, object] | None] | None pooler_output = pooler_output or [] if self._async_chunk and stage_sends_async_output(self.model_config): pooler_inter, pooler_client = partition_payload_list(pooler_output) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 03913aebe..80d55020f 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -78,6 +78,8 @@ def _filter_mrope_kwargs_for_model(model: object, kwargs: dict[str, Any]) -> dic class OmniGPUModelRunner(GPUModelRunner): + intermediate_tensors: IntermediateTensors | None + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.model_intermediate_buffer: dict[str, dict[str, Any]] = {} @@ -842,7 +844,9 @@ class OmniGPUModelRunner(GPUModelRunner): return None @torch.inference_mode() - def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict: + def extract_multimodal_outputs( + self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput + ) -> tuple[Any, Any]: if ( hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs @@ -1407,14 +1411,14 @@ class OmniGPUModelRunner(GPUModelRunner): def _process_additional_information_updates( self, hidden_states: torch.Tensor, - multimodal_outputs: object, + multimodal_outputs: Any, num_scheduled_tokens_np: np.ndarray, scheduler_output: "SchedulerOutput", combined_hidden_states: dict[str, torch.Tensor] | None = None, - combined_multimodal_outputs: dict[str, object] | None = None, + combined_multimodal_outputs: dict[str, Any] | None = None, req_ids_filter: set[str] | None = None, req_ids: list[str] | None = None, - query_start_loc_cpu: object | None = None, + query_start_loc_cpu: Any = 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 @@ -1476,9 +1480,9 @@ class OmniGPUModelRunner(GPUModelRunner): def _collect_additional_information_for_prefill( self, num_scheduled_tokens_np: np.ndarray, - ) -> dict[str, dict]: + ) -> None: """Overlay per-request prompt_embeds for the prefill portion and collect - additional_information slices for this step. Returns a map req_id -> dict.""" + additional_information slices for this step.""" for req_index, req_id in enumerate(self.input_batch.req_ids): req_state = self.requests[req_id] pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) @@ -1707,9 +1711,9 @@ class OmniGPUModelRunner(GPUModelRunner): # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only - decode_req_ids = [] - decode_start_offsets = [] - decode_batch_items = [] + decode_req_ids: list[str] = [] + decode_start_offsets: list[int] = [] + decode_batch_items: list[tuple[str, int, dict[str, Any]]] = [] batch_decode_preprocess = getattr(self.model, "preprocess_decode_batch", None) def flush_decode_batch() -> None: @@ -1717,6 +1721,7 @@ class OmniGPUModelRunner(GPUModelRunner): if not decode_batch_items: return + assert callable(batch_decode_preprocess) req_ids_b = [item[0] for item in decode_batch_items] start_offsets_b = [item[1] for item in decode_batch_items] req_infos_b = [item[2] for item in decode_batch_items] @@ -1761,7 +1766,7 @@ class OmniGPUModelRunner(GPUModelRunner): # mimo-audio check req_state = self.requests.get(req_id) - req_infos = self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id) + req_infos = cast(dict[str, Any], self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id)) start_offset = int(self.query_start_loc.cpu[req_index]) sched_tokens = int(num_scheduled_tokens_np[req_index])