anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Model] Accelerate Breeze-TTS-2 non-streaming inference (#7752)

Browse Source
This commit is contained in:
Yueqian Lin authored and GitHub committed 2026-09-22 17:10:15 -07:00
1 parent d2aa07159c
commit 3f473a358b
84 files changed
+8191 -5622
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/model_executor/models/breeze_tts_2/__init__.py
View File
Whitespace-only changes.
tests/model_executor/models/breeze_tts_2/test_audio_tokenizer.py
-128
View File
@@ -1,128 +0,0 @@
# 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
-95
View File
@@ -1,95 +0,0 @@
# 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
-68
View File
@@ -1,68 +0,0 @@
# 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
-113
View File
@@ -1,113 +0,0 @@
# 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
-157
View File
@@ -1,157 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
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 (
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"}))
@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_local_modules.py
-152
View File
@@ -1,152 +0,0 @@
# 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
-28
View File
@@ -1,28 +0,0 @@
# 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
-36
View File
@@ -1,36 +0,0 @@
# 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
-189
View File
@@ -1,189 +0,0 @@
# 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
-187
View File
@@ -1,187 +0,0 @@
# 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_repetition_penalty.py
-122
View File
@@ -1,122 +0,0 @@
# 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_stage_input_processor.py
-133
View File
@@ -1,133 +0,0 @@
# 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
-243
View File
@@ -1,243 +0,0 @@
# 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
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[str, list[torch.Tensor]] = 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 == []
@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
codec._num_codebooks = 4
codec._codebook_size = 8
codec._sample_rate = 24_000
codec._audio_tokenizer = SimpleNamespace(model=SimpleNamespace(decoder=decoder))
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.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": {"finished": count == 0}}
for count in token_counts
],
seq_token_counts=token_counts,
request_ids=request_ids,
)
# 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] == 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 == {}
tests/model_executor/models/breeze_tts_2/test_talker_state_machine.py
-243
View File
@@ -1,243 +0,0 @@
# 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
-73
View File
@@ -1,73 +0,0 @@
# 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
tests/model_executor/models/breeze_tts_2_helpers.py
+54
View File
@@ -0,0 +1,54 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Weight-free Breeze talker and request helpers shared by model tests."""
from dataclasses import dataclass
from unittest.mock import Mock
import torch
from vllm_omni.model_executor.models.breeze_tts_2.first_code_sampler import BreezeFirstCodeSampler
from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration
@dataclass
class TalkerConfig:
vocab_size: int
eos_token_id: int
@dataclass
class DepthStub:
generate_frames: Mock
def _small_talker():
model = BreezeForConditionalGeneration.__new__(BreezeForConditionalGeneration)
torch.nn.Module.__init__(model)
model._first_code_sampler = BreezeFirstCodeSampler()
model.config = TalkerConfig(vocab_size=8, eos_token_id=7)
model.num_codebooks, model.codebook_size, model.hidden_size = 3, 4, 2
model.lm_head = torch.nn.Linear(2, 8, bias=False)
with torch.no_grad():
model.lm_head.weight.zero_()
model.lm_head.weight[1] = torch.tensor([10.0, 0.0])
model.lm_head.weight[2] = torch.tensor([9.5, 0.0])
model.lm_head.weight[3] = torch.tensor([0.0, 10.0])
model.lm_head.weight[7] = torch.tensor([-10.0, -10.0])
model.depth_decoder = DepthStub(
generate_frames=Mock(side_effect=lambda hidden, first, **kwargs: first[:, None].repeat(1, 3))
)
return model
def _request_info(request_id):
return {
"global_request_id": [request_id],
"breeze_prompt": {"role": "cond", "guidance_scale": 1.0},
"breeze_sampling": {"temperature": 0.0, "top_k": 0, "top_p": 1.0, "repetition_penalty": 1.1},
"breeze_state": {
"generator": torch.Generator().manual_seed(42),
"history": torch.empty(1, 0, dtype=torch.long),
"current": torch.zeros(1, 3, dtype=torch.long),
},
}
tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py
+11 -3
View File
@@ -23,7 +23,7 @@ class _FakeDecoder(nn.Module):
def __init__(self, total_upsample: int = _TOTAL_UPSAMPLE):
super().__init__()
self.total_upsample = total_upsample
self.decode_calls: list[dict[str, int | tuple[int, ...]]] = []
self.decode_calls: list[dict[str, object]] = []
self.batched_decode_calls: list[dict[str, object]] = []
self.decode_codes: list[torch.Tensor] = []
self.cudagraph_calls: list[dict[str, int | torch.device]] = []
@@ -91,8 +91,8 @@ class _FakeDecoder(nn.Module):
wav_len = frames * self.total_upsample + 6
wav = torch.arange(wav_len, dtype=torch.float32).view(1, 1, -1)
offsets = torch.arange(batch, dtype=torch.float32).view(batch, 1, 1) * 1000
batched_wav = wav.expand(batch, 1, wav_len) + offsets
return [batched_wav[row, :, : lengths[row] * self.total_upsample] for row in range(batch)]
batched_outputs = wav.expand(batch, 1, wav_len) + offsets
return [batched_outputs[row, :, : lengths[row] * self.total_upsample] for row in range(batch)]
def enable_cudagraph(self, **kwargs):
self.cudagraph_calls.append(kwargs)
@@ -577,6 +577,7 @@ def test_decode_chunking_override_is_passed_to_cudagraph():
_load_weights_noop(model)
assert model.decoder.cudagraph_calls[-1] == {
"capture_modes": ("icl", "xvec"),
"capture_batch_sizes": None,
"stateless_capture_sizes": None,
"device": torch.device("cuda"),
@@ -607,6 +608,13 @@ def test_cudagraph_batch_config_is_preserved():
assert call["capture_batch_sizes"] == [1, 2, 4, 8]
def test_cudagraph_captures_only_modes_used_by_the_model():
model = _make_model(async_chunk=True, device=torch.device("cuda"))
model.decoder_cudagraph_modes = ("xvec",)
_load_weights_noop(model)
assert model.decoder.cudagraph_calls[-1]["capture_modes"] == ("xvec",)
def test_stateless_cudagraph_capture_sizes_are_passed_from_config():
model = _make_model(
async_chunk=False,
tests/model_executor/models/test_breeze_tts_2.py
+459
View File
@@ -0,0 +1,459 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Check depth attention against HF Llama and request-local sampling state."""
from collections import defaultdict
import numpy as np
import pytest
import torch
import torch.nn.functional as F
from transformers import LlamaConfig, MimiConfig
from transformers.models.llama.modeling_llama import LlamaModel
from transformers.models.mimi.modeling_mimi import MimiConv1d
from tests.model_executor.models.breeze_tts_2_helpers import _request_info, _small_talker
from vllm_omni.engine.serialization import deserialize_additional_information, serialize_additional_information
from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import (
BreezeDepthDecoder,
sample_graph_logits,
sample_logits,
)
from vllm_omni.model_executor.models.breeze_tts_2.prompt import CFG_UNCOND_SUFFIX, build_breeze_prompt
from vllm_omni.model_executor.models.breeze_tts_2.reference_encoder import BreezeReferenceConv
from vllm_omni.model_executor.stage_input_processors.breeze_tts_2 import expand_cfg_prompts, talker2code2wav_async_chunk
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize("parameter", ["temperature", "top_p", "repetition_penalty"])
@pytest.mark.parametrize("value", [float("nan"), float("inf")])
def test_offline_prompt_rejects_nonfinite_sampling(mocker, parameter, value):
with pytest.raises(ValueError, match="sampling"):
build_breeze_prompt(mocker.Mock(), "Hello.", **{parameter: value})
@pytest.mark.parametrize("stride", [2, 3])
def test_reference_convolution_batch_preserves_partial_final_frames(stride):
torch.manual_seed(7)
config = MimiConfig(use_causal_conv=True, pad_mode="constant")
first = MimiConv1d(config, 1, 4, kernel_size=2 * stride, stride=stride)
last = MimiConv1d(config, 4, 4, kernel_size=4, stride=2, pad_mode="replicate", bias=False)
batched_first = BreezeReferenceConv(first, stride)
batched_last = BreezeReferenceConv(last, 2 * stride)
lengths = torch.tensor([37, 62])
batched_first.sample_lengths = batched_last.sample_lengths = lengths
waves = [torch.randn(1, 1, int(length)) for length in lengths]
padded = torch.cat([F.pad(wave, (0, 96 - wave.shape[-1])) for wave in waves])
actual = batched_last(batched_first(padded))
for row, waveform in enumerate(waves):
expected = last(first(waveform))
torch.testing.assert_close(actual[row : row + 1, :, : expected.shape[-1]], expected)
assert not actual[row, :, expected.shape[-1] :].count_nonzero()
def test_reference_conditioning_survives_transport_and_cfg_keeps_true_prompt_lengths(mocker) -> None:
tokenizer = mocker.Mock()
tokenizer.encode.side_effect = lambda text, **kwargs: list(text.encode("utf-8"))
waveform = np.zeros(1921, dtype=np.float32)
prompt = build_breeze_prompt(
tokenizer, "Hello.", "Speak happily.", ref_audio=(waveform, 24000), ref_text="Reference.", guidance_scale=4
)
companion = expand_cfg_prompts(prompt, None)[0].prompt
assert len(prompt["prompt_token_ids"]) > len(companion["prompt_token_ids"])
positive = deserialize_additional_information(serialize_additional_information(prompt["additional_information"]))
negative = deserialize_additional_information(serialize_additional_information(companion["additional_information"]))
assert positive["breeze_prompt"]["target_ids"] == list(b"[S0]<ins_bos>Speak happily.<ins_eos>Hello.")
assert negative["breeze_prompt"]["target_ids"] == list(b"[S0]Hello.")
assert positive["breeze_prompt"]["reference_ids"] == negative["breeze_prompt"]["reference_ids"]
assert positive["breeze_prompt"]["reference_frames"] == 2
assert positive["cfg_group"]["role"] == "cond"
assert negative["cfg_group"]["role"] == "uncond"
torch.testing.assert_close(positive["reference_waveform"], torch.from_numpy(waveform), atol=0, rtol=0)
torch.testing.assert_close(negative["reference_waveform"], positive["reference_waveform"], atol=0, rtol=0)
@pytest.fixture
def depth() -> BreezeDepthDecoder:
torch.manual_seed(42)
config = LlamaConfig(
hidden_size=32,
intermediate_size=64,
num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=2,
head_dim=8,
vocab_size=35,
max_position_embeddings=16,
rope_theta=500000.0,
rope_scaling={
"rope_type": "llama3",
"factor": 32.0,
"low_freq_factor": 0.001953125,
"high_freq_factor": 0.0078125,
"original_max_position_embeddings": 16,
},
)
config.num_codebooks = 4
config.audio_embed_size = 48
config._attn_implementation = "eager"
model = BreezeDepthDecoder(config).eval()
torch.nn.init.normal_(model.codebooks_head, std=0.1)
return model
@pytest.fixture
def hf_depth(depth: BreezeDepthDecoder) -> LlamaModel:
reference = LlamaModel(depth.config).eval()
weights = {}
for name, parameter in depth.named_parameters():
if name.startswith("layers."):
if name.endswith("qkv.weight"):
layer = depth.layers[int(name.split(".")[1])]
for part, tensor in zip(
("q", "k", "v"), parameter.split([layer.q_size, layer.kv_size, layer.kv_size]), strict=True
):
weights[name.replace("qkv.weight", f"self_attn.{part}_proj.weight")] = tensor
elif name.endswith("gate_up.weight"):
for part, tensor in zip(("gate", "up"), parameter.chunk(2), strict=True):
weights[name.replace("gate_up.weight", f"mlp.{part}_proj.weight")] = tensor
else:
weights[name.replace(".o_proj.", ".self_attn.o_proj.").replace(".down_proj.", ".mlp.down_proj.")] = (
parameter
)
elif name == "norm.weight":
weights[name] = parameter
result = reference.load_state_dict(weights, strict=False)
assert result.missing_keys == ["embed_tokens.weight"]
assert result.unexpected_keys == []
return reference
@pytest.mark.parametrize("batch", [1, 2])
@torch.inference_mode()
def test_depth_cached_attention_matches_hf_full_sequence(
depth: BreezeDepthDecoder, hf_depth: LlamaModel, batch: int
) -> None:
"""Wrong cache positions or fused-QKV ordering change the reference hidden states."""
embeddings = torch.randn(batch, 4, depth.config.hidden_size)
expected = hf_depth(inputs_embeds=embeddings, use_cache=False).last_hidden_state
shape = (batch, 2, 4, 8)
caches = [(torch.zeros(shape), torch.zeros(shape)) for _ in depth.layers]
actual = []
for start, end in [(0, 2), (2, 3), (3, 4)]:
x = embeddings[:, start:end]
for layer, cache in zip(depth.layers, caches, strict=True):
x = layer(
x,
depth.positions[start:end],
cache,
depth.causal_mask[start:end, :4][None, None],
depth.rope_cos[start:end][None, None],
depth.rope_sin[start:end][None, None],
)
actual.append(depth.norm(x))
torch.testing.assert_close(torch.cat(actual, 1), expected, atol=2e-6, rtol=2e-5)
@pytest.mark.parametrize("guidance_scale", [0.5, 4.0])
@torch.inference_mode()
def test_cfg_depth_matches_hf_full_prefix_at_every_codebook(
depth: BreezeDepthDecoder, hf_depth: LlamaModel, guidance_scale: float
) -> None:
hidden = torch.randn(4, depth.config.audio_embed_size)
first = torch.tensor([7, 11])
prefix = torch.cat((hidden[:, None], depth.embed_tokens(first.repeat(2)[:, None])), dim=1)
expected = [first]
for codebook in range(depth.num_codebooks - 1):
states = hf_depth(inputs_embeds=depth.inputs_embeds_projector(prefix), use_cache=False).last_hidden_state
logits = F.linear(states[:, -1].float(), depth.codebooks_head[codebook].T)
cond, uncond = logits.chunk(2)
token = (uncond + guidance_scale * (cond - uncond))[:, : depth.vocab_size - 3].argmax(-1)
expected.append(token)
embed = depth.embed_tokens((token.repeat(2) + (codebook + 1) * depth.vocab_size)[:, None])
prefix = torch.cat((prefix, embed), dim=1)
actual = depth._generate_frame(
hidden, first, depth._allocate_cache(hidden), 0.0, 0, 1.0, torch.Generator(), guidance_scale=guidance_scale
)
torch.testing.assert_close(actual, torch.stack(expected, dim=-1), atol=0, rtol=0)
@torch.inference_mode()
def test_frame_generation_does_not_reuse_another_requests_cache(depth: BreezeDepthDecoder) -> None:
hidden = torch.randn(1, 48)
first = torch.tensor([7])
def generate(value: torch.Tensor) -> torch.Tensor:
return depth.generate_frame(
value, first, temperature=0.9, top_k=10, top_p=0.8, generator=torch.Generator().manual_seed(99)
)
expected = generate(hidden)
generate(hidden * -3)
torch.testing.assert_close(generate(hidden), expected, atol=0, rtol=0)
def test_top_k_keeps_tied_candidates_and_nucleus_keeps_crossing_token() -> None:
generator = torch.Generator().manual_seed(42)
logits = torch.tensor([[0.0, 0.0, -100.0]]).expand(128, -1)
sampled = sample_logits(logits, 1.0, 1, 1.0, generator)
assert set(sampled.tolist()) == {0, 1}
sampled = sample_logits(torch.tensor([[10.0, 0.0]]), 1.0, 0, 0.1, generator)
assert sampled.item() == 0
@pytest.mark.parametrize("parameters", [(0.0, 0, 1.0), (0.9, 50, 1.0), (0.7, 100, 0.8), (1.2, 0, 0.95)])
def test_graph_sampler_matches_scalar_filtering_with_identical_noise(parameters) -> None:
generator = torch.Generator().manual_seed(42)
logits = torch.randn(32, 2051, generator=generator)
logits[:, -3:] = -torch.inf
logits[:, 10:14] = 4.0
noise = torch.empty_like(logits).exponential_(generator=generator)
expected = sample_logits(logits, *parameters, generator, noise)
actual = sample_graph_logits(logits, torch.tensor(parameters), noise)
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
@pytest.mark.parametrize("top_p", [0.25, 0.5, 0.75])
@pytest.mark.parametrize("top_k", [0, 2])
def test_graph_sampler_excludes_token_after_exact_nucleus_boundary(top_p, top_k) -> None:
# Equal probabilities make the CDF boundary exact in float32. Bias the
# shared noise toward the first token outside that nucleus so retaining
# one extra candidate deterministically changes the sampled token.
logits = torch.tensor([[0.0, 0.0, 0.0, 0.0, -torch.inf, -torch.inf, -torch.inf]])
noise = torch.ones_like(logits)
first_excluded = int(top_p * 4)
noise[:, first_excluded] = 0.001
parameters = (1.0, top_k, top_p)
expected = sample_logits(logits, *parameters, torch.Generator().manual_seed(42), noise)
actual = sample_graph_logits(logits, torch.tensor(parameters), noise)
assert expected.item() != first_excluded
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
def test_codec_chunks_preserve_zero_frames_and_flush_each_request_once(mocker) -> None:
manager = mocker.Mock()
manager.connector.config = {"extra": {"codec_chunk_frames": 5, "initial_codec_chunk_frames": 5}}
manager.code_prompt_token_ids = defaultdict(list)
manager.request_payload = {}
first = mocker.Mock(external_req_id="first")
first.is_finished.return_value = False
second = mocker.Mock(external_req_id="second")
second.is_finished.return_value = False
zeros = {"codes": {"audio": torch.zeros(3, 16, dtype=torch.long)}}
assert talker2code2wav_async_chunk(manager, zeros, first) is None
assert talker2code2wav_async_chunk(manager, zeros, second) is None
codes = torch.arange(64).reshape(4, 16)
output = talker2code2wav_async_chunk(manager, {"codes": {"audio": codes}}, first)
expected = torch.cat((torch.zeros(3, 16, dtype=torch.long), codes))
torch.testing.assert_close(output.codes.audio.reshape(16, -1).T, expected)
final = talker2code2wav_async_chunk(manager, None, first, is_finished=True)
assert final.codes.audio.numel() == 0
assert final.meta.finished.item()
other = talker2code2wav_async_chunk(manager, None, second, is_finished=True)
torch.testing.assert_close(other.codes.audio, torch.zeros(48, dtype=torch.long))
def test_codec_chunk_ramp_keeps_frame_order_and_flushes_partial_tail(mocker) -> None:
manager = mocker.Mock()
manager.connector.config = {"extra": {"codec_chunk_frames": 5, "codec_chunk_ramp": [1, 2, 4, 5]}}
manager.code_prompt_token_ids = defaultdict(list)
manager.request_payload = {}
request = mocker.Mock(external_req_id="ramp")
request.is_finished.return_value = False
outputs = []
for index in range(15):
result = talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.full((1, 16), index)}}, request)
if result is not None:
outputs.append(result.codes.audio.reshape(16, -1).T)
tail = talker2code2wav_async_chunk(manager, None, request, is_finished=True)
outputs.append(tail.codes.audio.reshape(16, -1).T)
assert [len(part) for part in outputs] == [1, 2, 4, 5, 3]
torch.testing.assert_close(torch.cat(outputs), torch.arange(15)[:, None].expand(-1, 16))
assert tail.meta.finished.item()
@torch.inference_mode()
def test_talker_history_and_generators_follow_requests_across_reordering_and_new_admission():
model = _small_talker()
first, second = _request_info("first"), _request_info("second")
# Unequal spans must select the last hidden row of each request.
model.make_omni_output(
torch.tensor([[0.0, 1.0], [0.0, 1.0], [1.0, 0.0], [0.0, 1.0]]),
model_intermediate_buffer=[first, second],
request_token_spans=[(0, 3), (3, 4)],
)
reordered = model.make_omni_output(
torch.tensor([[0.0, 1.0], [1.0, 0.0]]),
model_intermediate_buffer=[second, first],
request_token_spans=[(0, 1), (1, 2)],
)
assert [frame[0, 0].item() for frame in reordered.multimodal_outputs["codes"]["audio"]] == [3, 2]
assert first["breeze_state"]["history"].tolist() == [[1, 2]]
assert second["breeze_state"]["history"].tolist() == [[3, 3]]
generators = model.depth_decoder.generate_frames.call_args.kwargs["generators"]
assert generators[0] is second["breeze_state"]["generator"]
assert generators[1] is first["breeze_state"]["generator"]
newcomer = _request_info("newcomer")
admitted = model.make_omni_output(
torch.tensor([[1.0, 0.0], [1.0, 0.0]]),
model_intermediate_buffer=[newcomer, first],
request_token_spans=[(0, 1), (1, 2)],
)
assert [frame[0, 0].item() for frame in admitted.multimodal_outputs["codes"]["audio"]] == [1, 1]
assert newcomer["breeze_state"]["history"].tolist() == [[1]]
assert first["breeze_state"]["history"].tolist() == [[1, 2, 1]]
assert second["breeze_state"]["history"].tolist() == [[3, 3]]
@torch.inference_mode()
def test_talker_eos_does_not_shift_live_state_or_poison_the_next_batch():
model = _small_talker()
finished, live = _request_info("finished"), _request_info("live")
output = model.make_omni_output(
torch.tensor([[-1.0, -1.0], [0.0, 1.0]]),
model_intermediate_buffer=[finished, live],
request_token_spans=[(0, 1), (1, 2)],
)
assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0
assert finished["breeze_state"]["history"].numel() == 0
assert live["breeze_state"]["current"].tolist() == [[3, 3, 3]]
assert model.compute_logits(torch.zeros(2, 2)).argmax(-1).tolist() == [7, 0]
model.make_omni_output(torch.tensor([[0.0, 1.0]]), model_intermediate_buffer=[live], request_token_spans=[(0, 1)])
assert live["breeze_state"]["history"].tolist() == [[3, 3]]
assert model.compute_logits(torch.zeros(1, 2)).argmax(-1).tolist() == [0]
@pytest.mark.parametrize("temperature", [0.0, 0.9])
@torch.inference_mode()
def test_talker_mixed_eos_and_cfg_matches_independent_request_state(temperature):
"""Batch EOS decisions must not mix CFG state or change a request's RNG stream."""
finished_companion = "finished_cfg" + CFG_UNCOND_SUFFIX
live_companion = "live_cfg" + CFG_UNCOND_SUFFIX
order = ["finished_cfg", "live_plain", finished_companion, "finished_plain", live_companion, "live_cfg"]
hidden_by_id = {
"finished_cfg": [-1.0, -1.0],
finished_companion: [-1.0, -1.0],
"finished_plain": [-1.0, -1.0],
"live_plain": [0.0, 1.0],
"live_cfg": [1.0, 0.0],
live_companion: [1.0, 0.0],
}
def make_infos():
infos = {request_id: _request_info(request_id) for request_id in order}
for index, (request_id, info) in enumerate(infos.items()):
info["breeze_sampling"].update(temperature=temperature, top_k=1)
if "cfg" not in request_id:
# Interleave distinct sampler groups while each request still
# matches its independently decoded state and RNG trajectory.
info["breeze_sampling"].update(temperature=temperature * 0.5, top_p=0.8)
info["breeze_state"]["generator"].manual_seed(42 + index)
info["breeze_state"]["history"] = torch.tensor([[3, 3]])
if "cfg" in request_id:
info["breeze_prompt"]["guidance_scale"] = 4.0
if request_id.endswith(CFG_UNCOND_SUFFIX):
info["breeze_prompt"]["role"] = "uncond"
return infos
def sample_depth(hidden, first, *, temperature, generators, **kwargs):
# Consume request-local RNG state in the downstream phase as well as
# codebook-0 sampling, without requiring depth CUDA graphs on CPU.
tail = (
torch.stack([torch.randint(0, 4, (2,), generator=generator) for generator in generators])
if temperature > 0
else first[:, None].expand(-1, 2)
)
return torch.cat((first[:, None], tail), dim=1)
batched, independent = _small_talker(), _small_talker()
batched.depth_decoder.generate_frames.side_effect = sample_depth
independent.depth_decoder.generate_frames.side_effect = sample_depth
batch_infos, independent_infos = make_infos(), make_infos()
initial_rng = {rid: info["breeze_state"]["generator"].get_state() for rid, info in batch_infos.items()}
saved_frames = []
# The second step drops both EOS requests and moves the live CFG companion
# ahead of its parent, exercising shrinking and reordered batches.
for step_ids in (order, [live_companion, "live_cfg", "live_plain"]):
expected_audio, expected_tokens = {}, {}
for request_id in step_ids:
if request_id.endswith(CFG_UNCOND_SUFFIX):
continue
pair = [request_id]
if independent_infos[request_id]["breeze_prompt"]["guidance_scale"] != 1.0:
pair.append(request_id + CFG_UNCOND_SUFFIX)
result = independent.make_omni_output(
torch.tensor([hidden_by_id[rid] for rid in pair]),
model_intermediate_buffer=[independent_infos[rid] for rid in pair],
request_token_spans=[(i, i + 1) for i in range(len(pair))],
)
tokens = independent.compute_logits(torch.zeros(len(pair), 2)).argmax(-1)
for index, rid in enumerate(pair):
expected_audio[rid] = result.multimodal_outputs["codes"]["audio"][index].clone()
expected_tokens[rid] = tokens[index]
result = batched.make_omni_output(
torch.tensor([hidden_by_id[rid] for rid in step_ids]),
model_intermediate_buffer=[batch_infos[rid] for rid in step_ids],
request_token_spans=[(i, i + 1) for i in range(len(step_ids))],
)
actual_tokens = batched.compute_logits(torch.zeros(len(step_ids), 2)).argmax(-1)
for index, rid in enumerate(step_ids):
actual = result.multimodal_outputs["codes"]["audio"][index]
torch.testing.assert_close(actual, expected_audio[rid], atol=0, rtol=0)
torch.testing.assert_close(actual_tokens[index], expected_tokens[rid], atol=0, rtol=0)
saved_frames.append((actual, actual.clone()))
for rid in order:
actual, expected = batch_infos[rid]["breeze_state"], independent_infos[rid]["breeze_state"]
for key in ("history", "current"):
torch.testing.assert_close(actual[key], expected[key], atol=0, rtol=0)
torch.testing.assert_close(actual["generator"].get_state(), expected["generator"].get_state())
# Unconditional companions never own sampling; an EOS request consumes
# only codebook-0 sampling, with no downstream depth draws.
for rid in (finished_companion, live_companion):
torch.testing.assert_close(batch_infos[rid]["breeze_state"]["generator"].get_state(), initial_rng[rid])
for rid in ("finished_cfg", "finished_plain"):
reference_generator = torch.Generator().set_state(initial_rng[rid])
sampling = batch_infos[rid]["breeze_sampling"]
sample_logits(
torch.zeros(1, batched.config.vocab_size),
sampling["temperature"],
sampling["top_k"],
sampling["top_p"],
reference_generator,
)
torch.testing.assert_close(
batch_infos[rid]["breeze_state"]["generator"].get_state(), reference_generator.get_state()
)
for rid in ("finished_cfg", finished_companion, "finished_plain", live_companion):
assert batch_infos[rid]["breeze_state"]["history"].tolist() == [[3, 3]]
assert batch_infos["live_cfg"]["breeze_state"]["history"].tolist() == [[3, 3, 1, 2]]
assert batch_infos["live_plain"]["breeze_state"]["history"].tolist() == [[3, 3, 3, 3]]
assert batch_infos["live_cfg"]["breeze_state"]["current"] is batch_infos[live_companion]["breeze_state"]["current"]
for original, snapshot in saved_frames:
torch.testing.assert_close(original, snapshot, atol=0, rtol=0)
@torch.inference_mode()
def test_talker_all_eos_preserves_history_and_skips_depth():
model = _small_talker()
infos = [_request_info("first"), _request_info("second")]
for info in infos:
info["breeze_state"]["history"] = torch.tensor([[1, 2]])
output = model.make_omni_output(
torch.tensor([[-1.0, -1.0], [-2.0, -2.0]]),
model_intermediate_buffer=infos,
request_token_spans=[(0, 1), (1, 2)],
)
model.depth_decoder.generate_frames.assert_not_called()
assert model.compute_logits(torch.zeros(2, 2)).argmax(-1).tolist() == [7, 7]
assert all(frame.numel() == 0 for frame in output.multimodal_outputs["codes"]["audio"])
for info in infos:
assert info["breeze_state"]["history"].tolist() == [[1, 2]]
assert info["breeze_state"]["current"].tolist() == [[0, 0, 0]]
tests/model_executor/models/test_breeze_tts_2_first_code_integration.py
+184
View File
@@ -0,0 +1,184 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Check only the caller's grouping seam, with the existing weight-free fixture."""
import pytest
import torch
import torch.nn.functional as F
from tests.model_executor.models.breeze_tts_2_helpers import _request_info, _small_talker
from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import sample_logits
from vllm_omni.model_executor.models.breeze_tts_2.prompt import CFG_UNCOND_SUFFIX
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@torch.inference_mode()
def test_first_code_groups_preserve_logits_request_order_cfg_and_eos(monkeypatch):
"""Grouping changes dispatch only; head/history arithmetic and routing stay exact."""
model = _small_talker()
cfg_a, cfg_b = "cfg_a" + CFG_UNCOND_SUFFIX, "cfg_b" + CFG_UNCOND_SUFFIX
order = ["plain_a", "plain_b", cfg_a, "finished_a", "cfg_a", "cfg_b", cfg_b, "plain_c"]
tuples = {"a": (0.9, 3, 1.0), "b": (0.7, 2, 0.8), "c": (0.0, 0, 1.0)}
chosen = {"plain_a": 1, "plain_b": 2, "finished_a": 7, "cfg_a": 3, "cfg_b": 0, "plain_c": 2}
hidden = {
"plain_a": [1.0, -0.5],
"plain_b": [-0.5, 1.0],
cfg_a: [0.25, 0.5],
"finished_a": [-1.0, -1.0],
"cfg_a": [0.75, -0.5],
"cfg_b": [0.5, 0.75],
cfg_b: [-0.25, 0.5],
"plain_c": [0.5, -0.75],
}
infos = {rid: _request_info(rid) for rid in order}
for index, (rid, info) in enumerate(infos.items()):
parent = rid.removesuffix(CFG_UNCOND_SUFFIX)
parameters = tuples[parent[-1]]
info["breeze_sampling"].update(
temperature=parameters[0],
top_k=parameters[1],
top_p=parameters[2],
repetition_penalty=2.0 if parent.startswith("cfg") else 1.25,
)
info["breeze_state"]["generator"].manual_seed(100 + index)
info["breeze_state"]["history"] = torch.tensor([[1, 3, 1]])
if parent.startswith("cfg"):
info["breeze_prompt"]["guidance_scale"] = 4.0 if parent == "cfg_a" else 2.0
if rid.endswith(CFG_UNCOND_SUFFIX):
info["breeze_prompt"]["role"] = "uncond"
original_states = {
rid: {key: value.clone() for key, value in info["breeze_state"].items() if key in ("history", "current")}
for rid, info in infos.items()
}
request_for_generator = {id(info["breeze_state"]["generator"]): rid for rid, info in infos.items()}
class RecordingSampler:
def __init__(self):
self.calls = []
def sample(self, logits, parameters, generators):
ids = [request_for_generator[id(generator)] for generator in generators]
self.calls.append((parameters, ids, [row.clone() for row in logits]))
return [torch.tensor([chosen[rid]], dtype=torch.long) for rid in ids]
recorder = RecordingSampler()
model._first_code_sampler = recorder
# Build the scalar pre-grouping reference before spying on actual head
# calls. Both positive and negative repeated logits exercise the penalty.
expected_logits = {}
for rid in chosen:
row = torch.tensor([hidden[rid]], dtype=torch.bfloat16).float()
scores = F.linear(row, model.lm_head.weight)
scale = infos[rid]["breeze_prompt"]["guidance_scale"]
if scale != 1.0:
uncond = F.linear(
torch.tensor([hidden[rid + CFG_UNCOND_SUFFIX]], dtype=torch.bfloat16).float(), model.lm_head.weight
)
scores = uncond + scale * (scores - uncond)
scores[:, model.codebook_size : model.config.eos_token_id] = -torch.inf
history = infos[rid]["breeze_state"]["history"]
selected = scores.gather(1, history)
penalty = infos[rid]["breeze_sampling"]["repetition_penalty"]
scores.scatter_(1, history, torch.where(selected < 0, selected * penalty, selected / penalty))
expected_logits[rid] = scores
actual_linear, head_inputs = F.linear, []
def scalar_head(inputs, weight, bias=None):
assert inputs.dtype == torch.float32 and inputs.shape == (1, model.hidden_size)
assert weight is model.lm_head.weight
head_inputs.append(inputs.clone())
return actual_linear(inputs, weight, bias)
# Prefill-like unequal spans must still select each physical row's last
# hidden state, including a CFG companion preceding its parent.
flat_hidden: list[list[float]] = []
spans: list[tuple[int, int]] = []
for index, rid in enumerate(order):
start = len(flat_hidden)
flat_hidden.extend([[-99.0, -99.0]] * (index % 3))
flat_hidden.append(hidden[rid])
spans.append((start, len(flat_hidden)))
values = torch.tensor(flat_hidden, dtype=torch.bfloat16)
with monkeypatch.context() as head_guard:
head_guard.setattr(F, "linear", scalar_head)
output = model.make_omni_output(
values, model_intermediate_buffer=[infos[rid] for rid in order], request_token_spans=spans
)
assert output.text_hidden_states is values
assert len(head_inputs) == len(chosen) + 2
assert [(parameters, ids) for parameters, ids, _ in recorder.calls] == [
(tuples["a"], ["plain_a", "finished_a", "cfg_a"]),
(tuples["b"], ["plain_b", "cfg_b"]),
(tuples["c"], ["plain_c"]),
]
for _, ids, rows in recorder.calls:
for rid, row in zip(ids, rows, strict=True):
assert row.dtype == torch.float32 and row.shape == (1, model.config.vocab_size)
torch.testing.assert_close(row, expected_logits[rid], atol=0, rtol=0)
# Depth group creation must follow original conditioned order, not the
# first-code sampler's grouped traversal (A,A,A,B,B,C).
live_order = ["plain_a", "plain_b", "cfg_a", "cfg_b", "plain_c"]
calls = model.depth_decoder.generate_frames.call_args_list
assert len(calls) == len(live_order)
for call, rid in zip(calls, live_order, strict=True):
depth_hidden, first = call.args
scale = infos[rid]["breeze_prompt"]["guidance_scale"]
expected_hidden = [hidden[rid]] + ([hidden[rid + CFG_UNCOND_SUFFIX]] if scale != 1 else [])
torch.testing.assert_close(depth_hidden, torch.tensor(expected_hidden, dtype=torch.bfloat16), atol=0, rtol=0)
assert first.tolist() == [chosen[rid]]
assert call.kwargs["generators"] == [infos[rid]["breeze_state"]["generator"]]
assert call.kwargs["guidance_scale"] == scale
assert tuple(call.kwargs[key] for key in ("temperature", "top_k", "top_p")) == tuples[rid[-1]]
assert model.compute_logits(torch.zeros(len(order), 2)).argmax(-1).tolist() == [0, 0, 0, 7, 0, 0, 0, 0]
for index, rid in enumerate(order):
state = infos[rid]["breeze_state"]
audio = output.multimodal_outputs["codes"]["audio"][index]
if rid in live_order:
assert state["history"].tolist() == [[1, 3, 1, chosen[rid]]]
assert audio.tolist() == [[chosen[rid]] * 3]
else:
assert audio.numel() == 0
torch.testing.assert_close(state["history"], original_states[rid]["history"], atol=0, rtol=0)
torch.testing.assert_close(
infos["finished_a"]["breeze_state"]["current"], original_states["finished_a"]["current"], atol=0, rtol=0
)
for parent, companion in (("cfg_a", cfg_a), ("cfg_b", cfg_b)):
assert infos[parent]["breeze_state"]["current"] is infos[companion]["breeze_state"]["current"]
@torch.inference_mode()
def test_first_code_shared_generator_preserves_order_across_parameter_groups(monkeypatch):
"""A shared RNG stream must retain A,B,A request order across tuple groups."""
model = _small_talker()
infos = [_request_info(rid) for rid in ("first_a", "middle_b", "last_a")]
parameters = [(0.9, 3, 1.0), (0.7, 2, 0.8), (0.9, 3, 1.0)]
actual_generator = torch.Generator().manual_seed(2**63 + 19)
expected_generator = torch.Generator().manual_seed(2**63 + 19)
hidden = torch.tensor([[0.2, 0.0], [0.0, 0.2], [0.1, 0.2]], dtype=torch.bfloat16)
expected = []
for index, (info, values) in enumerate(zip(infos, parameters, strict=True)):
info["breeze_sampling"].update(temperature=values[0], top_k=values[1], top_p=values[2])
info["breeze_state"]["generator"] = actual_generator
logits = F.linear(hidden[index : index + 1].float(), model.lm_head.weight)
logits[:, model.codebook_size : model.config.eos_token_id] = -torch.inf
expected.append(sample_logits(logits, *values, expected_generator))
assert all(token.item() != model.config.eos_token_id for token in expected)
sample, calls = model._first_code_sampler.sample, []
def record_sample(logits, settings, generators):
calls.append((settings, len(logits), [generator is actual_generator for generator in generators]))
return sample(logits, settings, generators)
monkeypatch.setattr(model._first_code_sampler, "sample", record_sample)
output = model.make_omni_output(
hidden, model_intermediate_buffer=infos, request_token_spans=[(0, 1), (1, 2), (2, 3)]
)
assert calls == [(values, 1, [True]) for values in parameters]
torch.testing.assert_close(actual_generator.get_state(), expected_generator.get_state(), atol=0, rtol=0)
for info, token, audio in zip(infos, expected, output.multimodal_outputs["codes"]["audio"], strict=True):
torch.testing.assert_close(audio, token[:, None].repeat(1, 3), atol=0, rtol=0)
torch.testing.assert_close(info["breeze_state"]["history"], token[:, None], atol=0, rtol=0)
assert info["breeze_state"]["generator"] is actual_generator
tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py
+307
View File
@@ -0,0 +1,307 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""First-code sampler token, RNG, graph-cache and fallback contracts."""
import pytest
import torch
from tests.helpers.mark import hardware_test
from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import sample_logits
from vllm_omni.model_executor.models.breeze_tts_2.first_code_sampler import BreezeFirstCodeSampler
pytestmark = [pytest.mark.core_model]
PARAMETERS = (0.9, 50, 1.0)
VOCAB = 2052
EOS = 2051
def _logits(count: int, *, device: str, seed: int = 100, vocab: int = VOCAB) -> list[torch.Tensor]:
source = torch.Generator(device="cpu").manual_seed(seed)
values = torch.randn((count, vocab), generator=source, dtype=torch.float32) * 3
if vocab == VOCAB:
values[:, 2048:EOS] = -torch.inf
return list(values.to(device).split(1))
def _generators(seeds: list[int], device: str) -> list[torch.Generator]:
return [torch.Generator(device=device).manual_seed(seed) for seed in seeds]
def _reference(rows, parameters, generators) -> list[torch.Tensor]:
return [sample_logits(row, *parameters, generator) for row, generator in zip(rows, generators, strict=True)]
def _assert_outputs(actual, expected, rows) -> None:
assert isinstance(actual, list) and len(actual) == len(expected)
for result, reference, row in zip(actual, expected, rows, strict=True):
assert result.shape == (1,) and result.dtype == torch.long and result.device == row.device
torch.testing.assert_close(result, reference, atol=0, rtol=0)
def _assert_states(actual, expected) -> None:
for left, right in zip(actual, expected, strict=True):
torch.testing.assert_close(left.get_state(), right.get_state(), atol=0, rtol=0)
@pytest.mark.cpu
def test_empty_and_cpu_sampling_preserve_eager_contract() -> None:
sampler = BreezeFirstCodeSampler()
assert sampler.sample([], PARAMETERS, []) == []
assert not sampler._graphs
rows = _logits(3, device="cpu")
before = [row.clone() for row in rows]
actual_generators = _generators([42, 2**63 + 19, 2**64 - 23], "cpu")
expected_generators = _generators([42, 2**63 + 19, 2**64 - 23], "cpu")
for parameters in (PARAMETERS, (0.7, 20, 0.8), (1.2, 0, 0.95)):
expected = _reference(rows, parameters, expected_generators)
_assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
_assert_states(actual_generators, expected_generators)
for row, original in zip(rows, before, strict=True):
torch.testing.assert_close(row, original, atol=0, rtol=0)
assert not sampler._graphs
@pytest.mark.cpu
def test_greedy_never_samples_or_advances_rng(monkeypatch) -> None:
sampler = BreezeFirstCodeSampler()
rows = _logits(3, device="cpu")
generators = _generators([42, 43, 44], "cpu")
original_states = [generator.get_state().clone() for generator in generators]
def unexpected_multinomial(*_args, **_kwargs):
raise AssertionError("Greedy sampling must not call multinomial")
monkeypatch.setattr(torch, "multinomial", unexpected_multinomial)
for parameters in ((0.0, 0, 1.0), (0.0, 1, 0.2)):
_assert_outputs(sampler.sample(rows, parameters, generators), [row.argmax(-1) for row in rows], rows)
for generator, state in zip(generators, original_states, strict=True):
torch.testing.assert_close(generator.get_state(), state, atol=0, rtol=0)
assert not sampler._graphs
@pytest.mark.cpu
@pytest.mark.parametrize(
"invalid", ["extra_generator", "missing_generator", "empty_rows", "multirow", "mixed_vocab", "mixed_dtype"]
)
def test_invalid_batch_fails_before_advancing_any_generator(invalid: str) -> None:
sampler = BreezeFirstCodeSampler()
rows = _logits(2, device="cpu")
generators = _generators([42, 43], "cpu")
if invalid == "extra_generator":
rows = rows[:1]
elif invalid == "missing_generator":
generators = generators[:1]
elif invalid == "empty_rows":
rows = []
elif invalid == "multirow":
rows[1] = rows[1].expand(2, -1)
elif invalid == "mixed_vocab":
rows[1] = rows[1][:, :-1]
else:
rows[1] = rows[1].double()
states = [generator.get_state().clone() for generator in generators]
with pytest.raises(ValueError):
sampler.sample(rows, PARAMETERS, generators)
for generator, state in zip(generators, states, strict=True):
torch.testing.assert_close(generator.get_state(), state, atol=0, rtol=0)
assert not sampler._graphs
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.inference_mode()
def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None:
sampler = BreezeFirstCodeSampler()
seeds = list(range(42, 77)) + [2**63 + 19, 2**64 - 23, 101]
actual_generators = _generators(seeds, "cuda")
expected_generators = _generators(seeds, "cuda")
# Torch also accepts an unspecified CUDA index. Exercise it alongside the
# indexed generators created from serving request tensors.
indexed_device = torch.device("cuda", torch.accelerator.current_device_index())
for pool in (actual_generators, expected_generators):
pool[1] = torch.Generator(device=indexed_device).manual_seed(seeds[1])
orders = (
[0],
[2, 0, 1],
list(range(16)),
[36, *range(16)],
[35, *range(30)],
list(reversed(range(32))),
[35, 1, 2],
)
retained: list[tuple[torch.Tensor, torch.Tensor]] = []
global_state = torch.cuda.get_rng_state().clone()
for sweep in range(2):
if sweep:
for pool in (actual_generators, expected_generators):
pool[2] = torch.Generator(device="cuda").manual_seed(2**64 - 101)
for case, order in enumerate(orders):
rows = _logits(len(order), device="cuda", seed=1000 + sweep * 10 + case)
if sweep and case == 3:
for index, row in enumerate(rows):
row.fill_(-torch.inf)
row[0, EOS if index % 2 else 2] = 0
original = [row.clone() for row in rows]
live_actual = [actual_generators[index] for index in order]
live_expected = [expected_generators[index] for index in order]
expected = _reference(rows, PARAMETERS, live_expected)
actual = sampler.sample(rows, PARAMETERS, live_actual)
_assert_outputs(actual, expected, rows)
_assert_states(actual_generators, expected_generators)
retained.extend((value, reference.clone()) for value, reference in zip(actual, expected, strict=True))
for row, before in zip(rows, original, strict=True):
torch.testing.assert_close(row, before, atol=0, rtol=0)
# Preserve the subsequent 15 independent depth draws, including
# skipping them for EOS. All inactive/companion generators remain
# in the full-pool comparison above and below.
for index, token in enumerate(torch.cat(expected).cpu().tolist()):
if token == EOS:
continue
left = torch.empty((15, 2051), device="cuda")
right = torch.empty_like(left)
for codebook in range(15):
left[codebook].exponential_(generator=live_actual[index])
right[codebook].exponential_(generator=live_expected[index])
torch.testing.assert_close(left, right, atol=0, rtol=0)
_assert_states(actual_generators, expected_generators)
# B1/3/16/17/31/32 must use exactly buckets 1/4/16/32.
assert len(sampler._graphs) == 4
rows = _logits(3, device="cuda", seed=2026)
expected = _reference(rows, PARAMETERS, expected_generators[:3])
def unexpected_multinomial(*_args, **_kwargs):
raise AssertionError("Warmed graph replay must not call Python multinomial")
with monkeypatch.context() as replay_guard:
replay_guard.setattr(torch, "multinomial", unexpected_multinomial)
_assert_outputs(sampler.sample(rows, PARAMETERS, actual_generators[:3]), expected, rows)
_assert_outputs(
sampler.sample(rows, (0.0, 1, 0.2), actual_generators[:3]), [row.argmax(-1) for row in rows], rows
)
_assert_states(actual_generators, expected_generators)
assert len(sampler._graphs) == 4
for value, expected in retained:
torch.testing.assert_close(value, expected, atol=0, rtol=0)
torch.testing.assert_close(torch.cuda.get_rng_state(), global_state, atol=0, rtol=0)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.inference_mode()
def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None:
sampler = BreezeFirstCodeSampler()
actual_generators = _generators([42, 43, 44], "cuda")
expected_generators = _generators([42, 43, 44], "cuda")
cases = ((0.9, 50, 1.0), (0.7, 100, 0.8), (1.2, 0, 0.95), (1.0, 1, 1.0), (1.0, 0, 0.5))
for index, parameters in enumerate(cases):
rows = _logits(3, device="cuda", seed=200 + index)
if index >= 3:
for row in rows:
row.fill_(-torch.inf)
row[0, :4] = 0 # top-k ties and exact top-p=0.5 boundary.
expected = _reference(rows, parameters, expected_generators)
_assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
_assert_states(actual_generators, expected_generators)
assert len(sampler._graphs) == len(cases)
rows = _logits(3, device="cuda", vocab=17)
expected = _reference(rows, PARAMETERS, expected_generators)
_assert_outputs(sampler.sample(rows, PARAMETERS, actual_generators), expected, rows)
_assert_states(actual_generators, expected_generators)
assert len(sampler._graphs) == len(cases) + 1
class UnavailableGraph:
def __init__(self, *_args, **_kwargs):
raise AssertionError("Unsupported graph API must select the eager fallback")
fallback = BreezeFirstCodeSampler()
expected = _reference(rows, PARAMETERS, expected_generators)
with monkeypatch.context() as api_guard:
api_guard.setattr(torch.cuda, "CUDAGraph", UnavailableGraph)
_assert_outputs(fallback.sample(rows, PARAMETERS, actual_generators), expected, rows)
_assert_states(actual_generators, expected_generators)
assert not fallback._graphs
# Validate every row before the first eager/captured draw, including an
# invalid device that appears only after an otherwise valid CUDA row.
invalid = [rows[0], rows[1].cpu(), rows[2]]
with pytest.raises(ValueError):
sampler.sample(invalid, PARAMETERS, actual_generators)
_assert_states(actual_generators, expected_generators)
# A late invalid generator must likewise fail before row zero samples.
cpu_generator = torch.Generator(device="cpu").manual_seed(77)
cpu_reference = torch.Generator(device="cpu").manual_seed(77)
mixed_generators = [actual_generators[0], cpu_generator, actual_generators[2]]
with pytest.raises(ValueError):
sampler.sample(rows, PARAMETERS, mixed_generators)
_assert_states(mixed_generators, [expected_generators[0], cpu_reference, expected_generators[2]])
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.inference_mode()
def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> None:
sampler = BreezeFirstCodeSampler()
rows = _logits(3, device="cuda")
# The duplicate-generator guard must precede lookup of an existing graph.
sampler.sample(rows, PARAMETERS, _generators([100, 101, 102], "cuda"))
cached = dict(sampler._graphs)
assert len(cached) == 1
actual = torch.Generator(device="cuda").manual_seed(2**64 - 23)
reference = torch.Generator(device="cuda").manual_seed(2**64 - 23)
multinomial = torch.multinomial
calls = []
def eager_spy(*args, **kwargs):
calls.append(kwargs.get("generator"))
return multinomial(*args, **kwargs)
for order in ([0, 1, 2], [2, 0, 1]):
ordered = [rows[index] for index in order]
expected = _reference(ordered, PARAMETERS, [reference] * 3)
calls.clear()
with monkeypatch.context() as fallback_guard:
fallback_guard.setattr(torch, "multinomial", eager_spy)
_assert_outputs(sampler.sample(ordered, PARAMETERS, [actual] * 3), expected, ordered)
assert len(calls) == 3 and all(generator is actual for generator in calls)
_assert_states([actual], [reference])
assert set(sampler._graphs) == set(cached)
assert all(sampler._graphs[key] is value for key, value in cached.items())
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.inference_mode()
def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch) -> None:
sampler = BreezeFirstCodeSampler()
rows = _logits(1, device="cuda")
actual_generators = _generators([42], "cuda")
expected_generators = _generators([42], "cuda")
for top_k in range(1, 33):
parameters = (0.9, top_k, 1.0)
expected = _reference(rows, parameters, expected_generators)
_assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
_assert_states(actual_generators, expected_generators)
cached = dict(sampler._graphs)
assert len(cached) == 32
parameters = (0.9, 33, 1.0)
expected = _reference(rows, parameters, expected_generators)
multinomial = torch.multinomial
calls = []
def eager_spy(*args, **kwargs):
calls.append(kwargs.get("generator"))
return multinomial(*args, **kwargs)
with monkeypatch.context() as fallback_guard:
fallback_guard.setattr(torch, "multinomial", eager_spy)
_assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows)
assert len(calls) == 1 and calls[0] is actual_generators[0]
_assert_states(actual_generators, expected_generators)
assert set(sampler._graphs) == set(cached)
assert all(sampler._graphs[key] is value for key, value in cached.items())
expected = _reference(rows, (0.9, 1, 1.0), expected_generators)
def unexpected_multinomial(*_args, **_kwargs):
raise AssertionError("Cache exhaustion must retain existing captured entries")
with monkeypatch.context() as replay_guard:
replay_guard.setattr(torch, "multinomial", unexpected_multinomial)
_assert_outputs(sampler.sample(rows, (0.9, 1, 1.0), actual_generators), expected, rows)
_assert_states(actual_generators, expected_generators)
tests/model_executor/models/test_breeze_tts_2_graphs.py
+405
View File
@@ -0,0 +1,405 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Check captured cache ownership and padded text masking against eager models."""
import pytest
import torch
from transformers import LlamaConfig, T5Gemma2TextConfig
from transformers.models.t5gemma2.modeling_t5gemma2 import T5Gemma2TextEncoder
from vllm.platforms import current_platform
from tests.helpers.mark import hardware_test
from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import BreezeDepthDecoder, sample_logits
from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration
from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import (
BreezeTextEncoderCompiled,
BreezeTextEncoderGraph,
)
pytestmark = [pytest.mark.core_model]
@pytest.fixture
def full_precision_matmul():
# Different padded shapes can select different TF32 kernels. Compare
# masking and replay semantics independently of that reduced precision.
original_precision = torch.get_float32_matmul_precision()
torch.set_float32_matmul_precision("highest")
try:
yield
finally:
torch.set_float32_matmul_precision(original_precision)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.inference_mode()
def test_text_graph_masks_padding_and_local_attention_across_replays(full_precision_matmul) -> None:
torch.manual_seed(42)
config = T5Gemma2TextConfig(
vocab_size=64,
hidden_size=32,
intermediate_size=64,
num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=2,
head_dim=8,
sliding_window=16,
layer_types=["sliding_attention", "full_attention"],
dropout_rate=0.0,
)
config._attn_implementation = "sdpa"
encoder = T5Gemma2TextEncoder(config).to("cuda").eval()
encoder.embed_tokens.eoi_token_index = 17
encoder.embed_tokens.eoi_embedding.fill_(0.4)
projection = torch.nn.Linear(32, 48, bias=False).to("cuda").eval()
graph = BreezeTextEncoderGraph(encoder, projection, 64)
for length in (31, 51, 17, 31):
prompt = torch.randint(0, 64, (1, length), device="cuda")
prompt[:, 3] = 17
expected = projection(encoder(input_ids=prompt).last_hidden_state)[0]
actual = graph.run(prompt)
torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5)
batched = BreezeTextEncoderGraph(encoder, projection, 64, batch_size=3)
prompts = [torch.randint(0, 64, (length,), device="cuda") for length in (31, 51, 17)]
outputs = batched.run_batch(prompts)
for prompt, actual in zip(prompts, outputs, strict=True):
expected = projection(encoder(input_ids=prompt[None]).last_hidden_state)[0]
torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5)
compiled = BreezeTextEncoderCompiled(encoder, projection)
long_prompts = [torch.randint(0, 64, (length,)) for length in (211, 173, 129)]
for batch in ([long_prompts[0]], long_prompts, [long_prompts[1]]):
outputs = compiled.run_batch(batch)
for prompt, actual in zip(batch, outputs, strict=True):
expected = projection(encoder(input_ids=prompt[None].to("cuda")).last_hidden_state)[0]
torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5)
def _make_depth_decoder(vocab_size: int = 35) -> BreezeDepthDecoder:
config = LlamaConfig(
hidden_size=32,
intermediate_size=64,
num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=2,
head_dim=8,
vocab_size=vocab_size,
max_position_embeddings=16,
rope_theta=500000.0,
rope_scaling={
"rope_type": "llama3",
"factor": 32.0,
"low_freq_factor": 0.001953125,
"high_freq_factor": 0.0078125,
"original_max_position_embeddings": 16,
},
)
config.num_codebooks = 16
config.audio_embed_size = 48
depth = BreezeDepthDecoder(config).to("cuda").eval()
torch.nn.init.normal_(depth.codebooks_head, std=0.1)
return depth
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.inference_mode()
def test_depth_graph_survives_interleaved_batches_and_workspace_allocations(full_precision_matmul, monkeypatch) -> None:
torch.manual_seed(42)
depth = _make_depth_decoder()
hidden = torch.randn(2, 48, device="cuda")
first = torch.tensor([7, 11], device="cuda")
generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in (42, 99)]
expected = torch.cat(
[
depth.generate_frame(hidden[i : i + 1], first[i : i + 1], temperature=0, top_k=0, top_p=1, generator=g)
for i, g in enumerate(generators)
]
)
for batch in (1, 2, 1, 2):
# Unowned allocations made before graph capture would be recycled here.
pressure = [torch.full((batch, 2, depth.num_codebooks, 8), float("nan"), device="cuda") for _ in range(32)]
actual = depth.generate_frames(
hidden[:batch], first[:batch], temperature=0, top_k=0, top_p=1, generators=generators[:batch]
)
torch.testing.assert_close(actual, expected[:batch], atol=0, rtol=0)
del pressure
def sample(batch: int) -> torch.Tensor:
return depth.generate_frames(
hidden[:batch],
first[:batch],
temperature=0.9,
top_k=10,
top_p=0.8,
generators=[torch.Generator(device="cuda").manual_seed(seed) for seed in (42, 99)[:batch]],
)
sampled = sample(1)
sampled_graph = depth._graphs[(1, 1, False)]
greedy_graph = depth._graphs[(1, 1, True)]
assert sampled_graph is not greedy_graph
sample(2)
torch.testing.assert_close(sample(1), sampled, atol=0, rtol=0)
for temperature, top_k, top_p in ((0.0, 0, 1.0), (0.7, 20, 0.95), (0.0, 1, 0.2), (1.2, 0, 0.8)):
rng_before = generators[0].get_state()
actual = depth.generate_frames(
hidden[:1], first[:1], temperature=temperature, top_k=top_k, top_p=top_p, generators=generators[:1]
)
if temperature == 0:
assert depth._graphs[(1, 1, True)] is greedy_graph
torch.testing.assert_close(actual, expected[:1], atol=0, rtol=0)
torch.testing.assert_close(generators[0].get_state(), rng_before, atol=0, rtol=0)
else:
assert depth._graphs[(1, 1, False)] is sampled_graph
torch.testing.assert_close(sample(1), sampled, atol=0, rtol=0)
graph = None
for scale in (4.0, 0.5, 4.0):
expected_cfg = depth._generate_frame(
hidden, first[:1], depth._allocate_cache(hidden), 0, 0, 1, generators[0], guidance_scale=scale
)
actual_cfg = depth.generate_frames(
hidden, first[:1], temperature=0, top_k=0, top_p=1, generators=generators[:1], guidance_scale=scale
)
torch.testing.assert_close(actual_cfg, expected_cfg, atol=0, rtol=0)
if graph is not None:
assert depth._graphs[(2, 2, True)] is graph
graph = depth._graphs[(2, 2, True)]
sample(2)
# Use a still-cold sampled CFG entry to exercise platforms without
# captured generator support, while reusing the existing compiled model.
fallback_key = (2, 2, False)
assert fallback_key not in depth._graphs
reference_generator = torch.Generator(device="cuda").set_state(generators[0].get_state())
inactive_rng = generators[1].get_state()
exponential = torch.Tensor.exponential_
observed_generators = []
def record_eager_noise(tensor, *args, **kwargs):
observed_generators.append(kwargs.get("generator"))
return exponential(tensor, *args, **kwargs)
fallback_graph = None
retained = []
with monkeypatch.context() as platform_guard:
platform_guard.setattr(current_platform, "is_cuda", lambda: False)
for scale in (4.0, 0.5, 4.0):
expected = depth._generate_frame(
hidden,
first[:1],
depth._allocate_cache(hidden),
0.9,
10,
0.8,
reference_generator,
guidance_scale=scale,
)
observed_generators.clear()
with monkeypatch.context() as noise_spy:
noise_spy.setattr(torch.Tensor, "exponential_", record_eager_noise)
actual = depth.generate_frames(
hidden,
first[:1],
temperature=0.9,
top_k=10,
top_p=0.8,
generators=generators[:1],
guidance_scale=scale,
)
entry = depth._graphs[fallback_key]
assert entry.noise_generators is None
assert len(observed_generators) == depth.num_codebooks - 1
assert all(generator is generators[0] for generator in observed_generators)
if fallback_graph is not None:
assert entry is fallback_graph
fallback_graph = entry
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
torch.testing.assert_close(generators[0].get_state(), reference_generator.get_state(), atol=0, rtol=0)
torch.testing.assert_close(generators[1].get_state(), inactive_rng, atol=0, rtol=0)
retained.append((actual, expected.clone()))
for actual, expected in retained:
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.inference_mode()
def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_matmul, monkeypatch) -> None:
def reject_eager_noise(*args, **kwargs):
raise AssertionError("Sampled depth replay must not issue eager RNG draws")
torch.manual_seed(17)
# Match the checkpoint's multinomial draw size, including reserved IDs.
depth = _make_depth_decoder(vocab_size=2051)
hidden = torch.randn(32, 48, device="cuda")
uncond = torch.randn_like(hidden)
first = torch.arange(32, device="cuda")
first_logits = torch.randn(32, depth.vocab_size, device="cuda")
first_logits[:, -3:] = -torch.inf
seeds = [42 + index for index in range(32)]
seeds[0], seeds[-1] = 2**63 + 17, 2**64 - 1
generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds]
reference_generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds]
expected_greedy = torch.cat(
[
depth.generate_frame(
hidden[index : index + 1],
first[index : index + 1],
temperature=0,
top_k=0,
top_p=1,
generator=reference_generators[index],
)
for index in range(32)
]
)
retained: list[tuple[torch.Tensor, torch.Tensor]] = []
for batch in (*range(1, 33), 17, 11, 3, 1):
rows = list(reversed(range(batch)))
actual = depth.generate_frames(
hidden[rows], first[rows], temperature=0, top_k=0, top_p=1, generators=[generators[row] for row in rows]
)
expected = expected_greedy[rows]
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
retained.append((actual, expected.clone()))
assert set(depth._graphs) == {(batch, 1, True) for batch in (1, 2, 4, 8, 16, 32)}
for generator, reference in zip(generators, reference_generators, strict=True):
torch.testing.assert_close(generator.get_state(), reference.get_state(), atol=0, rtol=0)
retired_generators: list[tuple[torch.Generator, torch.Tensor]] = []
for temperature, guidance_scale in ((0.9, 1.0), (0.0, 4.0), (0.9, 4.0)):
batches: tuple[int, ...] = (3, 5, 11, 7, 3, 1)
if guidance_scale == 1:
batches = (16, 17, 31, 32, 17, 31, 16, *batches)
for step, batch in enumerate(batches):
if guidance_scale == 1 and step == 4:
# New requests reuse rows in the same 32-request graph. Their
# seeds must replace the previous owners, not just offsets.
for row, seed in ((0, 2**64 - 1), (16, 2**63 + 321)):
retired_generators.append((generators[row], generators[row].get_state()))
generators[row] = torch.Generator(device="cuda").manual_seed(seed)
reference_generators[row] = torch.Generator(device="cuda").manual_seed(seed)
rows = list(reversed(range(batch)))
batch_hidden = hidden[rows]
if guidance_scale != 1:
batch_hidden = torch.cat((batch_hidden, uncond[rows]))
expected, first_codes = [], []
for row in rows:
# The talker and depth decoder share each request's generator.
# Exercise a codebook-0 sample immediately before graph replay.
row_first = sample_logits(first_logits[row : row + 1], temperature, 10, 0.8, generators[row])
reference_first = sample_logits(
first_logits[row : row + 1], temperature, 10, 0.8, reference_generators[row]
)
torch.testing.assert_close(row_first, reference_first, atol=0, rtol=0)
first_codes.append(row_first)
row_hidden = hidden[row : row + 1]
if guidance_scale != 1:
row_hidden = torch.cat((row_hidden, uncond[row : row + 1]))
expected.append(
depth._generate_frame(
row_hidden,
reference_first,
depth._allocate_cache(row_hidden),
temperature,
10,
0.8,
reference_generators[row],
guidance_scale=guidance_scale if guidance_scale != 1 else None,
)
)
request_bucket = 1 << (batch - 1).bit_length()
branches = 2 if guidance_scale != 1 else 1
graph_key = (request_bucket * branches, branches, temperature == 0)
with monkeypatch.context() as rng_guard:
if temperature > 0 and graph_key in depth._graphs:
# A warmed replay must generate noise within the CUDA
# graph, without dispatching new RNG kernels from Python.
rng_guard.setattr(torch.Tensor, "exponential_", reject_eager_noise)
actual = depth.generate_frames(
batch_hidden,
torch.cat(first_codes),
temperature=temperature,
top_k=10,
top_p=0.8,
generators=[generators[row] for row in rows],
guidance_scale=guidance_scale,
)
expected_frames = torch.cat(expected)
torch.testing.assert_close(actual, expected_frames, atol=0, rtol=0)
retained.append((actual, expected_frames.clone()))
# Check the next codebook-0 sample as well as the depth result: a
# replay must return the advanced RNG stream to its live request.
for row in rows:
actual_next = sample_logits(first_logits[row : row + 1], temperature, 10, 0.8, generators[row])
expected_next = sample_logits(
first_logits[row : row + 1], temperature, 10, 0.8, reference_generators[row]
)
torch.testing.assert_close(actual_next, expected_next, atol=0, rtol=0)
# Padding must not draw RNG, and inactive request streams must not advance.
for generator, reference in zip(generators, reference_generators, strict=True):
torch.testing.assert_close(generator.get_state(), reference.get_state(), atol=0, rtol=0)
for generator, state in retired_generators:
torch.testing.assert_close(generator.get_state(), state, atol=0, rtol=0)
expected_keys = {(batch, 1, True) for batch in (1, 2, 4, 8, 16, 32)}
expected_keys.update((batch, 1, False) for batch in (1, 4, 8, 16, 32))
expected_keys.update((batch, 2, greedy) for batch in (2, 8, 16, 32) for greedy in (False, True))
assert set(depth._graphs) == expected_keys
for (batch_bucket, branches, greedy), entry in depth._graphs.items():
if greedy:
assert entry.noise_generators is None
else:
assert entry.noise_generators is not None
assert len(entry.noise_generators) == batch_bucket // branches
assert {id(generator) for generator in entry.noise_generators}.isdisjoint(
id(generator) for generator in generators
)
for actual, expected in retained:
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@torch.inference_mode()
def test_batched_decode_embedding_matches_scalar_bfloat16() -> None:
# Use the checkpoint's RVQ table dimensions without loading its backbone.
model = BreezeForConditionalGeneration.__new__(BreezeForConditionalGeneration)
torch.nn.Module.__init__(model)
model.num_codebooks, model.hidden_size = 16, 2048
model.register_buffer("offsets", torch.arange(16, device="cuda") * 2051)
model.depth_decoder = torch.nn.Module()
torch.manual_seed(17)
model.depth_decoder.embed_tokens = torch.nn.Embedding(16 * 2051, 2048, device="cuda", dtype=torch.bfloat16)
for seed in (17, 42, 99):
generator = torch.Generator(device="cuda").manual_seed(seed)
frames = torch.randint(0, 2051, (32, 16), device="cuda", generator=generator)
frames[1] = frames[0]
for batch in (1, 2, 3, 4, 5, 6, 7, 8, 16, 32, 7, 1):
# Reverse physical rows and revisit smaller batches after expansion.
current_frames = [row[None] for row in frames[:batch].flip(0)]
if batch > 1:
# CFG branches share their generated audio frame.
current_frames[-2] = current_frames[-1]
infos = [
{
"_omni_is_prefill": False,
"breeze_state": {"current": current},
"global_request_id": [str(index)],
"breeze_prompt": {},
"breeze_sampling": {},
}
for index, current in enumerate(current_frames)
]
ids = torch.zeros(batch, device="cuda", dtype=torch.long)
expected = torch.cat(
[
model.preprocess(input_ids=ids[index : index + 1], input_embeds=None, **info)[1]
for index, info in enumerate(infos)
]
)
output_ids, actual, updates = model.preprocess_decode_batch(input_ids=ids, req_infos=infos)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
assert output_ids is ids
assert updates == [{} for _ in infos]
tests/model_executor/models/test_breeze_tts_2_preprocess.py
+119
View File
@@ -0,0 +1,119 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU contracts for Breeze's decode-only embedding preprocessing."""
import pytest
import torch
from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _embedding_talker(dtype):
model = BreezeForConditionalGeneration.__new__(BreezeForConditionalGeneration)
torch.nn.Module.__init__(model)
model.num_codebooks, model.hidden_size = 16, 8
model.register_buffer("offsets", torch.arange(model.num_codebooks) * 7)
model.depth_decoder = torch.nn.Module()
model.depth_decoder.embed_tokens = torch.nn.Embedding(16 * 7, 8, dtype=dtype)
with torch.no_grad():
model.depth_decoder.embed_tokens.weight.copy_(
torch.randn(16 * 7, 8, generator=torch.Generator().manual_seed(17))
)
return model
def _decode_info(request_id, current, seed):
return {
"global_request_id": [request_id],
"breeze_prompt": {"role": "cond", "guidance_scale": 1.0},
"breeze_sampling": {"temperature": 0.9, "top_k": 4, "top_p": 1.0, "repetition_penalty": 1.1},
"breeze_state": {
"current": current,
"history": torch.tensor([[1, 2, 3]], dtype=torch.long),
"generator": torch.Generator().manual_seed(seed),
},
"_omni_is_prefill": False,
}
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
@torch.inference_mode()
def test_decode_batch_matches_scalar_after_reordering_shrinking_and_cfg(dtype):
model = _embedding_talker(dtype)
shared_frame = torch.arange(16, dtype=torch.long).reshape(1, 16) % 7
infos = [
_decode_info("unconditional", shared_frame, 1),
_decode_info("other", (shared_frame + 1) % 7, 2),
_decode_info("conditioned", shared_frame, 3),
_decode_info("last", (shared_frame + 3) % 7, 4),
]
infos[0]["breeze_prompt"] = {"role": "uncond", "guidance_scale": 4.0}
infos[2]["breeze_prompt"]["guidance_scale"] = 4.0
saved = []
for indices in ([0, 1, 2, 3], [2, 0, 1], [1]):
selected = [infos[index] for index in indices]
# Placeholder IDs deliberately differ from the RVQ frame IDs.
input_ids = torch.arange(len(selected), dtype=torch.long) + 100
expected = torch.cat(
[
model.preprocess(input_ids=input_ids[row : row + 1], input_embeds=None, **info)[1]
for row, info in enumerate(selected)
]
)
state_before = [
(
info["breeze_state"]["current"].clone(),
info["breeze_state"]["history"].clone(),
info["breeze_state"]["generator"].get_state(),
)
for info in infos
]
output_ids, embeds, updates = model.preprocess_decode_batch(input_ids=input_ids, req_infos=selected)
assert output_ids is input_ids
torch.testing.assert_close(embeds, expected, rtol=0, atol=0)
assert updates == [{} for _ in selected]
for info, (current, history, rng) in zip(infos, state_before, strict=True):
state = info["breeze_state"]
assert torch.equal(state["current"], current)
assert torch.equal(state["history"], history)
assert torch.equal(state["generator"].get_state(), rng)
saved.append((embeds, embeds.clone()))
# Simulate the next generated frame. CFG branches share the new frame.
next_frame = (infos[2]["breeze_state"]["current"] + 2) % 7
infos[0]["breeze_state"]["current"] = next_frame
infos[2]["breeze_state"]["current"] = next_frame
infos[1]["breeze_state"]["current"] = (infos[1]["breeze_state"]["current"] + 1) % 7
assert infos[0]["breeze_state"]["current"] is infos[2]["breeze_state"]["current"]
for output, snapshot in saved:
assert torch.equal(output, snapshot)
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_decode_batch_empty_preserves_embedding_shape_and_dtype(dtype):
model = _embedding_talker(dtype)
ids = torch.empty(0, dtype=torch.long)
output_ids, embeds, updates = model.preprocess_decode_batch(input_ids=ids, req_infos=[])
assert output_ids is ids
assert embeds.shape == (0, 8)
assert embeds.dtype == dtype
assert embeds.device == model.depth_decoder.embed_tokens.weight.device
assert updates == []
@pytest.mark.parametrize("invalid", ["count", "state", "frame_shape", "prefill"])
def test_decode_batch_rejects_invalid_request_state(invalid):
model = _embedding_talker(torch.float32)
info = _decode_info("request", torch.zeros(1, 16, dtype=torch.long), 42)
ids = torch.tensor([0])
if invalid == "count":
ids = torch.tensor([0, 1])
elif invalid == "state":
del info["breeze_state"]
elif invalid == "frame_shape":
info["breeze_state"]["current"] = torch.zeros(2, 16, dtype=torch.long)
else:
info["_omni_is_prefill"] = True
with pytest.raises(ValueError, match="Breeze decode"):
model.preprocess_decode_batch(input_ids=ids, req_infos=[info])
tests/model_executor/models/test_breeze_tts_2_registration.py
+50
View File
@@ -0,0 +1,50 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The accelerated runtime preserves Breeze-TTS-2's model identifiers."""
from pathlib import Path
import pytest
import yaml
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.entrypoints.openai.tts_adapters import detect_tts_model_type, resolve_adapter
from vllm_omni.entrypoints.openai.tts_adapters.breeze_tts_2 import BreezeTTS2Adapter
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_tts_2_pipeline_keeps_checkpoint_and_stage_identifiers():
assert OMNI_PIPELINES["breeze"] is BREEZE_TTS_2_PIPELINE
assert BREEZE_TTS_2_PIPELINE.model_arch == "BreezeForConditionalGeneration"
assert BREEZE_TTS_2_PIPELINE.default_deploy_config_name == "breeze_tts_2.yaml"
assert [stage.model_stage for stage in BREEZE_TTS_2_PIPELINE.stages] == [
"breeze_tts_2",
"breeze_tts_2_codec",
]
def test_breeze_tts_2_adapter_remains_registered():
assert resolve_adapter("breeze_tts_2") is BreezeTTS2Adapter
@pytest.mark.parametrize(
"model_stage,model_arch",
[
("breeze_tts_2", None),
("breeze_tts_2_codec", None),
("breeze_tts_2", "BreezeForConditionalGeneration"),
(None, "BreezeForConditionalGeneration"),
],
)
def test_breeze_tts_2_detection_keeps_existing_identifiers(model_stage, model_arch):
assert detect_tts_model_type(model_stage, model_arch) == "breeze_tts_2"
@pytest.mark.parametrize("suffix", ["", "_throughput"])
def test_breeze_deploy_aliases_match_canonical_configuration(suffix):
deploy_dir = Path(__file__).parents[3] / "vllm_omni" / "deploy"
canonical = yaml.safe_load((deploy_dir / f"breeze_tts_2{suffix}.yaml").read_text())
alias = yaml.safe_load((deploy_dir / f"breeze_tts{suffix}.yaml").read_text())
assert canonical == alias
tests/model_executor/models/breeze_tts_2/test_sampling_eos.py → tests/model_executor/models/test_breeze_tts_2_sampling_eos.py
+1 -1
View File
@@ -19,7 +19,7 @@ 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_path = Path(__file__).parents[3] / "vllm_omni" / "deploy" / "breeze_tts_2.yaml"
deploy = yaml.safe_load(deploy_path.read_text(encoding="utf-8"))
params = SamplingParams(
**merge_sampling_constraints(