[Model] Add Breeze-TTS-2 two-stage AR TTS support (talker + codec, streaming) (#7084)

Signed-off-by: noname <liuzhiwei08@kuaishou.com>
Signed-off-by: noname <133508552+lzwnoname@users.noreply.github.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: noname <liuzhiwei08@kuaishou.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>

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` | ✅︎ | | | | — |

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

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.

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()

examples/online_serving/text_to_speech/README.md

@@ -14,6 +14,7 @@ For the full list of supported architectures across all modalities, see
 
 | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
 |---|---|---|---|---|---|
+| 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/).

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}"

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 |

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 |

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,
     }
 

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)

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",

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",

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

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

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

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)

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"}))

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))

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()

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

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

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)]

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

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

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

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

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()

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()

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",

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,

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

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)

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,

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"]

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"]

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",
+]

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"]

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",
+]

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"]

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"]

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"]

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 = "<ins_bos>"
+INSTRUCTION_EOS = "<ins_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"]

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",
+    ),
 }
 
 

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"]