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/models/supported_models.md b/docs/models/supported_models.md index de5d5e346..6b0952eac 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -97,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/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index da00ea44d..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. 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/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index a4cf2fdb9..b1221de63 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,6 +254,138 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \ --- +## Audio8 TTS Preview + +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 +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 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/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/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/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_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/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 6a6a9b7ed..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 @@ -2123,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/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/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/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/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 5586b956e..b711f0cfb 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -39,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, @@ -172,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/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/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 75a8a435e..88d9f090d 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -712,6 +712,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if self._tts_model_type in ( "cosyvoice3", "fish_tts", + "audio8_tts", "omnivoice", "moss_tts_nano", "glm_tts", @@ -721,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", 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/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/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/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/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/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", +]