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

[Perf][Ming-TTS] Cache speaker embeddings for uploaded reference audio (#5240)

Browse Source 
Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
This commit is contained in:
LHXuuu authored and GitHub committed 2026-07-23 11:22:51 +08:00
1 parent be667a8067
commit 00bdb5dcd7
7 files changed
+213 -37
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/entrypoints/openai_api/test_serving_speech.py
+152 -18
View File
@@ -45,6 +45,7 @@ from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
FISH_TEXT_ONLY_SYSTEM_PROMPT,
build_fish_voice_clone_prompt_ids,
)
from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -975,21 +976,25 @@ class TestTTSMethods:
)
def test_upload_ming_voice_embedding_wrong_dims_rejected_without_replacing_existing(self, speech_server):
"""Ming embedding uploads must be 192-dim before replacing stored voices."""
"""Ming embedding uploads must match the model dimension before replacement."""
speech_server._tts_model_type = "ming_tts"
invalid_dim = SPEAKER_EMBEDDING_DIM - 1
existing = {
"name": "bad_emb_voice",
"file_path": "/tmp/voice_samples/bad_emb_voice.safetensors",
"mime_type": "application/x-safetensors",
"embedding_source": "direct",
"embedding_dim": 192,
"embedding_dim": SPEAKER_EMBEDDING_DIM,
}
speech_server.uploaded_speakers = {"bad_emb_voice": existing.copy()}
with pytest.raises(ValueError, match="Ming speaker embedding must have 192 dims, got 191"):
with pytest.raises(
ValueError,
match=f"Ming speaker embedding must have {SPEAKER_EMBEDDING_DIM} dims, got {invalid_dim}",
):
asyncio.run(
speech_server.upload_voice_embedding(
embedding_json=json.dumps([0.0] * 191),
embedding_json=json.dumps([0.0] * invalid_dim),
consent="consent",
name="bad_emb_voice",
)
@@ -997,6 +1002,59 @@ class TestTTSMethods:
assert speech_server.uploaded_speakers["bad_emb_voice"] == existing
def test_upload_ming_audio_voice_defers_speaker_extraction(self, speech_server, mocker: MockerFixture, tmp_path):
"""Ming audio uploads stay model-agnostic and defer extraction to use."""
speech_server._tts_model_type = "ming_tts"
speech_server.uploaded_speakers_dir = tmp_path
speech_server.uploaded_speakers = {}
speech_server.supported_speakers = set()
speech_server._speaker_cache.clear()
left = np.full(32000, 0.5, dtype=np.float32)
right = np.full(32000, -0.25, dtype=np.float32)
samples = np.stack((left, right), axis=-1)
mocker.patch("soundfile.read", return_value=(samples, 16000))
mock_extract = mocker.patch.object(speech_server, "_extract_ming_speaker_embeddings_from_ref_audio")
saved: dict[str, object] = {}
def fake_save_file(tensors, path, metadata=None):
saved["tensors"] = tensors
saved["metadata"] = metadata
Path(path).touch()
mocker.patch("safetensors.torch.save_file", side_effect=fake_save_file)
audio_file = mocker.MagicMock()
audio_file.filename = "reference.wav"
audio_file.content_type = "audio/wav"
audio_file.file = io.BytesIO(b"fake audio")
audio_file.read = mocker.AsyncMock(return_value=b"fake audio")
asyncio.run(
speech_server.upload_voice(
audio_file,
consent="consent",
name="ming_audio_voice",
ref_text="Reference transcript.",
)
)
mock_extract.assert_not_called()
tensors = saved["tensors"]
metadata = saved["metadata"]
assert isinstance(tensors, dict)
assert isinstance(metadata, dict)
assert set(tensors) == {"audio"}
assert tensors["audio"].shape == (32000, 2)
assert metadata["embedding_source"] == "audio"
assert "embedding_dim" not in metadata
speaker_info = speech_server.uploaded_speakers["ming_audio_voice"]
cache_key = speech_server._speaker_cache.make_cache_key(
"ming_audio_voice",
model_type="ming_tts",
created_at=speaker_info["created_at"],
)
assert speech_server._speaker_cache.get(cache_key) is None
def test_base_task_requires_ref_audio_or_speaker_embedding(self, speech_server):
"""Base task without ref_audio or speaker_embedding is rejected."""
req = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
@@ -1591,10 +1649,10 @@ class TestTTSMethods:
"file_path": "/tmp/voice_samples/emb_voice.safetensors",
"mime_type": "application/x-safetensors",
"embedding_source": "direct",
"embedding_dim": 192,
"embedding_dim": SPEAKER_EMBEDDING_DIM,
}
}
fake_embedding = [0.1] * 192
fake_embedding = [0.1] * SPEAKER_EMBEDDING_DIM
mock_get_emb = mocker.patch.object(
speech_server, "_get_uploaded_speaker_embedding", return_value=fake_embedding
)
@@ -1614,25 +1672,36 @@ class TestTTSMethods:
mock_prompt.assert_called_once_with(req, ref_audio_data=None)
prompt_request = mock_prompt.call_args.args[0]
assert prompt_request.speaker_embedding == fake_embedding
assert len(prompt_request.speaker_embedding) == 192
assert len(prompt_request.speaker_embedding) == SPEAKER_EMBEDDING_DIM
assert prepared.prompt["additional_information"]["speaker_count"] == 1
assert prepared.model_type == "ming_tts"
def test_ming_adapter_uploaded_audio_path_preserves_ref_text(self, speech_server, mocker: MockerFixture):
"""Ming uploaded audio voices still resolve audio and stored ref_text."""
def test_ming_adapter_uploaded_audio_uses_cached_embedding(self, speech_server, mocker: MockerFixture):
"""Cached Ming audio voices skip repeated speaker extraction."""
speech_server._speaker_cache.clear()
speech_server.uploaded_speakers = {
"audio_voice": {
"name": "audio_voice",
"file_path": "/tmp/voice_samples/audio_voice.wav",
"file_path": "/tmp/voice_samples/audio_voice.safetensors",
"created_at": 123,
"mime_type": "audio/wav",
"embedding_source": "audio",
"ref_text": "Reference transcript.",
}
}
ref_audio_source = "data:audio/wav;base64,ZmFrZWF1ZGlv"
ref_audio_data = ([0.0, 0.1], 16000)
fake_embedding = [0.2] * 192
mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data", return_value=ref_audio_source)
mock_get_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding")
fake_embedding = [0.2] * SPEAKER_EMBEDDING_DIM
cache_key = speech_server._speaker_cache.make_cache_key(
"audio_voice",
model_type="ming_tts",
created_at=123,
)
speech_server._speaker_cache.put(
cache_key,
{"speaker_embedding": torch.tensor(fake_embedding)},
)
mocker.patch.object(speech_server, "_get_uploaded_audio_data", return_value=ref_audio_source)
mocker.patch.object(
speech_server,
"_resolve_ref_audio",
@@ -1641,7 +1710,6 @@ class TestTTSMethods:
mock_extract = mocker.patch.object(
speech_server,
"_extract_ming_speaker_embeddings_from_ref_audio",
return_value=[fake_embedding],
)
mock_prompt = mocker.patch.object(
speech_server,
@@ -1653,14 +1721,80 @@ class TestTTSMethods:
req = OpenAICreateSpeechRequest(input="Hello", voice="audio_voice")
prepared = asyncio.run(adapter.build(req, [], False))
mock_get_audio.assert_called_once_with("audio_voice")
mock_get_emb.assert_not_called()
mock_extract.assert_called_once_with([ref_audio_data])
mock_extract.assert_not_called()
mock_prompt.assert_called_once_with(req, ref_audio_data=ref_audio_data)
assert req.ref_text == "Reference transcript."
assert req.speaker_embedding == fake_embedding
assert req.speaker_embedding == pytest.approx(fake_embedding)
assert prepared.model_type == "ming_tts"
def test_ming_adapter_legacy_audio_only_voice_falls_back_and_caches(
self,
speech_server,
mocker: MockerFixture,
tmp_path,
):
"""Legacy audio-only profiles extract once instead of failing after upgrade."""
from safetensors.torch import save_file
speech_server._tts_model_type = "ming_tts"
speech_server.uploaded_speakers_dir = tmp_path
speech_server._speaker_cache.clear()
file_path = tmp_path / "legacy_audio_voice.safetensors"
save_file({"audio": torch.zeros(32000, dtype=torch.float32)}, str(file_path))
speech_server.uploaded_speakers = {
"legacy_audio_voice": {
"name": "legacy_audio_voice",
"file_path": str(file_path),
"created_at": 456,
"mime_type": "audio/wav",
"embedding_source": "audio",
"ref_text": "Reference transcript.",
}
}
ref_audio_source = "data:audio/wav;base64,ZmFrZWF1ZGlv"
ref_audio_data = ([0.0, 0.1], 16000)
fake_embedding = [0.4] * SPEAKER_EMBEDDING_DIM
mocker.patch.object(speech_server, "_get_uploaded_audio_data", return_value=ref_audio_source)
mocker.patch.object(
speech_server,
"_resolve_ref_audio",
new=mocker.AsyncMock(return_value=ref_audio_data),
)
mock_extract = mocker.patch.object(
speech_server,
"_extract_ming_speaker_embeddings_from_ref_audio",
return_value=[fake_embedding],
)
mock_to_thread = mocker.spy(asyncio, "to_thread")
mocker.patch.object(
speech_server,
"_build_ming_dense_prompt",
return_value={"additional_information": {"speaker_count": 1}},
)
adapter = MingTTSAdapter(SpeechServingContext(server=speech_server))
first_req = OpenAICreateSpeechRequest(input="Hello", voice="legacy_audio_voice")
first_prepared = asyncio.run(adapter.build(first_req, [], False))
second_req = OpenAICreateSpeechRequest(input="Hello again", voice="legacy_audio_voice")
second_prepared = asyncio.run(adapter.build(second_req, [], False))
mock_extract.assert_called_once_with([ref_audio_data])
mock_to_thread.assert_called_once_with(mock_extract, [ref_audio_data])
assert first_req.speaker_embedding == pytest.approx(fake_embedding)
assert second_req.speaker_embedding == pytest.approx(fake_embedding)
assert first_req.ref_text == "Reference transcript."
assert second_req.ref_text == "Reference transcript."
assert first_prepared.model_type == "ming_tts"
assert second_prepared.model_type == "ming_tts"
cache_key = speech_server._speaker_cache.make_cache_key(
"legacy_audio_voice",
model_type="ming_tts",
created_at=456,
)
cached = speech_server._speaker_cache.get(cache_key)
assert cached is not None
torch.testing.assert_close(cached["speaker_embedding"], torch.tensor(fake_embedding))
# ── regression: full flow from issue #1603 ──
def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server, mocker: MockerFixture):