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

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