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

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 == {}