[Bugfix] Fix Breeze-TTS-2 repeated silence and audio EOS handling

recipes/BreezeBlue/Breeze-TTS-2.md

@@ -30,10 +30,20 @@ Four prompt modes are selected automatically from the request fields:
 | `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
+Current scope: greedy sampling with `cfg_scale=1.0` and a default
+`repetition_penalty=1.1` on generated codebook-0 tokens; CFG ≠ 1.0 and
 `negative_prompt` are rejected with explicit client errors until companion
 request support lands.
 
+The [official inference example](https://github.com/breezeblue-ai/breeze-tts/blob/main/infer.py)
+also uses repetition penalty `1.1`, but its
+[runtime defaults](https://github.com/breezeblue-ai/breeze-tts/blob/main/breeze_infer/runtime.py)
+sample both the backbone and depth decoder with temperature `0.9` and top-k
+`50`. The [official voice-design examples](https://huggingface.co/BreezeBlue/Breeze-TTS-2#-voice-design)
+use CFG `4`. Those sampling and guidance settings are outside the current
+port's scope, so the examples below do not establish parity with the official
+voice-design outputs.
+
 ## Supported model contract
 
 | Item | Value |
@@ -42,9 +52,9 @@ request support lands.
 | 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` |
+| Output | 24 kHz mono; `wav` or `pcm`; raw audio streaming (`stream_format="audio"`) or SSE `speech.audio.delta` (`stream_format="sse"`) |
 | 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 |
+| Sampling | greedy (`temperature=0`) with `cfg_scale=1.0` and `repetition_penalty=1.1`; 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
@@ -106,7 +116,7 @@ Voice design (instruction only, no reference audio):
 curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
-        "input": "[laughs] Welcome to tonight's story time.",
+        "input": "(laugh) Welcome to our story time.",
         "instructions": "A warm young woman, clear voice, lively delivery.",
         "response_format": "wav"
     }' --output breeze_design.wav
@@ -139,7 +149,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     }' --output breeze_direction.wav
 ```
 
-Streaming PCM (SSE `speech.audio.delta` events):
+Streaming raw PCM (`stream_format="audio"`):
 
 ```bash
 curl -N -X POST http://localhost:8091/v1/audio/speech \
@@ -153,8 +163,9 @@ curl -N -X POST http://localhost:8091/v1/audio/speech \
     }' --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
+Expected results: finite, non-silent 24 kHz mono WAV files; the streaming example
+returns raw PCM bytes. Use `stream_format="sse"` for `speech.audio.delta`
+events containing base64-encoded audio. 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
@@ -166,6 +177,12 @@ warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time.
   supported yet.
 - **CFG**: `guidance_scale`/`cfg_scale` must be `1.0`. Non-1.0 values and
   `negative_prompt` return a client error.
+- **Repetition penalty**: the default `1.1` matches upstream's penalty on
+  generated codebook-0 tokens and discourages repeated silence. Override it
+  with `extra_params: {"repetition_penalty": 1.0}` to disable it, or another
+  finite positive value. Text and reference audio are excluded from its history.
+  The legacy `tools/breeze_tts_2_upstream_golden.py` uses penalty `1.0`; set
+  the same serving override when comparing against those golden frames.
 - **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.
@@ -181,7 +198,7 @@ warm-up) a few seconds of speech completes in roughly 3.4–3.8 s wall time.
 
 | Feature | Status | Notes |
 | ------- | ------ | ----- |
-| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks, SSE `speech.audio.delta` PCM |
+| Streaming output | ✓ | `async_chunk: true`, 8-frame inter-stage chunks; raw audio or SSE `speech.audio.delta` |
 | 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) |

tests/model_executor/models/breeze_tts_2/test_instruction.py

@@ -4,6 +4,7 @@
 from types import SimpleNamespace
 
 import pytest
+from vllm import SamplingParams
 
 from vllm_omni.entrypoints.openai.tts_adapters.breeze_tts_2 import BreezeTTS2Adapter
 from vllm_omni.model_executor.models.breeze_tts_2.prompt_builder import (
@@ -133,3 +134,24 @@ def test_non_greedy_sampling_overrides_are_rejected():
     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"}))
+
+
+@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf"), "invalid"])
+def test_invalid_repetition_penalties_are_rejected(value):
+    assert "repetition_penalty" in _adapter().validate(_request({"repetition_penalty": value}))
+
+
+@pytest.mark.parametrize("value", [0.9, 1.0, 1.1, "1.2"])
+@pytest.mark.parametrize("max_new_tokens", [None, 32])
+def test_repetition_penalty_override_reaches_standard_sampling_metadata(value, max_new_tokens):
+    adapter = _adapter()
+    request = _request({"repetition_penalty": value}, max_new_tokens=max_new_tokens)
+    original = [SamplingParams(max_tokens=128, repetition_penalty=1.1)]
+
+    assert adapter.validate(request) is None
+    updated = adapter.apply_sampling_overrides(original, request)
+
+    assert updated[0].repetition_penalty == float(value)
+    assert updated[0].max_tokens == (128 if max_new_tokens is None else max_new_tokens)
+    assert original[0].repetition_penalty == 1.1
+    assert original[0].max_tokens == 128

tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py

@@ -0,0 +1,122 @@
+# 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.v1.sample.logits_processor import LogitsProcessors
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from tests.model_executor.models.breeze_tts_2.test_talker_state_machine import _depth_codes, _talker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _model(first_score=11.0, second_score=10.5):
+    model = _talker()
+    model._async_chunk = True
+    logits = torch.full((1, 9), -100.0)
+    logits[0, 2:4] = torch.tensor([first_score, second_score])
+    model._main_head_logits = lambda hidden: logits.expand(hidden.shape[0], -1).clone()
+    model._generate_depth_codes = lambda hidden, code0: _depth_codes(model, hidden, code0)
+    return model
+
+
+def _step(model, infos, penalties, *, no_penalties=False):
+    return model.make_omni_output(
+        torch.zeros(len(infos), 16),
+        model_intermediate_buffer=infos,
+        request_token_spans=[(i, i + 1) for i in range(len(infos))],
+        sampling_metadata=SamplingMetadata(
+            temperature=None,
+            all_greedy=True,
+            all_random=False,
+            top_p=None,
+            top_k=None,
+            generators={},
+            max_num_logprobs=None,
+            no_penalties=no_penalties,
+            frequency_penalties=torch.zeros(len(infos)),
+            presence_penalties=torch.zeros(len(infos)),
+            repetition_penalties=torch.tensor(penalties),
+            # Text/reference conditioning must never be counted as generated
+            # codec history, even if the scheduler ids overlap the codec ids.
+            prompt_token_ids=torch.tensor([[2, 2]] * len(infos)),
+            output_token_ids=[list(info.get("breeze_generated_code0_ids", [])) for info in infos],
+            allowed_token_ids_mask=None,
+            bad_words_token_ids={},
+            logitsprocs=LogitsProcessors(),
+        ),
+    )
+
+
+@pytest.mark.parametrize("scores", [(11.0, 10.5), (-1.0, -1.05)])
+def test_penalty_applies_once_per_generated_token_and_respects_logit_sign(scores):
+    model = _model(*scores)
+    info = {"prompt_ids": torch.tensor([2]), "input_values": torch.tensor([[2, 3, 4, 5]])}
+    chosen = []
+
+    for _ in range(4):
+        output = _step(model, [info], [1.1])
+        chosen.append(output.multimodal_outputs["codes"]["audio"][0][0, 0].item())
+
+    assert chosen == [2, 3, 2, 2]
+
+
+def test_request_history_and_penalty_are_independent():
+    model = _model()
+    first: dict[str, Any] = {}
+    second: dict[str, Any] = {}
+    _step(model, [first], [1.1])
+    output = _step(model, [first, second], [1.1, 1.1])
+    assert [codes[0, 0].item() for codes in output.multimodal_outputs["codes"]["audio"]] == [3, 2]
+
+    # Reverse the batch order; a request's history follows its state, while
+    # each penalty follows the current sampling-metadata row.
+    output = _step(model, [second, first], [1.1, 1.0])
+    assert [codes[0, 0].item() for codes in output.multimodal_outputs["codes"]["audio"]] == [3, 2]
+
+
+def test_disabled_penalties_do_not_read_uninitialized_sampling_buffer():
+    model = _model()
+    info: dict[str, Any] = {}
+    for _ in range(2):
+        output = _step(model, [info], [float("nan")], no_penalties=True)
+        assert output.multimodal_outputs["codes"]["audio"][0][0, 0].item() == 2
+
+
+def test_scheduler_logits_preserve_selected_frame_under_further_sampling_transforms():
+    model = _model()
+    info: dict[str, Any] = {}
+    _step(model, [info], [1.1])
+    output = _step(model, [info], [1.1])
+    selected = output.multimodal_outputs["codes"]["audio"][0][0, 0].item()
+    logits = model.compute_logits(torch.zeros(1, 16))
+
+    assert selected == 3
+    assert logits.argmax(dim=-1).item() == selected
+    assert torch.isfinite(logits).sum().item() == 1
+    assert logits[0, selected].item() == 0.0
+    # The scheduler may apply repetition again and rescale temperature. A
+    # finite zero for the selected token preserves a valid point mass.
+    for temperature in (0.1, 0.9, 2.0):
+        probabilities = torch.softmax(logits / 1.1 / temperature, dim=-1)
+        assert torch.isfinite(probabilities).all()
+        assert probabilities[0, selected].item() == 1.0
+
+
+def test_repetition_penalty_can_select_eos_without_emitting_a_frame():
+    model = _model()
+    info: dict[str, Any] = {}
+    _step(model, [info], [1.1])
+    logits = torch.full((1, 9), -100.0)
+    logits[0, 2] = 11.0
+    logits[0, 8] = 10.5
+    model._main_head_logits = lambda hidden: logits.expand(hidden.shape[0], -1).clone()
+
+    output = _step(model, [info], [1.1])
+
+    assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0
+    assert info["breeze_generated_frames"] == 1
+    assert model.compute_logits(torch.zeros(1, 16)).argmax(dim=-1).item() == 8

tests/model_executor/models/breeze_tts_2/test_sampling_eos.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from pathlib import Path
+
+import pytest
+import torch
+import yaml
+from vllm import SamplingParams
+from vllm.v1.core.sched.utils import check_stop
+from vllm.v1.request import Request, RequestStatus
+from vllm.v1.sample.logits_processor import BatchUpdate, MinTokensLogitsProcessor
+
+from vllm_omni.config.stage_config import merge_sampling_constraints
+from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def talker_sampling_params():
+    deploy_path = Path(__file__).parents[4] / "vllm_omni" / "deploy" / "breeze_tts_2.yaml"
+    deploy = yaml.safe_load(deploy_path.read_text(encoding="utf-8"))
+    params = SamplingParams(
+        **merge_sampling_constraints(
+            deploy["stages"][0]["default_sampling_params"],
+            BREEZE_TTS_2_PIPELINE.stages[0].sampling_constraints,
+        )
+    )
+    # Mirror the input processor: Breeze's text tokenizer uses EOS 1,
+    # which is also an ordinary token in the generated codec vocabulary.
+    params.update_from_generation_config({"eos_token_id": [1]}, eos_token_id=1)
+    return params
+
+
+@pytest.mark.parametrize("codec_token, should_stop", [(0, False), (1, False), (2047, False), (2051, True)])
+def test_scheduler_stops_only_on_codec_eos(talker_sampling_params, codec_token, should_stop):
+    request = Request(
+        request_id="breeze-eos-test",
+        prompt_token_ids=[7, 8, 9],
+        sampling_params=talker_sampling_params,
+        pooling_params=None,
+    )
+    request.append_output_token_ids([42, codec_token])
+    request.status = RequestStatus.RUNNING
+
+    assert check_stop(request, max_model_len=4096) is should_stop
+    if should_stop:
+        assert request.status == RequestStatus.FINISHED_STOPPED
+        assert request.stop_reason == 2051
+    else:
+        assert request.status == RequestStatus.RUNNING
+
+
+@pytest.mark.parametrize("codec_token", [1, 2051])
+def test_first_step_keeps_codec_token_and_eos_logits(talker_sampling_params, codec_token, monkeypatch):
+    # vLLM's tensor helper uses a global flag rather than is_pin_memory,
+    # and pinned allocations require a CUDA device even for this CPU test.
+    monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False)
+    # ignore_eos still leaves text EOS in all_stop_token_ids. Exercise the
+    # actual processor so the first codec frame and immediate EOS stay valid.
+    processor = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False)
+    processor.update_state(
+        BatchUpdate(batch_size=1, removed=[], added=[(0, talker_sampling_params, None, [])], moved=[])
+    )
+    logits = torch.full((1, 2052), -float("inf"))
+    logits[0, codec_token] = 0.0
+
+    processed = processor.apply(logits.clone())
+
+    assert torch.equal(processed, logits)
+    assert processed.argmax(dim=-1).item() == codec_token
+
+
+def test_pipeline_disallows_minimum_tokens_that_mask_selected_codec_ids():
+    params = merge_sampling_constraints({"min_tokens": 10}, BREEZE_TTS_2_PIPELINE.stages[0].sampling_constraints)
+    assert params["min_tokens"] == 0

tests/model_executor/models/breeze_tts_2/test_streaming.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+import copy
 from collections import defaultdict
 from types import SimpleNamespace
 
@@ -20,7 +21,7 @@ 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)
+        self.code_prompt_token_ids: defaultdict[str, list[torch.Tensor]] = defaultdict(list)
 
 
 class _Request:
@@ -120,7 +121,8 @@ def test_stateful_codec_rejects_partial_frame_chunk():
     assert decoder.calls == []
 
 
-def test_stateful_codec_terminal_marker_pops_state_without_decoding():
+@pytest.mark.parametrize("token_counts", [(4, 0), (0, 4), (0, 4, 0, 8, 0)])
+def test_stateful_codec_terminal_marker_pops_state_without_decoding(token_counts):
     decoder = _Decoder()
     codec = object.__new__(BreezeTTS2MimiCodec)
     codec._async_chunk = True
@@ -128,24 +130,114 @@ def test_stateful_codec_terminal_marker_pops_state_without_decoding():
     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": {}}
+    request_ids = [f"scheduler-{index}" for index in range(len(token_counts))]
+    codec._decoder_state_cache = {request_id: {} for request_id in request_ids}
 
     output = codec.forward(
-        torch.tensor([0, 1, 2, 3], dtype=torch.long),
+        torch.arange(sum(token_counts), dtype=torch.long) % codec._codebook_size,
         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}},
+            {"meta": {"finished": count == 0}}
+            for count in token_counts
         ],
-        seq_token_counts=[4, 0],
-        request_ids=["scheduler-live", "scheduler-done"],
+        seq_token_counts=token_counts,
+        request_ids=request_ids,
     )
 
-    # The empty terminal marker decodes nothing and releases its state; the
-    # live request still decodes its own chunk.
+    # Empty terminal markers release their state without consuming an output
+    # slot belonging to another request, regardless of their batch position.
     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
+    assert decoder.calls[0][0][0] == sum(count > 0 for count in token_counts)
+    assert [item.numel() for item in output.multimodal_outputs["model_outputs"]] == [
+        count // codec._num_codebooks * 1920 for count in token_counts
+    ]
+    live_request_ids = [request_id for request_id, count in zip(request_ids, token_counts) if count]
+    assert list(codec._decoder_state_cache) == live_request_ids
+    assert all(codec._decoder_state_cache[request_id]["prefix_frames"] == 0 for request_id in live_request_ids)
+
+
+@pytest.mark.parametrize("configured_chunk_frames", [None, 4])
+def test_stateful_codec_advances_decoder_cache_after_sliding_window(monkeypatch, tmp_path, configured_chunk_frames):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
+        Qwen3TTSTokenizerV2DecoderConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+        Qwen3TTSTokenizerV2Decoder,
+    )
+
+    torch.manual_seed(0)
+    config = Qwen3TTSTokenizerV2DecoderConfig(
+        codebook_size=32,
+        hidden_size=16,
+        latent_dim=16,
+        codebook_dim=16,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        intermediate_size=32,
+        num_hidden_layers=1,
+        num_quantizers=2,
+        decoder_dim=32,
+        upsample_rates=(2,),
+        upsampling_ratios=(2,),
+        sliding_window=72,
+    )
+    decoder = Qwen3TTSTokenizerV2Decoder(config).eval()
+    # Codebooks initialize to zeros before checkpoint loading; populate them
+    # so different codec IDs produce distinguishable cached frame histories.
+    with torch.no_grad():
+        for name, parameter in decoder.quantizer.named_parameters():
+            if name.endswith("embedding_sum"):
+                parameter.normal_()
+    tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder, get_output_sample_rate=lambda: 24_000))
+    monkeypatch.setattr(Qwen3TTSTokenizer, "from_pretrained", lambda *args, **kwargs: tokenizer)
+    connector_extra = {} if configured_chunk_frames is None else {"breeze_codec_chunk_frames": configured_chunk_frames}
+    chunk_frames = configured_chunk_frames or 8
+    codec = object.__new__(BreezeTTS2MimiCodec)
+    codec.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_connector_config={"extra": connector_extra}),
+        device_config=SimpleNamespace(device="cpu"),
+    )
+    codec._tokenizer_path = tmp_path / "audio_tokenizer"
+    codec._audio_tokenizer = None
+    codec._async_chunk = True
+    codec._num_codebooks = config.num_quantizers
+    codec._codebook_size = config.codebook_size
+    codec._decoder_state_cache = {}
+    codec.load_weights(iter(()))
+
+    # At least two chunks after the decoder's rolling window expose a stale
+    # cache even when the first padded chunk still produces plausible audio.
+    total_frames = config.sliding_window + 2 * chunk_frames
+    tail_frames = 3
+    codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, total_frames + tail_frames))
+    for end in range(chunk_frames, total_frames + 1, chunk_frames):
+        flat = codes[0, :, end - chunk_frames : end].reshape(-1)
+        codec.forward(
+            flat,
+            runtime_additional_information=[{"meta": {"stream_finished": False}}],
+            seq_token_counts=[flat.numel()],
+            request_ids=["scheduler-live"],
+        )
+        state = codec._decoder_state_cache["scheduler-live"]
+        assert state["suffix_frames"] == min(end - chunk_frames, config.sliding_window) + chunk_frames
+        with torch.inference_mode():
+            expected_quantized = decoder.quantizer.decode(codes[..., max(0, end - config.sliding_window) : end])
+        torch.testing.assert_close(state["suffix_quantized"], expected_quantized)
+
+    # A short terminal chunk must emit only its real frames, then release
+    # state. Its waveform should agree with the unpadded single-request path.
+    with torch.inference_mode():
+        expected_tail = decoder.chunked_decode(codes[..., -tail_frames:], caches=copy.deepcopy(state))
+    flat_tail = codes[0, :, -tail_frames:].reshape(-1)
+    output = codec.forward(
+        flat_tail,
+        runtime_additional_information=[{"meta": {"stream_finished": True}}],
+        seq_token_counts=[flat_tail.numel()],
+        request_ids=["scheduler-live"],
+    )
+    waveform = output.multimodal_outputs["model_outputs"][0]
+    assert waveform.numel() == tail_frames * decoder.total_upsample
+    torch.testing.assert_close(waveform, expected_tail.reshape(-1), atol=1e-5, rtol=1e-4)
+    assert codec._decoder_state_cache == {}

vllm_omni/deploy/breeze_tts_2.yaml

@@ -32,8 +32,12 @@ stages:
       temperature: 0.0
       top_p: 1.0
       top_k: -1
+      # Match upstream's penalty on generated codebook-0 history.
+      repetition_penalty: 1.1
       max_tokens: 2048
-      min_tokens: 1
+      # min_tokens also masks the tokenizer's EOS 1 despite ignore_eos.
+      # Keep valid codec token 1 and an immediate backbone EOS available.
+      min_tokens: 0
 
   - stage_id: 1
     max_num_seqs: 8

vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py

@@ -4,6 +4,8 @@
 """Breeze-TTS-2 serving adapter for synchronous full-payload generation."""
 
 import asyncio
+import copy
+import math
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -49,6 +51,7 @@ class BreezeTTS2Adapter(ARTTSAdapter):
         return set()
 
     def _load_codec_frame_rate(self) -> float | None:
+        assert self.ctx.engine_client is not None
         config = self.ctx.engine_client.model_config.hf_config
         codec = getattr(config, "codec_config", None)
         if isinstance(codec, dict):
@@ -99,6 +102,14 @@ class BreezeTTS2Adapter(ARTTSAdapter):
         greedy_error = _validate_greedy_sampling(extra_params)
         if greedy_error is not None:
             return greedy_error
+        repetition_penalty = extra_params.get("repetition_penalty")
+        if repetition_penalty is not None:
+            try:
+                repetition_penalty = float(repetition_penalty)
+            except (TypeError, ValueError):
+                return "Breeze-TTS-2 repetition_penalty must be finite and positive"
+            if not math.isfinite(repetition_penalty) or repetition_penalty <= 0:
+                return "Breeze-TTS-2 repetition_penalty must be finite and positive"
         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}"
@@ -210,7 +221,16 @@ class BreezeTTS2Adapter(ARTTSAdapter):
         request_id: str | None = None,
     ) -> list:
         del prompt, request_id
-        return apply_max_new_tokens(sampling_params_list, request)
+        sampling_params_list = apply_max_new_tokens(sampling_params_list, request)
+        repetition_penalty = (request.extra_params or {}).get("repetition_penalty")
+        if repetition_penalty is not None:
+            # The generic speech handler copies extra_params to extra_args;
+            # promote this value so the worker receives sampling metadata.
+            # Only this scalar changes; preserve the caller's params without
+            # copying nested state or the unchanged codec-stage parameters.
+            sampling_params_list = [copy.copy(sampling_params_list[0]), *sampling_params_list[1:]]
+            sampling_params_list[0].repetition_penalty = float(repetition_penalty)
+        return sampling_params_list
 
 
 def _validate_greedy_sampling(extra_params: dict[str, Any]) -> str | None:

vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py

@@ -204,6 +204,7 @@ class BreezeTTS2MimiCodec(nn.Module):
                 f"Breeze codec request id count {len(request_ids)} does not match input request count {len(requests)}"
             )
         valid: list[tuple[str, torch.Tensor, bool]] = []
+        valid_indices: list[int] = []
         for index, request_ids_tensor in enumerate(requests):
             runtime_info = (
                 runtime_additional_information[index]
@@ -238,6 +239,7 @@ class BreezeTTS2MimiCodec(nn.Module):
                 )
             frames = int(flat.numel()) // self._num_codebooks
             valid.append((state_id, flat.reshape(self._num_codebooks, frames), finished))
+            valid_indices.append(index)
 
         waveforms = [empty.to(device=ids.device)] * len(requests)
         sample_rates = [sr] * len(requests)
@@ -247,7 +249,6 @@ class BreezeTTS2MimiCodec(nn.Module):
                 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):
@@ -258,7 +259,6 @@ class BreezeTTS2MimiCodec(nn.Module):
                 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,
@@ -353,6 +353,27 @@ class BreezeTTS2MimiCodec(nn.Module):
                     device_map=str(device),
                 )
                 self._sample_rate = int(self._audio_tokenizer.model.get_output_sample_rate())
+                if self._async_chunk:
+                    connector_config = getattr(self.vllm_config.model_config, "stage_connector_config", None)
+                    extra = (
+                        connector_config.get("extra", connector_config)
+                        if isinstance(connector_config, Mapping)
+                        else getattr(connector_config, "extra", None)
+                    )
+                    value = extra.get("breeze_codec_chunk_frames") if isinstance(extra, Mapping) else None
+                    try:
+                        chunk_frames = int(value) if value is not None else 8
+                    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}")
+                    # Qwen treats chunks shorter than its configured window as
+                    # terminal tails and does not advance their rolling cache.
+                    # Match the producer's chunk size instead of its default 25.
+                    decoder = self._audio_tokenizer.model.decoder
+                    decoder._initial_codec_chunk_frames = chunk_frames
+                    decoder._incremental_chunk_frames = chunk_frames
+                    decoder._incremental_chunk_ramp = []
             for _ in weights:
                 pass
             tokenizer_module = getattr(self._audio_tokenizer, "model", None)

vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py

@@ -13,6 +13,7 @@ generation order without importing HuggingFace GenerationMixin.
 from __future__ import annotations
 
 import hashlib
+import math
 import os
 from collections.abc import Iterable, Mapping, Sequence
 from pathlib import Path
@@ -367,10 +368,12 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
             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.
+        # ``make_omni_output`` already selected codebook 0 and completed its
+        # frame. Return that same decision to the scheduler, including EOS.
+        # A finite zero survives the sampler's repetition/temperature passes
+        # without applying the repetition penalty twice or changing the frame.
+        # These logits describe deterministic codec selection, not model
+        # confidence; sampled-token logprobs are consequently zero.
         states = self._batch_state or []
         if states:
             for state, row_start, row_end in _iter_request_rows(
@@ -378,10 +381,11 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
                 getattr(self, "_batch_state_spans", None),
                 int(logits.shape[0]),
             ):
-                if state.get("breeze_force_eos", False):
+                selected_code0 = state.get("breeze_selected_code0")
+                if selected_code0 is not None:
                     row = logits[row_start:row_end]
                     row.fill_(float("-inf"))
-                    row[:, self.codebook_vocab_size] = 0.0
+                    row[:, int(selected_code0)] = 0.0
         return logits
 
     def _main_head_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -443,6 +447,13 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
         self._batch_state = info_items
         self._batch_state_spans = spans
         per_request_codes: list[torch.Tensor] = []
+        # The runner passes the same per-request sampling metadata to this
+        # hook and compute_logits, including offline SamplingParams overrides.
+        # When no_penalties is true vLLM does not initialize this GPU buffer.
+        sampling_metadata = kwargs.get("sampling_metadata")
+        repetition_penalties = None
+        if sampling_metadata is not None and not sampling_metadata.no_penalties:
+            repetition_penalties = sampling_metadata.repetition_penalties.tolist()
 
         for index, info in enumerate(info_items):
             if spans is not None and index < len(spans):
@@ -458,11 +469,22 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
             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))
+            main_logits = self._main_head_logits(row_hidden.reshape(1, -1)).float()
+            penalty = float(repetition_penalties[index]) if repetition_penalties is not None else 1.0
+            if not math.isfinite(penalty) or penalty <= 0:
+                raise ValueError("Breeze repetition_penalty must be finite and positive")
+            # Match upstream: penalize each previously generated codebook-0
+            # id once, excluding the text prompt and reference audio codes.
+            history = info.setdefault("breeze_generated_code0_ids", [])
+            if penalty != 1.0 and history:
+                history_ids = torch.tensor(history, device=main_logits.device, dtype=torch.long)
+                scores = main_logits[0, history_ids]
+                main_logits[0, history_ids] = torch.where(scores < 0, scores * penalty, scores / penalty)
             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)
+            info["breeze_selected_code0"] = self.codebook_vocab_size if terminal else code0
             if terminal:
                 info["breeze_force_eos"] = True
                 existing = info.get("breeze_audio_codes")
@@ -476,6 +498,7 @@ class BreezeTTS2TalkerForGeneration(nn.Module):
                 raise RuntimeError(f"Breeze main head produced invalid codec id {code0}")
 
             frame = self._generate_depth_codes(row_hidden, code0)
+            history.append(code0)
             info["breeze_current_frame"] = frame[0].detach()
             info["breeze_generated_frames"] = generated_frames + 1
             info["breeze_force_eos"] = False

vllm_omni/model_executor/models/breeze_tts_2/pipeline.py

@@ -28,6 +28,12 @@ BREEZE_TTS_2_PIPELINE = PipelineConfig(
             custom_process_next_stage_input_func=f"{_PROC}.talker2codec_full_payload",
             sampling_constraints={
                 "detokenize": False,
+                # Text EOS 1 is a valid codec token; only the backbone's
+                # explicit codec EOS below may terminate audio generation.
+                "ignore_eos": True,
+                # The talker has already selected the emitted frame's token;
+                # min_tokens must not mask that token in the sampler.
+                "min_tokens": 0,
                 "stop_token_ids": [2051],
             },
         ),