diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 1f4093828..a60eb6670 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -71,6 +71,7 @@ th { | `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — | +| `BreezeForConditionalGeneration` | Breeze-TTS-2 (TTS, voice clone, voice direction) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) | | `DotsTTSForConditionalGeneration` | dots.tts | `dots-studio/dots.tts-soar` | ✅︎ | | | | — | | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — | | `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — | diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index f0658f278..38cca54d6 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -7,7 +7,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. - **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz. - **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz. -See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano. +See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, MOSS-TTS-Nano, and Breeze-TTS-2. !!! tip "Deployment recipes" TTS deployment recipes are published at @@ -783,6 +783,12 @@ hidden-state dtype before the output projection. Consequently, a float32 stage c | ------- | ------------- | | `OpenMOSS-Team/MOSS-TTS-Nano` | Voice cloning only. Requires `ref_audio` (or an uploaded `voice`); no built-in voice presets. `ref_text` is accepted but ignored — upstream's `voice_clone` mode does not consume a transcript. | +### Breeze-TTS-2 + +| Model | Description | +| ------- | ------------- | +| `BreezeBlue/Breeze-TTS-2` | Two-stage AR TTS (T5Gemma2 + Qwen3 talker with a depth decoder, bundled Qwen3-TTS codec) at 24 kHz. Four modes are selected from the request fields: plain (`input` + speaker tag `voice`, `S0`..`S9`), voice design (`instructions`), voice clone (`ref_audio` + `ref_text`, exactly one clip), and voice direction (reference + `instructions`). Greedy decoding only: `sample_rate` must be `24000`, `speed` must be `1.0`, and `guidance_scale`/`cfg_scale` other than `1.0`, `negative_prompt`, `temperature`/`top_p`/`top_k` overrides, `language`, and `speaker_embedding` are rejected. Streaming returns PCM `speech.audio.delta` events. See [`recipes/BreezeBlue/Breeze-TTS-2.md`](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md). | + ## Error Responses ### 400 Bad Request diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 89ff2dbce..954b96bd3 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -16,6 +16,7 @@ list of supported architectures across all modalities, see |---|---|---|---|---|---|---| | AuK | `tencent/AuK`, `tencent/AuK-Flash` (assembled with `tools/prepare_auk_checkpoint.py`) | 2 (thinker encoder + DiT) | ✓ | — | instruction-driven editing, enhancement, separation; no example script (use `vllm_omni.model_extras.auk`) | 24 kHz | | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz | +| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | 2 (talker + codec) | ✓ | ✓ (async chunk) | voice design / voice direction (`--instruction`) | 24 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 | @@ -662,6 +663,37 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). ### Notes + - `--num-prompts N` replicates the prompt for performance measurement. - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`. - Run `--help` for the full argument surface. + +--- + +## Breeze-TTS-2 + +Two-stage AR TTS (T5Gemma2 + Qwen3 talker with a depth decoder → bundled Qwen3-TTS codec) at 24 kHz mono. The prompt builder picks one of four templates automatically: plain, voice design (`--instruction`), clone (`--ref-audio` + `--ref-text`), or voice direction (reference + `--instruction`). + +### Quick start + +```bash +python examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py \ + --model BreezeBlue/Breeze-TTS-2 \ + --text "Hello, this is Breeze TTS 2 running on vLLM Omni." +``` + +### Voice direction + +```bash +python examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py \ + --text "We need to discuss what happened last night." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript of the reference audio." \ + --instruction "Speak slowly with a restrained, serious tone." +``` + +### Notes + +- Output: 24 kHz mono WAV; `--max-new-tokens` caps generated frames (80 ms each). +- Speaker tags `--voice S0`..`S9` (default `S0`) select the timbre in non-reference modes; with `--ref-audio` the timbre comes from the reference clip. +- Deploy config: `vllm_omni/deploy/breeze_tts_2.yaml` (auto-loaded by HF `model_type`); see [`recipes/BreezeBlue/Breeze-TTS-2.md`](../../../recipes/BreezeBlue/Breeze-TTS-2.md) for the online-serving recipe. diff --git a/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py b/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py new file mode 100644 index 000000000..942ead792 --- /dev/null +++ b/examples/offline_inference/text_to_speech/breeze_tts_2/end2end.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Offline Breeze-TTS-2 inference example (two-stage talker + codec pipeline). + +Uses the async-chunk deploy config (vllm_omni/deploy/breeze_tts_2.yaml). +Reference checkpoint: BreezeBlue/Breeze-TTS-2. +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path + +import soundfile as sf +import torch +from transformers import AutoConfig + +from vllm_omni import Omni +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) +from vllm_omni.utils.tracking_parser import TrackingArgumentParser + +REPO_ROOT = Path(__file__).resolve().parents[4] +SAMPLE_RATE = 24_000 + + +def parse_args(): + parser = TrackingArgumentParser(description="Offline Breeze-TTS-2 inference") + parser.add_argument( + "--model", + type=str, + default="BreezeBlue/Breeze-TTS-2", + help="Breeze-TTS-2 model path or HuggingFace repo ID.", + ) + parser.add_argument( + "--text", + type=str, + default="Hello, this is Breeze TTS 2 running on vLLM Omni.", + help="Text to synthesize.", + ) + parser.add_argument( + "--voice", + type=str, + default="S0", + help="Speaker tag such as S0..S9.", + ) + parser.add_argument( + "--instruction", + type=str, + default=None, + help="Optional style instruction (voice design / voice direction).", + ) + parser.add_argument( + "--ref-audio", + type=str, + default=None, + help="Optional reference wav path for cloning (requires --ref-text).", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Exact transcript of the reference audio.", + ) + parser.add_argument( + "--max-new-tokens", + type=int, + default=None, + help="Cap on generated codec frames (one frame is 80 ms of audio).", + ) + parser.add_argument( + "--output-dir", + type=str, + default="output_audio", + help="Directory for output WAV files.", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override the deploy config path. If unset, auto-loads " + "vllm_omni/deploy/breeze_tts_2.yaml based on the HF model_type.", + ) + return parser.parse_args() + + +def resolve_template(args) -> str: + if args.ref_audio is not None: + return "ref_edit_tata" if args.instruction else "ref_clone_tata" + return "tts_instruction" if args.instruction else "tts_plain" + + +def reference_payload(ref_audio_path: str | None) -> dict: + """Load the reference clip into the payload keys the prompt builder reads. + + ``BreezeTTS2PromptBuilder`` consumes ``ref_audio`` (a waveform) plus + ``ref_audio_sample_rate``; a bare path field is not read. + """ + if ref_audio_path is None: + return {} + waveform, sample_rate = sf.read(ref_audio_path, dtype="float32", always_2d=True) + # Down-mix to the single 1-D clip the reference encoder expects. + waveform = waveform.mean(axis=1) + return {"ref_audio": waveform, "ref_audio_sample_rate": int(sample_rate)} + + +def build_prompt(args) -> dict: + config = AutoConfig.from_pretrained(args.model, trust_remote_code=True) + reference_encoder = None + if args.ref_audio is not None: + from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( + BreezeReferenceAudioTokenizer, + ) + + reference_encoder = BreezeReferenceAudioTokenizer.from_pretrained( + args.model, + num_codebooks=int(getattr(config, "num_codebooks", 16)), + device_map="cpu", + ) + builder = BreezeTTS2PromptBuilder.from_pretrained( + args.model, + config, + reference_audio_encoder=reference_encoder, + ) + payload = { + "text": args.text, + "speaker": args.voice, + **reference_payload(args.ref_audio), + } + if args.instruction: + payload["instruction"] = args.instruction + if args.ref_text: + payload["ref_text"] = args.ref_text + prompt = builder.build(payload, template=resolve_template(args)) + if args.max_new_tokens is not None: + prompt["additional_information"]["breeze_max_new_frames"] = int(args.max_new_tokens) + return prompt + + +def extract_audio(multimodal_output: dict) -> torch.Tensor: + """Extract the final complete audio tensor from multimodal output. + + The output processor concatenates per-step delta tensors under + ``model_outputs``. Falls back to ``audio`` for backwards compat. + """ + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is None: + raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}") + + if isinstance(audio, list): + valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] + if not valid: + raise ValueError("Audio list is empty or all elements are None.") + return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] + + return torch.as_tensor(audio).float().cpu().reshape(-1) + + +def main(): + args = parse_args() + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + engine = Omni( + model=args.model, + deploy_config=args.deploy_config, + ) + + prompt = build_prompt(args) + + print(f"Model : {args.model}") + print(f"Text : {args.text}") + print(f"Template : {resolve_template(args)}") + print(f"Output dir : {output_dir}") + print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens") + + t_start = time.perf_counter() + outputs = engine.generate([prompt]) + elapsed = time.perf_counter() - t_start + + request_output = outputs[0] + mm = request_output.outputs[0].multimodal_output + audio = extract_audio(mm) + + duration = audio.numel() / SAMPLE_RATE + rtf = elapsed / duration if duration > 0 else float("inf") + + output_path = output_dir / "output.wav" + sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV") + + print(f"Saved : {output_path}") + print(f"Duration : {duration:.2f}s") + print(f"Inference : {elapsed:.2f}s") + print(f"RTF : {rtf:.3f}") + + +if __name__ == "__main__": + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + main() diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 5141a37c2..7ad939415 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -13,7 +13,8 @@ For the full list of supported architectures across all modalities, see ## Supported Models | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | -| --- | --- | --- | --- | --- | --- | +|---|---|---|---|---|---| +| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | speaker tags (`S0`..`S9`, default `S0`) | — | | 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) | — | ✓ | @@ -1025,3 +1026,63 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). - Voice cloning is gated upstream and may require a recent `mistral_common`. - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now. + +## Breeze-TTS-2 + +Two-stage AR TTS (T5Gemma2 + Qwen3 talker with a depth decoder → bundled Qwen3-TTS codec) at 24 kHz mono, from [BreezeBlue](https://huggingface.co/BreezeBlue/Breeze-TTS-2). Four modes are selected automatically from the request fields: plain (`input`), voice design (`instructions`), clone (`ref_audio`+`ref_text`), and voice direction (reference trio + `instructions`). + +### Launch + +```bash +vllm serve BreezeBlue/Breeze-TTS-2 --omni --port 8091 +# or: +./breeze_tts_2/run_server.sh +``` + +The deploy config at `vllm_omni/deploy/breeze_tts_2.yaml` auto-loads (async-chunk streaming, 8-frame inter-stage chunks). + +### Sending requests + +```bash +# Plain synthesis (speaker tags S0..S9, default S0) +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is Breeze TTS 2 on vLLM-Omni.", + "voice": "S0", + "response_format": "wav", + "sample_rate": 24000 + }' --output breeze_plain.wav + +# Voice direction (clone a reference, then steer the delivery) +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "We need to discuss what happened last night.", + "ref_audio": "file:///path/to/reference.wav", + "ref_text": "The exact transcript of the reference audio.", + "instructions": "Speak slowly with a restrained, serious tone.", + "response_format": "wav" + }' --output breeze_direction.wav +``` + +### Streaming PCM + +```bash +curl -N -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Streaming output from Breeze TTS 2.", + "stream": true, + "stream_format": "audio", + "response_format": "pcm", + "sample_rate": 24000 + }' --output breeze_stream.pcm +``` + +### Notes + +- Only `sample_rate=24000` is accepted; `speed` must be `1.0`. +- `guidance_scale`/`cfg_scale` must be `1.0`; `negative_prompt` is rejected (CFG companion support is a follow-up). +- `instructions` without `ref_audio` is voice design; with `ref_audio`+`ref_text` it is voice direction. +- See [`recipes/BreezeBlue/Breeze-TTS-2.md`](../../../recipes/BreezeBlue/Breeze-TTS-2.md) for the full recipe and the offline example under [`examples/offline_inference/text_to_speech/breeze_tts_2/`](../../offline_inference/text_to_speech/breeze_tts_2/). diff --git a/examples/online_serving/text_to_speech/breeze_tts_2/run_server.sh b/examples/online_serving/text_to_speech/breeze_tts_2/run_server.sh new file mode 100755 index 000000000..83ca27084 --- /dev/null +++ b/examples/online_serving/text_to_speech/breeze_tts_2/run_server.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# Launch Breeze-TTS-2 online serving with the async-chunk deploy config. + +set -euo pipefail + +MODEL="${1:-BreezeBlue/Breeze-TTS-2}" +PORT="${PORT:-8091}" + +exec vllm-omni serve "${MODEL}" \ + --deploy-config vllm_omni/deploy/breeze_tts_2.yaml \ + --omni --port "${PORT}" diff --git a/recipes/BreezeBlue/Breeze-TTS-2.md b/recipes/BreezeBlue/Breeze-TTS-2.md new file mode 100644 index 000000000..50e8858af --- /dev/null +++ b/recipes/BreezeBlue/Breeze-TTS-2.md @@ -0,0 +1,193 @@ +# Breeze-TTS-2 + +> Text-to-speech serving (plain / voice design / voice clone / voice direction) + +## Summary + +- Vendor: BreezeBlue +- Model: `BreezeBlue/Breeze-TTS-2` (`BreezeForConditionalGeneration`) +- Task: Text-to-speech with speaker tags, natural-language style instructions, single-reference voice cloning, and reference+instruction "voice direction" +- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API (streaming and non-streaming); offline `Omni` example +- Hardware: 1x NVIDIA L20 48GB (CUDA) +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want a known-good starting point for serving +Breeze-TTS-2 with vLLM-Omni. Breeze-TTS-2 is a two-stage AR TTS pipeline: + +| Stage | Components | Output | +| ----- | ---------- | ------ | +| 0 (talker, `LLM_AR`) | T5Gemma2 text encoder + Qwen3 backbone + 2052-way codebook-0 head + 12-layer depth decoder | 16-codebook codec frames | +| 1 (codec, `LLM_GENERATION`) | Bundled Qwen3-TTS audio tokenizer decoder (Mimi fallback) | 24 kHz mono waveform | + +Four prompt modes are selected automatically from the request fields: + +| Mode | Trigger | Fields | +| ---- | ------- | ------ | +| `tts_plain` | only `input` (default `voice=S0`) | `input`, `voice` | +| `tts_instruction` (voice design) | `input` + `instructions` | style/emotion/pace directives, not spoken | +| `ref_clone_tata` (clone) | `ref_audio` + `ref_text` + `input` | reference clip and its exact transcript | +| `ref_edit_tata` (voice direction) | reference trio + `instructions` | keep the reference timbre, direct the delivery | + +Current scope: greedy sampling with `cfg_scale=1.0`; CFG ≠ 1.0 and +`negative_prompt` are rejected with explicit client errors until companion +request support lands. + +## Supported model contract + +| Item | Value | +| ---- | ----- | +| Tasks | Text-to-speech through `/v1/audio/speech` (online) and the offline `Omni` example | +| Modes | plain, voice design (`instructions`), voice clone (`ref_audio` + `ref_text`), voice direction (reference + `instructions`) | +| Languages | English and Chinese (model card) | +| Reference audio | exactly one clip per request with its exact transcript; encoded by the bundled Qwen3-TTS tokenizer to 16 codebooks at 12.5 Hz | +| Output | 24 kHz mono; `wav` or `pcm`; streaming SSE `speech.audio.delta` | +| Length | prompt bounded by stage-0 `max_model_len=4096`; synthesis bounded by `max_new_tokens` codec frames (default 2048, about 164 s) | +| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0`; other guidance values and `negative_prompt` are rejected | +| Deployment profile | `vllm_omni/deploy/breeze_tts_2.yaml`: two stages on one GPU, async-chunk streaming with 8-frame codec chunks, stage 0 `gpu_memory_utilization=0.80`, stage 1 `0.15` | + +## References + +- Issue: [#6656 [New Model]: Breeze TTS 2](https://github.com/vllm-project/vllm-omni/issues/6656) +- Related examples under `examples/`: + [`examples/online_serving/text_to_speech/breeze_tts_2/`](../../examples/online_serving/text_to_speech/breeze_tts_2/), + [`examples/offline_inference/text_to_speech/breeze_tts_2/`](../../examples/offline_inference/text_to_speech/breeze_tts_2/) + +## Hardware + +- Accelerator model and per-device memory: 1x NVIDIA L20 48GB +- Number of devices: 1 (both stages are placed on device 0) +- Device interconnect: not applicable (single device) +- Host memory: no special requirement; per-request codec state is small +- Qualification scope: functional serving of all four prompt modes, streaming PCM, and 4 concurrent mixed-mode requests (greedy, `cfg_scale=1.0`). No other accelerator is qualified by this recipe. + +## Software environment + +- OS: Linux +- Python: 3.12.12 +- Driver / runtime: NVIDIA driver 580.82.07, CUDA 13.0 (PyTorch 2.13.0+cu130) +- vLLM version: 0.28.0 +- vLLM-Omni version or commit: the commit that introduced this recipe (branch based on `main` @ `67ef0b64`) +- Transformers: 5.14.1 + +## Command + +Start the server from the repository root. The deploy config +(`vllm_omni/deploy/breeze_tts_2.yaml`, async chunk streaming by default) +auto-loads from the checkpoint's `model_type`; pass it explicitly if you +customized it: + +```bash +vllm serve BreezeBlue/Breeze-TTS-2 \ + --deploy-config vllm_omni/deploy/breeze_tts_2.yaml \ + --omni --port 8091 +``` + +## Verification + +Plain synthesis (speaker tag `S0`..`S9`, default `S0`): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "BreezeBlue/Breeze-TTS-2", + "input": "Hello, this is Breeze TTS 2 running on vLLM-Omni.", + "voice": "S0", + "response_format": "wav", + "sample_rate": 24000 + }' --output breeze_plain.wav +``` + +Voice design (instruction only, no reference audio): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "[laughs] Welcome to tonight's story time.", + "instructions": "A warm young woman, clear voice, lively delivery.", + "response_format": "wav" + }' --output breeze_design.wav +``` + +Voice cloning (`ref_audio` and `ref_text` must be provided together): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "This is new text spoken in the cloned voice.", + "ref_audio": "file:///path/to/reference.wav", + "ref_text": "The exact transcript of the reference audio.", + "response_format": "wav" + }' --output breeze_clone.wav +``` + +Voice direction (reference + instruction): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "We need to discuss what happened last night.", + "ref_audio": "file:///path/to/reference.wav", + "ref_text": "The exact transcript of the reference audio.", + "instructions": "Speak slowly with a restrained, serious tone.", + "response_format": "wav" + }' --output breeze_direction.wav +``` + +Streaming PCM (SSE `speech.audio.delta` events): + +```bash +curl -N -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Streaming output from Breeze TTS 2.", + "stream": true, + "stream_format": "audio", + "response_format": "pcm", + "sample_rate": 24000 + }' --output breeze_stream.pcm +``` + +Expected results: finite, non-silent 24 kHz mono WAV files; streaming requests +return incremental PCM deltas. On 1x L20 (greedy, single request after +warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time. + +## Notes + +- **Sample rate**: only `24000` is accepted; other values fail validation + before inference. +- **Speed / seeds / multi-reference**: `speed` must be `1.0`; + `task_type=VoiceDesign`, `ref_audio_2`, and speaker embeddings are not + supported yet. +- **CFG**: `guidance_scale`/`cfg_scale` must be `1.0`. Non-1.0 values and + `negative_prompt` return a client error. +- **Length budget**: prompt length is bounded by stage 0 `max_model_len=4096`; + synthesis length by `max_new_tokens` (default 2048 frames ≈ 164 s at the + 12.5 Hz codec frame rate) — set `max_new_tokens` for shorter caps. +- **Checkpoint layout**: the checkpoint must contain the `audio_tokenizer/` + subdirectory (bundled Qwen3-TTS codec) for the default streaming path; the + Mimi fallback is used only when that directory is absent and does not + support async-chunk streaming. +- **License**: the reference inference code is Apache 2.0; the checkpoints are + distributed under the separate BreezeBlue Research and Non-Commercial + License. Commercial use requires written authorization from BreezeBlue. + +## Supported features + +| Feature | Status | Notes | +| ------- | ------ | ----- | +| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks, SSE `speech.audio.delta` PCM | +| Voice cloning | ✓ | one reference clip plus transcript per request | +| Voice design / voice direction | ✓ | `instructions` field, with or without a reference | +| Classifier-free guidance | ✗ | `cfg_scale` must be `1.0`; `negative_prompt` rejected (follow-up) | +| Non-greedy sampling | ✗ | greedy decoding only | +| Prefix caching | ✗ | disabled in the deploy config | +| Tensor / pipeline parallelism | untested | single-GPU profile only | +| Quantization | ✗ | not supported in this release | +| CUDA graphs | not tuned | depth decoder and stage-1 codec run eagerly | +| Platforms | CUDA | qualified on 1x NVIDIA L20 | diff --git a/recipes/README.md b/recipes/README.md index d9f4b6822..86eba6d87 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -45,6 +45,7 @@ recipes/ | [`black-forest-labs/FLUX.1-schnell.md`](./black-forest-labs/FLUX.1-schnell.md) | Offline text-to-image generation with CPU offload | 1x NVIDIA GeForce RTX 5090 32GB | | [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ | | [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB | +| [`BreezeBlue/Breeze-TTS-2.md`](./BreezeBlue/Breeze-TTS-2.md) | Online + offline TTS (plain / voice design / voice clone / voice direction, streaming PCM) | 1x NVIDIA L20 48GB | | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB | | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB | | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index d958bbf18..ce608dca3 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 27, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 63, + EnvironmentVariableCategory.MODEL_SPECIFIC: 64, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -181,7 +181,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): ModelEnvironmentVariableDisposition.PROMOTE: 38, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.EXTERNAL: 0, - ModelEnvironmentVariableDisposition.INTERNALIZE: 14, + ModelEnvironmentVariableDisposition.INTERNALIZE: 15, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5, } diff --git a/tests/e2e/online_serving/test_breeze_tts_2.py b/tests/e2e/online_serving/test_breeze_tts_2.py new file mode 100644 index 000000000..29cd43b68 --- /dev/null +++ b/tests/e2e/online_serving/test_breeze_tts_2.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +""" +E2E Online tests for Breeze-TTS-2 with text input and audio output. + +These tests verify the /v1/audio/speech endpoint works correctly with +actual model inference, not mocks. Breeze-TTS-2 plain mode needs only +``input`` and a speaker tag (``voice`` defaults to ``S0``); reference +cloning requires ``ref_audio`` + ``ref_text`` and is not exercised here. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "BreezeBlue/Breeze-TTS-2" +DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0 +MAX_CONCURRENT = 4 + +# ~0.5 s of 24 kHz mono PCM_16 in WAV (~24k payload + header). +_MIN_AUDIO_BYTES = 40_000 + + +def get_prompt(prompt_type="text"): + """English prompt for zero-shot TTS.""" + prompts = { + "text": "The weather is nice today, perfect for a walk in the park.", + } + return prompts.get(prompt_type, prompts["text"]) + + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("breeze_tts_2.yaml"), + server_args=["--trust-remote-code", "--disable-log-stats"], + ), + id="breeze_tts_2", + ) +] + + +@pytest.mark.slow +@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, online_client) -> None: + """ + Test plain text-to-audio via OpenAI API. + Deploy Setting: breeze_tts_2.yaml + Input Modal: text + Output Modal: audio + Input Setting: stream=False + Datasets: few requests (max_num_seqs=8) + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "voice": "S0", + "sample_rate": 24000, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + online_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) + + +@pytest.mark.slow +@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_002(omni_server, online_client) -> None: + """ + Test streaming text-to-audio via OpenAI API (async-chunk SSE deltas). + Deploy Setting: breeze_tts_2.yaml + Input Modal: text + Output Modal: audio + Input Setting: stream=True + Datasets: single request + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": True, + "stream_format": "audio", + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "voice": "S0", + "sample_rate": 24000, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + online_client.send_audio_speech_request(request_config) diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index a17a6ccbd..edd2b4cd4 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -65,6 +65,7 @@ EXPECTED_MODEL_TYPES = { "higgs_audio_v2", "higgs_audio_v3", "glm_tts", + "breeze_tts_2", "step_audio2", "indextts2", "indextts2_5", diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index d54ce36d5..6c7671659 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """TTS model detection: registry-driven stage -> model-type resolution. ``serving_speech.py`` used to carry a hand-written 20-branch ladder mapping @@ -138,6 +139,7 @@ _PIPELINE_STAGES = [ "audio_tokenizer", "audio_vae", "aura", + "breeze_tts_2", "code2wav", "cosyvoice3_code2wav", "cosyvoice3_talker", @@ -184,6 +186,7 @@ _STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"] _ARCHS = [ None, "VoxCPM2TalkerForConditionalGeneration", + "BreezeForConditionalGeneration", "MingTTSForConditionalGeneration", "CovoAudioForConditionalGeneration", "MyCovoAudioThing", diff --git a/tests/model_executor/models/breeze_tts_2/__init__.py b/tests/model_executor/models/breeze_tts_2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py b/tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py new file mode 100644 index 000000000..75dcbb1fb --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( + BreezeReferenceAudioTokenizer, + resolve_audio_tokenizer_path, +) + + +def _patch_hf_api(monkeypatch, snapshot_download): + """Point the resolver's tagged HF helper at a fake API object.""" + from vllm.transformers_utils import repo_utils + + monkeypatch.setattr(repo_utils, "hf_api", lambda: SimpleNamespace(snapshot_download=snapshot_download)) + + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Tokenizer: + def encode(self, audio, sr=None): + assert isinstance(audio, np.ndarray) + assert sr == 16000 + return {"audio_codes": [torch.arange(32).reshape(16, 2)]} + + +def test_reference_audio_tokenizer_normalizes_codebook_major_output(): + adapter = BreezeReferenceAudioTokenizer(_Tokenizer()) + + codes = adapter.encode(np.zeros(32, dtype=np.float32), 16000) + + assert tuple(codes.shape) == (2, 16) + assert codes.dtype == torch.int16 + assert codes.device.type == "cpu" + + +def test_reference_audio_tokenizer_rejects_waveform_without_sample_rate(): + adapter = BreezeReferenceAudioTokenizer(_Tokenizer()) + + try: + adapter.encode(np.zeros(32, dtype=np.float32)) + except ValueError as exc: + assert "sample_rate is required" in str(exc) + else: + raise AssertionError("missing sample rate should be rejected") + + +def test_resolve_prefers_local_directory_without_hub_access(tmp_path, monkeypatch): + bundled = tmp_path / "audio_tokenizer" + bundled.mkdir() + + def _no_hub(*_args, **_kwargs): + raise AssertionError("local directories must not touch the hub") + + _patch_hf_api(monkeypatch, _no_hub) + + assert resolve_audio_tokenizer_path(str(tmp_path)) == bundled + + +def test_resolve_local_directory_without_bundled_tokenizer_returns_none(tmp_path, monkeypatch): + def _no_hub(*_args, **_kwargs): + raise AssertionError("local directories must not touch the hub") + + _patch_hf_api(monkeypatch, _no_hub) + + assert resolve_audio_tokenizer_path(str(tmp_path)) is None + + +def test_resolve_repo_id_uses_hf_snapshot(tmp_path, monkeypatch): + snapshot_root = tmp_path / "snapshot" + (snapshot_root / "audio_tokenizer").mkdir(parents=True) + calls = [] + + def _fake_snapshot_download(repo_id, **kwargs): + calls.append((repo_id, kwargs)) + if kwargs.get("local_files_only"): + raise FileNotFoundError("cold cache") + return str(snapshot_root) + + _patch_hf_api(monkeypatch, _fake_snapshot_download) + + resolved = resolve_audio_tokenizer_path("BreezeBlue/Breeze-TTS-2") + + assert resolved == snapshot_root / "audio_tokenizer" + # The cold-cache probe runs first; the online fetch carries the pattern. + assert calls[0] == ( + "BreezeBlue/Breeze-TTS-2", + {"allow_patterns": ["audio_tokenizer/*"], "local_files_only": True}, + ) + assert calls[1] == ("BreezeBlue/Breeze-TTS-2", {"allow_patterns": ["audio_tokenizer/*"]}) + + +def test_resolve_repo_id_prefers_complete_cached_snapshot(tmp_path, monkeypatch): + snapshot_root = tmp_path / "snapshot" + (snapshot_root / "audio_tokenizer").mkdir(parents=True) + calls = [] + + def _fake_snapshot_download(repo_id, **kwargs): + calls.append((repo_id, kwargs)) + assert kwargs.get("local_files_only"), "cached snapshot must not hit the network" + return str(snapshot_root) + + _patch_hf_api(monkeypatch, _fake_snapshot_download) + + resolved = resolve_audio_tokenizer_path("BreezeBlue/Breeze-TTS-2") + + assert resolved == snapshot_root / "audio_tokenizer" + assert len(calls) == 1 + + +def test_resolve_repo_id_without_bundled_tokenizer_returns_none(tmp_path, monkeypatch): + snapshot_root = tmp_path / "snapshot" + snapshot_root.mkdir() + + def _fake_snapshot_download(repo_id, **kwargs): + if kwargs.get("local_files_only"): + raise FileNotFoundError("cold cache") + return str(snapshot_root) + + _patch_hf_api(monkeypatch, _fake_snapshot_download) + + assert resolve_audio_tokenizer_path("BreezeBlue/Breeze-TTS-2") is None diff --git a/tests/model_executor/models/breeze_tts_2/test_batching.py b/tests/model_executor/models/breeze_tts_2/test_batching.py new file mode 100644 index 000000000..41fc0b9ba --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_batching.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from typing import Any + +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_talker import ( + BreezeTTS2TalkerForGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _RequestHead: + def __call__(self, hidden: torch.Tensor) -> torch.Tensor: + # Hidden row 0 selects codec id 2; row 1 selects codec id 3. + code0 = 2 + int(hidden.reshape(-1)[0].item()) + logits = torch.zeros((hidden.shape[0], 9), dtype=hidden.dtype) + logits[:, code0] = 1.0 + return logits + + +def _model() -> BreezeTTS2TalkerForGeneration: + model = object.__new__(BreezeTTS2TalkerForGeneration) + model.num_codebooks = 4 + model.codebook_vocab_size = 8 + model.codebook_size = 8 + model.codebook_eos_token_id = 0 + model.audio_token_id = 100 + model.audio_eos_token_id = 101 + model.lm_head = _RequestHead() + model.logits_processor = lambda head, hidden: head(hidden) + model._codec_disallowed_mask = None + model._batch_state = None + model._batch_state_spans = None + model._async_chunk = True + model._golden_dump_dir = None + model._generate_depth_codes = lambda _hidden, code0: _frame(code0).unsqueeze(0) + return model + + +def _frame(code0: int) -> torch.Tensor: + return torch.tensor([code0, code0 + 1, code0 + 2, code0 + 3], dtype=torch.long) + + +def _step(model: BreezeTTS2TalkerForGeneration, infos: list[dict[str, Any]], hidden: torch.Tensor): + return model.make_omni_output( + hidden, + model_intermediate_buffer=infos, + request_token_spans=[(0, 1), (1, 2)], + ) + + +def test_two_requests_accumulate_and_emit_independent_frames(): + model = _model() + infos = [ + { + "breeze_generated_frames": 1, + "breeze_max_new_frames": 2, + "breeze_audio_codes": _frame(2).reshape(1, 4), + }, + {"breeze_generated_frames": 0, "breeze_max_new_frames": 3}, + ] + hidden = torch.tensor([[0.0], [1.0]], dtype=torch.float32) + + first = _step(model, infos, hidden) + assert [tuple(item.shape) for item in first.multimodal_outputs["codes"]["audio"]] == [(1, 4), (1, 4)] + assert torch.equal(first.multimodal_outputs["codes"]["audio"][0][0], _frame(2)) + assert torch.equal(first.multimodal_outputs["codes"]["audio"][1][0], _frame(3)) + + second = _step(model, infos, hidden) + # Request 0 reaches its two-frame budget and emits only a finish marker; + # request 1 still has budget and emits its next tail frame. + assert second.multimodal_outputs["codes"]["audio"][0].numel() == 0 + assert torch.equal(second.multimodal_outputs["codes"]["audio"][1][0], _frame(3)) + # Async streaming only maintains the per-frame tail; the cumulative + # buffer is deliberately not grown (request 0 keeps its pre-seeded + # entry untouched, request 1 never gains one). + assert infos[0]["breeze_audio_codes"].shape == (1, 4) + assert "breeze_audio_codes" not in infos[1] + assert infos[0]["breeze_force_eos"] is True + assert infos[1]["breeze_force_eos"] is False + + +def test_finished_requests_only_clear_model_owned_batch_mapping(): + model = _model() + model._batch_state = [{"id": "a"}, {"id": "b"}] + model._batch_state_spans = [(0, 1), (1, 2)] + + model.on_requests_finished(["a", "b"]) + + assert model._batch_state is None + assert model._batch_state_spans is None diff --git a/tests/model_executor/models/breeze_tts_2/test_codec_loading.py b/tests/model_executor/models/breeze_tts_2/test_codec_loading.py new file mode 100644 index 000000000..f558d067a --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_codec_loading.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from typing import Any + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_codec import ( + BreezeTTS2MimiCodec, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _BundledTokenizer: + def __init__(self) -> None: + self.model = nn.Linear(2, 2) + + def named_parameters(self): + for name, parameter in self.model.named_parameters(): + yield f"model.{name}", parameter + + def decode(self, _encoded: Any) -> tuple[list[torch.Tensor], int]: + return [torch.zeros(4)], 24_000 + + +def test_bundled_qwen_tokenizer_consumes_unused_root_fallback_weights(tmp_path): + bundled = tmp_path / "audio_tokenizer" + bundled.mkdir() + codec = object.__new__(BreezeTTS2MimiCodec) + codec.model_path = str(tmp_path) + codec.vllm_config = None + # load_weights reuses the path resolved once in __init__. + codec._tokenizer_path = bundled + codec._codec = None + codec._audio_tokenizer = _BundledTokenizer() + codec._loaded_local_weights = False + + weights = [ + ("backbone_model.layers.0.norm.weight", torch.ones(2)), + ("codec_model.decoder.layers.0.conv.weight", torch.ones(2)), + ("lm_head.weight", torch.ones(2, 2)), + ] + + loaded = codec.load_weights(iter(weights)) + + assert "_audio_tokenizer.model.weight" in loaded + assert "_audio_tokenizer.model.bias" in loaded + # The bundled tokenizer was preloaded, so load_weights must not try to copy + # the root Mimi fallback into the nonexistent ``self._codec`` module. + assert codec._codec is None + + +def test_codec_warmup_with_single_token_returns_empty_audio(): + codec = object.__new__(BreezeTTS2MimiCodec) + codec._num_codebooks = 4 + codec._codebook_size = 8 + codec._sample_rate = 24_000 + codec._audio_tokenizer = None + codec._codec = None + codec._async_chunk = False + + output = codec.forward(torch.tensor([1], dtype=torch.long), seq_token_counts=[1]) + + assert output.multimodal_outputs["model_outputs"][0].numel() == 0 + assert output.multimodal_outputs["sr"][0].item() == 24_000 diff --git a/tests/model_executor/models/breeze_tts_2/test_configuration.py b/tests/model_executor/models/breeze_tts_2/test_configuration.py new file mode 100644 index 000000000..950f420cf --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_configuration.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +from transformers.models.qwen3 import Qwen3Config + +from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( + BreezeTTS2Config, + BreezeTTS2DepthDecoderConfig, + BreezeTTS2TextEncoderConfig, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _config() -> BreezeTTS2Config: + return BreezeTTS2Config( + backbone_config={ + "model_type": "qwen3", + "hidden_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "vocab_size": 128, + "max_position_embeddings": 512, + }, + text_encoder_config={ + "model_type": "t5gemma2_text", + "hidden_size": 24, + "intermediate_size": 48, + "num_hidden_layers": 6, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + "vocab_size": 256, + }, + depth_decoder_config={ + "model_type": "breeze_depth_decoder_model", + "hidden_size": 32, + "backbone_hidden_size": 32, + "audio_embed_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "num_codebooks": 16, + "vocab_size": 2051, + "rope_scaling": { + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 0.001953125, + "high_freq_factor": 0.0078125, + "original_max_position_embeddings": 16, + }, + }, + codec_config={"model_type": "mimi", "codebook_size": 2048, "num_quantizers": 32}, + num_codebooks=16, + vocab_size=2051, + text_vocab_size=262158, + ) + + +def test_nested_configs_and_text_config_use_qwen_backbone(): + config = _config() + + assert isinstance(config.backbone_config, Qwen3Config) + assert config.get_text_config() is config.backbone_config + assert config.get_text_config().vocab_size == 128 + assert isinstance(config.text_encoder_config, BreezeTTS2TextEncoderConfig) + assert isinstance(config.depth_decoder_config, BreezeTTS2DepthDecoderConfig) + assert config.num_codebooks == config.depth_decoder_config.num_codebooks == 16 + assert config.codec_config["codebook_size"] == 2048 + + +def test_transformers_five_roundtrip_preserves_nested_configs_and_rope(): + config = _config() + + roundtrip = BreezeTTS2Config(**dict(config.to_dict())) + + assert isinstance(roundtrip.backbone_config, Qwen3Config) + assert isinstance(roundtrip.text_encoder_config, BreezeTTS2TextEncoderConfig) + assert isinstance(roundtrip.depth_decoder_config, BreezeTTS2DepthDecoderConfig) + assert roundtrip.depth_decoder_config.rope_parameters["rope_theta"] == 500_000.0 + + +def test_nested_autoconfig_registration_is_idempotent(): + import importlib + + module = importlib.import_module("vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2") + importlib.reload(module) + + assert module.BreezeTTS2Config.model_type == "breeze" + + +def test_text_encoder_defaults_to_upstream_layer_pattern(): + config = BreezeTTS2TextEncoderConfig(num_hidden_layers=6) + + assert config.layer_types == [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ] + + +def test_invalid_codebook_layout_is_rejected(): + with pytest.raises(ValueError, match="positive"): + BreezeTTS2DepthDecoderConfig(num_codebooks=0) diff --git a/tests/model_executor/models/breeze_tts_2/test_instruction.py b/tests/model_executor/models/breeze_tts_2/test_instruction.py new file mode 100644 index 000000000..3817ccb30 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_instruction.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.entrypoints.openai.tts_adapters.breeze_tts_2 import BreezeTTS2Adapter +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Tokenizer: + unk_token_id = -1 + + def convert_tokens_to_ids(self, token): + return {"<|AUDIO|>": 262144, "<|audio_eos|>": 262145}.get(token, self.unk_token_id) + + def __call__(self, texts, *, add_special_tokens=True, padding=False): + if isinstance(texts, str): + texts = [texts] + ids = [[1000 + ord(char) % 1000 for char in text] for text in texts] + if add_special_tokens: + ids = [[2, *row, 1] for row in ids] + return {"input_ids": ids} + + def decode(self, ids, *, skip_special_tokens=False): + del skip_special_tokens + return "".join(chr(int(token_id) - 1000) for token_id in ids if token_id not in (1, 2)) + + +def test_instruction_prompt_remains_one_independent_text_segment(): + builder = BreezeTTS2PromptBuilder( + _Tokenizer(), + SimpleNamespace( + audio_token_id=262144, + audio_eos_token_id=262145, + pad_token_id=0, + backbone_config=SimpleNamespace(vocab_size=100000), + num_codebooks=16, + codec_config={"codebook_size": 2048}, + ), + ) + + prompt = builder.build( + { + "template": "tts_instruction", + "speaker": "S2", + "instruction": "Use a bright and friendly delivery.", + "text": "Hello from Breeze.", + } + ) + + info = prompt["additional_information"] + assert info["template"] == "tts_instruction" + assert "input_values" not in info + assert bool(info["text_ids_mask"].all()) + assert info["text_ids_len"].tolist() == [len(info["prompt_ids"])] + + +def _request(extra_params=None, **overrides): + base = dict( + input="Hello.", + ref_text=None, + ref_audio=None, + task_type=None, + speed=1.0, + max_new_tokens=None, + language=None, + speaker_embedding=None, + x_vector_only_mode=None, + extra_params=extra_params or {}, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +def _adapter() -> BreezeTTS2Adapter: + adapter = object.__new__(BreezeTTS2Adapter) + adapter.ctx = SimpleNamespace(server=SimpleNamespace()) + return adapter + + +def test_cfg_scale_one_is_explicitly_supported_without_companion_request(): + adapter = _adapter() + + assert adapter.validate(_request({"guidance_scale": 1.0})) is None + assert adapter.validate(_request({"cfg_scale": 1.0})) is None + + +def test_non_unit_cfg_is_rejected_until_negative_branch_exists(): + adapter = _adapter() + + assert "guidance_scale=1.0" in adapter.validate(_request({"guidance_scale": 2.0})) + assert "guidance_scale=1.0" in adapter.validate(_request({"cfg_scale": 0.5})) + assert "negative_prompt" in adapter.validate(_request({"negative_prompt": "noisy"})) + + +def test_unsupported_conditioning_fields_are_rejected_not_ignored(): + adapter = _adapter() + + assert "language" in adapter.validate(_request(language="en")) + assert "speaker_embedding" in adapter.validate(_request(speaker_embedding=[0.1, 0.2])) + assert "x_vector_only_mode" in adapter.validate(_request(x_vector_only_mode=True)) + # Absent fields must keep the request valid. + assert adapter.validate(_request()) is None + assert adapter.validate(_request(x_vector_only_mode=False)) is None + + +def test_multiple_reference_clips_are_rejected_not_silently_truncated(): + adapter = _adapter() + adapter.ctx.server._validate_ref_audio_format = lambda _uri: None + + error = adapter.validate(_request(ref_audio=["a.wav", "b.wav"], ref_text="hello")) + assert "exactly one reference clip" in error + assert "2" in error + # A single-element list stays valid and reaches format validation. + assert adapter.validate(_request(ref_audio=["a.wav"], ref_text="hello")) is None + + +def test_non_greedy_sampling_overrides_are_rejected(): + adapter = _adapter() + + # Greedy values, including the deploy-config defaults, stay valid. + assert adapter.validate(_request({"temperature": 0.0, "top_p": 1.0, "top_k": -1})) is None + assert adapter.validate(_request({"temperature": 0})) is None + # The talker ignores the sampled code0, so any non-greedy override must be + # refused instead of silently producing greedy audio. + assert "greedy" in adapter.validate(_request({"temperature": 0.7})) + assert "greedy" in adapter.validate(_request({"top_p": 0.9})) + assert "greedy" in adapter.validate(_request({"top_k": 50})) + assert "temperature must be a number" in adapter.validate(_request({"temperature": "warm"})) diff --git a/tests/model_executor/models/breeze_tts_2/test_local_modules.py b/tests/model_executor/models/breeze_tts_2/test_local_modules.py new file mode 100644 index 000000000..2ce54b09a --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_local_modules.py @@ -0,0 +1,152 @@ +# 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.breeze_tts_2.configuration_breeze_tts_2 import ( + BreezeTTS2DepthDecoderConfig, + BreezeTTS2TextEncoderConfig, +) +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_depth import ( + BreezeDepthDecoderForCausalLM, +) +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_text import ( + BreezeTTS2TextEncoder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _depth_config() -> BreezeTTS2DepthDecoderConfig: + return BreezeTTS2DepthDecoderConfig( + hidden_size=16, + backbone_hidden_size=16, + audio_embed_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + max_position_embeddings=8, + num_codebooks=4, + vocab_size=8, + rope_scaling={ + "rope_type": "llama3", + "rope_theta": 100.0, + "factor": 1.0, + "low_freq_factor": 0.5, + "high_freq_factor": 1.0, + "original_max_position_embeddings": 4, + }, + ) + + +def test_depth_decoder_position_specific_logits_and_offsets(): + config = _depth_config() + decoder = BreezeDepthDecoderForCausalLM(config) + + output = decoder( + input_ids=torch.tensor([[0, 2]]), + backbone_last_hidden_state=torch.randn(1, 16), + use_cache=False, + ) + + assert output.logits.shape == (1, 1, 8) + with torch.no_grad(): + embedding = decoder.model.embed_tokens.weight + expected_first = embedding[2] @ decoder.model.inputs_embeds_projector.weight.T + # The first hidden row is replaced by the backbone state, so compare the + # second input row (codebook 0, id 2 + 0 * vocab_size). + actual_input = decoder.model.embed_tokens(torch.tensor([[2]]))[0, 0] + actual_projected = actual_input @ decoder.model.inputs_embeds_projector.weight.T + assert torch.allclose(expected_first, actual_projected) + + +def _text_config() -> BreezeTTS2TextEncoderConfig: + return BreezeTTS2TextEncoderConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=8, + vocab_size=32, + layer_types=["sliding_attention", "full_attention"], + rope_parameters={ + "sliding_attention": {"rope_type": "default", "rope_theta": 100.0}, + "full_attention": {"rope_type": "default", "rope_theta": 100.0}, + }, + ) + + +def test_text_encoder_masks_right_padding_independently(monkeypatch): + import vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_text as text_module + + class _ParallelEmbedding(torch.nn.Embedding): + def __init__(self, num_embeddings: int, embedding_dim: int): + super().__init__(num_embeddings, embedding_dim) + + class _QKVLinear(torch.nn.Module): + def __init__( + self, + hidden_size: int, + head_size: int, + total_num_heads: int, + total_num_kv_heads: int, + **_: object, + ): + super().__init__() + output_size = (total_num_heads + 2 * total_num_kv_heads) * head_size + self.linear = torch.nn.Linear(hidden_size, output_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: + return self.linear(hidden_states), None + + class _RowLinear(torch.nn.Module): + def __init__(self, input_size: int, output_size: int, **_: object): + super().__init__() + self.linear = torch.nn.Linear(input_size, output_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: + return self.linear(hidden_states), None + + class _MergedLinear(torch.nn.Module): + def __init__(self, input_size: int, output_sizes: list[int], **_: object): + super().__init__() + self.linear = torch.nn.Linear(input_size, sum(output_sizes), bias=False) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: + return self.linear(hidden_states), None + + monkeypatch.setattr(text_module, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(text_module, "VocabParallelEmbedding", _ParallelEmbedding) + monkeypatch.setattr(text_module, "QKVParallelLinear", _QKVLinear) + monkeypatch.setattr(text_module, "RowParallelLinear", _RowLinear) + monkeypatch.setattr(text_module, "MergedColumnParallelLinear", _MergedLinear) + + config = _text_config() + encoder = BreezeTTS2TextEncoder(config).eval() + ids = torch.tensor([[3, 4]]) + padded_ids = torch.tensor([[3, 4, 0, 0]]) + mask = torch.tensor([[True, True, False, False]]) + + with torch.inference_mode(): + short = encoder(ids) + padded = encoder(padded_ids, attention_mask=mask) + + assert short.shape == (1, 2, 16) + assert padded.shape == (1, 4, 16) + assert torch.allclose(short[0, :2], padded[0, :2], atol=1e-5) + + +def test_text_encoder_weight_loader_reports_fused_target_names(): + encoder = object.__new__(BreezeTTS2TextEncoder) + parameter = torch.nn.Parameter(torch.zeros(2, 2)) + parameter.weight_loader = lambda param, tensor, shard_id: param.data.copy_(tensor) + encoder.__dict__["named_parameters"] = lambda: [("layers.0.qkv_proj.weight", parameter)] + + loaded = encoder.load_weights([("layers.0.q_proj.weight", torch.ones(2, 2))]) + + assert loaded == {"layers.0.qkv_proj.weight"} + assert torch.equal(parameter.detach(), torch.ones(2, 2)) diff --git a/tests/model_executor/models/breeze_tts_2/test_model_registration.py b/tests/model_executor/models/breeze_tts_2/test_model_registration.py new file mode 100644 index 000000000..12a2b31f3 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_model_registration.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +from vllm.model_executor.models.registry import ModelRegistry + +from vllm_omni.model_executor.models.registry import OmniModelRegistry + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_breeze_architectures_resolve_to_omni_classes(): + assert "BreezeForConditionalGeneration" in OmniModelRegistry.get_supported_archs() + assert "BreezeTTS2MimiCodec" in OmniModelRegistry.get_supported_archs() + + talker = OmniModelRegistry._try_load_model_cls("BreezeForConditionalGeneration") + codec = OmniModelRegistry._try_load_model_cls("BreezeTTS2MimiCodec") + assert talker.__name__ == "BreezeTTS2TalkerForGeneration" + assert codec.__name__ == "BreezeTTS2MimiCodec" + + +def test_plugin_registration_publishes_architectures_to_upstream_registry(): + from vllm_omni.engine.arg_utils import register_omni_models_to_vllm + + register_omni_models_to_vllm() + + assert "BreezeForConditionalGeneration" in ModelRegistry.get_supported_archs() + assert "BreezeTTS2MimiCodec" in ModelRegistry.get_supported_archs() diff --git a/tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py b/tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py new file mode 100644 index 000000000..069296942 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_pipeline_registration.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from pathlib import Path + +import pytest +import yaml + +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_breeze_pipeline_is_registered_with_talker_and_codec_stages(): + assert OMNI_PIPELINES["breeze"] is BREEZE_TTS_2_PIPELINE + assert [stage.model_stage for stage in BREEZE_TTS_2_PIPELINE.stages] == [ + "breeze_tts_2", + "breeze_tts_2_codec", + ] + assert BREEZE_TTS_2_PIPELINE.stages[1].requires_full_payload_input is True + assert BREEZE_TTS_2_PIPELINE.stages[0].async_chunk_process_next_stage_input_func.endswith( + "talker2codec_async_chunk" + ) + assert BREEZE_TTS_2_PIPELINE.stages[1].retains_state_across_chunks is True + + +def test_breeze_default_deploy_is_streaming(): + deploy_path = Path(__file__).parents[4] / "vllm_omni" / "deploy" / "breeze_tts_2.yaml" + config = yaml.safe_load(deploy_path.read_text(encoding="utf-8")) + assert config["async_chunk"] is True + assert len(config["stages"]) == 2 + # Stage 0 must prefill each prompt in one step: the talker builds a codec + # frame from the last hidden row of every scheduled request. + assert config["stages"][0]["enable_chunked_prefill"] is False + assert config["connectors"]["connector_of_shared_memory"]["extra"]["breeze_codec_chunk_frames"] == 8 diff --git a/tests/model_executor/models/breeze_tts_2/test_prompt_builder.py b/tests/model_executor/models/breeze_tts_2/test_prompt_builder.py new file mode 100644 index 000000000..3bd26fce6 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_prompt_builder.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Tokenizer: + pad_token_id = 0 + eos_token_id = 1 + unk_token_id = -1 + _special = { + "<|AUDIO|>": 262144, + "<|audio_eos|>": 262145, + } + + def convert_tokens_to_ids(self, token): + return self._special.get(token, self.unk_token_id) + + def __call__(self, texts, *, add_special_tokens=True, padding=False): + if isinstance(texts, str): + texts = [texts] + ids = [] + for text in texts: + values = [1000 + ord(char) % 1000 for char in text] + if add_special_tokens: + values = [2, *values, 1] + ids.append(values) + return {"input_ids": ids} + + def decode(self, ids, *, skip_special_tokens=False): + del skip_special_tokens + rendered = [] + for token_id in ids: + if token_id == 262144: + rendered.append("<|AUDIO|>") + elif token_id == 262145: + rendered.append("<|audio_eos|>") + elif token_id in (2, 1): + continue + else: + rendered.append(chr(int(token_id) - 1000)) + return "".join(rendered) + + +class _RenderingTokenizer(_Tokenizer): + """Tiny tokenizer that preserves Breeze placeholders after decode.""" + + def __call__(self, texts, *, add_special_tokens=True, padding=False): + if isinstance(texts, str): + texts = [texts] + ids = [] + for text in texts: + row = [] + cursor = 0 + while cursor < len(text): + if text.startswith("<|AUDIO|>", cursor): + row.append(262144) + cursor += len("<|AUDIO|>") + elif text.startswith("<|audio_eos|>", cursor): + row.append(262145) + cursor += len("<|audio_eos|>") + else: + row.append(1000 + ord(text[cursor]) % 1000) + cursor += 1 + if add_special_tokens: + row = [2, *row, 1] + ids.append(row) + return {"input_ids": ids} + + +class _ReferenceEncoder: + def __init__(self): + self.calls = [] + + def encode(self, audio, sample_rate=None): + self.calls.append((audio, sample_rate)) + return torch.arange(32, dtype=torch.int16).reshape(2, 16) + + +def _config(): + return SimpleNamespace( + audio_token_id=262144, + audio_eos_token_id=262145, + pad_token_id=0, + backbone_config=SimpleNamespace(vocab_size=100000), + num_codebooks=16, + codec_config={"codebook_size": 2048}, + ) + + +def test_builder_accepts_missing_codec_config(): + config = _config() + config.codec_config = None + + builder = BreezeTTS2PromptBuilder(_Tokenizer(), config) + + assert builder.codebook_size == 2048 + + +def test_build_instruction_prompt(): + builder = BreezeTTS2PromptBuilder(_Tokenizer(), _config()) + + prompt = builder.build({"text": "hello", "instruction": "happy", "speaker": "S0"}) + + assert prompt["prompt_token_ids"] + info = prompt["additional_information"] + assert "input_values" not in info + assert set(prompt["prompt_token_ids"]) == {0} + assert info["prompt_ids"].numel() == len(prompt["prompt_token_ids"]) + assert info["text_ids_mask"].dtype == torch.bool + assert info["text_ids_len"].tolist() == [len(prompt["prompt_token_ids"])] + assert bool(info["text_ids_mask"].all()) + + +def test_build_reference_prompt_and_reuse_encoder(): + encoder = _ReferenceEncoder() + builder = BreezeTTS2PromptBuilder(_RenderingTokenizer(), _config(), encoder) + + prompt = builder.build( + { + "template": "ref_edit_tata", + "text": "target", + "instruction": "calm", + "ref_text": "reference", + "ref_audio": "ref.wav", + "ref_audio_sample_rate": 16000, + } + ) + + info = prompt["additional_information"] + codes = info["input_values"] + assert tuple(codes.shape) == (2, 16) + assert codes.dtype == torch.int16 + assert set(prompt["prompt_token_ids"]) == {0} + assert info["prompt_ids"].tolist().count(262144) == 2 + assert info["prompt_ids"].tolist().count(262145) == 1 + assert int((~info["text_ids_mask"]).sum()) == 3 + assert info["text_ids_len"].numel() == 2 + assert encoder.calls == [("ref.wav", 16000)] + + +def test_build_reference_prompt_preserves_special_tokens_after_rendering(): + encoder = _ReferenceEncoder() + builder = BreezeTTS2PromptBuilder(_RenderingTokenizer(), _config(), encoder) + + prompt = builder.build( + { + "template": "ref_clone_tata", + "text": "target", + "ref_text": "reference", + "ref_audio": "ref.wav", + "ref_audio_sample_rate": 16000, + } + ) + + info = prompt["additional_information"] + prompt_ids = info["prompt_ids"].tolist() + assert prompt_ids.count(262144) == 2 + assert prompt_ids.count(262145) == 1 + assert int((~info["text_ids_mask"]).sum()) == 3 + assert len(info["text_ids_len"]) == 2 + assert sum(info["text_ids_len"].tolist()) == int(info["text_ids_mask"].sum()) + + +def test_build_reference_prompt_accepts_preencoded_codes_mapping(): + builder = BreezeTTS2PromptBuilder(_RenderingTokenizer(), _config()) + codes = torch.arange(32, dtype=torch.int16).reshape(2, 16) + + prompt = builder.build( + { + "template": "ref_clone_tata", + "text": "target", + "ref_text": "reference", + "codes": {"ref": codes}, + } + ) + + info = prompt["additional_information"] + assert torch.equal(info["input_values"], codes) + assert info["ref_code_len"] == 2 diff --git a/tests/model_executor/models/breeze_tts_2/test_reference_voice.py b/tests/model_executor/models/breeze_tts_2/test_reference_voice.py new file mode 100644 index 000000000..d29b863c4 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_reference_voice.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import soundfile as sf +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_EXAMPLE_PATH = ( + Path(__file__).resolve().parents[4] + / "examples" + / "offline_inference" + / "text_to_speech" + / "breeze_tts_2" + / "end2end.py" +) + + +class _Tokenizer: + unk_token_id = -1 + + def convert_tokens_to_ids(self, token): + return { + "<|AUDIO|>": 262144, + "<|audio_eos|>": 262145, + }.get(token, self.unk_token_id) + + def __call__(self, texts, *, add_special_tokens=True, padding=False): + if isinstance(texts, str): + texts = [texts] + ids = [[1000 + ord(char) % 1000 for char in text] for text in texts] + for row_index, text in enumerate(texts): + if "<|AUDIO|>" in text or "<|audio_eos|>" in text: + cursor = 0 + row = [] + while cursor < len(text): + if text.startswith("<|AUDIO|>", cursor): + row.append(262144) + cursor += len("<|AUDIO|>") + elif text.startswith("<|audio_eos|>", cursor): + row.append(262145) + cursor += len("<|audio_eos|>") + else: + row.append(1000 + ord(text[cursor]) % 1000) + cursor += 1 + ids[row_index] = row + if add_special_tokens: + ids = [[2, *row, 1] for row in ids] + return {"input_ids": ids} + + def decode(self, ids, *, skip_special_tokens=False): + del skip_special_tokens + rendered = [] + for token_id in ids: + if token_id == 262144: + rendered.append("<|AUDIO|>") + elif token_id == 262145: + rendered.append("<|audio_eos|>") + elif token_id in (1, 2): + continue + else: + rendered.append(chr(int(token_id) - 1000)) + return "".join(rendered) + + +class _ReferenceEncoder: + def __init__(self) -> None: + self.calls: list[tuple[object, int | None]] = [] + + def encode(self, audio, sample_rate=None): + self.calls.append((audio, sample_rate)) + return torch.arange(32, dtype=torch.int16).reshape(2, 16) + + +def _builder() -> tuple[BreezeTTS2PromptBuilder, _ReferenceEncoder]: + encoder = _ReferenceEncoder() + builder = BreezeTTS2PromptBuilder( + _Tokenizer(), + SimpleNamespace( + audio_token_id=262144, + audio_eos_token_id=262145, + pad_token_id=0, + backbone_config=SimpleNamespace(vocab_size=100000), + num_codebooks=16, + codec_config={"codebook_size": 2048}, + ), + encoder, + ) + return builder, encoder + + +def test_reference_voice_uses_exact_transcript_and_audio_codes(): + builder, encoder = _builder() + + prompt = builder.build( + { + "template": "ref_clone_tata", + "speaker": "S1", + "ref_text": "exact reference transcript", + "ref_audio": "reference.wav", + "ref_audio_sample_rate": 24_000, + "text": "new target text", + } + ) + + info = prompt["additional_information"] + assert tuple(info["input_values"].shape) == (2, 16) + assert info["ref_code_len"] == 2 + assert info["prompt_ids"].tolist().count(262144) == 2 + assert info["text_ids_len"].numel() == 2 + assert int(info["text_ids_mask"].sum()) == int(info["text_ids_len"].sum()) + assert encoder.calls == [("reference.wav", 24_000)] + + +def test_reference_plus_instruction_composes_voice_direction_prompt(): + builder, _ = _builder() + + prompt = builder.build( + { + "template": "ref_edit_tata", + "speaker": "S0", + "ref_text": "exact reference transcript", + "ref_audio": "reference.wav", + "instruction": "speak calmly and warmly", + "text": "new target text", + } + ) + + info = prompt["additional_information"] + assert info["template"] == "ref_edit_tata" + assert tuple(info["input_values"].shape) == (2, 16) + assert info["prompt_ids"].tolist().count(262144) == 2 + assert info["text_ids_len"].numel() == 2 + + +def test_offline_example_reference_payload_feeds_prompt_builder(tmp_path): + """Smoke: the offline example's payload keys flow into the real builder. + + Loads ``end2end.py`` (voxtral-style importlib), reads a real stereo wav + through its ``reference_payload`` helper, and feeds the result into the + real ``BreezeTTS2PromptBuilder`` — guarding the example↔builder field + contract that previously broke via an unread ``ref_audio_path`` key. + """ + spec = importlib.util.spec_from_file_location("breeze_tts_2_end2end", _EXAMPLE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + sample_rate = 16_000 + t = np.arange(sample_rate, dtype=np.float32) / sample_rate + tone = np.sin(2 * np.pi * 220.0 * t) + ref_path = tmp_path / "reference.wav" + sf.write(str(ref_path), np.stack([tone, tone], axis=1), sample_rate) + + payload = module.reference_payload(str(ref_path)) + + # The helper emits exactly the keys the prompt builder consumes. + assert set(payload) == {"ref_audio", "ref_audio_sample_rate"} + assert payload["ref_audio_sample_rate"] == sample_rate + assert payload["ref_audio"].ndim == 1 # stereo down-mixed to one clip + + builder, encoder = _builder() + prompt = builder.build( + { + "template": "ref_clone_tata", + "speaker": "S0", + "ref_text": "exact reference transcript", + "text": "new target text", + **payload, + } + ) + + info = prompt["additional_information"] + assert tuple(info["input_values"].shape) == (2, 16) + assert info["ref_code_len"] == 2 + assert info["prompt_ids"].tolist().count(262144) == 2 + # The builder received the example's waveform and sample rate verbatim. + assert encoder.calls == [(payload["ref_audio"], sample_rate)] diff --git a/tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py b/tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py new file mode 100644 index 000000000..b9795f798 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_stage_input_processor.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.breeze_tts_2 import ( + talker2codec, + talker2codec_full_payload, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_talker2codec_flattens_codebook_major(): + codes = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int16) + output = SimpleNamespace(multimodal_output={"codes": {"audio": codes}}) + + prompts = talker2codec([output]) + + assert prompts[0]["prompt_token_ids"] == [1, 3, 5, 2, 4, 6] + assert torch.equal(prompts[0]["multi_modal_data"]["codes"]["audio"], codes.transpose(0, 1)) + + +def test_talker2codec_keeps_empty_request_terminal(): + output = SimpleNamespace(multimodal_output={"codes": {}}) + + prompts = talker2codec([output]) + + assert prompts == [{"prompt_token_ids": []}] + + +def test_talker2codec_accepts_engine_core_output_wrapper(): + codes = torch.tensor([[7, 8]], dtype=torch.int16) + output = SimpleNamespace(outputs=[SimpleNamespace(multimodal_output={"codes": {"audio": codes}})]) + + prompts = talker2codec([output]) + + assert prompts[0]["prompt_token_ids"] == [7, 8] + + +def test_talker2codec_unwraps_per_request_audio_container(): + codes = torch.tensor([[7, 8], [9, 10]], dtype=torch.int16) + output = SimpleNamespace(multimodal_output={"codes": {"audio": [codes]}}) + + prompts = talker2codec([output]) + + assert prompts[0]["prompt_token_ids"] == [7, 9, 8, 10] + + +def test_talker2codec_accepts_flat_full_payload(): + output = {"codes.audio": torch.tensor([1, 3, 2, 4], dtype=torch.int16)} + + prompts = talker2codec([output]) + + assert prompts[0]["prompt_token_ids"] == [1, 3, 2, 4] + assert torch.equal( + prompts[0]["multi_modal_data"]["codes"]["audio"], + torch.tensor([1, 3, 2, 4], dtype=torch.int16), + ) + + +def test_talker2codec_full_payload_transposes_accumulated_frames(): + payload = talker2codec_full_payload( + transfer_manager=None, + pooling_output={"codes.audio": torch.tensor([[1, 2], [3, 4], [5, 6]])}, + request=None, + ) + + assert payload is not None + assert torch.equal( + payload.codes.audio, + torch.tensor([1, 3, 5, 2, 4, 6]), + ) + + +def test_talker2codec_full_payload_accepts_runner_pooling_output(): + payload = talker2codec_full_payload( + pooling_output={"codes.audio": torch.tensor([[9, 8]], dtype=torch.int16)}, + request=None, + ) + + assert payload is not None + assert torch.equal(payload.codes.audio, torch.tensor([9, 8])) + + +def test_full_payload_zero_frame_finish_sends_explicit_terminal_marker(): + payload = talker2codec_full_payload( + pooling_output={"codes.audio": torch.empty(0, 16, dtype=torch.long)}, + request=None, + is_finished=True, + ) + + assert payload is not None + assert payload.codes.audio.numel() == 0 + assert payload.meta.code_flat_numel == 0 + assert bool(payload.meta.finished.item()) is True + + +def test_full_payload_missing_codes_finish_sends_explicit_terminal_marker(): + payload = talker2codec_full_payload( + pooling_output={}, + request=None, + is_finished=True, + ) + + assert payload is not None + assert payload.codes.audio.numel() == 0 + assert bool(payload.meta.finished.item()) is True + + +def test_full_payload_empty_without_finish_stays_silent(): + payload = talker2codec_full_payload( + pooling_output={"codes.audio": torch.empty(0, 16, dtype=torch.long)}, + request=None, + is_finished=False, + ) + + assert payload is None + + +def test_full_payload_carries_finished_flag_on_nonempty_frames(): + payload = talker2codec_full_payload( + pooling_output={"codes.audio": torch.tensor([[1, 2], [3, 4]], dtype=torch.long)}, + request=None, + is_finished=True, + ) + + assert payload is not None + assert payload.codes.audio.numel() == 4 + assert bool(payload.meta.finished.item()) is True diff --git a/tests/model_executor/models/breeze_tts_2/test_streaming.py b/tests/model_executor/models/breeze_tts_2/test_streaming.py new file mode 100644 index 000000000..045444ec3 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_streaming.py @@ -0,0 +1,151 @@ +# 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.models.breeze_tts_2.modeling_breeze_tts_2_codec import ( + BreezeTTS2MimiCodec, +) +from vllm_omni.model_executor.stage_input_processors.breeze_tts_2 import ( + talker2codec_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _TransferManager: + def __init__(self, chunk_frames: int = 4): + self.connector = SimpleNamespace(config={"breeze_codec_chunk_frames": chunk_frames}) + self.code_prompt_token_ids = defaultdict(list) + + +class _Request: + def __init__(self, request_id: str, finished: bool = False): + self.external_req_id = request_id + self._finished = finished + + def is_finished(self): + return self._finished + + +def test_async_processor_sends_only_unemitted_tail_and_flushes_finish(): + manager = _TransferManager(chunk_frames=4) + request = _Request("req-1") + + for code0 in range(3): + frame = torch.tensor([code0, code0 + 4, code0 + 8, code0 + 12], dtype=torch.long) + assert talker2codec_async_chunk(manager, {"codes": {"audio": frame}}, request) is None + + payload = talker2codec_async_chunk( + manager, + {"codes": {"audio": torch.tensor([3, 7, 11, 15], dtype=torch.long)}}, + request, + ) + assert payload is not None + assert payload.codes.audio.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] + assert bool(payload.meta.finished.item()) is False + assert manager.code_prompt_token_ids["req-1"] == [] + + assert talker2codec_async_chunk(manager, {"codes": {"audio": torch.arange(4)}}, request) is None + payload = talker2codec_async_chunk(manager, None, _Request("req-1", finished=True)) + assert payload.codes.audio.tolist() == [0, 1, 2, 3] + assert bool(payload.meta.finished.item()) is True + assert bool(payload.meta.stream_finished.item()) is True + assert payload.meta.codec_streaming is True + + +class _Decoder: + def __init__(self): + self.calls = [] + + def batched_chunked_decode(self, codes, lengths, caches=None, **kwargs): + self.calls.append((codes.shape, tuple(lengths), tuple(caches), kwargs)) + return [torch.ones(length * 1920, dtype=torch.float32) for length in lengths] + + +def test_stateful_codec_batches_request_chunks_and_releases_finished_state(): + decoder = _Decoder() + codec = object.__new__(BreezeTTS2MimiCodec) + codec._async_chunk = True + codec._num_codebooks = 4 + codec._codebook_size = 8 + codec._sample_rate = 24_000 + codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) + codec._decoder_state_cache = {} + + input_ids = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], dtype=torch.long) + output = codec.forward( + input_ids, + runtime_additional_information=[ + # ``stream_finished`` is the model-visible flag: the connector + # receiver strips ``meta.finished`` while merging runtime info. + {"meta": {"request_id": "live", "stream_finished": False}}, + {"meta": {"request_id": "done", "stream_finished": True}}, + ], + seq_token_counts=[4, 4], + request_ids=["scheduler-live", "scheduler-done"], + ) + + assert decoder.calls[0][0] == (2, 4, 1) + assert [len(item) for item in output.multimodal_outputs["model_outputs"]] == [1920, 1920] + assert set(codec._decoder_state_cache) == {"scheduler-live"} + + codec.on_requests_finished(["scheduler-live"]) + assert codec._decoder_state_cache == {} + + +def test_stateful_codec_rejects_partial_frame_chunk(): + decoder = _Decoder() + codec = object.__new__(BreezeTTS2MimiCodec) + codec._async_chunk = True + codec._num_codebooks = 4 + codec._codebook_size = 8 + codec._sample_rate = 24_000 + codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) + codec._decoder_state_cache = {} + + # 5 codes cannot form a whole frame; a silent state reset would corrupt + # the rest of the stream, so the codec must fail loudly. + with pytest.raises(ValueError, match="not divisible"): + codec.forward( + torch.tensor([0, 1, 2, 3, 4], dtype=torch.long), + runtime_additional_information=[{"meta": {"request_id": "bad", "finished": False}}], + seq_token_counts=[5], + request_ids=["scheduler-bad"], + ) + assert decoder.calls == [] + + +def test_stateful_codec_terminal_marker_pops_state_without_decoding(): + decoder = _Decoder() + codec = object.__new__(BreezeTTS2MimiCodec) + codec._async_chunk = True + codec._num_codebooks = 4 + codec._codebook_size = 8 + codec._sample_rate = 24_000 + codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder)) + codec._decoder_state_cache = {"scheduler-live": {}, "scheduler-done": {}} + + output = codec.forward( + torch.tensor([0, 1, 2, 3], dtype=torch.long), + runtime_additional_information=[ + # Plain ``finished`` covers the fallback for direct in-process + # callers that do not go through the connector receiver. + {"meta": {"request_id": "live", "finished": False}}, + {"meta": {"request_id": "done", "finished": True}}, + ], + seq_token_counts=[4, 0], + request_ids=["scheduler-live", "scheduler-done"], + ) + + # The empty terminal marker decodes nothing and releases its state; the + # live request still decodes its own chunk. + assert len(decoder.calls) == 1 + assert decoder.calls[0][0][0] == 1 + assert [item.numel() for item in output.multimodal_outputs["model_outputs"]] == [1920, 0] + assert list(codec._decoder_state_cache) == ["scheduler-live"] + assert codec._decoder_state_cache["scheduler-live"]["prefix_frames"] == 0 diff --git a/tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py b/tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py new file mode 100644 index 000000000..2f4a13df1 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from typing import Any + +import pytest +import torch + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_talker import ( + BreezeTTS2TalkerForGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Head: + def __init__(self) -> None: + self.code0 = 2 + + def __call__(self, hidden: torch.Tensor) -> torch.Tensor: + logits = torch.zeros((hidden.shape[0], 9), dtype=hidden.dtype) + logits[:, self.code0] = 1.0 + return logits + + +class _Embedding: + def __init__(self) -> None: + self.weight = torch.randn(4 * 8, 16) + + def __call__(self, ids: torch.Tensor) -> torch.Tensor: + return self.weight[ids.reshape(-1)].reshape(*ids.shape, -1) + + +def _talker() -> BreezeTTS2TalkerForGeneration: + model = object.__new__(BreezeTTS2TalkerForGeneration) + model.num_codebooks = 4 + model.codebook_vocab_size = 8 + model.codebook_size = 8 + model.hidden_size = 16 + model.codebook_eos_token_id = 0 + model.audio_token_id = 100 + model.audio_eos_token_id = 101 + # Bypass nn.Module registration; this fake is created with object.__new__. + model.__dict__["embed_audio_tokens"] = _Embedding() + model.audio_token_offsets = torch.arange(4) * 8 + model.lm_head = _Head() + model.logits_processor = lambda head, hidden: head(hidden) + model._codec_disallowed_mask = None + model._batch_state = None + model._batch_state_spans = None + model._async_chunk = False + model._golden_dump_dir = None + return model + + +def _depth_codes(_model: Any, _hidden: torch.Tensor, code0: int) -> torch.Tensor: + if code0 >= 8: + return torch.empty(0, dtype=torch.long) + return torch.tensor([code0, code0 + 1, code0 + 2, code0 + 3], dtype=torch.long).unsqueeze(0) + + +def test_prefill_hidden_generates_and_stores_first_complete_frame(): + model = _talker() + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"_omni_is_prefill": True, "breeze_generated_frames": 0, "breeze_max_new_frames": 3} + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + + codes = output.multimodal_outputs["codes"]["audio"][0] + assert torch.equal(codes, torch.tensor([[2, 3, 4, 5]])) + assert torch.equal(info["breeze_current_frame"], codes[0]) + assert info["breeze_generated_frames"] == 1 + assert info["breeze_force_eos"] is False + + +def test_first_decode_consumes_full_frame_instead_of_scalar_eos_fallback(): + model = _talker() + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"_omni_is_prefill": True, "breeze_generated_frames": 0, "breeze_max_new_frames": 3} + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + frame = output.multimodal_outputs["codes"]["audio"][0][0] + decode_info = dict(info) + decode_info.pop("_omni_is_prefill", None) + + safe_id, decode_embeds, update = model.preprocess( + torch.tensor([2], dtype=torch.long), + None, + **decode_info, + ) + + assert safe_id.item() == 2 + assert update == {} + expected = model.embed_input_ids(frame.reshape(1, 1, -1)).reshape(1, -1) + assert torch.equal(decode_embeds, expected) + + +def test_real_decode_rejects_missing_current_frame(): + model = _talker() + info = { + "breeze_generated_frames": 1, + "breeze_audio_codes": torch.tensor([[2, 3, 4, 5]]), + } + + try: + model.preprocess(torch.tensor([2], dtype=torch.long), None, **info) + except RuntimeError as exc: + assert "breeze_current_frame" in str(exc) + assert "codebooks 1..15" in str(exc) + else: + raise AssertionError("missing breeze_current_frame must fail explicitly") + + +def test_embed_input_ids_rejects_scalar_codec_id(): + model = _talker() + + try: + model.embed_input_ids(torch.tensor([[2]])) + except ValueError as exc: + assert "complete codec frames" in str(exc) + else: + raise AssertionError("scalar codec ids must not be embedded") + + +def test_decode_accumulates_frames_then_natural_eos_stops_without_fake_frame(): + model = _talker() + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"breeze_generated_frames": 1, "breeze_max_new_frames": 3} + info["breeze_audio_codes"] = torch.tensor([[2, 3, 4, 5]]) + model.lm_head.code0 = 3 + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + assert output.multimodal_outputs["codes"]["audio"][0].shape == (2, 4) + + # Main-head class 8 is the Breeze backbone EOS (one past vocab_size=8). + model.lm_head.code0 = 8 + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + assert output.multimodal_outputs["codes"]["audio"][0].shape == (2, 4) + assert info["breeze_force_eos"] is True + assert info["breeze_force_eos"] is True + + logits = model.compute_logits(hidden) + assert torch.equal(logits[0, :8], torch.full((8,), float("-inf"))) + assert logits[0, 8].item() == 0.0 + + +def test_frame_budget_forces_eos_on_followup_step(): + model = _talker() + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"breeze_generated_frames": 2, "breeze_max_new_frames": 2} + info["breeze_audio_codes"] = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 0]]) + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + + assert output.multimodal_outputs["codes"]["audio"][0].shape == (2, 4) + assert info["breeze_generated_frames"] == 2 + assert info["breeze_force_eos"] is True + + +def test_async_chunk_mode_skips_full_accumulation(): + model = _talker() + model._async_chunk = True + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"breeze_generated_frames": 0, "breeze_max_new_frames": 8} + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + for _ in range(3): + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 4) + + # Streaming keeps only the per-frame tail; the cumulative buffer (and its + # per-step torch.cat) must not exist without the opt-in golden dump. + assert "breeze_audio_codes" not in info + assert info["breeze_generated_frames"] == 3 + assert torch.equal(info["breeze_current_frame"], torch.tensor([2, 3, 4, 5])) + + +def test_async_chunk_mode_keeps_accumulation_for_golden_dump(tmp_path): + model = _talker() + model._async_chunk = True + model._golden_dump_dir = str(tmp_path) + model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0) + info = {"breeze_generated_frames": 0, "breeze_max_new_frames": 8} + hidden = torch.arange(16, dtype=torch.float32).reshape(1, 16) + + for _ in range(3): + output = model.make_omni_output( + hidden, + model_intermediate_buffer=[info], + request_token_spans=[(0, 1)], + ) + assert output.multimodal_outputs["codes"]["audio"][0].shape == (1, 4) + + assert info["breeze_audio_codes"].shape == (3, 4) + assert torch.equal(info["breeze_audio_codes"][-1], torch.tensor([2, 3, 4, 5])) + + +def test_opt_in_golden_dump_writes_terminal_frames(tmp_path): + model = _talker() + model._golden_dump_dir = str(tmp_path) + info = { + "template": "tts_plain", + "prompt_ids": torch.tensor([1, 2, 3]), + "breeze_audio_codes": torch.tensor([[1, 2, 3, 4]]), + "breeze_generated_frames": 1, + "breeze_force_eos": True, + } + + model._dump_golden_frames(info) + + dump_path = next(tmp_path.glob("breeze_*.pt")) + payload = torch.load(dump_path, weights_only=False) + assert payload["template"] == "tts_plain" + assert torch.equal(payload["prompt_ids"], torch.tensor([1, 2, 3])) + assert torch.equal(payload["codes"], torch.tensor([[1, 2, 3, 4]])) + assert payload["finished"] is True diff --git a/tests/model_executor/models/breeze_tts_2/test_weight_loading.py b/tests/model_executor/models/breeze_tts_2/test_weight_loading.py new file mode 100644 index 000000000..f4c0da429 --- /dev/null +++ b/tests/model_executor/models/breeze_tts_2/test_weight_loading.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_talker import ( + BreezeTTS2TalkerForGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _PaddedHead(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(2_112, 32), requires_grad=False) + + +class _DepthDecoder: + def __init__(self) -> None: + self.loaded: list[tuple[str, torch.Tensor]] = [] + + def load_weights(self, weights): + self.loaded.extend(weights) + return {"model.embed_tokens.weight"} + + +class _TalkerHost(nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace(tie_codebooks_embeddings=True) + self.codebook_vocab_size = 2_051 + self.hidden_size = 32 + self.lm_head = _PaddedHead() + self.embed_audio_tokens = nn.Embedding(4 * 2_051, 32) + self.depth_decoder = _DepthDecoder() + + +def _padded_loader(parameter: nn.Parameter, weight: torch.Tensor) -> None: + assert parameter.shape != weight.shape + parameter[: weight.shape[0]].copy_(weight) + + +def test_lm_head_loader_accepts_vocab_padding(): + host = _TalkerHost() + host.lm_head.weight.weight_loader = _padded_loader + checkpoint_weight = torch.full((2_052, 32), 1.5) + + loaded = BreezeTTS2TalkerForGeneration.load_weights( + host, + iter([("lm_head.weight", checkpoint_weight)]), + ) + + assert "lm_head.weight" in loaded + assert torch.equal(host.lm_head.weight[:2_052], checkpoint_weight) + assert torch.count_nonzero(host.lm_head.weight[2_052:]).item() == 0 + + +def test_tied_audio_embedding_reports_both_strict_loading_aliases(): + host = _TalkerHost() + checkpoint_weight = torch.full((4 * 2_051, 32), 2.0) + + loaded = BreezeTTS2TalkerForGeneration.load_weights( + host, + iter([("depth_decoder.model.embed_tokens.weight", checkpoint_weight)]), + ) + + assert "embed_audio_tokens.weight" in loaded + assert "depth_decoder.model.embed_tokens.weight" in loaded diff --git a/tools/breeze_tts_2_compare_golden.py b/tools/breeze_tts_2_compare_golden.py new file mode 100644 index 000000000..23306896e --- /dev/null +++ b/tools/breeze_tts_2_compare_golden.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Compare upstream and vLLM-Omni Breeze-TTS-2 golden artifacts.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import soundfile as sf +import torch + +TEMPLATES = { + "plain": "tts_plain", + "instruction": "tts_instruction", + "reference_edit": "ref_edit_tata", +} + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _audio_metrics(path: Path) -> dict[str, object]: + data, sample_rate = sf.read(path, dtype="float32") + return { + "path": str(path), + "bytes": path.stat().st_size, + "sha256": _sha256(path), + "samples": int(data.shape[0]), + "sample_rate": int(sample_rate), + "channels": int(1 if data.ndim == 1 else data.shape[1]), + "duration_s": round(float(data.shape[0]) / sample_rate, 6), + "finite": bool(torch.isfinite(torch.from_numpy(data)).all().item()), + "peak": round(float(abs(data).max()), 6), + "rms": round(float((data**2).mean() ** 0.5), 6), + } + + +def _vector_metrics(upstream: torch.Tensor, actual: torch.Tensor) -> dict[str, float]: + difference = upstream.double() - actual.double() + return { + "mae": round(float(difference.abs().mean()), 8), + "rmse": round(float(difference.pow(2).mean().sqrt()), 8), + "max_abs_error": round(float(difference.abs().max()), 8), + "relative_l2": round(float(difference.norm() / upstream.norm()), 8), + "cosine": round( + float(torch.dot(upstream.double(), actual.double()) / (upstream.norm() * actual.norm())), + 8, + ), + } + + +def _load_vllm_payload(vllm_dir: Path, template: str) -> dict[str, object]: + matches = [] + for path in vllm_dir.glob("breeze_*.pt"): + payload = torch.load(path, weights_only=False) + if payload.get("template") == template: + matches.append(payload) + if len(matches) != 1: + raise ValueError(f"expected one vLLM dump for {template!r}, found {len(matches)}") + return matches[0] + + +def _longest_exact_frame_prefix(upstream: torch.Tensor, actual: torch.Tensor) -> int: + count = min(int(upstream.shape[0]), int(actual.shape[0])) + equal = (upstream[:count] == actual[:count]).all(dim=-1) + prefix = 0 + for is_equal in equal.tolist(): + if not is_equal: + break + prefix += 1 + return prefix + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--upstream-dir", type=Path, required=True) + parser.add_argument("--vllm-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--plain-audio", type=Path) + parser.add_argument("--instruction-audio", type=Path) + parser.add_argument("--reference-audio", type=Path) + args = parser.parse_args() + + audio_paths = { + "plain": args.plain_audio, + "instruction": args.instruction_audio, + "reference_edit": args.reference_audio, + } + report: dict[str, object] = { + "upstream_dir": str(args.upstream_dir.resolve()), + "vllm_dir": str(args.vllm_dir.resolve()), + "cases": [], + } + for name, template in TEMPLATES.items(): + upstream_codes = torch.load(args.upstream_dir / name / "codes.pt", weights_only=False) + upstream_prompt = torch.load(args.upstream_dir / name / "prompt_ids.pt", weights_only=False) + actual_payload = _load_vllm_payload(args.vllm_dir, template) + actual_codes = actual_payload["codes"] + actual_prompt = actual_payload["prompt_ids"] + common_frames = min(int(upstream_codes.shape[0]), int(actual_codes.shape[0])) + code_equal = upstream_codes[:common_frames] == actual_codes[:common_frames] + frame_equal = code_equal.all(dim=-1) + prefill_metrics = None + upstream_hidden_path = args.upstream_dir / name / "prefill_hidden.pt" + if upstream_hidden_path.exists(): + prefill_files = [ + path + for path in args.vllm_dir.glob("prefill_*.pt") + if torch.load(path, weights_only=False).get("template") == template + ] + if len(prefill_files) == 1: + actual_hidden = torch.load(prefill_files[0], weights_only=False)["hidden"] + prefill_metrics = _vector_metrics( + torch.load(upstream_hidden_path, weights_only=False), + actual_hidden, + ) + entry: dict[str, object] = { + "name": name, + "template": template, + "prompt": { + "upstream_tokens": int(upstream_prompt.numel()), + "vllm_tokens": int(actual_prompt.numel()), + "exact": bool(torch.equal(upstream_prompt, actual_prompt)), + }, + "frames": { + "upstream": int(upstream_codes.shape[0]), + "vllm": int(actual_codes.shape[0]), + "codebooks": int(upstream_codes.shape[-1]), + "exact_sequence": bool(torch.equal(upstream_codes, actual_codes)), + "common_frames": common_frames, + "exact_common_frames": int(frame_equal.sum()), + "longest_exact_prefix": _longest_exact_frame_prefix(upstream_codes, actual_codes), + "exact_codes": int(code_equal.sum()), + "total_common_codes": int(code_equal.numel()), + "code_accuracy": round(float(code_equal.float().mean()), 8), + "codebook0_accuracy": round( + float((upstream_codes[:common_frames, 0] == actual_codes[:common_frames, 0]).float().mean()), + 8, + ), + "range_valid": bool( + upstream_codes.min() >= 0 + and upstream_codes.max() < 2048 + and actual_codes.min() >= 0 + and actual_codes.max() < 2048 + ), + }, + "prefill_hidden": prefill_metrics, + } + if audio_paths[name] is not None and Path(audio_paths[name]).exists(): # type: ignore[arg-type] + entry["audio"] = _audio_metrics(Path(audio_paths[name])) # type: ignore[arg-type] + report["cases"].append(entry) + + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/breeze_tts_2_upstream_golden.py b/tools/breeze_tts_2_upstream_golden.py new file mode 100644 index 000000000..337a0444f --- /dev/null +++ b/tools/breeze_tts_2_upstream_golden.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Generate deterministic Breeze-TTS-2 golden frames with the upstream runtime. + +This tool is intentionally separate from vLLM-Omni: it imports the unmodified +``breezeblue-ai/breeze-tts`` reference implementation, fixes greedy decoding, +and writes both generated codec frames and decoded audio for offline parity +checks. It is a diagnostic tool, not a serving path. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import soundfile as sf +import torch + + +def _plain_template(upstream_templates): + """The upstream CLI does not register tts_plain; reuse its segment builder.""" + return SimpleNamespace( + name="tts_plain", + required_fields=("text",), + build_segments=upstream_templates._tts_plain_segments, + build_negative_segments=None, + build_dual_branches=None, + ) + + +def _greedy(model) -> None: + model.generation_config.do_sample = False + model.generation_config.temperature = None + model.generation_config.top_p = None + model.generation_config.top_k = None + depth = model.depth_decoder.generation_config + depth.depth_decoder_do_sample = False + depth.do_sample = False + depth.temperature = None + depth.top_p = None + depth.top_k = None + + +def _valid_frames(sequences: torch.Tensor, pad_token_id: int) -> torch.Tensor: + frames = sequences[0].detach().cpu().to(torch.long) + if frames.ndim == 1: + frames = frames.unsqueeze(0) + is_pad = (frames == pad_token_id).all(dim=-1) + if bool(is_pad.any()): + frames = frames[: int(torch.argmax(is_pad.int()).item())] + return frames.contiguous() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--upstream", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--reference-audio", type=Path) + parser.add_argument("--max-new-tokens", type=int, default=32) + args = parser.parse_args() + + args.output_dir.mkdir(parents=True, exist_ok=True) + upstream_posix = str(args.upstream.resolve()) + if upstream_posix not in sys.path: + sys.path.insert(0, upstream_posix) + + from breeze_infer.runtime import load_runtime, resolve_device + from breeze_infer.templates import get_template, prepare_inputs + + device = resolve_device() + # ``torch.manual_seed`` seeds CUDA RNGs as well, so no device-specific + # seeding call is needed here. + torch.manual_seed(0) + tokenizer, model, audio_tokenizer = load_runtime( + args.model, + device=device, + attn_implementation="eager", + ) + # Do not enable Transformers 4.x's ``fix_mistral_regex`` compatibility + # path: vLLM-Omni runs with Transformers 5.x, and direct tokenization checks + # showed both runtimes produce the same token IDs without that patch. + _greedy(model) + model.eval() + first_backbone_hidden: list[torch.Tensor] = [] + + def _capture_first_backbone_hidden(_module, _inputs, output): + if first_backbone_hidden: + return + hidden = getattr(output, "last_hidden_state", None) + if isinstance(hidden, torch.Tensor): + first_backbone_hidden.append(hidden[0, -1].detach().cpu().to(torch.float32).contiguous()) + + model.backbone_model.register_forward_hook(_capture_first_backbone_hidden) + + cases = [ + { + "name": "plain", + "template": _plain_template(sys.modules["breeze_infer.templates"]), + "request": { + "id": "golden-plain", + "text": "Golden plain alignment check.", + "speaker": "S0", + }, + }, + { + "name": "instruction", + "template": get_template("tts_instruction"), + "request": { + "id": "golden-instruction", + "text": "Golden instruction alignment check.", + "instruction": "Speak calmly, warmly, and clearly.", + "speaker": "S0", + }, + }, + ] + if args.reference_audio is not None: + cases.append( + { + "name": "reference_edit", + "template": get_template("ref_edit_tata"), + "request": { + "id": "golden-reference-edit", + "text": "Golden reference alignment check.", + "instruction": "Speak calmly and clearly.", + "ref_audio_path": str(args.reference_audio), + "ref_text": "Hello from Breeze.", + "speaker": "S0", + }, + } + ) + + manifest = { + "runtime": "upstream breezeblue-ai/breeze-tts", + "model": str(args.model.resolve()), + "upstream": upstream_posix, + "device": device, + "seed": 0, + "max_new_tokens": args.max_new_tokens, + "decoding": "greedy", + "cfg_scale": 1.0, + "cases": [], + } + for case in cases: + first_backbone_hidden.clear() + inputs = prepare_inputs( + tokenizer, + audio_tokenizer, + model, + [case["request"]], + case["template"], + guidance_scale=1.0, + guidance_scale_ref=None, + guidance_scale_ins=None, + ) + output = model.generate( + **inputs, + max_new_tokens=args.max_new_tokens, + do_sample=False, + return_dict_in_generate=True, + output_audio=True, + audio_tokenizer=audio_tokenizer, + ) + frames = _valid_frames(output.sequences, model.config.codebook_pad_token_id) + audio = output.audio[0].detach().cpu().to(torch.float32).reshape(-1) + case_dir = args.output_dir / case["name"] + case_dir.mkdir(parents=True, exist_ok=True) + torch.save(frames, case_dir / "codes.pt") + prompt_ids = inputs["input_ids"].detach().cpu().to(torch.long).reshape(-1) + torch.save(prompt_ids, case_dir / "prompt_ids.pt") + if first_backbone_hidden: + torch.save(first_backbone_hidden[0], case_dir / "prefill_hidden.pt") + sf.write(case_dir / "upstream.wav", audio.numpy(), 24_000, subtype="PCM_16") + entry = { + "name": case["name"], + "request": case["request"], + "prompt_tokens": int(prompt_ids.numel()), + "frames": int(frames.shape[0]), + "codebooks": int(frames.shape[-1]) if frames.ndim == 2 else 0, + "min_code": int(frames.min().item()) if frames.numel() else None, + "max_code": int(frames.max().item()) if frames.numel() else None, + "audio_samples": int(audio.numel()), + "audio_sha256": __import__("hashlib").sha256((case_dir / "upstream.wav").read_bytes()).hexdigest(), + } + manifest["cases"].append(entry) + print(json.dumps(entry), flush=True) + + (args.output_dir / "manifest.json").write_text( + json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index fd7617144..b5214720a 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -201,6 +201,7 @@ _MODEL_INTERNALIZE = ( "MING_TTS_STAGE1_FINAL_LOG", "MOSS_TTS_DEBUG_STOP", "NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", + "VLLM_OMNI_BREEZE_GOLDEN_DIR", "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS", "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS_LOG_EVERY", "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS", diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 2378ce5b9..92aad9f4d 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -47,6 +47,7 @@ from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_SINGLE_STAGE_PIPELINE, BAGEL_THINK_PIPELINE, ) +from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE @@ -128,6 +129,7 @@ PipelineResolverFunc: TypeAlias = Callable[[PretrainedConfig | None], PipelineCo # --- Multi-stage omni pipelines (LLM-centric; audio / video I/O) --- OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "aura_omni": AURA_OMNI_PIPELINE, + "breeze": BREEZE_TTS_2_PIPELINE, "joyai_vl_interaction": JOYAI_VL_INTERACTION_PIPELINE, "qwen2_5_omni": QWEN2_5_OMNI_PIPELINE, "qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE, diff --git a/vllm_omni/deploy/breeze_tts_2.yaml b/vllm_omni/deploy/breeze_tts_2.yaml new file mode 100644 index 000000000..1d1b85391 --- /dev/null +++ b/vllm_omni/deploy/breeze_tts_2.yaml @@ -0,0 +1,51 @@ +# Breeze-TTS-2: talker -> stateful Qwen3-TTS codec decoder. +# Stage 0 batches codec frames and sends incremental codebook-major chunks. +async_chunk: true +trust_remote_code: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + breeze_codec_chunk_frames: 8 + +stages: + - stage_id: 0 + max_num_seqs: 8 + gpu_memory_utilization: 0.80 + enable_prefix_caching: false + # The talker derives one codec frame from the last hidden row of every + # scheduled request. A partially scheduled prompt would yield a frame from + # a mid-prompt hidden state, so keep each prefill in a single step. + enable_chunked_prefill: false + async_scheduling: true + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + min_tokens: 1 + + - stage_id: 1 + max_num_seqs: 8 + gpu_memory_utilization: 0.15 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 20319de86..d1e53775a 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -54,6 +54,9 @@ def _register_omni_hf_configs() -> None: try: from transformers import AutoConfig + from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( + BreezeTTS2Config, + ) from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import ( IndexTTS2Config, IndexTTS25Config, @@ -105,6 +108,7 @@ def _register_omni_hf_configs() -> None: ("glm_tts", GLMTTSConfig), ("omnivoice", OmniVoiceConfig), ("voxcpm2", VoxCPM2Config), + ("breeze", BreezeTTS2Config), ]: try: AutoConfig.register(model_type, config_cls) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index 39882274d..8ed54afd8 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -105,6 +105,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 audex, audex_tta, audio8_tts, + breeze_tts_2, cosyvoice3, covo_audio, dots_tts, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py new file mode 100644 index 000000000..d2851e32c --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze-TTS-2 serving adapter for synchronous full-payload generation.""" + +import asyncio +from typing import TYPE_CHECKING, Any + +import numpy as np + +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.breeze_tts_2.audio_tokenizer import ( + BreezeReferenceAudioTokenizer, +) +from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import ( + BreezeTTS2PromptBuilder, +) + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + + +@register_tts_adapter +class BreezeTTS2Adapter(ARTTSAdapter): + """Build Breeze prompts before scheduler submission. + + The adapter owns only request-facing concerns. Model workers still own the + stage-0 embeddings and stage-1 waveform decode. + """ + + stage_keys = frozenset({"breeze_tts_2"}) + model_archs = frozenset({"BreezeForConditionalGeneration"}) + name = "breeze_tts_2" + # Breeze's architecture and stage key are unique, so this detector cannot + # collide with another adapter. Lower priority runs first; 9 keeps it + # ahead of the default-priority (100) detectors. + detect_priority = 9 + supported_output_sample_rates = frozenset({24000}) + + def _load_supported_speakers(self) -> set[str]: + # Breeze accepts arbitrary speaker tags such as S0/S1 in its prompt; + # there is no finite checkpoint speaker inventory to advertise. + return set() + + def _load_codec_frame_rate(self) -> float | None: + config = self.ctx.engine_client.model_config.hf_config + codec = getattr(config, "codec_config", None) + if isinstance(codec, dict): + frame_rate = codec.get("_frame_rate") + if frame_rate is not None: + return float(frame_rate) + return 12.5 + + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: + if not request.input or not request.input.strip(): + return "Breeze-TTS-2 input text cannot be empty" + if request.ref_text is not None and not request.ref_text.strip(): + return "Breeze-TTS-2 ref_text cannot be empty" + if request.ref_text is not None and request.ref_audio is None: + return "Breeze-TTS-2 ref_text requires ref_audio" + if request.ref_audio is not None and (request.ref_text is None or not request.ref_text.strip()): + return "Breeze-TTS-2 ref_audio requires ref_text" + if isinstance(request.ref_audio, list) and len(request.ref_audio) != 1: + # The prompt path conditions on exactly one reference waveform; + # extra clips would be silently discarded otherwise. + return f"Breeze-TTS-2 supports exactly one reference clip, got {len(request.ref_audio)}" + if request.task_type == "Base" and request.ref_audio is None: + return "Breeze-TTS-2 Base task requires ref_audio and ref_text" + if request.task_type == "VoiceDesign": + return "Breeze-TTS-2 does not support task_type=VoiceDesign" + if request.speed is not None and request.speed != 1.0: + return "Breeze-TTS-2 does not support speed adjustment" + if request.language is not None: + return "Breeze-TTS-2 does not support 'language'; the prompt language follows the input text" + if request.speaker_embedding is not None: + return "Breeze-TTS-2 does not support 'speaker_embedding' (voice-cloning helper field)" + if request.x_vector_only_mode: + return "Breeze-TTS-2 does not support 'x_vector_only_mode' (voice-cloning helper field)" + extra_params = request.extra_params or {} + guidance_scale = extra_params.get("guidance_scale", extra_params.get("cfg_scale", 1.0)) + try: + unsupported_guidance = guidance_scale is not None and float(guidance_scale) != 1.0 + except (TypeError, ValueError): + return "Breeze-TTS-2 guidance_scale must be a number" + if unsupported_guidance: + return "Breeze-TTS-2 currently supports only guidance_scale=1.0" + if extra_params.get("negative_prompt"): + return "Breeze-TTS-2 does not support negative_prompt/CFG yet" + # The talker selects codebook 0 by argmax inside make_omni_output and + # ignores the scheduler's sampled token, so non-greedy overrides would + # not change the audio while still letting the sampler draw an early + # EOS. Reject them until the talker consumes the sampled code. + greedy_error = _validate_greedy_sampling(extra_params) + if greedy_error is not None: + return greedy_error + 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}" + if request.ref_audio is not None: + ref_audio = request.ref_audio[0] if isinstance(request.ref_audio, list) else request.ref_audio + return self.ctx.server._validate_ref_audio_format(ref_audio) + return None + + async def build( + self, + request: "OpenAICreateSpeechRequest", + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + del has_inline_ref_audio + server = self.ctx.server + builder = await self._get_builder() + template = self._resolve_template(request) + payload: dict[str, Any] = { + "text": request.input, + "template": template, + "speaker": request.voice or "S0", + } + if request.instructions: + payload["instruction"] = request.instructions + if request.ref_text: + payload["ref_text"] = request.ref_text + + if request.ref_audio is not None: + ref_audio = request.ref_audio[0] if isinstance(request.ref_audio, list) else request.ref_audio + wav_list, sr, cache_key = await server._resolve_ref_audio(ref_audio) + payload["ref_audio"] = np.asarray(wav_list, dtype=np.float32) + payload["ref_audio_sample_rate"] = int(sr) + else: + cache_key = None + + prompt = await asyncio.to_thread(builder.build, payload, template=template) + # The AR scheduler's completion budget is measured in codebook-0 + # frames for Breeze. Preserve it in the mutable request metadata so + # the talker can emit its own EOS before a scheduler length cutoff. + max_new_frames = request.max_new_tokens + if max_new_frames is None and sampling_params_list: + max_new_frames = getattr(sampling_params_list[0], "max_tokens", None) + if max_new_frames is not None: + prompt["additional_information"]["breeze_max_new_frames"] = int(max_new_frames) + tts_params = { + "template": [template], + "text": [request.input], + } + if cache_key: + tts_params["ref_audio_cache_key"] = cache_key + prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) + return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name) + + async def _get_builder(self) -> BreezeTTS2PromptBuilder: + server = self.ctx.server + cached = getattr(server, "_breeze_tts_2_prompt_builder", None) + if cached is not None: + return cached + # Multiple speech requests can arrive while the first tokenizer load + # is still in progress. Serialize the one-time CPU initialization so + # model files are not opened and decoded repeatedly. + lock = getattr(server, "_breeze_tts_2_prompt_builder_lock", None) + if lock is None: + lock = asyncio.Lock() + server._breeze_tts_2_prompt_builder_lock = lock + async with lock: + cached = getattr(server, "_breeze_tts_2_prompt_builder", None) + if cached is not None: + return cached + model_path = server.engine_client.model_config.model + config = server.engine_client.model_config.hf_config + audio_tokenizer = await asyncio.to_thread( + BreezeReferenceAudioTokenizer.from_pretrained, + model_path, + num_codebooks=int(getattr(config, "num_codebooks", 16)), + codebook_size=self._codebook_size(config), + device_map="cpu", + ) + builder = await asyncio.to_thread( + BreezeTTS2PromptBuilder.from_pretrained, + model_path, + config, + reference_audio_encoder=audio_tokenizer, + ) + server._breeze_tts_2_prompt_builder = builder + return builder + + @staticmethod + def _codebook_size(config: Any) -> int: + codec = getattr(config, "codec_config", None) + if isinstance(codec, dict): + return int(codec.get("codebook_size", 2048)) + return int(getattr(codec, "codebook_size", 2048)) + + @staticmethod + def _resolve_template(request: "OpenAICreateSpeechRequest") -> str: + if request.ref_audio is not None: + return "ref_edit_tata" if request.instructions else "ref_clone_tata" + return "tts_instruction" if request.instructions else "tts_plain" + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + del prompt, request_id + return apply_max_new_tokens(sampling_params_list, request) + + +def _validate_greedy_sampling(extra_params: dict[str, Any]) -> str | None: + """Return an error unless the sampling overrides keep decoding greedy.""" + temperature = extra_params.get("temperature") + if temperature is not None: + try: + if float(temperature) != 0.0: + return "Breeze-TTS-2 currently supports only greedy decoding (temperature=0)" + except (TypeError, ValueError): + return "Breeze-TTS-2 temperature must be a number" + top_p = extra_params.get("top_p") + if top_p is not None: + try: + if float(top_p) != 1.0: + return "Breeze-TTS-2 currently supports only greedy decoding (top_p=1.0)" + except (TypeError, ValueError): + return "Breeze-TTS-2 top_p must be a number" + top_k = extra_params.get("top_k") + if top_k is not None: + try: + if int(top_k) not in (-1, 0, 1): + return "Breeze-TTS-2 currently supports only greedy decoding (top_k=-1)" + except (TypeError, ValueError): + return "Breeze-TTS-2 top_k must be an integer" + return None + + +__all__ = ["BreezeTTS2Adapter"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/__init__.py b/vllm_omni/model_executor/models/breeze_tts_2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py b/vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py new file mode 100644 index 000000000..ea2796a53 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/audio_tokenizer.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Reference-audio encoding helpers for Breeze-TTS-2. + +The Breeze prompt uses a Qwen3-TTS tokenizer only for reference audio. This +module deliberately does not contain text tokenization or waveform decoding; +the latter belongs to the pipeline's codec stage. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import torch + + +def resolve_audio_tokenizer_path(model_path: str) -> Path | None: + """Resolve the bundled ``audio_tokenizer`` dir for a local path or HF repo id. + + Local directories (dev checkouts) resolve directly: present returns the + subdirectory, absent returns ``None`` so callers can fall back. Repo ids + resolve through the vLLM-tagged HF helpers: an already-complete cached + snapshot is preferred (offline deployments keep working), and a partial + cache falls through to a targeted online fetch restricted to the bundled + tokenizer subtree. Returns ``None`` when the checkpoint carries no + ``audio_tokenizer`` subdirectory. + """ + root = Path(model_path) + if root.is_dir(): + bundled = root / "audio_tokenizer" + return bundled if bundled.is_dir() else None + + from vllm.transformers_utils.repo_utils import hf_api + + kwargs: dict[str, object] = {"allow_patterns": ["audio_tokenizer/*"]} + try: + cached = Path(hf_api().snapshot_download(model_path, local_files_only=True, **kwargs)) + except Exception: + cached = None + if cached is not None and (cached / "audio_tokenizer").is_dir(): + return cached / "audio_tokenizer" + online = Path(hf_api().snapshot_download(model_path, **kwargs)) + bundled = online / "audio_tokenizer" + return bundled if bundled.is_dir() else None + + +class BreezeReferenceAudioTokenizer: + """Small adapter that normalizes Qwen3-TTS encoder output to ``(T, Q)``. + + The underlying tokenizer is created once per serving worker and is reused + for all requests. Keeping this wrapper independent from the + prompt builder makes it straightforward to inject a fake encoder in tests + and to replace the upstream Qwen3 implementation later. + """ + + def __init__(self, tokenizer: Any, *, num_codebooks: int = 16, codebook_size: int = 2048) -> None: + self.tokenizer = tokenizer + self.num_codebooks = int(num_codebooks) + self.codebook_size = int(codebook_size) + if self.num_codebooks <= 0 or self.codebook_size <= 0: + raise ValueError("num_codebooks and codebook_size must be positive") + + @classmethod + def from_pretrained( + cls, + model_path: str, + *, + num_codebooks: int = 16, + codebook_size: int = 2048, + **kwargs: Any, + ) -> BreezeReferenceAudioTokenizer: + """Load the bundled ``audio_tokenizer`` with the existing Qwen wrapper.""" + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import ( + Qwen3TTSTokenizer, + ) + + resolved = resolve_audio_tokenizer_path(model_path) + if resolved is None: + raise FileNotFoundError( + "Breeze reference audio tokenizer was not found for " + f"{model_path} (expected an audio_tokenizer/ subdirectory)" + ) + tokenizer = Qwen3TTSTokenizer.from_pretrained(str(resolved), **kwargs) + return cls(tokenizer, num_codebooks=num_codebooks, codebook_size=codebook_size) + + @torch.inference_mode() + def encode(self, audio: Any, sample_rate: int | None = None) -> torch.Tensor: + """Encode one waveform or audio source into ``(T, Q)`` codes.""" + if isinstance(audio, torch.Tensor): + audio = audio.detach().cpu().numpy() + elif isinstance(audio, list) and audio and all(isinstance(x, (int, float)) for x in audio): + audio = np.asarray(audio, dtype=np.float32) + if isinstance(audio, np.ndarray) and audio.ndim > 1: + # Accept both common layouts: (samples, channels) and + # (channels, samples). Qwen3TTSTokenizer expects one 1-D clip. + channel_first = audio.shape[0] <= 8 and audio.shape[-1] > audio.shape[0] + audio = audio.mean(axis=0 if channel_first else -1) + + if isinstance(audio, np.ndarray) and sample_rate is None: + raise ValueError("sample_rate is required when ref_audio is a numpy waveform") + + encoded = self.tokenizer.encode(audio, sr=sample_rate) + return self._normalize_codes(encoded) + + def _normalize_codes(self, encoded: Any) -> torch.Tensor: + codes = getattr(encoded, "audio_codes", None) + if codes is None and isinstance(encoded, Mapping): + codes = encoded.get("audio_codes") + if codes is None and isinstance(encoded, (tuple, list)): + codes = encoded[0] if encoded else None + if codes is None: + raise ValueError("reference audio tokenizer returned no audio_codes") + + if isinstance(codes, (tuple, list)) and codes: + # Qwen3TTSTokenizer returns one tensor per input waveform. + if len(codes) != 1: + raise ValueError(f"only one reference waveform is supported, got {len(codes)}") + codes = codes[0] + codes = torch.as_tensor(codes) + if codes.ndim == 3: + if codes.shape[0] != 1: + raise ValueError(f"only one reference waveform is supported, got {tuple(codes.shape)}") + codes = codes[0] + if codes.ndim != 2: + raise ValueError(f"reference audio codes must be 2D, got {tuple(codes.shape)}") + if codes.shape[0] == 0: + raise ValueError("reference audio tokenizer returned no frames") + if codes.shape[-1] != self.num_codebooks: + if codes.shape[0] == self.num_codebooks: + codes = codes.transpose(0, 1) + else: + raise ValueError( + f"reference audio codes must have {self.num_codebooks} codebooks, got {tuple(codes.shape)}" + ) + codes = codes.to(device="cpu", dtype=torch.long).contiguous() + if codes.numel() and (int(codes.min()) < 0 or int(codes.max()) >= self.codebook_size): + raise ValueError( + f"reference audio code outside [0, {self.codebook_size}): " + f"min={int(codes.min())}, max={int(codes.max())}" + ) + # int16 is enough for Mimi/Qwen3-TTS codes and halves IPC/cache memory. + return codes.to(dtype=torch.int16) + + +__all__ = ["BreezeReferenceAudioTokenizer", "resolve_audio_tokenizer_path"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py b/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py new file mode 100644 index 000000000..2c04768a9 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/configuration_breeze_tts_2.py @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Configuration classes for Breeze-TTS-2.""" + +from __future__ import annotations + +from typing import Any + +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig +from transformers.models.qwen3 import Qwen3Config + + +class BreezeTTS2Config(PretrainedConfig): + """Configuration for ``BreezeForConditionalGeneration``.""" + + model_type = "breeze" + + def __init__( + self, + backbone_config: dict | Qwen3Config | None = None, + text_encoder_config: dict | BreezeTTS2TextEncoderConfig | None = None, + depth_decoder_config: dict | BreezeTTS2DepthDecoderConfig | None = None, + codec_config: dict | None = None, + architectures: list[str] | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + audio_embed_size: int = 2048, + audio_eos_token_id: int = 262145, + audio_num_codebooks: int = 16, + audio_token_id: int = 262144, + audio_vocab_size: int = 2051, + backbone_flavor: str = "llama-1B", + backbone_model_name_or_path: str | None = None, + backbone_model_type: str = "qwen3", + bos_token_id: int = 2, + codebook_eos_token_id: int = 0, + codebook_pad_token_id: int = 2050, + decoder_flavor: str = "llama-100M", + depth_header_loss_weight: float = 1.0, + dtype: str = "bfloat16", + eos_token_id: int = 1, + head_dim: int = 128, + hidden_act: str = "silu", + hidden_size: int = 2048, + initializer_range: float = 0.02, + intermediate_size: int = 6144, + max_position_embeddings: int = 2048, + mlp_bias: bool = False, + num_attention_heads: int = 16, + num_codebooks: int = 16, + num_hidden_layers: int = 28, + num_key_value_heads: int = 8, + pad_token_id: int = 0, + rms_norm_eps: float = 1e-5, + rope_scaling: dict | None = None, + rope_theta: float = 500_000.0, + text_encoder_bucket_max_length_ratio: float = 4.0, + text_encoder_lora_config: dict | None = None, + text_encoder_proj_type: str = "linear", + text_encoder_special_tokens_config: dict | None = None, + text_vocab_size: int = 262158, + tie_codebooks_embeddings: bool = True, + tie_word_embeddings: bool = False, + use_cache: bool = True, + vocab_size: int = 2051, + sampling_rate: int = 24000, + codec_model_name_or_path: str = "kyutai/mimi", + **kwargs: Any, + ) -> None: + serialized_rope_parameters = kwargs.pop("rope_parameters", None) + if rope_scaling is None and serialized_rope_parameters is not None: + rope_scaling = serialized_rope_parameters + # Breeze stores the Qwen3 backbone in backbone_config. This follows the + # MOSS-TTS pattern and deliberately does not declare sub_configs. + if backbone_config is None: + backbone_config = {} + if isinstance(backbone_config, dict): + backbone_config = dict(backbone_config) + backbone_config.pop("model_type", None) + self.backbone_config = Qwen3Config(**backbone_config) + else: + self.backbone_config = backbone_config + + if text_encoder_config is None: + text_encoder_config = {} + if isinstance(text_encoder_config, dict): + text_encoder_config = dict(text_encoder_config) + text_encoder_config.pop("model_type", None) + self.text_encoder_config = BreezeTTS2TextEncoderConfig(**text_encoder_config) + else: + self.text_encoder_config = text_encoder_config + + if depth_decoder_config is None: + depth_decoder_config = {} + if isinstance(depth_decoder_config, dict): + depth_decoder_config = dict(depth_decoder_config) + depth_decoder_config.pop("model_type", None) + self.depth_decoder_config = BreezeTTS2DepthDecoderConfig(**depth_decoder_config) + else: + self.depth_decoder_config = depth_decoder_config + + # Nested configs must exist before the parent constructor: Transformers + # 5.x may call get_text_config() during config validation. + super().__init__( + architectures=architectures, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.audio_embed_size = audio_embed_size + self.audio_eos_token_id = audio_eos_token_id + self.audio_num_codebooks = audio_num_codebooks + self.audio_token_id = audio_token_id + self.audio_vocab_size = audio_vocab_size + # Keep Mimi's nested config as a plain checkpoint dictionary. Stage 1 + # owns codec construction; stage 0 only needs codebook_size here. + self.codec_config = codec_config + self.backbone_flavor = backbone_flavor + self.backbone_model_name_or_path = backbone_model_name_or_path + self.backbone_model_type = backbone_model_type + self.codebook_eos_token_id = codebook_eos_token_id + self.codebook_pad_token_id = codebook_pad_token_id + self.decoder_flavor = decoder_flavor + self.depth_header_loss_weight = depth_header_loss_weight + self.dtype = dtype + self.head_dim = head_dim + self.hidden_act = hidden_act + self.hidden_size = hidden_size + self.initializer_range = initializer_range + self.intermediate_size = intermediate_size + self.max_position_embeddings = max_position_embeddings + self.mlp_bias = mlp_bias + self.num_attention_heads = num_attention_heads + self.num_codebooks = num_codebooks + self.num_hidden_layers = num_hidden_layers + self.num_key_value_heads = num_key_value_heads + self.rms_norm_eps = rms_norm_eps + if rope_scaling is not None: + # Preserve the upstream legacy field while also satisfying the + # normalized ``rope_parameters`` validation added in Transformers 5. + rope_scaling = dict(rope_scaling) + rope_scaling.setdefault("rope_type", rope_scaling.get("type", "default")) + rope_scaling.setdefault("rope_theta", rope_theta) + self.rope_scaling = rope_scaling + self.rope_theta = rope_theta + self.text_encoder_bucket_max_length_ratio = text_encoder_bucket_max_length_ratio + self.text_encoder_lora_config = text_encoder_lora_config or {} + self.text_encoder_proj_type = text_encoder_proj_type + self.text_encoder_special_tokens_config = text_encoder_special_tokens_config or {} + self.text_vocab_size = text_vocab_size + self.tie_codebooks_embeddings = tie_codebooks_embeddings + self.use_cache = use_cache + self.vocab_size = vocab_size + self.sampling_rate = sampling_rate + self.codec_model_name_or_path = codec_model_name_or_path + + if not getattr(self, "architectures", None): + self.architectures = ["BreezeForConditionalGeneration"] + + def get_text_config(self, **_: Any) -> Qwen3Config: + return self.backbone_config + + +class BreezeTTS2TextEncoderConfig(PretrainedConfig): + """Configuration for Breeze's T5Gemma2 text encoder.""" + + model_type = "t5gemma2_text" + + def __init__( + self, + hidden_size: int = 1152, + intermediate_size: int = 6912, + num_hidden_layers: int = 26, + num_attention_heads: int = 4, + num_key_value_heads: int = 1, + head_dim: int = 256, + max_position_embeddings: int = 32768, + hidden_activation: str = "gelu_pytorch_tanh", + rms_norm_eps: float = 1e-6, + rope_theta: float = 10_000.0, + sliding_window: int = 512, + query_pre_attn_scalar: int = 256, + eoi_token_index: int = 256000, + vocab_size: int = 262158, + layer_types: list[str] | None = None, + rope_parameters: dict | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + dropout_rate: float = 0.0, + initializer_range: float = 0.02, + tie_word_embeddings: bool = True, + use_bidirectional_attention: bool = False, + use_cache: bool = True, + requires_grad: bool = False, + preferred_attn_implementation: str = "flash_attention_2", + dtype: str = "bfloat16", + is_encoder_decoder: bool = False, + bos_token_id: int = 2, + eos_token_id: int = 1, + pad_token_id: int = 0, + **kwargs: Any, + ) -> None: + kwargs.pop("model_type", None) + + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + is_encoder_decoder=is_encoder_decoder, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.max_position_embeddings = max_position_embeddings + self.hidden_activation = hidden_activation + self.hidden_act = hidden_activation + self.rms_norm_eps = rms_norm_eps + self.rope_theta = rope_theta + self.sliding_window = sliding_window + self.query_pre_attn_scalar = query_pre_attn_scalar + self.eoi_token_index = eoi_token_index + self.vocab_size = vocab_size + if layer_types is None: + # T5Gemma2 uses sliding attention except every sixth layer. + layer_types = [ + "full_attention" if (index + 1) % 6 == 0 else "sliding_attention" for index in range(num_hidden_layers) + ] + self.layer_types = list(layer_types) + if rope_parameters is None: + rope_parameters = { + "full_attention": { + "factor": 8.0, + "rope_theta": 1_000_000.0, + "rope_type": "linear", + }, + "sliding_attention": { + "rope_theta": 10_000.0, + "rope_type": "default", + }, + } + self.rope_parameters = dict(rope_parameters) + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.dropout_rate = dropout_rate + self.initializer_range = initializer_range + self.use_bidirectional_attention = use_bidirectional_attention + self.use_cache = use_cache + self.requires_grad = requires_grad + self.preferred_attn_implementation = preferred_attn_implementation + self.dtype = dtype + + +class BreezeTTS2DepthDecoderConfig(PretrainedConfig): + """Configuration for Breeze's per-frame depth decoder.""" + + model_type = "breeze_depth_decoder_model" + + def __init__( + self, + hidden_size: int = 1024, + backbone_hidden_size: int = 2048, + audio_embed_size: int = 2048, + intermediate_size: int = 8192, + num_hidden_layers: int = 12, + num_attention_heads: int = 8, + num_key_value_heads: int = 2, + head_dim: int = 128, + max_position_embeddings: int = 33, + hidden_act: str = "silu", + rms_norm_eps: float = 1e-5, + rope_theta: float = 500_000.0, + rope_scaling: dict | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + mlp_bias: bool = False, + initializer_range: float = 0.02, + use_cache: bool = True, + num_codebooks: int = 16, + vocab_size: int = 2051, + codebook_loss_weights: list[int] | None = None, + dtype: str = "bfloat16", + tie_word_embeddings: bool = False, + pad_token_id: int | None = 0, + **kwargs: Any, + ) -> None: + kwargs.pop("model_type", None) + # Transformers 5 serializes the normalized field as + # ``rope_parameters`` while the upstream Breeze checkpoint calls it + # ``rope_scaling``. Normalize before ``PretrainedConfig`` validates + # so ``to_dict() -> BreezeTTS2DepthDecoderConfig(**dict)`` round-trips. + serialized_rope_parameters = kwargs.pop("rope_parameters", None) + if rope_scaling is None: + rope_scaling = { + "factor": 32.0, + "high_freq_factor": 0.0078125, + "low_freq_factor": 0.001953125, + "original_max_position_embeddings": 16, + "rope_type": "llama3", + } + if codebook_loss_weights is None: + codebook_loss_weights = [3] + if num_codebooks <= 0 or vocab_size <= 0: + raise ValueError("num_codebooks and vocab_size must be positive") + rope_parameters = dict(serialized_rope_parameters or rope_scaling) + rope_parameters.setdefault("rope_type", rope_parameters.get("type", "default")) + rope_parameters.setdefault("rope_theta", rope_theta) + + super().__init__( + vocab_size=vocab_size, + tie_word_embeddings=tie_word_embeddings, + pad_token_id=pad_token_id, + **kwargs, + ) + # The vLLM depth decoder is a plain ``nn.Module``, so it does not get + # Transformers' PreTrainedModel default attention initialization. + self._attn_implementation = getattr(self, "_attn_implementation", None) or "eager" + self.hidden_size = hidden_size + self.backbone_hidden_size = backbone_hidden_size + self.audio_embed_size = audio_embed_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.max_position_embeddings = max_position_embeddings + self.hidden_act = hidden_act + self.rms_norm_eps = rms_norm_eps + self.rope_theta = rope_theta + self.rope_parameters = dict(getattr(self, "rope_parameters", None) or rope_parameters) + self.rope_parameters.setdefault("rope_theta", rope_theta) + # Transformers 4 exposes this as a plain attribute; Transformers 5 + # implements ``rope_scaling`` as a property backed by rope_parameters. + if not hasattr(type(self), "rope_scaling"): + self.rope_scaling = self.rope_parameters + self.rope_type = rope_scaling.get("rope_type", "default") + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.mlp_bias = mlp_bias + self.initializer_range = initializer_range + self.use_cache = use_cache + self.num_codebooks = num_codebooks + self.vocab_size = vocab_size + self.codebook_loss_weights = list(codebook_loss_weights) + self.dtype = dtype + + +AutoConfig.register("breeze", BreezeTTS2Config, exist_ok=True) + + +__all__ = [ + "BreezeTTS2Config", + "BreezeTTS2DepthDecoderConfig", + "BreezeTTS2TextEncoderConfig", +] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py new file mode 100644 index 000000000..7addb3934 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py @@ -0,0 +1,404 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze-TTS-2 stage-1 codec (codec codes -> waveform). + +The talker emits a frame-major ``(T, 16)`` matrix. The stage input +processor flattens it codebook-major; this module restores the ``(T, 16)`` +layout and delegates decoding to Breeze's bundled Qwen3-TTS tokenizer. A +Mimi fallback remains available for checkpoints that omit ``audio_tokenizer``. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from typing import Any + +import torch +from torch import nn +from vllm.config import VllmConfig + +from vllm_omni.model_executor.models.breeze_tts_2.audio_tokenizer import ( + resolve_audio_tokenizer_path, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + + +def _split_flat_ids(ids: torch.Tensor, counts: list[int] | None) -> list[torch.Tensor]: + if counts is None or len(counts) <= 1: + return [ids] + boundaries = [0] + for count in counts: + boundaries.append(boundaries[-1] + int(count)) + return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(counts))] + + +class BreezeTTS2MimiCodec(nn.Module): + """Synchronous Breeze codec stage used by ``GenerationModelRunner``.""" + + input_modalities = "audio" + requires_raw_input_tokens = True + # Stateful chunk decoding needs the scheduler-side request id, which is + # also used by on_requests_finished to release decoder state. + requires_request_ids = True + have_multimodal_outputs = True + has_preprocess = False + has_postprocess = False + enable_update_additional_information = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + del prefix + self.vllm_config = vllm_config + self.config = vllm_config.model_config.hf_config + self.model_path = vllm_config.model_config.model + self._async_chunk = bool(getattr(vllm_config.model_config, "async_chunk", False)) + # Resolve once: a local directory works directly, an HF repo id + # resolves to its cached snapshot. ``None`` means the checkpoint has + # no bundled tokenizer and the Mimi fallback applies. + self._tokenizer_path = resolve_audio_tokenizer_path(self.model_path) + if self._async_chunk and self._tokenizer_path is None: + raise RuntimeError( + "Breeze-TTS-2 async_chunk requires the bundled Qwen3-TTS audio_tokenizer; " + "the stateless Mimi fallback cannot preserve chunk decoder state" + ) + + self._audio_tokenizer = None + codec_config = getattr(self.config, "codec_config", None) + if isinstance(codec_config, Mapping): + codec_kwargs = dict(codec_config) + elif codec_config is not None: + codec_kwargs = codec_config.to_dict() + else: + codec_kwargs = {} + codec_kwargs.pop("model_type", None) + codec_kwargs.pop("architectures", None) + + # Breeze checkpoints normally carry a Qwen3-TTS decoder in this + # subdirectory. Avoid constructing a second, unused Mimi network in + # that case; this saves both startup time and worker memory. + self._codec = None + if self._tokenizer_path is None: + # Import lazily: stage 0 does not need the Mimi implementation. + from transformers import MimiConfig, MimiModel + + self._codec = MimiModel(MimiConfig(**codec_kwargs)).eval() + for parameter in self._codec.parameters(): + parameter.requires_grad_(False) + + self._num_codebooks = int(getattr(self.config, "num_codebooks", 16)) + self._codebook_size = int(codec_kwargs.get("codebook_size", 2048)) + self._sample_rate = int(codec_kwargs.get("sampling_rate", getattr(self.config, "sampling_rate", 24000))) + self._decoder_state_cache: dict[str, dict[str, Any]] = {} + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + """Return stable dummy embeddings; this stage never samples.""" + 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: + del hidden_states, sampling_metadata + return None + + @torch.inference_mode() + def forward( + self, + input_ids: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + empty = torch.empty(0, dtype=torch.float32) + sr = torch.tensor(self._sample_rate, dtype=torch.int32) + if input_ids is None or input_ids.numel() == 0: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr]}, + ) + # The generation runner's profiling/dummy invocation has no request ids + # or runtime payload. Real requests always provide both. + if runtime_additional_information is None and kwargs.get("request_ids") is None: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr]}, + ) + + ids = input_ids.reshape(-1).to(dtype=torch.long) + counts = kwargs.get("seq_token_counts") + if counts is not None: + counts = [int(item) for item in counts] + requests = _split_flat_ids(ids, counts) + waveforms: list[torch.Tensor] = [] + sample_rates: list[torch.Tensor] = [] + if self._async_chunk: + return self._forward_stateful_chunks( + ids=ids, + requests=requests, + runtime_additional_information=runtime_additional_information, + empty=empty, + sr=sr, + **kwargs, + ) + for index, request_ids in enumerate(requests): + runtime_info = ( + runtime_additional_information[index] + if runtime_additional_information is not None and index < len(runtime_additional_information) + else None + ) + flat = self._payload_codes(request_ids, runtime_info) + if flat.numel() == 0: + waveforms.append(empty.to(device=ids.device)) + sample_rates.append(sr) + continue + if flat.numel() % self._num_codebooks != 0: + raise ValueError( + f"Breeze codec input length {int(flat.numel())} is not divisible by " + f"num_codebooks {self._num_codebooks}" + ) + if int(flat.min()) < 0 or int(flat.max()) >= self._codebook_size: + raise ValueError( + f"Breeze codec id outside [0, {self._codebook_size}): min={int(flat.min())}, max={int(flat.max())}" + ) + frames = flat.numel() // self._num_codebooks + codes = flat.reshape(self._num_codebooks, frames).unsqueeze(0) + waveform = self._decode_codes(codes) + if waveform.ndim == 3: + waveform = waveform[0, 0] + elif waveform.ndim == 2: + waveform = waveform[0] + if waveform.ndim != 1: + raise ValueError(f"Breeze Mimi waveform must be 1D, got {tuple(waveform.shape)}") + # Serving/output serialization consumes CPU tensors. Move only + # the final waveform off the worker device; codec inference itself + # remains on the accelerator. + waveforms.append(waveform.detach().to(device="cpu", dtype=torch.float32).contiguous()) + sample_rates.append(sr) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, + ) + + def _forward_stateful_chunks( + self, + *, + ids: torch.Tensor, + requests: list[torch.Tensor], + runtime_additional_information: list[dict[str, Any]] | None, + empty: torch.Tensor, + sr: torch.Tensor, + **kwargs: Any, + ) -> OmniOutput: + """Decode incremental codebook-major chunks with per-request state.""" + decoder = self._audio_tokenizer.model.decoder + scheduler_request_ids = kwargs.get("request_ids") + if scheduler_request_ids is None: + raise ValueError("Breeze async_chunk codec inputs require scheduler request_ids") + request_ids = [str(request_id) for request_id in scheduler_request_ids] + if len(request_ids) != len(requests): + raise ValueError( + f"Breeze codec request id count {len(request_ids)} does not match input request count {len(requests)}" + ) + valid: list[tuple[str, torch.Tensor, bool]] = [] + for index, request_ids_tensor in enumerate(requests): + runtime_info = ( + runtime_additional_information[index] + if runtime_additional_information is not None and index < len(runtime_additional_information) + else {} + ) + state_id = request_ids[index] + # The connector receiver strips ``meta.finished`` while merging + # runtime metadata, so the stream-terminated signal arrives on + # the model-visible ``stream_finished`` field (the same contract + # CosyVoice3 / Audex use). Fall back to ``finished`` for direct + # in-process callers that still set only the transport flag. + finished = self._meta_bool(runtime_info, "stream_finished") or self._meta_bool(runtime_info, "finished") + flat = self._payload_codes(request_ids_tensor, runtime_info) + if flat.numel() == 0: + # Terminal marker: the stage-0 processor sends an empty + # payload with ``finished=True`` once generation ends. + self._decoder_state_cache.pop(state_id, None) + continue + if flat.numel() % self._num_codebooks != 0: + # A partial frame means the transport payload is corrupt. + # Silently resetting decoder state would corrupt the rest of + # the stream; fail loudly instead. The scheduler aborts the + # request and ``on_requests_finished`` releases the state. + raise ValueError( + f"Breeze codec chunk length {int(flat.numel())} is not divisible by " + f"num_codebooks {self._num_codebooks}" + ) + if int(flat.min()) < 0 or int(flat.max()) >= self._codebook_size: + raise ValueError( + f"Breeze codec id outside [0, {self._codebook_size}): min={int(flat.min())}, max={int(flat.max())}" + ) + frames = int(flat.numel()) // self._num_codebooks + valid.append((state_id, flat.reshape(self._num_codebooks, frames), finished)) + + waveforms = [empty.to(device=ids.device)] * len(requests) + sample_rates = [sr] * len(requests) + if not valid: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, + ) + states: list[dict[str, Any]] = [] + valid_indices: list[int] = [] + max_frames = max(codes.shape[-1] for _, codes, _ in valid) + batched_codes = valid[0][1].new_zeros((len(valid), self._num_codebooks, max_frames)) + for row, (state_id, row_codes, _) in enumerate(valid): + batched_codes[row, :, : row_codes.shape[-1]].copy_(row_codes) + state = self._decoder_state_cache.get(state_id) + if state is None: + state = {} + self._decoder_state_cache[state_id] = state + state.setdefault("prefix_frames", 0) + states.append(state) + valid_indices.append(row) + + decoded = decoder.batched_chunked_decode( + batched_codes, + [int(codes.shape[-1]) for _, codes, _ in valid], + caches=states, + chunk_size=300, + left_context_size=25, + ) + if len(decoded) != len(valid): + raise RuntimeError(f"Breeze codec returned {len(decoded)} chunks for {len(valid)} requests") + for request_index, waveform in zip(valid_indices, decoded, strict=True): + if waveform.ndim == 2 and waveform.shape[0] == 1: + waveform = waveform[0] + elif waveform.ndim != 1: + waveform = waveform.reshape(-1) + waveforms[request_index] = waveform.detach().to(device="cpu", dtype=torch.float32).contiguous() + + for state_id, _, finished in valid: + if finished: + self._decoder_state_cache.pop(state_id, None) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": waveforms, "sr": sample_rates}, + ) + + @staticmethod + def _meta_bool(info: Mapping[str, Any] | None, key: str) -> bool: + if not isinstance(info, Mapping): + return False + meta = info.get("meta", {}) + if not isinstance(meta, Mapping): + return False + value = meta.get(key, False) + if isinstance(value, torch.Tensor): + return bool(value.reshape(-1)[0].item()) if value.numel() else False + return bool(value) + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + for request_id in finished_req_ids: + self._decoder_state_cache.pop(str(request_id), None) + + def _decode_codes(self, codes: torch.Tensor) -> torch.Tensor: + """Decode one ``(1, Q, T)`` sequence using Breeze's bundled codec.""" + if self._audio_tokenizer is not None: + wavs, sample_rate = self._audio_tokenizer.decode({"audio_codes": [codes[0].transpose(0, 1).contiguous()]}) + self._sample_rate = int(sample_rate) + if not wavs: + raise RuntimeError("Breeze Qwen3-TTS tokenizer returned no waveform") + return torch.as_tensor(wavs[0], device=codes.device, dtype=torch.float32).reshape(-1) + + if self._codec is None: + raise RuntimeError("Breeze codec is not initialized; load_weights() was not completed") + decoded = self._codec.decode(codes) + waveform = getattr(decoded, "audio_values", decoded) + if isinstance(waveform, Mapping): + waveform = waveform.get("audio_values") + if waveform is None: + raise RuntimeError("MimiModel.decode returned no audio_values") + return torch.as_tensor(waveform, device=codes.device, dtype=torch.float32) + + @staticmethod + def _payload_codes(input_ids: torch.Tensor, runtime_info: Mapping[str, Any] | None) -> torch.Tensor: + if isinstance(runtime_info, Mapping): + codes = runtime_info.get("codes") + if isinstance(codes, Mapping): + audio = codes.get("audio") + if audio is not None: + return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1) + # Full-payload transport may flatten nested OmniPayload keys at + # the wire boundary and restore them only after this hook runs. + audio = runtime_info.get("codes.audio") + if audio is not None: + return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1) + return input_ids + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load Breeze's bundled Qwen3-TTS codec, or fallback Mimi weights.""" + if self._tokenizer_path is not None: + # The root checkpoint contains both the talker weights and a Mimi + # ``codec_model.*`` fallback. When the bundled Qwen3-TTS tokenizer + # is present, it owns decoding; consume the root iterator without + # copying the unused fallback into worker memory. + if self._audio_tokenizer is None: + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import ( + Qwen3TTSTokenizer, + ) + + device = self.vllm_config.device_config.device + self._audio_tokenizer = Qwen3TTSTokenizer.from_pretrained( + str(self._tokenizer_path), + device_map=str(device), + ) + self._sample_rate = int(self._audio_tokenizer.model.get_output_sample_rate()) + for _ in weights: + pass + tokenizer_module = getattr(self._audio_tokenizer, "model", None) + named_parameters = getattr(tokenizer_module, "named_parameters", None) + if callable(named_parameters): + return {f"_audio_tokenizer.model.{name}" for name, _ in named_parameters()} + return set() + + params = dict(self._codec.named_parameters()) if self._codec is not None else {} + loaded: set[str] = set() + for name, tensor in weights: + if name.startswith("codec_model."): + target_name = name[len("codec_model.") :] + else: + continue + parameter = params.get(target_name) + if parameter is None: + raise ValueError(f"Unknown Breeze Mimi weight: {target_name}") + loader = getattr(parameter, "weight_loader", None) + if loader is None: + parameter.data.copy_(tensor.to(device=parameter.device, dtype=parameter.dtype)) + else: + loader(parameter, tensor) + loaded.add(target_name) + + if loaded: + return {f"_codec.{name}" for name in loaded} + + # Breeze inference uses the bundled Qwen3-TTS tokenizer for both + # reference encoding and generated-code decoding. Load it once per + # Stage-1 worker; unlike the request-side wrapper this instance is + # placed on the worker device and reused for the entire batch. + # A few Breeze distributions omit codec tensors and point to the + # standalone kyutai/mimi repository instead. Load it once, never per + # request, and fail loudly if it cannot be resolved. + codec_name = getattr(self.config, "codec_model_name_or_path", None) or "kyutai/mimi" + from transformers import AutoModel + + self._codec = AutoModel.from_pretrained(codec_name).eval() + for parameter in self._codec.parameters(): + parameter.requires_grad_(False) + self._sample_rate = int(getattr(self._codec.config, "sampling_rate", self._sample_rate)) + self._codebook_size = int( + getattr(getattr(self._codec.config, "quantizer", None), "cardinality", self._codebook_size) + ) + return {f"_codec.{name}" for name, _ in self._codec.named_parameters()} + + +__all__ = ["BreezeTTS2MimiCodec"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py new file mode 100644 index 000000000..559c88b07 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_depth.py @@ -0,0 +1,523 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections.abc import Callable +from inspect import signature +from typing import Any + +import torch +from torch import nn +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from transformers.integrations import use_kernel_forward_from_hub +from transformers.masking_utils import create_causal_mask +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs, logging +from transformers.utils.deprecation import deprecate_kwarg + +logger = logging.get_logger(__name__) + + +def _create_depth_causal_mask( + *, + config: Any, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.LongTensor | None, + cache_position: torch.LongTensor, +) -> torch.Tensor | None: + """Bridge Transformers 4.57 and 5.x ``create_causal_mask`` contracts.""" + parameters = signature(create_causal_mask).parameters + kwargs: dict[str, Any] = { + "config": config, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + if "inputs_embeds" in parameters: + kwargs["inputs_embeds"] = inputs_embeds + else: + kwargs["input_embeds"] = inputs_embeds + kwargs["cache_position"] = cache_position + return create_causal_mask(**kwargs) + + +class BreezeDepthDecoderModel(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + if hasattr(config, "audio_embed_size") and config.audio_embed_size: + self.audio_embed_size = config.audio_embed_size + else: + self.audio_embed_size = config.backbone_hidden_size + self.embed_tokens = nn.Embedding((config.num_codebooks * config.vocab_size), self.audio_embed_size) + self.layers = nn.ModuleList( + [BreezeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = BreezeRotaryEmbedding(config=config) + self.gradient_checkpointing = False + self.inputs_embeds_projector = nn.Linear(self.audio_embed_size, config.hidden_size, bias=False) + + if config.backbone_hidden_size != self.audio_embed_size: + self.backbone_hidden_state_projector = nn.Linear( + config.backbone_hidden_size, self.audio_embed_size, bias=False + ) + else: + self.backbone_hidden_state_projector = None + + # Weights are loaded from the Breeze checkpoint by the parent talker. + # Avoid PreTrainedModel.post_init here because this module is a plain + # nn.Module used for eager per-frame inference. + + def forward( + self, + input_ids: torch.LongTensor = None, + backbone_last_hidden_state: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPast: + r""" + backbone_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, backbone_hidden_size)`, *optional*): + The last hidden state of the backbone model. This is required when + the first codebook token (generated by the backbone model) is + provided in the `input_ids` argument. + """ + if position_ids is not None and not torch.compiler.is_compiling(): + logger.warning_once( + "Custom `position_ids` were provided but will be ignored. " + "Breeze depth decoder derives them from `cache_position`; " + "they must be identical across the batch." + ) + position_ids = None + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds.") + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + inputs_seq_length = inputs_embeds.shape[1] if inputs_embeds is not None else input_ids.shape[1] + device = inputs_embeds.device if inputs_embeds is not None else input_ids.device + cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_seq_length, device=device) + + if inputs_embeds is None: + codebook_idxs = torch.clamp(cache_position - 1, min=0) + offset = codebook_idxs * self.vocab_size + inputs_embeds = self.embed_tokens(input_ids + offset) + + input_ids_are_first_codebook = cache_position[0] == 0 + if backbone_last_hidden_state is not None: + if self.backbone_hidden_state_projector is not None: + inputs_embeds[:, 0] = self.backbone_hidden_state_projector(backbone_last_hidden_state) + else: + inputs_embeds[:, 0] = backbone_last_hidden_state + else: + if not torch.compiler.is_compiling() and input_ids_are_first_codebook: + logger.warning( + "When the first codebook token is provided, " + "`backbone_last_hidden_state` is required for correct inference." + ) + + inputs_embeds = self.inputs_embeds_projector(inputs_embeds) + + causal_mask = _create_depth_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + cache_position=cache_position, + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_ids = cache_position.unsqueeze(0) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + +@use_kernel_forward_from_hub("RMSNorm") +class BreezeRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + BreezeRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class BreezeRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict): + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class BreezeMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class BreezeAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim, + bias=config.attention_bias, + ) + self.k_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.v_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + + @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class BreezeDecoderLayer(nn.Module): + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = BreezeAttention(config=config, layer_idx=layer_idx) + + self.mlp = BreezeMLP(config) + self.input_layernorm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + cache_position: torch.LongTensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, # necessary, but kept here for BC + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class BreezeCodebooksHead(nn.Module): + """Position-specific heads for codebooks 1..15.""" + + def __init__(self, hidden_size: int, num_codebooks: int, vocab_size: int) -> None: + super().__init__() + self.num_codebooks = int(num_codebooks) + self.weight = nn.Parameter(torch.empty(self.num_codebooks - 1, hidden_size, vocab_size)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # ``hidden_states`` contains the placeholder at position 0 followed + # by codebook-0..N tokens. Return logits for codebook-1..N. + depth_states = hidden_states[:, 1 : self.num_codebooks, :] + if depth_states.shape[1] > self.num_codebooks - 1: + raise ValueError(f"depth decoder received too many input positions: {hidden_states.shape[1]}") + return torch.einsum("bkh,khv->bkv", depth_states, self.weight[: depth_states.shape[1]]) + + +class BreezeDepthDecoderForCausalLM(nn.Module): + """Eager depth decoder used by the vLLM stage-0 talker.""" + + def __init__(self, config) -> None: + super().__init__() + self.config = config + self.model = BreezeDepthDecoderModel(config) + self.codebooks_head = BreezeCodebooksHead( + config.hidden_size, + config.num_codebooks, + config.vocab_size, + ) + + def forward( + self, + input_ids: torch.Tensor, + backbone_last_hidden_state: torch.Tensor, + *, + use_cache: bool = False, + **kwargs, + ) -> CausalLMOutputWithPast: + outputs = self.model( + input_ids=input_ids, + backbone_last_hidden_state=backbone_last_hidden_state, + use_cache=use_cache, + **kwargs, + ) + logits = self.codebooks_head(outputs.last_hidden_state) + return CausalLMOutputWithPast( + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=getattr(outputs, "hidden_states", None), + attentions=getattr(outputs, "attentions", None), + ) + + def load_weights(self, weights): + params_dict = dict(self.named_parameters()) + loaded: set[str] = set() + for name, tensor in weights: + param = params_dict.get(name) + if param is None: + continue + if name == "codebooks_head.weight": + expected = ( + int(self.config.num_codebooks) - 1, + int(self.config.hidden_size), + int(self.config.vocab_size), + ) + actual = tuple(int(dim) for dim in tensor.shape) + if actual != expected: + raise ValueError(f"Breeze depth weight {name!r} has shape {actual}, expected {expected}") + loader = getattr(param, "weight_loader", None) + if loader is None: + param.data.copy_(tensor.to(device=param.device, dtype=param.dtype)) + else: + loader(param, tensor) + loaded.add(name) + return loaded + + +__all__ = [ + "BreezeCodebooksHead", + "BreezeDepthDecoderForCausalLM", + "BreezeDepthDecoderModel", +] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py new file mode 100644 index 000000000..beaa390c8 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py @@ -0,0 +1,618 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze-TTS-2 stage-0 talker for vLLM-Omni. + +The scheduler samples Breeze codebook-0 tokens one at a time. The other 15 +codebooks are kept in the request-local Omni buffer: after each backbone step, +the eager depth decoder completes the frame and the next backbone step +consumes the summed 16-codebook embedding. This preserves the upstream Breeze +generation order without importing HuggingFace GenerationMixin. +""" + +from __future__ import annotations + +import hashlib +import os +from collections.abc import Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + +import torch +from torch import nn +from torch.nn.utils.rnn import pad_sequence +from vllm.config import VllmConfig +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.qwen3 import Qwen3Model +from vllm.sequence import IntermediateTensors +from vllm.v1.sample.metadata import SamplingMetadata + +from vllm_omni.model_executor.models.breeze_tts_2.configuration_breeze_tts_2 import ( + BreezeTTS2Config, +) +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_depth import ( + BreezeDepthDecoderForCausalLM, +) +from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze_tts_2_text import ( + BreezeTTS2TextEncoder, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + + +def _prefixed(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + + +def _normalize_request_infos(infos: Any) -> list[dict[str, Any]] | None: + """Normalize runner metadata to one mutable dict per request.""" + if not isinstance(infos, Sequence) or isinstance(infos, (str, bytes, bytearray)): + return None + if not all(isinstance(info, dict) for info in infos): + return None + return infos + + +def _iter_request_rows( + states: Sequence[dict[str, Any]], + spans: Sequence[tuple[int, int]] | None, + num_rows: int, +) -> Iterable[tuple[dict[str, Any], int, int]]: + """Map request state to logits rows for full or sampled hidden layouts.""" + if num_rows <= 0 or not states: + return + if spans is not None and len(spans) == len(states): + if all(0 <= int(start) <= int(end) <= num_rows for start, end in spans): + for state, (start, end) in zip(states, spans, strict=True): + if int(start) < int(end): + yield state, int(start), int(end) + return + # The runner passes full-sequence spans to make_omni_output but only + # one sampled row per request to compute_logits. + if num_rows == len(states): + for index, state in enumerate(states): + yield state, index, index + 1 + return + if num_rows == len(states): + for index, state in enumerate(states): + yield state, index, index + 1 + return + raise RuntimeError( + "Breeze talker cannot align request state with logits rows: " + f"states={len(states)}, rows={num_rows}, spans={spans}" + ) + + +def _check_weight_shape(name: str, tensor: torch.Tensor, expected: tuple[int, ...]) -> None: + """Fail early for Breeze-owned, non-sharded checkpoint tensors.""" + actual = tuple(int(dim) for dim in tensor.shape) + if actual != expected: + raise ValueError(f"Breeze weight {name!r} has shape {actual}, expected {expected}") + + +class BreezeTTS2TalkerForGeneration(nn.Module): + """Stage-0 Breeze talker: text/audio conditioning -> codec frames.""" + + have_multimodal_outputs = True + has_preprocess = True + has_postprocess = True + # Stage 0 emits a tiny per-frame codec payload. Let the AR runner snapshot + # it through the async-output path instead of the final-only client path. + use_async_omni_output = True + eager_omni_postprocess_before_async_output = True + omni_pooler_payload_include_hidden = False + input_modalities = "text" + gpu_resident_buffer_keys: set[tuple[str, str]] = { + ("hidden_states", "last"), + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config: BreezeTTS2Config = vllm_config.model_config.hf_config + self._async_chunk = bool(vllm_config.model_config.async_chunk) + self._golden_dump_dir = os.environ.get("VLLM_OMNI_BREEZE_GOLDEN_DIR") + + backbone_config = self.config.backbone_config + backbone_vllm_config = vllm_config.with_hf_config( + backbone_config, + architectures=["Qwen3ForCausalLM"], + ) + backbone_vllm_config.model_config.hf_text_config = backbone_config + self.model = Qwen3Model( + vllm_config=backbone_vllm_config, + prefix=_prefixed(prefix, "model"), + ) + # Upstream's Qwen3 adapter keeps only decoder layers/norm/RoPE and + # replaces the standard text embedding with Breeze's summed audio/text + # tables. Breeze always supplies ``inputs_embeds`` through preprocess, + # so remove the unused Qwen embedding instead of allocating and then + # reporting an uninitialized 151936x2048 parameter at load time. + self.model.embed_tokens = nn.Identity() + + self.hidden_size = int(backbone_config.hidden_size) + self.num_codebooks = int(self.config.num_codebooks) + self.codebook_vocab_size = int(self.config.vocab_size) + codec_config = self.config.codec_config + self.codebook_size = ( + int(codec_config.get("codebook_size", self.codebook_vocab_size)) + if isinstance(codec_config, Mapping) + else self.codebook_vocab_size + ) + self.codebook_eos_token_id = int(self.config.codebook_eos_token_id) + self.audio_token_id = int(self.config.audio_token_id) + self.audio_eos_token_id = int(self.config.audio_eos_token_id) + self.text_pad_token_id = int(self.config.pad_token_id) + + self.lm_head = ParallelLMHead( + self.codebook_vocab_size + 1, + self.hidden_size, + bias=False, + prefix=_prefixed(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(self.codebook_vocab_size + 1) + + # Upstream Breeze uses embedding(code_i + i * vocab_size) and sums + # over codebooks before entering the Qwen3 backbone. + self.embed_audio_tokens = nn.Embedding( + self.num_codebooks * self.codebook_vocab_size, + self.hidden_size, + ) + # Keep Breeze's standalone text table even though the normal path + # uses T5Gemma2 outputs. It is part of the checkpoint contract and is + # also the correct fallback when a checkpoint omits text_encoder. + self.embed_text_tokens = nn.Embedding(int(self.config.text_vocab_size), self.hidden_size) + self.register_buffer( + "audio_token_offsets", + torch.arange(self.num_codebooks, dtype=torch.long) * self.codebook_vocab_size, + persistent=False, + ) + + self.text_encoder = BreezeTTS2TextEncoder(self.config.text_encoder_config) + self.text_encoder_proj = nn.Linear( + int(self.config.text_encoder_config.hidden_size), + self.hidden_size, + bias=False, + ) + self.depth_decoder = BreezeDepthDecoderForCausalLM(self.config.depth_decoder_config) + if bool(self.config.tie_codebooks_embeddings): + # The upstream HF model ties these tables. Sharing the Parameter + # also prevents a tied checkpoint from leaving the depth input + # embedding randomly initialized when only one key is serialized. + self.depth_decoder.model.embed_tokens.weight = self.embed_audio_tokens.weight + + self._codec_disallowed_mask: torch.Tensor | None = None + self._batch_state: list[dict[str, Any]] | None = None + self._batch_state_spans: Sequence[tuple[int, int]] | None = None + + # ------------------------------------------------------------------ + # Embedding and backbone hooks + # ------------------------------------------------------------------ + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + """Embed complete generated codec frames; scalar ids are invalid.""" + if input_ids.ndim < 1 or input_ids.shape[-1] != self.num_codebooks: + raise ValueError( + "Breeze embed_input_ids requires complete codec frames with shape " + f"(..., {self.num_codebooks}), got {tuple(input_ids.shape)}" + ) + codes = input_ids.to(dtype=torch.long) + offsets = self.audio_token_offsets.to(device=codes.device) + ids = codes.clamp(0, self.codebook_vocab_size - 1) + offsets + return self.embed_audio_tokens(ids).sum(dim=-2) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + # ------------------------------------------------------------------ + # Prompt/text conditioning + # ------------------------------------------------------------------ + + def _encode_text_segments( + self, + input_ids: torch.Tensor, + text_mask: torch.Tensor, + text_lengths: torch.Tensor, + ) -> torch.Tensor: + text_ids = input_ids[text_mask.to(device=input_ids.device, dtype=torch.bool)].to(torch.long) + lengths = [int(x) for x in text_lengths.reshape(-1).tolist() if int(x) > 0] + if int(text_ids.numel()) != sum(lengths): + raise ValueError(f"Breeze text mask/length mismatch: {text_ids.numel()} != {sum(lengths)}") + if not lengths: + return input_ids.new_empty((0, self.hidden_size), dtype=self.text_encoder.dtype) + segments = list(torch.split(text_ids, lengths)) + padded = pad_sequence(segments, batch_first=True, padding_value=0) + attn = torch.zeros_like(padded, dtype=torch.bool) + for row, length in enumerate(lengths): + attn[row, :length] = True + hidden = self.text_encoder(padded, attention_mask=attn) + projected = self.text_encoder_proj(hidden) + return torch.cat([projected[row, :length] for row, length in enumerate(lengths)], dim=0) + + def _build_prompt_embeddings(self, prompt_ids: torch.Tensor, info: Mapping[str, Any]) -> torch.Tensor: + device = next(self.parameters()).device + prompt_ids = prompt_ids.reshape(-1).to(device=device, dtype=torch.long) + text_mask = torch.as_tensor(info.get("text_ids_mask"), device=device, dtype=torch.bool).reshape(-1) + text_lengths = torch.as_tensor(info.get("text_ids_len"), device=device, dtype=torch.long) + if text_mask.numel() != prompt_ids.numel(): + raise ValueError("Breeze text_ids_mask length does not match prompt length") + output = torch.zeros( + (prompt_ids.numel(), self.hidden_size), + device=device, + dtype=self.embed_text_tokens.weight.dtype, + ) + text_embeds = self._encode_text_segments(prompt_ids, text_mask, text_lengths) + output[text_mask] = text_embeds.to(dtype=output.dtype) + + audio_mask = prompt_ids == self.audio_token_id + input_values = info.get("input_values") + if bool(audio_mask.any()): + if input_values is None: + raise ValueError("Breeze audio placeholders require input_values") + codes = torch.as_tensor(input_values, device=device, dtype=torch.long) + if codes.ndim != 2 or codes.shape[-1] != self.num_codebooks: + raise ValueError(f"Breeze input_values must be (T, {self.num_codebooks})") + if int(codes.shape[0]) != int(audio_mask.sum()): + raise ValueError("Breeze reference code count does not match audio placeholders") + output[audio_mask] = self.embed_input_ids(codes).to(dtype=output.dtype) + + eos_mask = prompt_ids == self.audio_eos_token_id + if bool(eos_mask.any()): + eos_codes = torch.full( + (1, self.num_codebooks), + self.codebook_eos_token_id, + dtype=torch.long, + device=device, + ) + output[eos_mask] = self.embed_input_ids(eos_codes).expand(int(eos_mask.sum()), -1) + return output + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build prefill embeddings and consume one completed frame per decode.""" + del input_embeds + info = dict(info_dict) + span_len = int(input_ids.shape[0]) + is_prefill = bool(info.get("_omni_is_prefill", span_len > 1)) + if is_prefill: + full_ids = info.get("prompt_ids") + if not isinstance(full_ids, torch.Tensor): + raise RuntimeError("Breeze prefill is missing prompt_ids") + full_embeds = info.get("breeze_prompt_embeds") + if not isinstance(full_embeds, torch.Tensor): + full_embeds = self._build_prompt_embeddings(full_ids, info).detach().cpu().contiguous() + offset = int(info.get("breeze_prefill_offset", info.get("_omni_num_computed_tokens", 0)) or 0) + take = full_embeds[offset : offset + span_len].to(device=input_ids.device) + if take.shape[0] != span_len: + raise RuntimeError("Breeze prompt embedding slice is shorter than scheduled span") + return ( + torch.full_like(input_ids, self.text_pad_token_id), + take, + { + "breeze_prompt_embeds": full_embeds, + "breeze_prefill_offset": offset + span_len, + "breeze_generated_frames": int(info.get("breeze_generated_frames", 0) or 0), + "breeze_max_new_frames": int(info.get("breeze_max_new_frames", -1)), + }, + ) + + frame = info.get("breeze_current_frame") + if isinstance(frame, torch.Tensor) and frame.numel() == self.num_codebooks: + frame = frame.to(device=input_ids.device, dtype=torch.long).reshape(1, 1, -1) + embeds = self.embed_input_ids(frame).reshape(1, -1) + safe_id = frame[..., 0].reshape_as(input_ids).to(dtype=input_ids.dtype) + return safe_id, embeds, {} + has_breeze_state = any( + key in info + for key in ( + "breeze_current_frame", + "breeze_audio_codes", + "breeze_generated_frames", + "breeze_prompt_embeds", + ) + ) + if has_breeze_state: + # A real decode must consume the full frame produced from the + # preceding backbone hidden state. A scalar fallback would silently + # fill codebooks 1..15 with EOS and corrupt AR state. + raise RuntimeError( + "Breeze decode is missing breeze_current_frame; refusing scalar fallback " + "because it would fill codebooks 1..15 with EOS" + ) + # Runner-only dummy execution has no Breeze codec state. It must not + # invent a semantically invalid audio frame; a stable zero embedding is + # sufficient for shape/profiling warmup. + dummy_embeds = torch.zeros( + (*input_ids.shape, self.hidden_size), + device=input_ids.device, + dtype=self.embed_audio_tokens.weight.dtype, + ) + return input_ids, dummy_embeds, {} + + 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()}} + + # ------------------------------------------------------------------ + # Main head, depth decoder and Omni output + # ------------------------------------------------------------------ + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: SamplingMetadata | None = None, + ) -> torch.Tensor | None: + del sampling_metadata + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + logits = self._main_head_logits(hidden_states) + + # ``make_omni_output`` runs immediately before this hook. It records + # requests whose previous hidden state selected EOS or reached the + # frame budget, so the scheduler receives a real stop token instead + # of depending on an arbitrary deploy-level max_tokens timeout. + states = self._batch_state or [] + if states: + for state, row_start, row_end in _iter_request_rows( + states, + getattr(self, "_batch_state_spans", None), + int(logits.shape[0]), + ): + if state.get("breeze_force_eos", False): + row = logits[row_start:row_end] + row.fill_(float("-inf")) + row[:, self.codebook_vocab_size] = 0.0 + return logits + + def _main_head_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is None: + raise RuntimeError("Breeze main head returned no logits") + if self._codec_disallowed_mask is None or self._codec_disallowed_mask.device != logits.device: + mask = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) + if self.codebook_size < self.codebook_vocab_size: + mask[self.codebook_size : self.codebook_vocab_size] = True + self._codec_disallowed_mask = mask + return logits.masked_fill(self._codec_disallowed_mask, float("-inf")) + + @torch.no_grad() + def _generate_depth_codes(self, backbone_hidden: torch.Tensor, code0: int) -> torch.Tensor: + sequence = torch.tensor([[0, code0]], dtype=torch.long, device=backbone_hidden.device) + generated = [int(code0)] + for _ in range(self.num_codebooks - 1): + depth_out = self.depth_decoder( + input_ids=sequence, + backbone_last_hidden_state=backbone_hidden.reshape(1, -1), + use_cache=False, + ) + logits = depth_out.logits[:, -1, :].clone() + if self.codebook_size < logits.shape[-1]: + logits[:, self.codebook_size :] = float("-inf") + next_code = int(torch.argmax(logits, dim=-1).item()) + generated.append(next_code) + sequence = torch.cat([sequence, torch.tensor([[next_code]], device=sequence.device)], dim=1) + return torch.tensor(generated, dtype=torch.long, device=backbone_hidden.device).unsqueeze(0) + + @staticmethod + def _accumulate_codes(info: dict[str, Any], frame: torch.Tensor) -> torch.Tensor: + """Append one ``(1, Q)`` frame to the request's cumulative codes.""" + accumulated = info.get("breeze_audio_codes") + accumulated = ( + torch.cat([accumulated.to(frame.device), frame], dim=0) + if isinstance(accumulated, torch.Tensor) and accumulated.numel() + else frame + ) + info["breeze_audio_codes"] = accumulated + return accumulated + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + self._batch_state = None + self._batch_state_spans = None + return model_outputs + hidden = model_outputs + infos = kwargs.get("model_intermediate_buffer", []) + spans = kwargs.get("request_token_spans") + info_items = _normalize_request_infos(infos) + if info_items is None: + self._batch_state = None + self._batch_state_spans = None + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + # The runner invokes compute_logits after this method, so retain the + # same request ordering and spans for per-request EOS masking. + self._batch_state = info_items + self._batch_state_spans = spans + per_request_codes: list[torch.Tensor] = [] + + for index, info in enumerate(info_items): + if spans is not None and index < len(spans): + row_start = max(0, int(spans[index][0])) + row_end = min(int(spans[index][1]), int(hidden.shape[0])) + elif len(info_items) == 1: + row_start, row_end = 0, int(hidden.shape[0]) + else: + raise RuntimeError("Breeze make_omni_output requires request_token_spans for batched decode") + if row_end <= row_start: + per_request_codes.append(hidden.new_empty((0, self.num_codebooks), dtype=torch.long)) + continue + row_hidden = hidden[row_end - 1] + if self._golden_dump_dir and bool(info.get("_omni_is_prefill", False)): + self._dump_golden_prefill(info, row_hidden) + main_logits = self._main_head_logits(row_hidden.reshape(1, -1)) + code0 = int(torch.argmax(main_logits, dim=-1).item()) + generated_frames = int(info.get("breeze_generated_frames", 0) or 0) + max_new_frames = int(info.get("breeze_max_new_frames", -1)) + terminal = code0 == self.codebook_vocab_size or (max_new_frames > 0 and generated_frames >= max_new_frames) + if terminal: + info["breeze_force_eos"] = True + existing = info.get("breeze_audio_codes") + per_request_codes.append( + existing + if isinstance(existing, torch.Tensor) and existing.numel() and not self._async_chunk + else hidden.new_empty((0, self.num_codebooks), dtype=torch.long) + ) + continue + if code0 < 0 or code0 >= self.codebook_size: + raise RuntimeError(f"Breeze main head produced invalid codec id {code0}") + + frame = self._generate_depth_codes(row_hidden, code0) + info["breeze_current_frame"] = frame[0].detach() + info["breeze_generated_frames"] = generated_frames + 1 + info["breeze_force_eos"] = False + if self._async_chunk: + # Streaming consumers only ever read the per-frame tail. The + # cumulative buffer exists for full-payload transport and the + # opt-in golden dump; maintaining it here would also cost an + # O(T^2) tensor-cat per request. + if self._golden_dump_dir: + self._accumulate_codes(info, frame) + per_request_codes.append(frame) + else: + per_request_codes.append(self._accumulate_codes(info, frame)) + if self._golden_dump_dir: + for info in info_items: + self._dump_golden_frames(info) + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={"codes": {"audio": per_request_codes}}) + + def _dump_golden_frames(self, info: dict[str, Any]) -> None: + """Write terminal codec frames for opt-in golden-parity diagnostics.""" + codes = info.get("breeze_audio_codes") + if not info.get("breeze_force_eos", False) or not isinstance(codes, torch.Tensor) or codes.numel() == 0: + return + prompt_ids = info.get("prompt_ids") + if not isinstance(prompt_ids, torch.Tensor): + raise RuntimeError("Breeze golden dump is missing prompt_ids") + digest = hashlib.sha256(prompt_ids.detach().cpu().numpy().tobytes()).hexdigest()[:16] + output_dir = Path(self._golden_dump_dir) # type: ignore[arg-type] + output_dir.mkdir(parents=True, exist_ok=True) + torch.save( + { + "template": info.get("template"), + "prompt_ids": prompt_ids.cpu(), + "codes": codes.detach().cpu().to(torch.long).contiguous(), + "generated_frames": int(info.get("breeze_generated_frames", 0) or 0), + "finished": True, + }, + output_dir / f"breeze_{digest}.pt", + ) + + def _dump_golden_prefill(self, info: dict[str, Any], hidden: torch.Tensor) -> None: + """Write the terminal-prompt hidden row for opt-in parity diagnostics.""" + prompt_ids = info.get("prompt_ids") + if not isinstance(prompt_ids, torch.Tensor): + raise RuntimeError("Breeze golden dump is missing prompt_ids") + digest = hashlib.sha256(prompt_ids.detach().cpu().numpy().tobytes()).hexdigest()[:16] + output_dir = Path(self._golden_dump_dir) # type: ignore[arg-type] + output_dir.mkdir(parents=True, exist_ok=True) + torch.save( + { + "template": info.get("template"), + "prompt_ids": prompt_ids.cpu(), + "hidden": hidden.detach().reshape(-1).cpu().to(torch.float32).contiguous(), + }, + output_dir / f"prefill_{digest}.pt", + ) + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + """Drop cross-step references retained for the finished AR requests.""" + del finished_req_ids + # Codec/frame state lives in request-local additional_information and is + # freed with the Request. Only the short logits-to-state mapping is + # model-owned and must not leak into the next micro-batch. + self._batch_state = None + self._batch_state_spans = None + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded: set[str] = set() + params = dict(self.named_parameters()) + backbone_weights: list[tuple[str, torch.Tensor]] = [] + text_weights: list[tuple[str, torch.Tensor]] = [] + depth_weights: list[tuple[str, torch.Tensor]] = [] + for name, tensor in weights: + # This table belongs to Breeze's custom audio embedding module, + # although its checkpoint key is nested under backbone_model. + # Handle it before forwarding the remaining backbone weights to + # Qwen3Model. + if name == "backbone_model.embed_tokens.embed_audio_tokens.weight": + target = "embed_audio_tokens.weight" + param = params.get(target) + if param is None: + raise ValueError(f"Missing Breeze parameter: {target}") + _check_weight_shape(name, tensor, tuple(param.shape)) + default_weight_loader(param, tensor) + loaded.add(target) + continue + if name.startswith("backbone_model."): + # Qwen3Model.load_weights() is called on the nested module and + # therefore expects names relative to that module (the same + # convention used by MOSS-TTS), not the outer ``model.`` + # namespace used by ``named_parameters()``. + backbone_weights.append((name[len("backbone_model.") :], tensor)) + continue + if name.startswith("text_encoder."): + text_name = name[len("text_encoder.") :] + text_weights.append((text_name, tensor)) + continue + if name.startswith("depth_decoder."): + depth_weights.append((name[len("depth_decoder.") :], tensor)) + continue + if name in {"lm_head.weight", "embed_text_tokens.weight"} or name.startswith("text_encoder_proj."): + target = name + else: + continue + param = params.get(target) + if param is None: + continue + # ParallelLMHead stores a TP-local shard while checkpoint tensors + # are global; validate against the global Breeze head shape. + expected_shape = ( + (self.codebook_vocab_size + 1, self.hidden_size) if name == "lm_head.weight" else tuple(param.shape) + ) + _check_weight_shape(name, tensor, expected_shape) + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, tensor) + loaded.add(target) + + if backbone_weights: + loaded.update("model." + key for key in self.model.load_weights(backbone_weights)) + if text_weights: + loaded.update("text_encoder." + key for key in self.text_encoder.load_weights(text_weights)) + if depth_weights: + loaded.update("depth_decoder." + key for key in self.depth_decoder.load_weights(depth_weights)) + if bool(getattr(self.config, "tie_codebooks_embeddings", False)): + # The checkpoint serializes only one side of the tied table. + # ``named_parameters(remove_duplicate=True)`` exposes it as the + # talker-owned ``embed_audio_tokens.weight``, so report both + # aliases to the strict weight tracker. + if "depth_decoder.model.embed_tokens.weight" in loaded: + loaded.add("embed_audio_tokens.weight") + if "embed_audio_tokens.weight" in loaded: + loaded.add("depth_decoder.model.embed_tokens.weight") + return loaded + + +__all__ = ["BreezeTTS2TalkerForGeneration"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py new file mode 100644 index 000000000..0276e8278 --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_text.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze T5Gemma2 text encoder. + +The Breeze checkpoint contains a T5Gemma2 encoder, not the older T5Gemma +implementation used by some diffusion models. This module keeps the encoder +local to Breeze so checkpoint names and attention semantics stay explicit, +while using vLLM tensor-parallel layers for the expensive projections. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.activation import get_act_fn +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + + +class _RMSNorm(nn.Module): + """T5Gemma2 RMSNorm, whose checkpoint stores a zero-centered weight.""" + + def __init__(self, dim: int, eps: float) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(dim)) + self.eps = float(eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + normalized = hidden_states.float() * torch.rsqrt( + hidden_states.float().pow(2).mean(dim=-1, keepdim=True) + self.eps + ) + return (normalized * (1.0 + self.weight.float())).to(hidden_states.dtype) + + +class _RotaryEmbedding(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + self.config = config + for layer_type in dict.fromkeys(config.layer_types): + params = dict(config.rope_parameters.get(layer_type, {})) + base = float(params.get("rope_theta", 10_000.0)) + factor = float(params.get("factor", 1.0)) + dim = int(config.head_dim) + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if params.get("rope_type") == "linear": + inv_freq = inv_freq / factor + elif params.get("rope_type", "default") != "default": + raise ValueError(f"Unsupported T5Gemma2 rope type: {params.get('rope_type')!r}") + self.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + layer_type: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = getattr(self, f"{layer_type}_inv_freq") + expanded = inv_freq[None, :, None].to(device=hidden_states.device) + positions = position_ids[:, None, :].float() + freqs = (expanded.float() @ positions).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos().to(hidden_states.dtype), emb.sin().to(hidden_states.dtype) + + +def _rotate_half(value: torch.Tensor) -> torch.Tensor: + first, second = value[..., : value.shape[-1] // 2], value[..., value.shape[-1] // 2 :] + return torch.cat((-second, first), dim=-1) + + +def _apply_rotary( + query: torch.Tensor, + key: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return (query * cos) + (_rotate_half(query) * sin), (key * cos) + (_rotate_half(key) * sin) + + +def _repeat_kv(hidden_states: torch.Tensor, repeats: int) -> torch.Tensor: + if repeats == 1: + return hidden_states + batch, heads, seq_len, head_dim = hidden_states.shape + return ( + hidden_states[:, :, None, :, :] + .expand(batch, heads, repeats, seq_len, head_dim) + .reshape(batch, heads * repeats, seq_len, head_dim) + ) + + +def _attention_mask( + padding_mask: torch.Tensor | None, + *, + batch_size: int, + seq_len: int, + layer_type: str, + sliding_window: int | None, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor | None: + if layer_type == "full_attention" and padding_mask is None: + return None + valid = ( + torch.ones((batch_size, seq_len), device=device, dtype=torch.bool) + if padding_mask is None + else padding_mask.to(device=device, dtype=torch.bool) + ) + allowed = valid[:, None, None, :].expand(batch_size, 1, seq_len, seq_len) + if layer_type == "sliding_attention": + if sliding_window is None: + raise ValueError("T5Gemma2 sliding_attention requires sliding_window") + q_idx = torch.arange(seq_len, device=device)[:, None] + kv_idx = torch.arange(seq_len, device=device)[None, :] + left = (int(sliding_window) + 1) // 2 + right = int(sliding_window) // 2 + 1 + distance = q_idx - kv_idx + local = ((distance >= 0) & (distance < left)) | ((distance < 0) & (-distance < right)) + allowed = allowed & local[None, None, :, :] + mask = torch.zeros((batch_size, 1, seq_len, seq_len), device=device, dtype=dtype) + mask.masked_fill_(~allowed, torch.finfo(dtype).min) + return mask + + +class _SelfAttention(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + tp_size = get_tensor_model_parallel_world_size() + if config.num_attention_heads % tp_size != 0: + raise ValueError("T5Gemma2 attention heads must be divisible by tensor parallel size") + self.layer_type = config.layer_type + self.head_dim = int(config.head_dim) + self.num_heads = int(config.num_attention_heads) // tp_size + self.num_kv_heads = max(1, int(config.num_key_value_heads) // tp_size) + self.num_key_value_groups = self.num_heads // self.num_kv_heads + self.scaling = float(config.query_pre_attn_scalar) ** -0.5 + self.qkv_proj = QKVParallelLinear( + hidden_size=config.hidden_size, + head_size=self.head_dim, + total_num_heads=config.num_attention_heads, + total_num_kv_heads=config.num_key_value_heads, + bias=config.attention_bias, + ) + self.o_proj = RowParallelLinear( + input_size=config.num_attention_heads * self.head_dim, + output_size=config.hidden_size, + bias=config.attention_bias, + input_is_parallel=True, + ) + self.q_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) + self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q_size = self.num_heads * self.head_dim + kv_size = self.num_kv_heads * self.head_dim + query, key, value = qkv.split((q_size, kv_size, kv_size), dim=-1) + shape = (*hidden_states.shape[:-1], -1, self.head_dim) + query = self.q_norm(query.view(shape).transpose(1, 2)) + key = self.k_norm(key.view(shape).transpose(1, 2)) + value = value.view(shape).transpose(1, 2) + query, key = _apply_rotary(query, key, *position_embeddings) + key = _repeat_kv(key, self.num_key_value_groups) + value = _repeat_kv(value, self.num_key_value_groups) + # SDPA scales by head_dim**-0.5; adjust Q to T5Gemma2's scalar. + query = query * (self.scaling / (self.head_dim**-0.5)) + output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + ) + output = output.transpose(1, 2).contiguous().reshape(*hidden_states.shape[:-1], -1) + output, _ = self.o_proj(output) + return output + + +class _MLP(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[config.intermediate_size, config.intermediate_size], + bias=False, + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + input_is_parallel=True, + ) + self.act_fn = get_act_fn(config.hidden_activation) + self.dropout = nn.Dropout(float(config.dropout_rate)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(hidden_states) + gate, up = gate_up.chunk(2, dim=-1) + output, _ = self.down_proj(self.dropout(self.act_fn(gate) * up)) + return output + + +class _EncoderLayer(nn.Module): + def __init__(self, config: Any, layer_idx: int) -> None: + super().__init__() + layer_config = type("LayerConfig", (), dict(vars(config)))() + layer_config.layer_type = config.layer_types[layer_idx] + self.attention_type = layer_config.layer_type + self.pre_self_attn_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.self_attn = _SelfAttention(layer_config) + self.post_self_attn_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.pre_feedforward_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.mlp = _MLP(config) + self.post_feedforward_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.dropout = nn.Dropout(float(config.dropout_rate)) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.pre_self_attn_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_embeddings, attention_mask) + hidden_states = residual + self.dropout(self.post_self_attn_layernorm(hidden_states)) + residual = hidden_states + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + self.dropout(self.post_feedforward_layernorm(hidden_states)) + return hidden_states + + +class BreezeTTS2TextEncoder(nn.Module): + """Frozen T5Gemma2 encoder returning ``(B, S, hidden_size)`` states.""" + + def __init__(self, config: Any) -> None: + super().__init__() + self.config = config + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.embed_scale = float(config.hidden_size**0.5) + self.eoi_token_index = int(getattr(config, "eoi_token_index", 256000)) + # T5Gemma2 stores this learned replacement vector under the input + # embedding module (``embed_tokens.eoi_embedding``), so keep the same + # parameter path for checkpoint loading and tied-weight inspection. + self.embed_tokens.eoi_embedding = nn.Parameter(torch.zeros(config.hidden_size)) + self.layers = nn.ModuleList([_EncoderLayer(config, i) for i in range(config.num_hidden_layers)]) + self.norm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.rotary_emb = _RotaryEmbedding(config) + self.dropout = nn.Dropout(float(config.dropout_rate)) + + @property + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if input_ids.ndim != 2: + raise ValueError(f"Breeze text encoder expects 2D input_ids, got {tuple(input_ids.shape)}") + if attention_mask is not None and attention_mask.shape != input_ids.shape: + raise ValueError( + "Breeze text encoder attention_mask must match input_ids: " + f"{tuple(attention_mask.shape)} != {tuple(input_ids.shape)}" + ) + hidden_states = self.embed_tokens(input_ids) * self.embed_scale + eoi_mask = input_ids == self.eoi_token_index + if bool(eoi_mask.any()): + hidden_states = torch.where( + eoi_mask.unsqueeze(-1), + self.embed_tokens.eoi_embedding.to(dtype=hidden_states.dtype), + hidden_states, + ) + hidden_states = self.dropout(hidden_states) + positions = torch.arange(input_ids.shape[1], device=input_ids.device, dtype=torch.long).unsqueeze(0) + position_embeddings = { + layer_type: self.rotary_emb(hidden_states, positions, layer_type) + for layer_type in set(self.config.layer_types) + } + masks = { + layer_type: _attention_mask( + attention_mask, + batch_size=input_ids.shape[0], + seq_len=input_ids.shape[1], + layer_type=layer_type, + sliding_window=self.config.sliding_window, + dtype=hidden_states.dtype, + device=input_ids.device, + ) + for layer_type in set(self.config.layer_types) + } + for layer in self.layers: + hidden_states = layer( + hidden_states, + position_embeddings[layer.attention_type], + masks[layer.attention_type], + ) + return self.dropout(self.norm(hidden_states)) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + 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(self.named_parameters()) + loaded: set[str] = set() + for name, tensor in weights: + target = name + shard_id: str | int | None = None + for packed_name, source_name, source_id in stacked_params_mapping: + marker = f".{source_name}." + if marker in name: + target = name.replace(marker, f".{packed_name}.") + shard_id = source_id + break + parameter = params.get(target) + if parameter is None: + continue + loader = getattr(parameter, "weight_loader", default_weight_loader) + if shard_id is None: + loader(parameter, tensor) + else: + loader(parameter, tensor, shard_id) + loaded.add(target) + return loaded + + +__all__ = ["BreezeTTS2TextEncoder"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py new file mode 100644 index 000000000..d1aaeb13e --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Synchronous Breeze-TTS-2 pipeline: talker -> Mimi codec.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.breeze_tts_2" + +BREEZE_TTS_2_PIPELINE = PipelineConfig( + model_type="breeze", + default_deploy_config_name="breeze_tts_2.yaml", + model_arch="BreezeForConditionalGeneration", + hf_architectures=("BreezeForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="breeze_tts_2", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + async_chunk_process_next_stage_input_func=f"{_PROC}.talker2codec_async_chunk", + custom_process_next_stage_input_func=f"{_PROC}.talker2codec_full_payload", + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [2051], + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="breeze_tts_2_codec", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="BreezeTTS2MimiCodec", + sync_process_input_func=f"{_PROC}.talker2codec", + sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, + retains_state_across_chunks=True, + ), + ), +) + +__all__ = ["BREEZE_TTS_2_PIPELINE"] diff --git a/vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py b/vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py new file mode 100644 index 000000000..32c8145bb --- /dev/null +++ b/vllm_omni/model_executor/models/breeze_tts_2/prompt_builder.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Breeze-TTS-2 request-to-prompt conversion. + +This module is intentionally limited to request preparation. It creates the +real prompt token ids needed by the scheduler and stores the text/audio masks +and reference codes in ``additional_information``. The stage-0 model later +turns those values into embeddings; this keeps tokenization independent from +the Qwen3/text-encoder implementation and avoids changing sequence length in +the model ``preprocess`` hook. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +from transformers import AutoTokenizer + +from vllm_omni.inputs.data import OmniTokensPrompt + +from .audio_tokenizer import BreezeReferenceAudioTokenizer + +AUDIO_TAG = "<|AUDIO|>" +AUDIO_EOS = "<|audio_eos|>" +INSTRUCTION_BOS = "" +INSTRUCTION_EOS = "" + + +@dataclass(frozen=True) +class _TextSegment: + text: str + + +@dataclass(frozen=True) +class _AudioSegment: + source: Any + sample_rate: int | None + + +class BreezeTTS2PromptBuilder: + """Build one Breeze prompt without depending on model modules. + + ``tokenizer`` is a normal HuggingFace tokenizer. ``reference_audio_encoder`` + is injected so the builder can use the serving worker's cached audio + encoder and never reload it per request. The returned ``OmniTokensPrompt`` + contains scheduler token ids and compact CPU metadata for stage-0 + ``preprocess``. + """ + + def __init__( + self, + tokenizer: Any, + config: Any, + reference_audio_encoder: BreezeReferenceAudioTokenizer | None = None, + ) -> None: + self.tokenizer = tokenizer + self.audio_encoder = reference_audio_encoder + self.audio_token_id = int(config.audio_token_id) + self.audio_eos_token_id = int(config.audio_eos_token_id) + # The Breeze tokenizer vocabulary is larger than the Qwen3 backbone + # vocabulary. ``prompt_token_ids`` are scheduler bookkeeping only, so + # keep them in the backbone vocabulary and retain the real Breeze ids + # in ``additional_information['prompt_ids']`` for the talker. + self.scheduler_token_id = int(config.pad_token_id) + scheduler_vocab_size = int(config.backbone_config.vocab_size) + if self.scheduler_token_id >= scheduler_vocab_size: + raise ValueError( + "prompt_pad_token_id must be smaller than the backbone vocabulary: " + f"{self.scheduler_token_id} >= {scheduler_vocab_size}" + ) + self.num_codebooks = int(config.num_codebooks) + codec_config = getattr(config, "codec_config", None) + if isinstance(codec_config, Mapping): + self.codebook_size = int(codec_config.get("codebook_size", 2048)) + elif codec_config is not None: + self.codebook_size = int(getattr(codec_config, "codebook_size", 2048)) + else: + # ``codec_config`` is optional in BreezeTTS2Config; a minimal + # checkpoint may omit it. Breeze's codec uses 2048-entry + # codebooks, the same fallback the talker and codec stages apply. + self.codebook_size = 2048 + if self.num_codebooks <= 0 or self.codebook_size <= 0: + raise ValueError("num_codebooks and codec codebook_size must be positive") + self._validate_audio_special_tokens() + + @classmethod + def from_pretrained( + cls, + model_path: str, + config: Any, + *, + reference_audio_encoder: BreezeReferenceAudioTokenizer | None = None, + tokenizer_kwargs: Mapping[str, Any] | None = None, + ) -> BreezeTTS2PromptBuilder: + """Load the text tokenizer once and construct a reusable builder.""" + kwargs = dict(tokenizer_kwargs or {}) + kwargs.setdefault("use_fast", True) + kwargs.setdefault("trust_remote_code", True) + tokenizer = AutoTokenizer.from_pretrained(model_path, **kwargs) + return cls(tokenizer, config, reference_audio_encoder) + + def build( + self, + request: Mapping[str, Any], + *, + template: str | None = None, + ) -> OmniTokensPrompt: + """Convert one request to a vLLM-Omni token prompt. + + Supported templates are ``tts_plain``, ``tts_instruction``, + ``ref_clone_tata`` and ``ref_edit_tata``. When omitted, the template + is inferred from the presence of instruction/reference fields. + """ + template_name = template or request.get("template") + if not isinstance(template_name, str) or not template_name: + template_name = _infer_template(request) + segments = self._build_segments(template_name, request) + + # Breeze's reference implementation first adds tokenizer special + # tokens to every text segment, decodes those segments back to text, + # inserts the audio placeholders, and tokenizes the complete string + # once. This matters at segment boundaries: tokenizing each raw text + # fragment independently can produce different boundary tokens. + rendered_segments: list[tuple[str, bool, int]] = [] + reference_codes: list[torch.Tensor] = [] + + for segment in segments: + if isinstance(segment, _TextSegment): + ids = self._encode_one_text(segment.text, add_special_tokens=True) + if not ids: + raise ValueError("Breeze text tokenizer returned an empty segment") + rendered = str(self.tokenizer.decode(ids, skip_special_tokens=False)) + rendered_segments.append((rendered, True, 0)) + continue + + codes = self._resolve_reference_codes(segment) + reference_codes.append(codes) + audio_text = AUDIO_TAG * int(codes.shape[0]) + AUDIO_EOS + rendered_segments.append((audio_text, False, int(codes.shape[0]) + 1)) + + prompt_ids: list[int] = [] + text_mask: list[bool] = [] + text_lengths: list[int] = [] + rendered_prompt = "".join(item[0] for item in rendered_segments) + prompt_ids = self._encode_one_text(rendered_prompt, add_special_tokens=False) + offset = 0 + for rendered, is_text, expected_len in rendered_segments: + segment_ids = self._encode_one_text(rendered, add_special_tokens=False) + actual_len = len(segment_ids) + if not is_text and actual_len != expected_len: + raise ValueError( + "Breeze audio placeholder token length changed after rendering: " + f"expected {expected_len}, got {actual_len}" + ) + if not is_text and segment_ids != [ + *([self.audio_token_id] * (expected_len - 1)), + self.audio_eos_token_id, + ]: + raise ValueError("Breeze tokenizer did not preserve audio placeholder ids") + text_mask.extend([is_text] * actual_len) + if is_text: + text_lengths.append(actual_len) + offset += actual_len + if offset != len(prompt_ids) or len(text_mask) != len(prompt_ids): + raise RuntimeError("Breeze rendered prompt length accounting error") + if not prompt_ids: + raise ValueError("Breeze prompt is empty") + + input_values = torch.cat(reference_codes, dim=0) if reference_codes else None + if input_values is not None and int(input_values.shape[0]) != prompt_ids.count(self.audio_token_id): + raise RuntimeError("reference code count does not match audio placeholder count") + + additional_information: dict[str, Any] = { + # The model runner may invoke preprocess on a chunk of a long + # prompt. Keep one compact CPU copy so the talker can reconstruct + # the full embedding buffer before slicing that chunk. + "prompt_ids": torch.tensor(prompt_ids, dtype=torch.long), + "attention_mask": torch.ones(len(prompt_ids), dtype=torch.bool), + "text_ids_mask": torch.tensor(text_mask, dtype=torch.bool), + "text_ids_len": torch.tensor(text_lengths, dtype=torch.long), + "template": template_name, + } + if input_values is not None: + # ``input_values`` is the name used by the upstream Breeze model. + # Keep the tensor unbatched here: Omni transports one request at a + # time and the talker adds the batch dimension after collation. + additional_information["input_values"] = input_values + additional_information["ref_code_len"] = int(input_values.shape[0]) + + if self.scheduler_token_id < 0: + raise ValueError("prompt_pad_token_id must be non-negative") + + return OmniTokensPrompt( + # Never expose Breeze's high text/audio ids to the Qwen3 scheduler + # or its token-range validation. The real ids remain available in + # additional_information and are consumed by ``preprocess``. + prompt_token_ids=[self.scheduler_token_id] * len(prompt_ids), + additional_information=additional_information, + ) + + def _build_segments(self, template: str, request: Mapping[str, Any]) -> list[_TextSegment | _AudioSegment]: + speaker = _speaker_prefix(request.get("speaker", "S0")) + text = _required_text(request, "text") + instruction = request.get("instruction") + ref_text = request.get("ref_text") + + if template == "tts_plain": + return [_TextSegment(f"{speaker}{text}")] + if template == "tts_instruction": + if not isinstance(instruction, str) or not instruction: + raise ValueError("tts_instruction requires a non-empty instruction") + return [_TextSegment(f"{speaker}{INSTRUCTION_BOS}{instruction}{INSTRUCTION_EOS}{text}")] + if template in ("ref_clone_tata", "ref_edit_tata"): + if not isinstance(ref_text, str) or not ref_text: + raise ValueError(f"{template} requires a non-empty ref_text") + audio = self._audio_segment(request) + if template == "ref_clone_tata": + return [_TextSegment(f"{speaker}{ref_text}"), audio, _TextSegment(f"{speaker}{text}")] + if not isinstance(instruction, str) or not instruction: + raise ValueError("ref_edit_tata requires a non-empty instruction") + return [ + _TextSegment(f"{speaker}{ref_text}"), + audio, + _TextSegment(f"{speaker}{INSTRUCTION_BOS}{instruction}{INSTRUCTION_EOS}{text}"), + ] + raise ValueError( + f"unknown Breeze template {template!r}; expected tts_plain, tts_instruction, " + "ref_clone_tata or ref_edit_tata" + ) + + def _audio_segment(self, request: Mapping[str, Any]) -> _AudioSegment: + source = request.get("ref_audio") + if source is None: + source = request.get("ref_audio_codes") + if source is None: + raw_codes = request.get("codes") + source = raw_codes.get("ref") if isinstance(raw_codes, Mapping) else None + if source is None: + raise ValueError("reference template requires ref_audio or ref_audio_codes") + sample_rate = request.get("ref_audio_sample_rate") + return _AudioSegment(source, int(sample_rate) if sample_rate is not None else None) + + def _resolve_reference_codes(self, segment: _AudioSegment) -> torch.Tensor: + source = segment.source + if _looks_like_codes(source, self.num_codebooks): + codes = torch.as_tensor(source) + return _normalize_codes(codes, self.num_codebooks, self.codebook_size) + if self.audio_encoder is None: + raise RuntimeError("reference audio was provided but no reference_audio_encoder is configured") + codes = self.audio_encoder.encode(source, segment.sample_rate) + return _normalize_codes(codes, self.num_codebooks, self.codebook_size) + + def _encode_one_text(self, text: str, *, add_special_tokens: bool) -> list[int]: + encoded = self.tokenizer(text, add_special_tokens=add_special_tokens, padding=False) + ids = encoded["input_ids"] if isinstance(encoded, Mapping) else encoded.input_ids + if isinstance(ids, torch.Tensor): + ids = ids.tolist() + return _as_int_list(ids) + + def _validate_audio_special_tokens(self) -> None: + convert = self.tokenizer.convert_tokens_to_ids + for token, expected in ((AUDIO_TAG, self.audio_token_id), (AUDIO_EOS, self.audio_eos_token_id)): + actual = convert(token) + unk = getattr(self.tokenizer, "unk_token_id", None) + if actual is None or (unk is not None and actual == unk) or int(actual) != expected: + raise ValueError(f"tokenizer id for {token!r} is {actual}, expected config value {expected}") + + +def _infer_template(request: Mapping[str, Any]) -> str: + has_ref = ( + request.get("ref_audio") is not None + or request.get("ref_audio_codes") is not None + or bool(request.get("ref_text")) + ) + has_instruction = bool(request.get("instruction")) + if has_ref and has_instruction: + return "ref_edit_tata" + if has_ref: + return "ref_clone_tata" + return "tts_instruction" if has_instruction else "tts_plain" + + +def _required_text(request: Mapping[str, Any], key: str) -> str: + value = request.get(key) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Breeze request requires non-empty {key}") + return value + + +def _speaker_prefix(value: Any) -> str: + if value in (None, ""): + return "" + value = str(value) + return value if value.startswith("[") and value.endswith("]") else f"[{value}]" + + +def _as_int_list(value: Any) -> list[int]: + if isinstance(value, torch.Tensor): + value = value.tolist() + elif isinstance(value, np.ndarray): + value = value.tolist() + if not isinstance(value, list): + raise ValueError("tokenizer input_ids must be a list") + if value and isinstance(value[0], (list, tuple, np.ndarray, torch.Tensor)): + if len(value) != 1: + raise ValueError("tokenizer input_ids contains multiple sequences") + return _as_int_list(value[0]) + return [int(item) for item in value] + + +def _looks_like_codes(value: Any, num_codebooks: int) -> bool: + if isinstance(value, (torch.Tensor, np.ndarray)): + return value.ndim == 2 and value.shape[-1] == num_codebooks + return False + + +def _normalize_codes(codes: Any, num_codebooks: int, codebook_size: int) -> torch.Tensor: + codes = torch.as_tensor(codes) + if codes.ndim != 2: + raise ValueError(f"reference audio codes must be 2D, got {tuple(codes.shape)}") + if codes.shape[0] == 0: + raise ValueError("reference audio codes must contain at least one frame") + if codes.shape[-1] != num_codebooks: + raise ValueError(f"expected {num_codebooks} codebooks, got {tuple(codes.shape)}") + codes = codes.to(device="cpu", dtype=torch.long).contiguous() + if int(codes.min()) < 0 or int(codes.max()) >= codebook_size: + raise ValueError(f"reference code outside [0, {codebook_size})") + return codes.to(dtype=torch.int16) + + +__all__ = ["BreezeTTS2PromptBuilder"] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index febda68af..e5d7b4115 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -492,6 +492,16 @@ _OMNI_MODELS = { "acoustic", "MiniMaxMusic3AcousticForConditionalGeneration", ), + "BreezeForConditionalGeneration": ( + "breeze_tts_2", + "modeling_breeze_tts_2_talker", + "BreezeTTS2TalkerForGeneration", + ), + "BreezeTTS2MimiCodec": ( + "breeze_tts_2", + "modeling_breeze_tts_2_codec", + "BreezeTTS2MimiCodec", + ), } diff --git a/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py b/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py new file mode 100644 index 000000000..ae393948a --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/breeze_tts_2.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Stage-0 to stage-1 input conversion for Breeze-TTS-2. + +The prompt builder and the reference-audio tokenizer live under the Breeze +model directory. This module intentionally contains only the inter-stage +contract: flatten the generated ``(T, Q)`` codec matrix into the codebook-major +token sequence consumed by the stage-1 codec model. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct +from vllm_omni.inputs.data import OmniTokensPrompt + +logger = init_logger(__name__) + +# ``BreezeTTS2TalkerForGeneration`` emits the complete sequence accumulated +# so far on every decode step. The generic full-payload accumulator normally +# concatenates rank-2 tensors, which would duplicate all earlier frames. +# Replace the latest complete sequence instead. +_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"codes.audio"}) +_DEFAULT_CODEC_CHUNK_FRAMES = 8 + + +def _extract_audio_codes(output: Any) -> torch.Tensor | None: + """Extract ``codes.audio`` from the supported Omni output shapes.""" + # Depending on the engine path, source_outputs contains either the + # per-request output itself or an EngineCoreOutput with ``outputs[0]``. + if isinstance(output, Mapping): + multimodal = output.get("multimodal_output", output.get("multimodal_outputs")) + if multimodal is None and "codes" in output: + multimodal = output + if multimodal is None and "codes.audio" in output: + audio = output["codes.audio"] + if isinstance(audio, (list, tuple)): + if len(audio) != 1: + raise ValueError(f"expected one Breeze audio-code tensor, got {len(audio)}") + audio = audio[0] + audio = torch.as_tensor(audio) + if audio.ndim not in (1, 2): + raise ValueError(f"Breeze audio codes must be rank 1 or 2, got {tuple(audio.shape)}") + return audio.to(device="cpu", dtype=torch.long).contiguous() + else: + multimodal = getattr(output, "multimodal_output", None) + if multimodal is None: + multimodal = getattr(output, "multimodal_outputs", None) + if multimodal is None: + candidates = getattr(output, "outputs", None) + if isinstance(candidates, Sequence) and candidates: + return _extract_audio_codes(candidates[0]) + if not isinstance(multimodal, Mapping): + return None + codes = multimodal.get("codes") + audio = codes.get("audio") if isinstance(codes, Mapping) else multimodal.get("codes.audio") + if audio is None: + return None + # ``OmniOutput`` stores one tensor per request under ``codes.audio``. + # ``process_engine_inputs`` invokes this function once per source output, + # so unwrap that single-item container before tensor conversion. + if isinstance(audio, (list, tuple)): + if len(audio) != 1: + raise ValueError(f"expected one Breeze audio-code tensor, got {len(audio)}") + audio = audio[0] + audio = torch.as_tensor(audio) + if audio.ndim == 3 and audio.shape[0] == 1: + audio = audio[0] + if audio.ndim not in (1, 2): + raise ValueError( + f"Breeze audio codes must have shape (T, Q) or codebook-major (Q*T,), got {tuple(audio.shape)}" + ) + return audio.to(device="cpu", dtype=torch.long).contiguous() + + +def talker2codec( + source_outputs: Sequence[Any], + prompt: Any = None, + _requires_multimodal_data: bool = False, +) -> list[OmniTokensPrompt]: + """Build synchronous stage-1 prompts from completed stage-0 outputs. + + Breeze's Mimi decoder receives one flat codebook-major sequence. The + generated matrix is therefore transposed from ``(T, Q)`` to ``(Q, T)`` + before flattening. No audio decoding, text tokenization, or streaming + state is handled here. + """ + del prompt, _requires_multimodal_data + results: list[OmniTokensPrompt] = [] + for source_index, output in enumerate(source_outputs): + codes = _extract_audio_codes(output) + if codes is None or codes.numel() == 0: + logger.warning("Breeze stage output %d contains no audio codes", source_index) + results.append(OmniTokensPrompt(prompt_token_ids=[])) + continue + # Connector full-payload builders send a codebook-major flat tensor; + # direct in-process stage outputs carry frame-major ``(T, Q)``. + if codes.ndim == 1: + codebook_major_data = codes.contiguous() + codebook_major = codebook_major_data + else: + codebook_major_data = codes.transpose(0, 1).contiguous() + codebook_major = codebook_major_data.reshape(-1) + results.append( + OmniTokensPrompt( + prompt_token_ids=codebook_major.tolist(), + multi_modal_data={"codes": {"audio": codebook_major_data}}, + ) + ) + return results + + +def talker2codec_full_payload( + transfer_manager: Any = None, + pooling_output: Mapping[str, Any] | None = None, + request: Any = None, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Build the final Stage-1 payload from accumulated talker codes. + + ``pooling_output`` is the AR runner's flattened per-request payload. Mimi + consumes codebook-major IDs, so the frame-major talker matrix is transposed + before flattening. A request that finishes without any codec frame still + gets an explicit empty payload with ``finished=True``: returning ``None`` + would leave Stage 1 without a terminal marker, and the connector receiver + requeues marker-less empty chunks forever. + """ + del transfer_manager, request + audio = None + if isinstance(pooling_output, Mapping): + audio = pooling_output.get("codes.audio") + if audio is None: + nested = pooling_output.get("codes") + if isinstance(nested, Mapping): + audio = nested.get("audio") + + codes = torch.empty(0, dtype=torch.long) if audio is None else torch.as_tensor(audio, dtype=torch.long) + if codes.numel() == 0: + if not is_finished: + return None + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct( + code_flat_numel=0, + finished=torch.tensor(True, dtype=torch.bool), + ), + ) + if codes.ndim == 2: + flat = codes.transpose(0, 1).contiguous().reshape(-1) + elif codes.ndim == 1: + flat = codes.contiguous() + else: + raise ValueError(f"Breeze audio codes must be rank 1 or 2, got {tuple(codes.shape)}") + return OmniPayloadStruct( + codes=CodesStruct(audio=flat), + meta=MetaStruct( + code_flat_numel=int(flat.numel()), + finished=torch.tensor(bool(is_finished), dtype=torch.bool), + ), + ) + + +def _extract_current_audio_frame( + multimodal_output: Mapping[str, Any] | None, +) -> torch.Tensor | None: + """Extract the single ``(Q,)`` frame emitted by async Stage 0.""" + if not isinstance(multimodal_output, Mapping): + return None + codes = multimodal_output.get("codes") + audio = codes.get("audio") if isinstance(codes, Mapping) else None + if audio is None: + return None + frame = torch.as_tensor(audio, dtype=torch.long).reshape(-1) + if frame.numel() == 0: + return None + return frame.cpu().contiguous() + + +def _codec_chunk_frames(transfer_manager: Any) -> int: + connector = getattr(transfer_manager, "connector", None) + # SharedMemoryConnector stores its ConnectorSpec.extra directly in + # ``connector.config``; there is no second ``extra`` wrapper at runtime. + config = getattr(connector, "config", {}) or {} + value = config.get("breeze_codec_chunk_frames") + try: + chunk_frames = int(value) if value is not None else _DEFAULT_CODEC_CHUNK_FRAMES + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid Breeze codec_chunk_frames={value!r}") from exc + if chunk_frames <= 0: + raise ValueError(f"Breeze codec_chunk_frames must be positive, got {chunk_frames}") + return chunk_frames + + +def talker2codec_async_chunk( + transfer_manager: Any, + multimodal_output: Mapping[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Batch Stage-0 frames and send only the un-emitted codebook-major tail.""" + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + frame = _extract_current_audio_frame(multimodal_output) + if frame is not None: + transfer_manager.code_prompt_token_ids[request_id].append(frame) + + pending = transfer_manager.code_prompt_token_ids[request_id] + chunk_frames = _codec_chunk_frames(transfer_manager) + if not pending and not finished: + return None + if len(pending) < chunk_frames and not finished: + return None + + if pending: + # Stage 0 supplies frame-major ``(T, Q)`` rows. The Qwen3-TTS decoder + # consumes ``(Q, T)``, so transpose before flattening the emitted tail. + tail = torch.stack(pending, dim=0) + flat = tail.transpose(0, 1).contiguous().reshape(-1) + else: + flat = torch.empty(0, dtype=torch.long) + pending.clear() + + return OmniPayloadStruct( + codes=CodesStruct(audio=flat), + meta=MetaStruct( + request_id=request_id, + code_flat_numel=int(flat.numel()), + codec_streaming=True, + finished=torch.tensor(finished, dtype=torch.bool), + # The chunk adapter strips ``meta.finished`` before it reaches the + # stage-1 model; ``stream_finished`` survives and is what the + # codec uses to release per-request decoder state. + stream_finished=torch.tensor(finished, dtype=torch.bool), + ), + ) + + +__all__ = ["talker2codec", "talker2codec_async_chunk", "talker2codec_full_payload"]