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

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>

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):

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
 

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)"

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)

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
 
 

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

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)