From 5aadcbb0abae2e7065966483fcb90ac5ebcbd701 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:40:12 -0700 Subject: [PATCH] [Bugfix] Fix Breeze-TTS-2 repeated silence and audio EOS handling --- recipes/BreezeBlue/Breeze-TTS-2.md | 33 +++-- .../models/breeze_tts_2/test_instruction.py | 22 ++++ .../breeze_tts_2/test_repetition_penalty.py | 122 ++++++++++++++++++ .../models/breeze_tts_2/test_sampling_eos.py | 77 +++++++++++ .../models/breeze_tts_2/test_streaming.py | 120 +++++++++++++++-- vllm_omni/deploy/breeze_tts_2.yaml | 6 +- .../openai/tts_adapters/breeze_tts_2.py | 22 +++- .../modeling_breeze_tts_2_codec.py | 25 +++- .../modeling_breeze_tts_2_talker.py | 37 +++++- .../models/breeze_tts_2/pipeline.py | 6 + 10 files changed, 437 insertions(+), 33 deletions(-) create mode 100644 tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py create mode 100644 tests/model_executor/models/breeze_tts_2/test_sampling_eos.py diff --git a/recipes/BreezeBlue/Breeze-TTS-2.md b/recipes/BreezeBlue/Breeze-TTS-2.md index 50e8858af..fb26ce91c 100644 --- a/recipes/BreezeBlue/Breeze-TTS-2.md +++ b/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) | diff --git a/tests/model_executor/models/breeze_tts_2/test_instruction.py b/tests/model_executor/models/breeze_tts_2/test_instruction.py index 3817ccb30..c7f2fef66 100644 --- a/tests/model_executor/models/breeze_tts_2/test_instruction.py +++ b/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 diff --git a/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py b/tests/model_executor/models/breeze_tts_2/test_repetition_penalty.py new file mode 100644 index 000000000..0edaabc6f --- /dev/null +++ b/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 diff --git a/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py b/tests/model_executor/models/breeze_tts_2/test_sampling_eos.py new file mode 100644 index 000000000..8cd389d17 --- /dev/null +++ b/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 diff --git a/tests/model_executor/models/breeze_tts_2/test_streaming.py b/tests/model_executor/models/breeze_tts_2/test_streaming.py index 045444ec3..a630b90bc 100644 --- a/tests/model_executor/models/breeze_tts_2/test_streaming.py +++ b/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 == {} diff --git a/vllm_omni/deploy/breeze_tts_2.yaml b/vllm_omni/deploy/breeze_tts_2.yaml index 1d1b85391..fd75dd69e 100644 --- a/vllm_omni/deploy/breeze_tts_2.yaml +++ b/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 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py b/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py index d2851e32c..03f0d3c4b 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/breeze_tts_2.py +++ b/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: diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py index 7addb3934..860d64c44 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_codec.py @@ -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) diff --git a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py index beaa390c8..765df5fbf 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py +++ b/vllm_omni/model_executor/models/breeze_tts_2/modeling_breeze_tts_2_talker.py @@ -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 diff --git a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py b/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py index d1aaeb13e..2b538df01 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/pipeline.py +++ b/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], }, ),