diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 8faca1b78..aab1acfbc 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -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): diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index f07ee7e10..ca3b8c21f 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -67,6 +67,7 @@ from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( create_instruction as ming_create_instruction, ) +from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.speaker_cache import ( get_speaker_cache, @@ -1396,8 +1397,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): emb_dim = len(embedding) if self._tts_model_type == "ming_tts": - if emb_dim != 192: - raise ValueError(f"Ming speaker embedding must have 192 dims, got {emb_dim}") + if emb_dim != SPEAKER_EMBEDDING_DIM: + raise ValueError(f"Ming speaker embedding must have {SPEAKER_EMBEDDING_DIM} dims, got {emb_dim}") else: dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim) if dim_err is not None: @@ -2148,8 +2149,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): waveform = torch.as_tensor(wav_samples, dtype=torch.float32).reshape(1, -1) embedding = extractor.extract_from_waveform(waveform, int(sr)) flat = embedding.detach().reshape(-1).to(torch.float32).cpu() - if int(flat.numel()) != 192: - raise ValueError(f"Ming speaker extractor returned {int(flat.numel())} dims; expected 192") + if int(flat.numel()) != SPEAKER_EMBEDDING_DIM: + raise ValueError( + f"Ming speaker extractor returned {int(flat.numel())} dims; expected {SPEAKER_EMBEDDING_DIM}" + ) embeddings.append(flat.tolist()) return embeddings diff --git a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py index c736afb51..dff392c28 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py @@ -1,16 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 """Ming-TTS (dense) serving adapter.""" +import asyncio from typing import TYPE_CHECKING +import torch +from vllm.logger import init_logger + from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest -from vllm.logger import init_logger - logger = init_logger(__name__) @@ -33,6 +36,7 @@ class MingTTSAdapter(ARTTSAdapter): ) -> PreparedRequest: server = self.ctx.server ref_audio_source = request.ref_audio + speaker_cache_key: tuple[str, str, int] | None = None voice_lower = request.voice.lower() if isinstance(request.voice, str) else None if ref_audio_source is None and voice_lower in server.uploaded_speakers: speaker_info = server.uploaded_speakers[voice_lower] @@ -42,6 +46,21 @@ class MingTTSAdapter(ARTTSAdapter): if request.speaker_embedding is None: raise ValueError(f"Speaker embedding for uploaded voice '{request.voice}' is missing") else: + if request.speaker_embedding is None: + speaker_cache_key = server._speaker_cache.make_cache_key( + voice_lower, + model_type=self.name, + created_at=int(speaker_info.get("created_at") or 0), + ) + cached = server._speaker_cache.get(speaker_cache_key) + if cached is not None: + cached_embedding = cached.get("speaker_embedding") + if isinstance(cached_embedding, torch.Tensor): + flat = cached_embedding.detach().reshape(-1).to(torch.float32).cpu() + if int(flat.numel()) == SPEAKER_EMBEDDING_DIM: + request.speaker_embedding = flat.tolist() + speaker_cache_key = None + logger.debug("Speaker cache HIT for Ming-TTS speaker '%s'", voice_lower) ref_audio_source = server._get_uploaded_audio_data(request.voice) if not ref_audio_source: raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing") @@ -51,12 +70,26 @@ class MingTTSAdapter(ARTTSAdapter): if isinstance(ref_audio_source, list): ref_audio_data = await server._resolve_ref_audio_many(ref_audio_source) if request.speaker_embedding is None: - request.speaker_embedding = server._extract_ming_speaker_embeddings_from_ref_audio(ref_audio_data) + request.speaker_embedding = await asyncio.to_thread( + server._extract_ming_speaker_embeddings_from_ref_audio, + ref_audio_data, + ) elif ref_audio_source is not None and isinstance(ref_audio_source, str): wav_list, sr = await server._resolve_ref_audio(ref_audio_source) ref_audio_data = (wav_list, sr) if request.speaker_embedding is None: - request.speaker_embedding = server._extract_ming_speaker_embeddings_from_ref_audio([ref_audio_data])[0] + embeddings = await asyncio.to_thread( + server._extract_ming_speaker_embeddings_from_ref_audio, + [ref_audio_data], + ) + request.speaker_embedding = embeddings[0] + if speaker_cache_key is not None and request.speaker_embedding is not None: + embedding = torch.as_tensor(request.speaker_embedding, dtype=torch.float32).detach().reshape(-1).cpu() + server._speaker_cache.put( + speaker_cache_key, + {"speaker_embedding": embedding}, + ) + logger.debug("Speaker cache STORE for Ming-TTS speaker '%s'", voice_lower) prompt = server._build_ming_dense_prompt(request, ref_audio_data=ref_audio_data) tts_params = prompt.get("additional_information", {}) # Ming stop-token / max_tokens sampling stays in the orchestrator tail. @@ -73,11 +106,12 @@ class MingTTSAdapter(ARTTSAdapter): if not request.speaker_embedding: return "'speaker_embedding' must be a non-empty list of floats" emb_len = len(request.speaker_embedding) - if emb_len != 192: + if emb_len != SPEAKER_EMBEDDING_DIM: logger.warning( - "speaker_embedding has %d dimensions; Ming dense expects 192. " + "speaker_embedding has %d dimensions; Ming dense expects %d. " "Wrong dimensions will likely fail or degrade output.", emb_len, + SPEAKER_EMBEDDING_DIM, ) voice_lower = request.voice.lower() if isinstance(request.voice, str) else None @@ -127,8 +161,10 @@ class MingTTSAdapter(ARTTSAdapter): ) if embeddings and isinstance(embeddings[0], list): for item in embeddings: - if len(item) != 192: - return "Podcast-style Ming speaker embeddings must each have 192 dimensions" + if len(item) != SPEAKER_EMBEDDING_DIM: + return ( + f"Podcast-style Ming speaker embeddings must each have {SPEAKER_EMBEDDING_DIM} dimensions" + ) if request.instructions and len(request.instructions) > server._max_instructions_length: return f"Instructions too long (max {server._max_instructions_length} characters)" diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index e7b305b1e..8868e9d2a 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -27,6 +27,7 @@ from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig +from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig @@ -120,8 +121,7 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin) ) self.aggregator = Aggregator(llm_input_dim=self.hidden_size, **config.aggregator) self.stop_head = nn.Linear(self.hidden_size, 2, bias=True) - # CAMPPlus 192-dim -> hidden - self.spk_head = nn.Linear(192, self.hidden_size, bias=True) + self.spk_head = nn.Linear(SPEAKER_EMBEDDING_DIM, self.hidden_size, bias=True) # AudioVAE self.audio_vae, self._vae_weight_source = self._init_audio_vae(config, self.talker_dir, model_path) diff --git a/vllm_omni/model_executor/models/ming_tts/audio_prep.py b/vllm_omni/model_executor/models/ming_tts/audio_prep.py index 378cd222d..055948932 100644 --- a/vllm_omni/model_executor/models/ming_tts/audio_prep.py +++ b/vllm_omni/model_executor/models/ming_tts/audio_prep.py @@ -16,6 +16,7 @@ from .config_ming_tts import ( VAE_PATCH_SIZE, MingTTSConfig, ) +from .constants import SPEAKER_EMBEDDING_DIM def pad_prompt_waveform( @@ -56,7 +57,7 @@ def coerce_prompt_waveform(value: Any) -> torch.Tensor: def coerce_speaker_embeddings(value: Any, *, use_zero_spk_emb: bool = False) -> list[torch.Tensor] | None: if value is None: - return [torch.zeros((192,), dtype=torch.float32)] if use_zero_spk_emb else None + return [torch.zeros((SPEAKER_EMBEDDING_DIM,), dtype=torch.float32)] if use_zero_spk_emb else None if isinstance(value, torch.Tensor): tensor = value.detach() if tensor.ndim == 1: @@ -78,10 +79,10 @@ def coerce_speaker_embeddings(value: Any, *, use_zero_spk_emb: bool = False) -> else: return coerce_speaker_embeddings(torch.as_tensor(value), use_zero_spk_emb=use_zero_spk_emb) if not items: - return [torch.zeros((192,), dtype=torch.float32)] if use_zero_spk_emb else None + return [torch.zeros((SPEAKER_EMBEDDING_DIM,), dtype=torch.float32)] if use_zero_spk_emb else None for item in items: - if int(item.numel()) != 192: - raise ValueError(f"Ming speaker embedding must have 192 dims, got {int(item.numel())}") + if int(item.numel()) != SPEAKER_EMBEDDING_DIM: + raise ValueError(f"Ming speaker embedding must have {SPEAKER_EMBEDDING_DIM} dims, got {int(item.numel())}") return items diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py index e0508bda7..3db8ff3b8 100644 --- a/vllm_omni/model_executor/models/ming_tts/constants.py +++ b/vllm_omni/model_executor/models/ming_tts/constants.py @@ -43,6 +43,7 @@ LLM_VOCAB_SIZE = 151936 AGGREGATOR_HIDDEN_SIZE = 1024 VAE_PATCH_SIZE = 4 SAMPLE_RATE = 44100 +SPEAKER_EMBEDDING_DIM = 192 # CAMPPlus output width and speaker projection input width # AudioVAE frame/hop geometry (confirmed) AUDIO_FRAME_HOP = 882 # enc input_dim / hop_size / dec output_dim diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py index 39ddd7b9f..909befb3a 100644 --- a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py @@ -33,6 +33,7 @@ from .config_ming_tts import ( KEY_TEXT_MODE, MingTTSConfig, ) +from .constants import SPEAKER_EMBEDDING_DIM from .flowloss_head import FlowLoss from .patch_emission import ( MING_STOP_REASON_CODES, @@ -91,7 +92,7 @@ class MingLLMModel(nn.Module): **self.ming_config.ditar_config, ) self.stop_head = nn.Linear(self.ming_config.llm_hidden_size, 2, bias=True) - self.spk_head = nn.Linear(192, self.ming_config.llm_hidden_size, bias=True) + self.spk_head = nn.Linear(SPEAKER_EMBEDDING_DIM, self.ming_config.llm_hidden_size, bias=True) self.flowloss.to(dtype=self.fm_dtype) self.linear_proj_audio.to(dtype=self.fm_dtype) self.stop_head.to(dtype=self.fm_dtype)