diff --git a/benchmarks/tts/bench_tts.py b/benchmarks/tts/bench_tts.py index f8c1676be..322aa437e 100644 --- a/benchmarks/tts/bench_tts.py +++ b/benchmarks/tts/bench_tts.py @@ -1,4 +1,7 @@ #!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Universal TTS benchmark CLI for vllm-omni. Runs ``vllm bench serve --omni`` with model-aware defaults loaded from @@ -87,6 +90,7 @@ def build_bench_args( served_model_name: str | None = None, num_warmups: int = 2, request_seed: int | None = None, + duration_seconds: float | None = None, ) -> list[str]: """Build the ``vllm bench serve --omni`` command for one (task, concurrency) run.""" dataset_name = _TASK_TO_DATASET[task] @@ -95,7 +99,10 @@ def build_bench_args( task_extra_body: dict[str, Any] = dict((model_cfg.get("task_extra_body") or {}).get(task) or {}) if request_seed is not None: task_extra_body["seed"] = request_seed - + if duration_seconds is not None: + if not math.isfinite(duration_seconds) or duration_seconds <= 0: + raise ValueError("duration_seconds must be finite and positive") + task_extra_body["duration_seconds"] = duration_seconds # Resolve dataset path if dataset_path: resolved_dataset_path = dataset_path @@ -194,7 +201,7 @@ def print_summary_table(results: list[dict[str, Any]]) -> None: if not results: return header = ( - f"{'Task':<16} {'Concurrency':>11} {'RTF mean':>10} " + f"{'Task':<16} {'Task Type':<12} {'Concurrency':>11} {'RTF mean':>10} " f"{'TTFP (ms)':>10} {'Throughput':>12} {'WER':>7} {'SIM':>7} {'UTMOS':>7}" ) print(f"\n{'=' * len(header)}") @@ -204,6 +211,7 @@ def print_summary_table(results: list[dict[str, Any]]) -> None: print("-" * len(header)) for r in results: task = r.get("_task", "?") + task_type = r.get("_task_type", "-") conc = r.get("_concurrency", "?") rtf = r.get("mean_audio_rtf", float("nan")) ttfp = r.get("mean_audio_ttfp_ms", float("nan")) @@ -216,7 +224,7 @@ def print_summary_table(results: list[dict[str, Any]]) -> None: return f"{v:.{digits}f}" if not math.isnan(v) else " n/a" print( - f"{task:<16} {str(conc):>11} {fmt(rtf):>10} {fmt(ttfp, 0):>10} " + f"{task:<16} {task_type:<12} {str(conc):>11} {fmt(rtf):>10} {fmt(ttfp, 0):>10} " f"{fmt(throughput):>12} {fmt(wer):>7} {fmt(sim):>7} {fmt(utmos):>7}" ) print("=" * len(header)) @@ -249,6 +257,12 @@ def main() -> None: help="Number of prompts per run. If one value, applied to all concurrency levels.", ) parser.add_argument("--num-warmups", type=int, default=2, help="Warmup requests before each measured run") + parser.add_argument( + "--duration-seconds", + type=float, + default=None, + help="Target audio duration per request (AuK defaults to 5 seconds). Independent of --output-len.", + ) parser.add_argument( "--request-seed", type=int, @@ -277,6 +291,8 @@ def main() -> None: if args.num_warmups < 0: parser.error("--num-warmups cannot be negative") + if args.duration_seconds is not None and (not math.isfinite(args.duration_seconds) or args.duration_seconds <= 0): + parser.error("--duration-seconds must be finite and positive") model_configs = load_model_configs(Path(args.model_configs)) if args.model not in model_configs: @@ -333,11 +349,16 @@ def main() -> None: served_model_name=args.served_model_name, num_warmups=args.num_warmups, request_seed=args.request_seed, + duration_seconds=args.duration_seconds, ) result = run_one_benchmark(cmd) if result is not None: result["_task"] = task result["_concurrency"] = concurrency + task_body = model_cfg.get("task_extra_body", {}).get(task, {}) + task_type = task_body.get("task_type") + if task_type is not None: + result["_task_type"] = task_type all_results.append(result) # Persist the metadata so plot_results.py can pick it up. if args.output_dir and result_filename: diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index b8921c1d3..ffe0e259a 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -7,6 +7,22 @@ # is tracked here. models: + tencent/AuK: &auk + supported_tasks: [default_voice, voice_clone] + backend: openai-audio-speech + endpoint: /v1/audio/speech + trust_remote_code: true + task_extra_body: + default_voice: + duration_seconds: 5.0 + voice: default + task_type: CustomVoice + voice_clone: + duration_seconds: 5.0 + voice: default + task_type: Base + tencent/AuK-Flash: *auk + BreezeBlue/Breeze-TTS-2: supported_tasks: [default_voice, voice_design, voice_clone] backend: openai-audio-speech 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 c1107ad60..acd7076d4 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | | --- | --- | --- | --- | --- | --- | +| AuK / AuK-Flash | assembled local bundle | ✓ (`ref_audio`) | HTTP audio after full generation | — | — | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | | Gepard-1.0 | `nineninesix/gepard-1.0` | — (zero-shot default voice) | ✓ (PCM / WAV stream) | `default` only | — | @@ -99,6 +100,41 @@ For full request-shape documentation (all parameters, response formats, error co --- +## AuK / AuK-Flash + +AuK serves speech generation and reference-voice synthesis at 24 kHz. Assemble a local bundle before serving: + +```bash +python tools/prepare_auk_checkpoint.py \ + --auk-dir ckpts/AuK --qwen-dir ckpts/Qwen2.5-Omni-3B --out ckpts/auk-omni +MODEL=ckpts/auk-omni bash examples/online_serving/text_to_speech/auk/run_server.sh +``` + +AuK accepts the complete instruction after you fill in the [local recipe template](../../../../recipes/Tencent/AuK-H100.md#supported-model-contract) or the [official AuK cookbook](https://github.com/Tencent-Hunyuan/AuK/blob/main/docs/COOKBOOK.md). Pass it in `instructions` and leave `input` empty. For instruct TTS, fill in `Generate speech based on the following description: "{voice description}". The content to speak is: "{text}".`; for zero-shot TTS, editing, enhancement and separation, use the corresponding complete instruction from the cookbook. Set `duration_seconds` for text-only generation; with `ref_audio`, omitting it keeps the source length. + +For compatibility with clients that can only send `input`, AuK treats an input-only request as the complete instruction and logs a recommendation to move it to `instructions` with `input=""`. + +`extra_params` accepts `num_inference_steps`, `guidance_scale`, `sway`, `t_grid`, and `vae_sample`; `seed` is a top-level field. + +```bash +python examples/online_serving/text_to_speech/auk/speech_client.py \ + --model ckpts/auk-omni \ + --instructions 'Generate speech based on the following description: "A calm young woman speaking warmly and slowly.". The content to speak is: "Welcome back, how was your day?".' \ + --duration-seconds 3 --seed 7 +``` + +Complete-instruction editing also does not need a separate text argument: + +```bash +python examples/online_serving/text_to_speech/auk/speech_client.py \ + --model ckpts/auk-omni --instructions "Keep pure speech voice, remove noise and reverberation." \ + --ref-audio noisy.wav --seed 7 --output clean.wav +``` + +The client accepts local reference files, URLs, and data URLs. `--stream` writes signed 16-bit 24 kHz PCM after the complete waveform is ready. See `speech_client.py --help` for editing and sampling options. + +--- + ## CosyVoice3 2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets. diff --git a/examples/online_serving/text_to_speech/auk/run_server.sh b/examples/online_serving/text_to_speech/auk/run_server.sh new file mode 100644 index 000000000..9dbb43f8c --- /dev/null +++ b/examples/online_serving/text_to_speech/auk/run_server.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# MODEL=/models/auk-omni bash examples/online_serving/text_to_speech/auk/run_server.sh +set -euo pipefail + +: "${MODEL:?Set MODEL to the assembled AuK or AuK-Flash checkpoint directory}" +exec vllm serve "$MODEL" --omni --host "${HOST:-127.0.0.1}" --port "${PORT:-8091}" "$@" diff --git a/examples/online_serving/text_to_speech/auk/speech_client.py b/examples/online_serving/text_to_speech/auk/speech_client.py new file mode 100644 index 000000000..ff16c33d9 --- /dev/null +++ b/examples/online_serving/text_to_speech/auk/speech_client.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK speech generation and audio editing through /v1/audio/speech. + +See the online TTS user guide for checkpoint preparation and serving. Local +reference audio is sent as a data URL; it need not exist on the server. +""" + +from __future__ import annotations + +import argparse +import base64 +import mimetypes +import os +from pathlib import Path + +import httpx + + +def reference_audio(value: str) -> str: + if value.startswith(("http://", "https://", "data:")): + return value + path = Path(value).expanduser() + mime = mimetypes.guess_type(path.name)[0] or "audio/wav" + return f"data:{mime};base64,{base64.b64encode(path.read_bytes()).decode('ascii')}" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--api-base", default="http://localhost:8091") + parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY", "EMPTY")) + parser.add_argument("--model", required=True, help="Model name or assembled bundle path used by the server") + parser.add_argument( + "--instructions", + required=True, + help=( + "Complete AuK instruction after filling the cookbook template " + "(https://github.com/Tencent-Hunyuan/AuK/blob/main/docs/COOKBOOK.md)" + ), + ) + parser.add_argument("--ref-audio", help="Reference/source audio: local file, HTTP URL or data URL") + parser.add_argument("--duration-seconds", type=float, help="Target duration; defaults to source length with audio") + parser.add_argument("--seed", type=int) + parser.add_argument("--num-inference-steps", type=int) + parser.add_argument("--guidance-scale", type=float) + parser.add_argument("--sway", type=float) + parser.add_argument("--t-grid", type=float, nargs="+", help="Explicit increasing Euler time grid") + parser.add_argument("--vae-sample", action="store_true", help="Sample the VAE posterior instead of taking its mean") + parser.add_argument("--stream", action="store_true", help="Receive raw PCM over HTTP after full generation") + parser.add_argument("--response-format", choices=["wav", "pcm", "flac", "mp3", "opus"], default="wav") + parser.add_argument("--output", type=Path, help="Output file; defaults to output.wav or output.pcm") + parser.add_argument("--timeout", type=float, default=300) + args = parser.parse_args() + if args.duration_seconds is None and args.ref_audio is None: + parser.error("--duration-seconds is required without --ref-audio") + + response_format = "pcm" if args.stream else args.response_format + payload = { + "model": args.model, + # OpenAI's shared Speech schema still declares input, while AuK's + # canonical request carries all content in the complete instruction. + "input": "", + "voice": "default", + "response_format": response_format, + } + for name in ("duration_seconds", "seed"): + value = getattr(args, name) + if value is not None: + payload[name] = value + payload["instructions"] = args.instructions + if args.ref_audio: + payload["ref_audio"] = reference_audio(args.ref_audio) + extra = { + name: getattr(args, name) + for name in ("num_inference_steps", "guidance_scale", "sway", "t_grid") + if getattr(args, name) is not None + } + if args.vae_sample: + extra["vae_sample"] = True + if extra: + payload["extra_params"] = extra + if args.stream: + payload.update(stream=True, stream_format="audio") + + output = args.output or Path(f"output.{response_format}") + headers = {"Authorization": f"Bearer {args.api_key}"} + with httpx.Client(timeout=args.timeout) as client: + with client.stream( + "POST", f"{args.api_base.rstrip('/')}/v1/audio/speech", json=payload, headers=headers + ) as response: + if response.is_error: + response.read() + raise RuntimeError(f"Speech API returned {response.status_code}: {response.text}") + with output.open("wb") as handle: + for chunk in response.iter_bytes(): + handle.write(chunk) + print(f"Saved {output} ({output.stat().st_size} bytes). Native audio: 24000 Hz, mono.") + if response_format == "pcm": + print("PCM format: signed 16-bit little-endian, 24000 Hz, mono.") + + +if __name__ == "__main__": + main() diff --git a/tests/benchmarks/test_bench_tts_cli.py b/tests/benchmarks/test_bench_tts_cli.py index e32d94698..64c17a26b 100644 --- a/tests/benchmarks/test_bench_tts_cli.py +++ b/tests/benchmarks/test_bench_tts_cli.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for the universal benchmarks/tts/bench_tts.py CLI.""" from __future__ import annotations @@ -51,6 +54,73 @@ def test_load_model_configs(model_configs_path: Path) -> None: assert configs["test/ModelA"]["supported_tasks"] == ["voice_clone", "default_voice"] +@pytest.mark.parametrize("model", ["tencent/AuK", "tencent/AuK-Flash"]) +@pytest.mark.parametrize("task", ["default_voice", "voice_clone"]) +def test_auk_benchmark_duration_and_local_bundle(model, task): + config = bench_tts.load_model_configs(bench_tts._DEFAULT_MODEL_CONFIGS)[model] + cmd = bench_tts.build_bench_args( + host="localhost", + port=8000, + model=model, + task=task, + model_cfg=config, + locale="en", + num_prompts=20, + concurrency=8, + dataset_path="/data/seed-tts", + wer_eval=False, + output_dir=None, + result_filename=None, + extra_cli_args=[], + served_model_name="/models/auk", + duration_seconds=3.5, + request_seed=7, + ) + assert cmd[cmd.index("--model") + 1] == "/models/auk" + body = json.loads(cmd[cmd.index("--extra-body") + 1]) + assert body == { + "voice": "default", + "duration_seconds": 3.5, + "seed": 7, + "task_type": "CustomVoice" if task == "default_voice" else "Base", + } + assert config["task_extra_body"][task]["duration_seconds"] == 5.0 + + +@pytest.mark.parametrize( + ("task", "expected_task_type"), + [ + ("default_voice", "CustomVoice"), + ("voice_clone", "Base"), + ], +) +def test_auk_task_type_survives_served_model_alias(task, expected_task_type): + model = "tencent/AuK-Flash" + served_model = "models/flash" + config = bench_tts.load_model_configs(bench_tts._DEFAULT_MODEL_CONFIGS)[model] + cmd = bench_tts.build_bench_args( + host="localhost", + port=8000, + model=model, + task=task, + model_cfg=config, + locale="en", + num_prompts=1, + concurrency=1, + dataset_path="/data/seed-tts", + wer_eval=False, + output_dir=None, + result_filename=None, + extra_cli_args=[], + served_model_name=served_model, + ) + + assert cmd[cmd.index("--model") + 1] == served_model + extra_body = json.loads(cmd[cmd.index("--extra-body") + 1]) + assert extra_body["task_type"] == expected_task_type + assert "_vllm_omni_benchmark_model" not in extra_body + + def test_indextts25_is_registered_in_shared_model_configs() -> None: config_path = Path(bench_tts.__file__).with_name("model_configs.yaml") config = bench_tts.load_model_configs(config_path)["IndexTeam/IndexTTS-2.5"] @@ -173,3 +243,20 @@ def test_unsupported_task_exits(model_configs_path: Path, capsys: pytest.Capture ) with pytest.raises(SystemExit): bench_tts.main() + + +def test_summary_table_includes_task_type(capsys: pytest.CaptureFixture) -> None: + bench_tts.print_summary_table( + [ + { + "_task": "voice_clone", + "_task_type": "Base", + "_concurrency": 1, + } + ] + ) + + output = capsys.readouterr().out + assert "Task Type" in output + assert "voice_clone" in output + assert "Base" in output diff --git a/tests/diffusion/models/auk/test_auk_cudagraph.py b/tests/diffusion/models/auk/test_auk_cudagraph.py new file mode 100644 index 000000000..436434235 --- /dev/null +++ b/tests/diffusion/models/auk/test_auk_cudagraph.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Single-request eager and CUDA-graph AuK DiT sampling equivalence.""" + +import pytest +import torch + +from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, sample_latents +from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_dit(device: str) -> AuKTransformer: + torch.manual_seed(12) + return ( + AuKTransformer( + dim=32, + heads=2, + dim_head=16, + ff_mult=2, + latent_dim=4, + text_hidden_dim=8, + num_layers=2, + num_single_layers=2, + ) + .eval() + .to(device) + ) + + +def _sample_inputs(device: str, offset: float = 0.0) -> dict[str, torch.Tensor]: + generator = torch.Generator(device=device).manual_seed(42) + return { + "text": torch.randn(1, 7, 8, generator=generator, device=device) + offset, + "c_mask": torch.ones(1, 7, dtype=torch.bool, device=device), + "ref": torch.randn(1, 4, 4, generator=generator, device=device) - offset, + "ref_mask": torch.ones(1, 4, dtype=torch.bool, device=device), + } + + +@torch.inference_mode() +@pytest.mark.parametrize("cfg_strength", [0.0, 2.0]) +def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: float, mocker) -> None: + dit = _make_dit("cpu") + wrapper = AuKCUDAGraphWrapper(dit) + run_spy = mocker.spy(wrapper, "_run") + run_cfg_spy = mocker.spy(wrapper, "_run_cfg") + capture_spy = mocker.spy(wrapper, "_capture") + inputs = _sample_inputs("cpu") + common = dict( + **inputs, + gen_frames=9, + t_grid=[0.0, 0.4, 1.0], + cfg_strength=cfg_strength, + ) + eager = sample_latents(dit, **common, generator=torch.Generator().manual_seed(7)) + graph = sample_latents(dit, **common, generator=torch.Generator().manual_seed(7), sampler=wrapper) + torch.testing.assert_close(graph, eager) + + if cfg_strength >= 1e-5: + run_spy.assert_not_called() + assert run_cfg_spy.call_count == 2 + else: + assert run_spy.call_count == 2 + run_cfg_spy.assert_not_called() + + capture_spy.assert_not_called() + assert not wrapper._cache + + +def test_graph_inputs_use_bounded_length_buckets() -> None: + wrapper = AuKCUDAGraphWrapper(_make_dit("cpu")) + x = torch.ones(1, 65, 4) + text = torch.ones(1, 65, 8) + c_mask = torch.ones(1, 65, dtype=torch.bool) + ref = torch.ones(1, 51, 4) + ref_mask = torch.ones(1, 51, dtype=torch.bool) + + padded = wrapper._bucket_inputs(x, text, c_mask, ref, ref_mask) + + assert padded[0].shape == (1, 128, 4) + assert padded[1].shape == (1, 128) + assert padded[2].shape == (1, 128, 8) + assert padded[3].shape == (1, 128) + assert padded[4].shape == (1, 100, 4) + assert padded[5].shape == (1, 100) + assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51] + assert wrapper._key(padded[0], padded[2], padded[4], False) == (128, 128, 100, False) + assert wrapper.max_graphs == 32 + + +def test_full_graph_cache_retires_as_one_generation() -> None: + wrapper = AuKCUDAGraphWrapper(_make_dit("cpu"), max_graphs=2) + first = (64, 64, 50, False) + second = (128, 64, 50, False) + wrapper._cache[first] = object() + + wrapper._retire_graph_generation_if_full() + assert list(wrapper._cache) == [first] + + wrapper._cache[second] = object() + wrapper._retire_graph_generation_if_full() + assert not wrapper._cache + + +@torch.inference_mode() +@pytest.mark.parametrize("cfg_strength", [0.0, 2.0]) +def test_bucket_padding_preserves_real_frame_outputs(cfg_strength: float) -> None: + dit = _make_dit("cpu") + wrapper = AuKCUDAGraphWrapper(dit) + inputs = _sample_inputs("cpu") + x = torch.randn(1, 9, 4) + timestep = torch.tensor(0.4) + cfg = torch.tensor(cfg_strength) + if cfg_strength >= 1e-5: + eager = wrapper._run_cfg( + x, + None, + inputs["text"], + inputs["c_mask"], + inputs["ref"], + inputs["ref_mask"], + timestep, + cfg_strength=cfg, + ) + else: + eager = wrapper._run( + x, + None, + inputs["text"], + inputs["c_mask"], + inputs["ref"], + inputs["ref_mask"], + timestep, + ) + dit.clear_cache() + + bucketed = wrapper._bucket_inputs(x, inputs["text"], inputs["c_mask"], inputs["ref"], inputs["ref_mask"]) + if cfg_strength >= 1e-5: + padded = wrapper._run_cfg(*bucketed, timestep, cfg_strength=cfg) + else: + padded = wrapper._run(*bucketed, timestep) + dit.clear_cache() + + torch.testing.assert_close(padded[:, : x.shape[1]], eager) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") +@torch.inference_mode() +@pytest.mark.parametrize("cfg_strength", [0.0, 2.0]) +def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_strength: float) -> None: + dit = _make_dit("cuda") + wrapper = AuKCUDAGraphWrapper(dit) + + for offset in (0.0, 0.25): + inputs = _sample_inputs("cuda", offset) + common = dict( + **inputs, + gen_frames=9, + t_grid=[0.0, 0.4, 1.0], + cfg_strength=cfg_strength, + ) + eager = sample_latents(dit, **common, generator=torch.Generator(device="cuda").manual_seed(7)) + graph = sample_latents(dit, **common, generator=torch.Generator(device="cuda").manual_seed(7), sampler=wrapper) + torch.testing.assert_close(graph, eager, atol=3e-6, rtol=3e-5) + + bucketed = wrapper._bucket_inputs( + torch.empty(1, 9, 4, device="cuda"), + inputs["text"], + inputs["c_mask"], + inputs["ref"], + inputs["ref_mask"], + ) + key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5) + assert key in wrapper._cache + + entry = wrapper._cache[key] + torch.testing.assert_close(entry.static_x_mask, bucketed[1]) + torch.testing.assert_close(entry.static_text, bucketed[2]) + torch.testing.assert_close(entry.static_c_mask, bucketed[3]) + torch.testing.assert_close(entry.static_ref, bucketed[4]) + torch.testing.assert_close(entry.static_ref_mask, bucketed[5]) + torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda")) + + assert len(wrapper._cache) == 1 + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") +@torch.inference_mode() +def test_graph_capture_failure_is_propagated(mocker) -> None: + dit = _make_dit("cuda") + wrapper = AuKCUDAGraphWrapper(dit) + capture = mocker.patch.object(wrapper, "_capture", side_effect=RuntimeError("capture failed")) + eager = mocker.spy(wrapper, "_run") + inputs = _sample_inputs("cuda") + common = dict( + **inputs, + gen_frames=9, + t_grid=[0.0, 0.4, 1.0], + cfg_strength=0.0, + ) + + with pytest.raises(RuntimeError, match="capture failed"): + sample_latents( + dit, + **common, + generator=torch.Generator(device="cuda").manual_seed(7), + sampler=wrapper, + ) + + result = sample_latents( + dit, + **common, + generator=torch.Generator(device="cuda").manual_seed(7), + sampler=wrapper, + ) + + assert result.shape == (1, 9, 4) + assert capture.call_count == 1 + assert eager.call_count == 2 + assert not wrapper.enabled + assert not wrapper._cache diff --git a/tests/diffusion/models/auk/test_pipeline_auk.py b/tests/diffusion/models/auk/test_pipeline_auk.py index 954491e9e..df1953ed5 100644 --- a/tests/diffusion/models/auk/test_pipeline_auk.py +++ b/tests/diffusion/models/auk/test_pipeline_auk.py @@ -227,13 +227,16 @@ def build_pipeline(tmp_path, monkeypatch): monkeypatch.setattr(pipeline_auk, "AuKTransformer", _StubTransformer) monkeypatch.setattr(pipeline_auk, "AuKVAE", _StubVAE) - def _build(variant: str = "base") -> tuple[AuKPipeline, list[dict[str, Any]]]: + def _build( + variant: str = "base", *, model_config: dict[str, Any] | None = None + ) -> tuple[AuKPipeline, list[dict[str, Any]]]: calls: list[dict[str, Any]] = [] monkeypatch.setattr(pipeline_auk, "sample_latents", _stub_sampler(calls)) od_config = OmniDiffusionConfig( model=str(_write_checkpoint(tmp_path, variant)), dtype=torch.float32, model_class_name="AuKPipeline", + model_config=model_config or {}, ) return AuKPipeline(od_config=od_config), calls @@ -274,6 +277,20 @@ def _silence(seconds: float, sample_rate: int = SAMPLE_RATE) -> tuple[np.ndarray @pytest.mark.core_model @pytest.mark.cpu class TestRequestParsing: + @pytest.mark.parametrize( + "model_config, expected_slots", + [({}, 32), ({"max_dit_graphs": 1}, 1), ({"max_dit_graphs": 3}, 3), ({"max_dit_graphs": 64}, 64)], + ) + def test_max_dit_graphs_reaches_lazy_cache(self, build_pipeline, model_config, expected_slots): + pipeline, _ = build_pipeline(model_config=model_config) + assert pipeline.cudagraph_wrapper.max_graphs == expected_slots + assert not pipeline.cudagraph_wrapper._cache + + @pytest.mark.parametrize("max_dit_graphs", [0, -1, True, False, 1.5, "3", None]) + def test_invalid_max_dit_graphs_is_rejected(self, build_pipeline, max_dit_graphs): + with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"): + build_pipeline(model_config={"max_dit_graphs": max_dit_graphs}) + def test_pipeline_declares_audio_output(self, build_pipeline): pipeline, _ = build_pipeline() @@ -299,6 +316,7 @@ class TestRequestParsing: assert calls[0]["ref"].shape == (1, 100, LATENT_DIM) assert calls[0]["ref_mask"].shape == (1, 100) assert calls[0]["c_mask"].shape == (1, 8) + assert calls[0]["sampler"] is pipeline.cudagraph_wrapper assert outputs[0].output.shape == (300 * HOP,) assert outputs[0].output.dtype is torch.float32 diff --git a/tests/e2e/online_serving/test_auk.py b/tests/e2e/online_serving/test_auk.py new file mode 100644 index 000000000..1b0cfaadc --- /dev/null +++ b/tests/e2e/online_serving/test_auk.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK Speech API E2E coverage for the encoder -> diffusion deployment.""" + +from __future__ import annotations + +import io +import os + +import numpy as np +import pytest +import requests +import soundfile as sf + +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("VLLM_OMNI_AUK_MODEL_DIR") +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.skipif(not MODEL, reason="set VLLM_OMNI_AUK_MODEL_DIR to an assembled AuK checkpoint"), +] + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("reference", [False, True], ids=["instruct", "voice-clone"]) +@pytest.mark.parametrize("stream", [False, True], ids=["wav", "streaming-pcm"]) +@pytest.mark.parametrize( + "omni_server", + [ + pytest.param( + OmniServerParams( + model=MODEL or ".", + stage_config_path=get_deploy_config_path("auk.yaml"), + server_args=["--trust-remote-code", "--disable-log-stats"], + ), + id="auk", + ) + ], + indirect=True, +) +def test_auk_speech_api_request(omni_server, reference, stream) -> None: + """The adapter accepts Speech API requests and returns bounded 24 kHz WAVs.""" + + def generate(duration: float) -> int: + extra = {} + if reference: + extra["ref_audio"] = get_asset_path("cosyvoice3/zero_shot_prompt.wav", as_data_url=True) + # Match the shared seed-tts benchmark request envelope. + extra["ref_text"] = "Reference transcript is not required by AuK." + response = requests.post( + f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech", + json={ + "model": omni_server.model, + "input": "", + "instructions": ( + "Say the following with the same voice: 'Hello, this is an AuK speech API test.'" + if reference + else 'Generate speech based on the following description: "A clear, natural voice.". ' + 'The content to speak is: "Hello, this is an AuK speech API test.".' + ), + "voice": "default", + "duration_seconds": duration, + "seed": 7, + "response_format": "pcm" if stream else "wav", + "stream": stream, + "stream_format": "audio" if stream else None, + **extra, + }, + timeout=300, + ) + assert response.status_code == 200, response.text + if stream: + waveform = np.frombuffer(response.content, dtype=" 0 + return int(np.asarray(waveform).reshape(-1).shape[0]) + + assert generate(2.0) == 100 * 480 diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 1ec32a95e..6504359e4 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -5367,6 +5367,9 @@ class TestTTSAsyncOffloading: adapter_model_type = "adapter_dummy_tts" class FakeAdapter: + def normalize(self, request): + return + def validate(self, request): return None diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index babf8a2b9..92f13462e 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -8,11 +8,22 @@ Pure-Python registry/resolution logic; no model or GPU resources are loaded. import asyncio from types import SimpleNamespace from typing import Any +from unittest.mock import AsyncMock import pytest +import torch from vllm.sampling_params import SamplingParams +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient +from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.stage_pool import StagePool from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.tts_adapters import ( TTS_ADAPTER_REGISTRY, ARTTSAdapter, @@ -22,6 +33,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( detect_tts_model_type, resolve_adapter, ) +from vllm_omni.entrypoints.openai.tts_adapters.auk import AuKAdapter from vllm_omni.entrypoints.openai.tts_adapters.base import TTSCapabilities, resolve_stage_model_path from vllm_omni.entrypoints.openai.tts_adapters.covo_audio import CovoAudioAdapter from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter @@ -42,10 +54,12 @@ from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import ( ) from vllm_omni.entrypoints.openai.tts_adapters.step_audio2 import StepAudio2Adapter from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.models.indextts2 import prompt_utils from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import ( INDEXTTS25_TOKENIZER_FILE, ) +from vllm_omni.model_executor.stage_input_processors.auk import encoder2dit pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -70,6 +84,7 @@ EXPECTED_MODEL_TYPES = { "step_audio2", "indextts2", "indextts2_5", + "auk", "gepard", } @@ -99,6 +114,314 @@ def test_resolve_unknown_returns_none(): assert resolve_adapter(None) is None +@pytest.fixture +def auk_adapter(mocker): + server = mocker.Mock(spec=OmniOpenAIServingSpeech) + server._max_instructions_length = 4096 + server._validate_ref_audio_format.return_value = None + server._resolve_ref_audio = AsyncMock(return_value=([0.1] * 480, 24000, "key")) + return AuKAdapter(SpeechServingContext(server=server)) + + +def test_auk_adapter_detection(mocker): + assert detect_tts_model_type("encoder", "AuKForConditionalGeneration") == "auk" + # ``encoder`` is shared with MiniMax H3, so AuK must not claim the stage key. + assert AuKAdapter.stage_keys == frozenset() + assert AuKAdapter.arch_identifies_entry_stage + assert detect_tts_model_type("encoder", "MiniMaxH3Encoder") is None + assert detect_tts_model_type("encoder", None) is None + + server = mocker.Mock(spec=OmniOpenAIServingSpeech) + server._diffusion_mode = False + server._tts_model_type = detect_tts_model_type("encoder", "AuKForConditionalGeneration") + server._adapter = None + server.engine_client = SimpleNamespace() + + adapter = OmniOpenAIServingSpeech._get_tts_adapter(server) + assert isinstance(adapter, AuKAdapter) + server._drop_shadowing_uploads.assert_called_once_with() + + +def test_auk_source_length_default_and_complete_instruction(auk_adapter): + request = OpenAICreateSpeechRequest( + input="", + ref_audio="reference.wav", + ref_text="Accepted transcript", + instructions="Change the pitch by one semitone.", + extra_params={"sway": -0.5, "t_grid": [0, 0.25, 1], "vae_sample": True}, + ) + assert auk_adapter.validate(request) is None + prepared = asyncio.run(auk_adapter.build(request, [], True)) + assert "Change the pitch by one semitone." in prepared.prompt["prompt"] + assert "content to speak" not in prepared.prompt["prompt"] + knobs = prepared.prompt["additional_information"]["auk"] + assert knobs["gen_seconds"] is None + assert knobs["vae_sample"] is True + assert knobs["t_grid"] == [0, 0.25, 1] + auk_adapter.ctx.server._resolve_ref_audio.assert_awaited_once_with("reference.wav") + + +@pytest.mark.parametrize( + "t_grid", + [ + [], + [0.0], + [0.0, 0.0], + [0.5, 0.25], + [0.0, float("nan")], + [0.0, float("inf")], + [0.0, "invalid"], + "0,1", + ], +) +def test_auk_rejects_invalid_t_grid_before_dispatch(auk_adapter, t_grid): + request = OpenAICreateSpeechRequest( + input="Say the following: 'target text'", + duration_seconds=2, + extra_params={"t_grid": t_grid}, + ) + + error = auk_adapter.validate(request) + + assert error == "AuK extra_params.t_grid must contain at least two finite, strictly increasing values" + auk_adapter.ctx.server._resolve_ref_audio.assert_not_awaited() + + +@pytest.mark.parametrize( + ("task_type", "expected_instruction"), + [ + ("CustomVoice", 'Say the following: "target text"'), + ("Base", 'Say the following with the same voice: "target text"'), + ], +) +def test_auk_task_type_normalizes_benchmark_text(auk_adapter, task_type, expected_instruction, mocker): + warning_once = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning_once") + request = OpenAICreateSpeechRequest( + input="target text", + task_type=task_type, + duration_seconds=2, + ref_audio="reference.wav" if task_type == "Base" else None, + ) + + server = mocker.Mock(spec=OmniOpenAIServingSpeech) + server._validate_speech_sample_rate.return_value = None + server._get_tts_adapter.return_value = auk_adapter + assert OmniOpenAIServingSpeech._validate_tts_request(server, request) is None + + assert request.input == "" + assert request.instructions == expected_instruction + assert request.task_type is None + warning_once.assert_called_once() + assert "prefer a complete `instructions` prompt" in warning_once.call_args.args[0] + + auk_adapter.normalize(request) + + assert request.instructions == expected_instruction + warning_once.assert_called_once() + + +def test_auk_task_type_quotes_embedded_delimiters(auk_adapter): + request = OpenAICreateSpeechRequest( + input='It\'s called "AuK"\\Flash', + task_type="CustomVoice", + duration_seconds=2, + ) + + auk_adapter.normalize(request) + + assert request.instructions == 'Say the following: "It\'s called \\"AuK\\"\\\\Flash"' + assert request.task_type is None + + +def test_auk_task_type_preserves_explicit_instructions(auk_adapter, mocker): + warning_once = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning_once") + warning = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning") + request = OpenAICreateSpeechRequest( + input="target text", + task_type="Base", + instructions="Use this complete AuK instruction.", + duration_seconds=2, + ) + + server = mocker.Mock(spec=OmniOpenAIServingSpeech) + server._validate_speech_sample_rate.return_value = None + server._get_tts_adapter.return_value = auk_adapter + assert OmniOpenAIServingSpeech._validate_tts_request(server, request) is None + + assert request.input == "" + assert request.instructions == "Use this complete AuK instruction." + assert request.task_type is None + warning_once.assert_called_once() + warning.assert_called_once() + assert "without applying another task template" in warning.call_args.args[0] + + +def test_auk_task_type_does_not_double_wrap_preformatted_input(auk_adapter, mocker): + warning_once = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning_once") + warning = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning") + instruction = "Say the following with the same voice: 'target text'" + request = OpenAICreateSpeechRequest( + input=instruction, + task_type="Base", + duration_seconds=2, + ref_audio="reference.wav", + ) + server = mocker.Mock(spec=OmniOpenAIServingSpeech) + server._validate_speech_sample_rate.return_value = None + server._get_tts_adapter.return_value = auk_adapter + + assert OmniOpenAIServingSpeech._validate_tts_request(server, request) is None + + assert request.input == "" + assert request.instructions == instruction + assert request.task_type is None + warning_once.assert_called_once() + warning.assert_called_once() + assert "already contains a complete task instruction" in warning.call_args.args[0] + + +def test_auk_sampling_overrides_reach_stage1_pipeline(auk_adapter, mocker): + """Exercise one request from the Speech adapter through stage-1 admission.""" + + class RecordingDiffusionStage: + stage_type = "diffusion" + final_output = True + engine_input_source = [0] + requires_multimodal_data = False + custom_process_input_func = staticmethod(encoder2dit) + + def __init__(self) -> None: + self.request: OmniDiffusionRequest | None = None + + async def add_request_async(self, request_id, prompt, sampling_params, **_kwargs) -> None: + # Match the stage-client wire boundary and the receiving diffusion + # process, where params are serialized then reconstructed before + # creating the request consumed by the pipeline. + wire_params = StageDiffusionClient._sampling_params_to_dict(sampling_params) + self.request = OmniDiffusionRequest( + prompt=prompt, + sampling_params=StageDiffusionProc._reconstruct_sampling_params( + object.__new__(StageDiffusionProc), wire_params + ), + request_id=request_id, + ) + + defaults = [SamplingParams(max_tokens=1), OmniDiffusionSamplingParams(seed=0, num_inference_steps=32)] + request = OpenAICreateSpeechRequest( + input="", + instructions='Generate speech based on the following description: "A calm voice". The content to speak is: "Hello".', + duration_seconds=2, + seed=9, + extra_params={ + "num_inference_steps": 4, + "guidance_scale": 1.5, + "sway": -0.5, + "t_grid": [0.0, 0.25, 1.0], + "vae_sample": True, + }, + ) + prepared = asyncio.run(auk_adapter.build(request, defaults, False)) + updated = auk_adapter.apply_sampling_overrides(defaults, request) + assert defaults[1].num_inference_steps == 32 + assert defaults[1].seed == 0 + assert updated[0].max_tokens == 1 + + source_stage = SimpleNamespace( + stage_type="llm", + final_output=False, + get_kv_sender_info=lambda: None, + ) + diffusion_stage = RecordingDiffusionStage() + orchestrator = object.__new__(Orchestrator) + orchestrator.stage_pools = [StagePool(0, source_stage), StagePool(1, diffusion_stage)] + orchestrator._cfg_tracker = CfgCompanionTracker() + orchestrator.duplex_control_plane = None + orchestrator._running_counter = None + + req_state = OrchestratorRequestState( + request_id="auk-sampling-override", + prompt=prepared.prompt, + sampling_params_list=updated, + final_stage_id=1, + ) + encoder_output = SimpleNamespace( + request_id=req_state.request_id, + finished=True, + multimodal_output={"hidden_states": {"output": torch.ones(3, 4)}}, + outputs=[], + ) + asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, req_state.request_id, 0, encoder_output, req_state)) + + received = diffusion_stage.request + assert received is not None + assert received.prompt["prompt_embeds"].shape == (3, 4) + assert received.prompt["additional_information"]["auk"] == { + "gen_seconds": 2, + "sway": -0.5, + "t_grid": [0.0, 0.25, 1.0], + "vae_sample": True, + "has_audio": False, + } + assert received.sampling_params.num_inference_steps == 4 + assert received.sampling_params.guidance_scale == 1.5 + assert received.sampling_params.guidance_scale_provided is True + assert received.sampling_params.seed == 9 + + runner = mocker.Mock(spec=DiffusionModelRunner) + runner.device = torch.device("cpu") + DiffusionModelRunner._initialize_generator(runner, received.sampling_params) + assert received.sampling_params.generator.initial_seed() == 9 + + pipeline_params = DiffusionRequestBatch(requests=[received]).sampling_params + assert pipeline_params is received.sampling_params + assert pipeline_params.generator.initial_seed() == 9 + + +def test_auk_instructions_take_precedence_over_input(auk_adapter, mocker): + warning = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning") + request = OpenAICreateSpeechRequest( + input="ignored input", + instructions='Generate speech based on the following description: "Speak warmly". The content to speak is: "Hello".', + duration_seconds=2, + ) + assert auk_adapter.validate(request) is None + prepared = asyncio.run(auk_adapter.build(request, [], False)) + prompt = prepared.prompt + assert 'Generate speech based on the following description: "Speak warmly".' in prompt["prompt"] + assert 'The content to speak is: "Hello".' in prompt["prompt"] + assert "||" in prompt["prompt"] + assert prompt["additional_information"]["auk"]["gen_seconds"] == 2 + warning.assert_called_once() + assert "using the complete `instructions`" in warning.call_args.args[0] + + +def test_auk_input_only_is_treated_as_complete_instruction(auk_adapter, mocker): + warning = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning") + request = OpenAICreateSpeechRequest( + input="Say the following with the same voice: 'Hello world.'", + ref_audio="reference.wav", + duration_seconds=2, + ) + assert auk_adapter.validate(request) is None + prepared = asyncio.run(auk_adapter.build(request, [], True)) + assert "Say the following with the same voice: 'Hello world.'" in prepared.prompt["prompt"] + assert warning.call_count == 1 + assert 'Prefer `input=""`' in warning.call_args.args[0] + + +def test_auk_complete_instruction_does_not_require_spoken_text(auk_adapter): + request = OpenAICreateSpeechRequest( + input="", + instructions="Replace 'morning' with 'evening' in the source recording.", + ref_audio="reference.wav", + ) + assert auk_adapter.validate(request) is None + prepared = asyncio.run(auk_adapter.build(request, [], True)) + assert prepared.prompt["prompt"].count("Replace 'morning'") == 1 + assert "Generate speech based on" not in prepared.prompt["prompt"] + assert prepared.prompt["additional_information"]["auk"]["gen_seconds"] is None + + def test_stage_model_path_prefers_typed_override(): engine_client = SimpleNamespace( stage_configs=[ diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index cc88969d1..3ef55f673 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -182,6 +182,7 @@ _PIPELINE_STAGES = [ "token2text", "token2wav", "tts", + "encoder", ] _STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"] @@ -283,14 +284,22 @@ def test_pipeline_stage_list_is_complete(): assert not missing, f"stage keys claimed by adapters but absent from _PIPELINE_STAGES: {sorted(missing)}" -def test_entry_stage_archs_is_ming_only(): - """Only Ming dense identifies its entry stage by architecture. +def test_entry_stage_archs_are_explicit(): + """Only Ming dense and AuK identify their entry stage by architecture. - Widening this set would change which stage ``_find_tts_stage`` selects in - mixed deployments — notably VoxCPM2, which declares ``model_archs`` but is - found by its ``latent_generator`` stage key. + Both use a generic ``model_stage`` (``llm`` and ``encoder``) that other + pipelines share. Widening this set would change which stage + ``_find_tts_stage`` selects in mixed deployments — notably VoxCPM2, which + declares ``model_archs`` but is found by its ``latent_generator`` stage key. """ - assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration"}) + assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration", "AuKForConditionalGeneration"}) + + +def test_minimax_h3_encoder_is_not_a_tts_stage(): + """MiniMax H3 shares AuK's ``encoder`` stage key but is not a speech model.""" + assert "encoder" not in all_tts_stage_keys() + assert detect_tts_model_type("encoder", "MiniMaxH3Encoder") is None + assert detect_tts_model_type("encoder", "AuKForConditionalGeneration") == "auk" def test_every_detected_type_is_adapter_backed(): diff --git a/tests/helpers/client.py b/tests/helpers/client.py index 5991ab09d..8637e510d 100644 --- a/tests/helpers/client.py +++ b/tests/helpers/client.py @@ -1475,6 +1475,7 @@ class OnlineOmniClient: "max_new_tokens", "seed", "instructions", + "duration_seconds", "speed", "sample_rate", "extra_params", diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 5e07706d2..ef9712eb9 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -214,6 +214,9 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc ) class Adapter: + def normalize(self, request): + return + def validate(self, request): return None diff --git a/tools/pre_commit/check_tts_adapter.py b/tools/pre_commit/check_tts_adapter.py index 199c46a30..752e5940e 100644 --- a/tools/pre_commit/check_tts_adapter.py +++ b/tools/pre_commit/check_tts_adapter.py @@ -29,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 = 6 +MAX_MODEL_TYPE_BRANCHES = 5 _GUIDANCE = """ Move the per-model behaviour into its adapter under diff --git a/vllm_omni/deploy/auk.yaml b/vllm_omni/deploy/auk.yaml index 43c403bf3..0268ea97b 100644 --- a/vllm_omni/deploy/auk.yaml +++ b/vllm_omni/deploy/auk.yaml @@ -43,8 +43,8 @@ stages: # Not eager: the codec decode is compiled into bucketed CUDA graphs at # startup (about a minute more on a cold Inductor cache, ~30 s warm, and # about 1 GB of graph pools); clips longer than the largest bucket are - # decoded in overlapping tiles of it. The DiT stays eager unless - # diffusion_compile_granularity is set to full. Buckets, tile size and the + # decoded in overlapping tiles of it. Each DiT denoise step replays a lazily + # captured per-shape CUDA graph. Buckets, tile size and the # plain-graph cache are tunable through model_config: # model_config: {auk_vae_compile_shapes: [128, 256, 512], auk_vae_tile_frames: 512, auk_vae_max_graphs: 32} - stage_id: 1 @@ -52,6 +52,10 @@ stages: max_num_seqs: 1 gpu_memory_utilization: 0.35 enforce_eager: false + model_config: + # Maximum number of cached DiT graphs. When the cache reaches this limit, + # capturing a new shape retires all existing graphs before capture. + max_dit_graphs: 32 default_sampling_params: seed: 0 num_inference_steps: 32 diff --git a/vllm_omni/diffusion/models/auk/auk_transformer.py b/vllm_omni/diffusion/models/auk/auk_transformer.py index 5aca8f39b..4383fe971 100644 --- a/vllm_omni/diffusion/models/auk/auk_transformer.py +++ b/vllm_omni/diffusion/models/auk/auk_transformer.py @@ -28,13 +28,13 @@ Deliberate differences from the reference, all inference-only: """ import math -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping import torch import torch.nn.functional as F from torch import nn -__all__ = ["AuKTransformer", "dit_state_dict", "sample_latents"] +__all__ = ["AuKTransformer", "build_time_grid", "dit_state_dict", "sample_latents"] def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor: @@ -86,14 +86,17 @@ class Rotary(nn.Module): exponents = torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim return 1.0 / (self.base**exponents) - def forward(self, seq_len: int) -> torch.Tensor: - """Return frequencies ``[1, 1, seq_len, dim]`` for positions ``0..seq_len-1``.""" + def forward(self, seq_len: int, mask: torch.Tensor | None = None) -> torch.Tensor: + """Return rotary frequencies, compressing positions across padding.""" inv_freq = self.inv_freq if inv_freq.dtype != torch.float32: inv_freq = self._frequencies(inv_freq.device) - pos = torch.arange(seq_len, device=inv_freq.device, dtype=torch.float32) - freqs = torch.outer(pos, inv_freq) - return torch.stack((freqs, freqs), dim=-1).flatten(-2)[None, None] + if mask is None: + pos = torch.arange(seq_len, device=inv_freq.device, dtype=torch.float32)[None] + else: + pos = (mask.to(torch.int32).cumsum(dim=1) - 1).clamp_min(0).to(torch.float32) + freqs = pos.unsqueeze(-1) * inv_freq + return torch.stack((freqs, freqs), dim=-1).flatten(-2).unsqueeze(1) class TimeEmbedding(nn.Module): @@ -541,15 +544,15 @@ class AuKTransformer(nn.Module): seq_len = x.shape[1] text_len = c.shape[1] - rope_audio = self.rotary_embed(seq_len) - rope_text = self.rotary_embed(text_len) + rope_audio = self.rotary_embed(seq_len, audio_mask) + rope_text = self.rotary_embed(text_len, c_mask) for block in self.transformer_blocks: c, x = block(x, c, t, mask=audio_mask, rope=rope_audio, c_rope=rope_text, c_mask=c_mask) x = torch.cat([c, x], dim=1) - rope = self.rotary_embed(text_len + seq_len) single_mask = None if audio_mask is None else torch.cat([c_mask, audio_mask], dim=1) + rope = self.rotary_embed(text_len + seq_len, single_mask) for block in self.single_transformer_blocks: x = block(x, t, mask=single_mask, rope=rope) @@ -571,6 +574,31 @@ def dit_state_dict( return {name[len(prefix) :]: tensor for name, tensor in items if name.startswith(prefix)} +def build_time_grid( + *, + nfe: int, + sway_sampling_coef: float | None, + t_grid: list[float] | None, + device: torch.device | str, +) -> torch.Tensor: + """Build and validate the Euler schedule before a graph replay.""" + if t_grid is not None: + timesteps = torch.tensor(t_grid, device="cpu", dtype=torch.float32) + else: + timesteps = torch.linspace(0, 1, nfe + 1, device="cpu", dtype=torch.float32) + if sway_sampling_coef is not None: + timesteps = timesteps + sway_sampling_coef * (torch.cos(math.pi / 2 * timesteps) - 1 + timesteps) + if ( + timesteps.ndim != 1 + or timesteps.numel() < 2 + or not bool(torch.isfinite(timesteps).all() & torch.all(timesteps[1:] > timesteps[:-1])) + ): + raise ValueError( + f"AuK sampling needs a strictly increasing time grid with at least two points; got {timesteps.tolist()}" + ) + return timesteps.to(device) + + @torch.no_grad() def sample_latents( dit: AuKTransformer, @@ -589,6 +617,7 @@ def sample_latents( latent_dim: int | None = None, device: torch.device | str | None = None, dtype: torch.dtype | None = None, + sampler: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor: """Integrate the flow from noise to audio latents with explicit Euler steps. @@ -640,25 +669,61 @@ def sample_latents( .unsqueeze(0) ) - if t_grid is not None: - t = torch.tensor(t_grid, device=device, dtype=torch.float32) - else: - t = torch.linspace(0, 1, nfe + 1, device=device, dtype=torch.float32) - if sway_sampling_coef is not None: - t = t + sway_sampling_coef * (torch.cos(math.pi / 2 * t) - 1 + t) + t = build_time_grid( + nfe=nfe, + sway_sampling_coef=sway_sampling_coef, + t_grid=t_grid, + device=device, + ) + if sampler is not None: + try: + for i in range(t.shape[0] - 1): + velocity = sampler( + x=x, + text=text, + c_mask=c_mask, + ref=ref, + ref_mask=ref_mask, + timestep=t[i], + cfg_strength=cfg_strength, + ) + x = x + (t[i + 1] - t[i]) * velocity + finally: + dit.clear_cache() + return x + return _sample_latents( + dit, + initial_latents=x, + text=text, + c_mask=c_mask, + ref=ref, + ref_mask=ref_mask, + timesteps=t, + cfg_strength=cfg_strength, + ) + +def _sample_latents( + dit: AuKTransformer, + *, + initial_latents: torch.Tensor, + text: torch.Tensor, + c_mask: torch.Tensor | None, + ref: torch.Tensor, + ref_mask: torch.Tensor | None, + timesteps: torch.Tensor, + cfg_strength: float, +) -> torch.Tensor: + """Euler integration shared by eager sampling and CUDA graph capture.""" + x = initial_latents guided = cfg_strength >= 1e-5 - if t.numel() < 2 or not bool(torch.all(t[1:] > t[:-1])): - raise ValueError( - f"AuK sampling needs a strictly increasing time grid with at least two points; got {t.tolist()}" - ) try: - for i in range(t.shape[0] - 1): + for i in range(timesteps.shape[0] - 1): if guided: pred = dit( x, text, - t[i], + timesteps[i], c_mask=c_mask, ref=ref, ref_mask=ref_mask, @@ -668,8 +733,8 @@ def sample_latents( v_cond, v_uncond = pred.chunk(2, dim=0) v = v_cond + (v_cond - v_uncond) * cfg_strength else: - v = dit(x, text, t[i], c_mask=c_mask, ref=ref, ref_mask=ref_mask) - x = x + (t[i + 1] - t[i]) * v + v = dit(x, text, timesteps[i], c_mask=c_mask, ref=ref, ref_mask=ref_mask) + x = x + (timesteps[i + 1] - timesteps[i]) * v finally: # The cached text projections belong to this request only; a failed # step must not leak them into the next one. diff --git a/vllm_omni/diffusion/models/auk/cudagraph_wrapper.py b/vllm_omni/diffusion/models/auk/cudagraph_wrapper.py new file mode 100644 index 000000000..d9226baf1 --- /dev/null +++ b/vllm_omni/diffusion/models/auk/cudagraph_wrapper.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Per-shape CUDA graph for one AuK DiT denoise step.""" + +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from vllm.logger import init_logger +from vllm.utils.math_utils import round_up + +from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer + +logger = init_logger(__name__) + + +@dataclass +class _GraphEntry: + graph: torch.cuda.CUDAGraph + static_x: torch.Tensor + static_x_mask: torch.Tensor + static_text: torch.Tensor + static_c_mask: torch.Tensor + static_ref: torch.Tensor + static_ref_mask: torch.Tensor + static_timestep: torch.Tensor + static_cfg: torch.Tensor | None + static_out: torch.Tensor + + +class AuKCUDAGraphWrapper: + """Replay one DiT denoise step and leave Euler scheduling to the caller. + + The graph is keyed by the target, text and reference sequence lengths plus + whether the CFG branch is enabled. Timestep and the CFG strength are + mutable scalar buffers, so all Euler steps and CFG values within one path + reuse one graph. + """ + + _TARGET_ALIGNMENT = 64 + _TEXT_ALIGNMENT = 64 + _REF_ALIGNMENT = 50 + + def __init__(self, dit: AuKTransformer, *, enabled: bool = True, max_graphs: int = 32) -> None: + self.dit = dit + self.enabled = bool(enabled) + self.max_graphs = max(1, int(max_graphs)) + self._cache: OrderedDict[tuple, _GraphEntry] = OrderedDict() + self._pool_handle: int | None = None + logger.info("Initialized AuK DiT lazy CUDA graph cache: enabled=%s slots=%d", self.enabled, self.max_graphs) + + @classmethod + def _bucket_inputs( + cls, + x: torch.Tensor, + text: torch.Tensor, + c_mask: torch.Tensor, + ref: torch.Tensor, + ref_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + target_bucket = round_up(x.shape[1], cls._TARGET_ALIGNMENT) + text_bucket = round_up(text.shape[1], cls._TEXT_ALIGNMENT) + ref_bucket = round_up(ref.shape[1], cls._REF_ALIGNMENT) + x_mask = torch.ones(x.shape[:2], dtype=torch.bool, device=x.device) + return ( + F.pad(x, (0, 0, 0, target_bucket - x.shape[1])), + F.pad(x_mask, (0, target_bucket - x_mask.shape[1]), value=False), + F.pad(text, (0, 0, 0, text_bucket - text.shape[1])), + F.pad(c_mask, (0, text_bucket - c_mask.shape[1]), value=False), + F.pad(ref, (0, 0, 0, ref_bucket - ref.shape[1])), + F.pad(ref_mask, (0, ref_bucket - ref_mask.shape[1]), value=False), + ) + + @staticmethod + def _key( + x: torch.Tensor, + text: torch.Tensor, + ref: torch.Tensor, + uses_cfg: bool, + ) -> tuple[int, int, int, bool]: + return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg) + + def _retire_graph_generation_if_full(self) -> None: + """Retire all graphs together so none outlive shared workspaces.""" + if len(self._cache) >= self.max_graphs: + self._cache.clear() + + @torch.no_grad() + def __call__( + self, + *, + x: torch.Tensor, + text: torch.Tensor, + c_mask: torch.Tensor, + ref: torch.Tensor, + ref_mask: torch.Tensor, + timestep: torch.Tensor, + cfg_strength: float, + ) -> torch.Tensor: + uses_cfg = cfg_strength >= 1e-5 + if not self.enabled or x.device.type != "cuda" or torch.cuda.is_current_stream_capturing(): + if uses_cfg: + return self._run_cfg(x, None, text, c_mask, ref, ref_mask, timestep, cfg_strength=cfg_strength) + return self._run(x, None, text, c_mask, ref, ref_mask, timestep) + + target_frames = x.shape[1] + x, x_mask, text, c_mask, ref, ref_mask = self._bucket_inputs(x, text, c_mask, ref, ref_mask) + inputs = (x, x_mask, text, c_mask, ref, ref_mask, timestep) + key = self._key(x, text, ref, uses_cfg) + entry = self._cache.get(key) + if entry is None: + # Same Retirement as #6587, details discussed in #7469 + self._retire_graph_generation_if_full() + try: + entry = self._capture(*inputs, cfg_strength=cfg_strength, uses_cfg=uses_cfg) + except Exception: + self._cache.clear() + self.enabled = False + logger.exception( + "Disabling AuK DiT CUDA graphs after capture failure for key=%s; " + "subsequent requests will use eager execution.", + key, + ) + raise + self._cache[key] = entry + else: + self._cache.move_to_end(key) + + entry.static_x.copy_(x) + entry.static_x_mask.copy_(x_mask) + entry.static_text.copy_(text) + entry.static_c_mask.copy_(c_mask) + entry.static_ref.copy_(ref) + entry.static_ref_mask.copy_(ref_mask) + entry.static_timestep.copy_(timestep) + if entry.static_cfg is not None: + entry.static_cfg.fill_(cfg_strength) + entry.graph.replay() + return entry.static_out[:, :target_frames].clone() + + def _run( + self, + x: torch.Tensor, + x_mask: torch.Tensor | None, + text: torch.Tensor, + c_mask: torch.Tensor, + ref: torch.Tensor, + ref_mask: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + return self.dit(x, text, timestep, mask=x_mask, c_mask=c_mask, ref=ref, ref_mask=ref_mask) + + def _run_cfg( + self, + x: torch.Tensor, + x_mask: torch.Tensor | None, + text: torch.Tensor, + c_mask: torch.Tensor, + ref: torch.Tensor, + ref_mask: torch.Tensor, + timestep: torch.Tensor, + *, + cfg_strength: torch.Tensor | float, + ) -> torch.Tensor: + pred = self.dit( + x, + text, + timestep, + mask=x_mask, + c_mask=c_mask, + ref=ref, + ref_mask=ref_mask, + cfg_infer=True, + cache=True, + ) + conditional, unconditional = pred.chunk(2, dim=0) + return conditional + (conditional - unconditional) * cfg_strength + + def _capture( + self, + *inputs: torch.Tensor, + cfg_strength: float, + uses_cfg: bool, + ) -> _GraphEntry: + static_inputs = tuple(value.clone() for value in inputs) + static_cfg = None + if uses_cfg: + static_cfg = torch.empty((), device=static_inputs[0].device, dtype=torch.float32) + static_cfg.fill_(cfg_strength) + try: + for _ in range(3): + if uses_cfg: + assert static_cfg is not None + self._run_cfg(*static_inputs, cfg_strength=static_cfg) + else: + self._run(*static_inputs) + # CFG warm-up populates AuKTransformer's Python-side projected-text + # cache. Clear it before capture so the graph includes + # project_text(static_text), rather than closing over the first + # request's projection and ignoring later static_text updates. + self.dit.clear_cache() + if static_cfg is not None: + static_cfg.fill_(cfg_strength) + if self._pool_handle is None: + self._pool_handle = torch.cuda.graph_pool_handle() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self._pool_handle): + # Capture the actual CFG path. The scalar value remains a + # mutable buffer, so cfg=2 and cfg=3 share this graph. + if uses_cfg: + assert static_cfg is not None + static_out = self._run_cfg(*static_inputs, cfg_strength=static_cfg) + else: + static_out = self._run(*static_inputs) + finally: + self.dit.clear_cache() + + logger.info( + "Captured AuK DiT single-step CUDA graph: target_frames=%d text_tokens=%d ref_frames=%d cfg=%s", + inputs[0].shape[1], + inputs[2].shape[1], + inputs[4].shape[1], + cfg_strength, + ) + return _GraphEntry( + graph=graph, + static_x=static_inputs[0], + static_x_mask=static_inputs[1], + static_text=static_inputs[2], + static_c_mask=static_inputs[3], + static_ref=static_inputs[4], + static_ref_mask=static_inputs[5], + static_timestep=static_inputs[6], + static_cfg=static_cfg, + static_out=static_out, + ) + + +__all__ = ["AuKCUDAGraphWrapper"] diff --git a/vllm_omni/diffusion/models/auk/pipeline_auk.py b/vllm_omni/diffusion/models/auk/pipeline_auk.py index b4b379894..87dabb316 100644 --- a/vllm_omni/diffusion/models/auk/pipeline_auk.py +++ b/vllm_omni/diffusion/models/auk/pipeline_auk.py @@ -27,6 +27,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE +from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph from vllm_omni.diffusion.models.interface import ( SupportAudioInput, @@ -139,6 +140,9 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp super().__init__() del prefix # Weights are not namespaced: one checkpoint, one pipeline. self.od_config = od_config + max_dit_graphs = od_config.model_config.get("max_dit_graphs", 32) + if isinstance(max_dit_graphs, bool) or not isinstance(max_dit_graphs, int) or max_dit_graphs < 1: + raise ValueError("AuK max_dit_graphs must be a positive integer") self.device = get_local_device() self.dtype = getattr(od_config, "dtype", None) or torch.bfloat16 @@ -183,6 +187,9 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp self.dit.load_state_dict(_read_dit_weights(model_dir, self.dtype), strict=True) self.dit = self.dit.to(device=self.device).eval() self.dit.requires_grad_(False) + self.cudagraph_wrapper = AuKCUDAGraphWrapper( + self.dit, enabled=not od_config.enforce_eager, max_graphs=max_dit_graphs + ) # The compiled decode buckets are warmed by setup_compile(), which the # model runner calls at startup unless the stage is enforce_eager. The # bucket list and the plain-graph cache size come from the stage's @@ -462,6 +469,7 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp device=self.device, dtype=torch.float32, generator=generator, + sampler=self.cudagraph_wrapper, ) latents = latents.float() diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 02a6c8cf6..b6c495a64 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1173,6 +1173,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): adapter = self._get_tts_adapter() if adapter is not None: + adapter.normalize(request) return adapter.validate(request) adapter_cls = resolve_adapter("qwen3_tts") @@ -1183,7 +1184,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): server=self, engine_client=self.engine_client, ) - return adapter_cls(ctx).validate(request) + adapter = adapter_cls(ctx) + adapter.normalize(request) + return adapter.validate(request) def _validate_speech_sample_rate(self, request: OpenAICreateSpeechRequest) -> str | None: if request.sample_rate is None: @@ -2000,6 +2003,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type: str | None = None has_inline_ref_audio = (request.ref_audio is not None) if has_inline_ref_audio is None else has_inline_ref_audio if (adapter := self._get_tts_adapter()) is not None: + adapter.normalize(request) validation_error = adapter.validate(request) if validation_error: raise ValueError(validation_error) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index 2b7f93a62..f98fa87df 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -105,6 +105,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 audex, audex_tta, audio8_tts, + auk, breeze_tts_2, cosyvoice3, covo_audio, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/auk.py b/vllm_omni/entrypoints/openai/tts_adapters/auk.py new file mode 100644 index 000000000..4ba11bf7c --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/auk.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK Speech API adapter for the current multi-stage serving contract.""" + +from __future__ import annotations + +import copy +import json +import math +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 +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.model_extras.auk import DEFAULT_SWAY, auk_prompt + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + + +@register_tts_adapter +class AuKAdapter(ARTTSAdapter): + """Translate Speech API fields into the existing AuK encoder/DiT prompt.""" + + name = "auk" + # AuK's entry stage uses the generic ``model_stage="encoder"``, which other + # pipelines (MiniMax H3) also declare, so both stage discovery and + # model-type detection go through the architecture instead. + model_archs = frozenset({"AuKForConditionalGeneration"}) + arch_identifies_entry_stage = True + + def normalize(self, request: OpenAICreateSpeechRequest) -> None: + if request.task_type is None: + return + if request.task_type in ("CustomVoice", "Base"): + logger.warning_once( + "AuK `task_type` is a compatibility shortcut; prefer a complete `instructions` prompt. " + 'CustomVoice applies `Say the following: ""`; Base applies ' + '`Say the following with the same voice: ""`.' + ) + if request.instructions and request.instructions.strip(): + logger.warning( + "AuK request contains both `task_type` and a complete `instructions` prompt; " + "using `instructions` without applying another task template." + ) + request.input = "" + request.task_type = None + return + instruction = request.input.strip() + if instruction.startswith(("Say the following:", "Say the following with the same voice:")): + logger.warning( + "AuK request `input` already contains a complete task instruction; " + "using it without applying another task template." + ) + request.instructions = instruction + request.input = "" + request.task_type = None + return + if not instruction: + return + quoted_instruction = json.dumps(instruction, ensure_ascii=False) + if request.task_type == "Base": + request.instructions = f"Say the following with the same voice: {quoted_instruction}" + elif request.task_type == "CustomVoice": + request.instructions = f"Say the following: {quoted_instruction}" + else: + return + request.input = "" + # task_type is a compatibility shortcut for AuK, not a model-side + # parameter. Consuming it makes normalization idempotent when the + # shared batch path validates and later prepares the same request. + request.task_type = None + + def validate(self, request: OpenAICreateSpeechRequest) -> str | None: + if not request.input.strip() and not (request.instructions and request.instructions.strip()): + return "AuK requires input text or a complete instructions prompt" + if request.duration_seconds is not None: + if not math.isfinite(request.duration_seconds) or request.duration_seconds <= 0: + return "AuK requires duration_seconds to be a positive finite value" + if request.duration_seconds is None and request.ref_audio is None: + return "AuK requires duration_seconds for requests without ref_audio" + if request.instructions is not None: + if len(request.instructions) > self.ctx.server._max_instructions_length: + return f"instructions exceeds max length {self.ctx.server._max_instructions_length}" + if request.ref_audio is not None: + if not isinstance(request.ref_audio, str): + return "AuK supports one ref_audio string per request" + error = self.ctx.server._validate_ref_audio_format(request.ref_audio) + if error: + return error + extra = request.extra_params or {} + for key in extra: + if key not in {"num_inference_steps", "guidance_scale", "sway", "t_grid", "vae_sample"}: + return f"AuK does not support extra_params.{key}" + t_grid = extra.get("t_grid") + if t_grid is not None: + t_grid_error = "AuK extra_params.t_grid must contain at least two finite, strictly increasing values" + if not isinstance(t_grid, (list, tuple)) or len(t_grid) < 2: + return t_grid_error + try: + values = [float(value) for value in t_grid] + except (TypeError, ValueError): + return t_grid_error + if not all(math.isfinite(value) for value in values) or any( + right <= left for left, right in zip(values, values[1:]) + ): + return t_grid_error + return None + + async def build( + self, + request: OpenAICreateSpeechRequest, + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + del sampling_params_list, has_inline_ref_audio + reference: tuple[np.ndarray, int] | None = None + if request.ref_audio is not None: + samples, sample_rate, _ = await self.ctx.server._resolve_ref_audio(request.ref_audio) + reference = (np.asarray(samples, dtype=np.float32), int(sample_rate)) + + extra = request.extra_params or {} + # The caller supplies the complete cookbook prompt in instructions; + # this adapter deliberately does not invent or compose a template. + if request.instructions and request.instructions.strip(): + if request.input.strip(): + logger.warning( + "AuK request contains both non-empty `input` and `instructions`; " + "using the complete `instructions` prompt and ignoring `input`." + ) + instruction = request.instructions.strip() + else: + logger.warning( + "AuK request only provides `input`; treating it as the complete AuK instruction. " + 'Prefer `input=""` with the instruction in `instructions`.' + ) + instruction = request.input.strip() + prompt = auk_prompt( + instruction, + reference, + gen_seconds=request.duration_seconds, + sway=float(extra.get("sway", DEFAULT_SWAY)), + t_grid=extra.get("t_grid"), + vae_sample=extra.get("vae_sample", False), + ) + return PreparedRequest(prompt=prompt, tts_params={}, model_type=self.name) + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: OpenAICreateSpeechRequest, + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + if len(sampling_params_list) < 2: + return sampling_params_list + updated = copy.deepcopy(sampling_params_list) + stage1 = updated[1] + extra = request.extra_params or {} + if "num_inference_steps" in extra: + stage1.num_inference_steps = int(extra["num_inference_steps"]) + if "guidance_scale" in extra: + stage1.guidance_scale = float(extra["guidance_scale"]) + if request.seed is not None: + stage1.seed = int(request.seed) + return updated + + +__all__ = ["AuKAdapter"]