[Refactor][Ming-TTS] Move speaker extraction and caching to Stage-0 (#5351)

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
Co-authored-by: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com>

tests/entrypoints/openai_api/test_serving_speech.py

@@ -45,7 +45,13 @@ 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.model_executor.models.ming_tts.constants import (
+    KEY_SPEAKER_EMBEDDING,
+    KEY_SPEAKER_SAMPLE_RATES,
+    KEY_SPEAKER_WAVEFORM,
+    KEY_SPEAKER_WAVEFORM_LENGTHS,
+    SPEAKER_EMBEDDING_DIM,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -1014,7 +1020,6 @@ class TestTTSMethods:
         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):
@@ -1038,7 +1043,6 @@ class TestTTSMethods:
             )
         )
 
-        mock_extract.assert_not_called()
         tensors = saved["tensors"]
         metadata = saved["metadata"]
         assert isinstance(tensors, dict)
@@ -1669,16 +1673,20 @@ class TestTTSMethods:
 
         mock_get_emb.assert_called_once_with("emb_voice")
         mock_get_audio.assert_not_called()
-        mock_prompt.assert_called_once_with(req, ref_audio_data=None)
+        mock_prompt.assert_called_once_with(
+            req,
+            ref_audio_data=None,
+            voice_name=None,
+            voice_created_at=0,
+        )
         prompt_request = mock_prompt.call_args.args[0]
         assert prompt_request.speaker_embedding == fake_embedding
         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_uses_cached_embedding(self, speech_server, mocker: MockerFixture):
-        """Cached Ming audio voices skip repeated speaker extraction."""
-        speech_server._speaker_cache.clear()
+    def test_ming_adapter_uploaded_audio_forwards_worker_cache_metadata(self, speech_server, mocker: MockerFixture):
+        """Ming audio voices defer extraction and cache ownership to Stage-0."""
         speech_server.uploaded_speakers = {
             "audio_voice": {
                 "name": "audio_voice",
@@ -1691,26 +1699,12 @@ class TestTTSMethods:
         }
         ref_audio_source = "data:audio/wav;base64,ZmFrZWF1ZGlv"
         ref_audio_data = ([0.0, 0.1], 16000)
-        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",
             new=mocker.AsyncMock(return_value=ref_audio_data),
         )
-        mock_extract = mocker.patch.object(
-            speech_server,
-            "_extract_ming_speaker_embeddings_from_ref_audio",
-        )
         mock_prompt = mocker.patch.object(
             speech_server,
             "_build_ming_dense_prompt",
@@ -1721,13 +1715,82 @@ class TestTTSMethods:
         req = OpenAICreateSpeechRequest(input="Hello", voice="audio_voice")
         prepared = asyncio.run(adapter.build(req, [], False))
 
-        mock_extract.assert_not_called()
-        mock_prompt.assert_called_once_with(req, ref_audio_data=ref_audio_data)
+        mock_prompt.assert_called_once_with(
+            req,
+            ref_audio_data=ref_audio_data,
+            voice_name="audio_voice",
+            voice_created_at=123,
+        )
         assert req.ref_text == "Reference transcript."
-        assert req.speaker_embedding == pytest.approx(fake_embedding)
+        assert req.speaker_embedding is None
         assert prepared.model_type == "ming_tts"
 
-    def test_ming_adapter_legacy_audio_only_voice_falls_back_and_caches(
+    def test_build_ming_prompt_forwards_audio_segments_and_worker_cache_metadata(self, speech_server):
+        class FakeTokenizer:
+            unk_token_id = -1
+
+            def encode(self, text):
+                return [777] if text == "<|vision_pad|>" else [1]
+
+            def convert_tokens_to_ids(self, token):
+                return 888 if token == "<audioPatch>" else self.unk_token_id
+
+        speech_server._tts_tokenizer = FakeTokenizer()
+        request = OpenAICreateSpeechRequest(
+            input="Hello",
+            ref_text="Two reference speakers.",
+        )
+        ref_audio_data = [
+            ([0.1, 0.2, 0.3], 16000),
+            ([0.4, 0.5], 24000),
+        ]
+
+        prompt = speech_server._build_ming_dense_prompt(
+            request,
+            ref_audio_data=ref_audio_data,
+            voice_name="uploaded_voice",
+            voice_created_at=123,
+        )
+
+        info = prompt["additional_information"]
+        assert KEY_SPEAKER_EMBEDDING not in info
+        torch.testing.assert_close(
+            info[KEY_SPEAKER_WAVEFORM],
+            torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5]], dtype=torch.float32),
+        )
+        assert info[KEY_SPEAKER_WAVEFORM_LENGTHS].tolist() == [3, 2]
+        assert info[KEY_SPEAKER_SAMPLE_RATES].tolist() == [16000, 24000]
+        assert info["prompt_waveform"].shape[-1] > info[KEY_SPEAKER_WAVEFORM].shape[-1]
+        assert info["voice_name"] == "uploaded_voice"
+        assert info["voice_created_at"] == 123
+        assert prompt["prompt_token_ids"].count(777) == 2
+
+    def test_build_ming_prompt_forwards_speaker_waveform_without_ref_text(self, speech_server):
+        class FakeTokenizer:
+            unk_token_id = -1
+
+            def encode(self, text):
+                return [1]
+
+            def convert_tokens_to_ids(self, token):
+                return 888 if token == "<audioPatch>" else self.unk_token_id
+
+        speech_server._tts_tokenizer = FakeTokenizer()
+        request = OpenAICreateSpeechRequest(input="Hello")
+
+        prompt = speech_server._build_ming_dense_prompt(
+            request,
+            ref_audio_data=([0.1, 0.2, 0.3], 16000),
+        )
+
+        info = prompt["additional_information"]
+        assert info[KEY_SPEAKER_WAVEFORM].shape == (1, 3)
+        assert info[KEY_SPEAKER_WAVEFORM_LENGTHS].tolist() == [3]
+        assert info[KEY_SPEAKER_SAMPLE_RATES].tolist() == [16000]
+        assert "prompt_waveform" not in info
+        assert KEY_SPEAKER_EMBEDDING not in info
+
+    def test_ming_adapter_legacy_audio_only_voice_forwards_audio_to_worker(
         self,
         speech_server,
         mocker: MockerFixture,
@@ -1753,20 +1816,13 @@ class TestTTSMethods:
         }
         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(
+        mock_prompt = mocker.patch.object(
             speech_server,
             "_build_ming_dense_prompt",
             return_value={"additional_information": {"speaker_count": 1}},
@@ -1778,22 +1834,25 @@ class TestTTSMethods:
         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.speaker_embedding is None
+        assert second_req.speaker_embedding is None
         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,
+        assert mock_prompt.call_count == 2
+        mock_prompt.assert_any_call(
+            first_req,
+            ref_audio_data=ref_audio_data,
+            voice_name="legacy_audio_voice",
+            voice_created_at=456,
+        )
+        mock_prompt.assert_any_call(
+            second_req,
+            ref_audio_data=ref_audio_data,
+            voice_name="legacy_audio_voice",
+            voice_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 ──
 

tests/model_executor/models/ming_tts/test_ming_shared_modules.py

@@ -12,17 +12,24 @@ from vllm_omni.model_executor.models.common.ming.fm import Solver
 from vllm_omni.model_executor.models.ming_tts.constants import (
     AGGREGATOR_HIDDEN_SIZE,
     HISTORY_PATCH_SIZE,
+    KEY_SPEAKER_EMBEDDING,
+    KEY_SPEAKER_SAMPLE_RATES,
+    KEY_SPEAKER_WAVEFORM,
+    KEY_SPEAKER_WAVEFORM_LENGTHS,
     LATENT_DIM,
     LLM_HIDDEN_SIZE,
     LLM_VOCAB_SIZE,
     PATCH_SIZE,
     SAMPLE_RATE,
+    SPEAKER_EMBEDDING_DIM,
     VAE_PATCH_SIZE,
 )
 from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss
 from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel
+from vllm_omni.model_executor.models.ming_tts.speaker_extractor import _resolve_speaker_embeddings
 from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
 
@@ -99,6 +106,122 @@ def test_ming_instruction_parser_preserves_dense_and_flash_defaults():
     assert flash_fields == {"风格": "calm"}
 
 
+class _FakeMingSpeakerExtractor:
+    def __init__(self):
+        self.calls: list[tuple[torch.Tensor, int]] = []
+
+    def extract_from_waveform(self, waveform, sample_rate):
+        waveform = torch.as_tensor(waveform).detach().clone()
+        self.calls.append((waveform, int(sample_rate)))
+        return torch.full((SPEAKER_EMBEDDING_DIM,), float(len(self.calls)), dtype=torch.float32)
+
+
+def _make_ming_speaker_wrapper():
+    return SimpleNamespace(
+        _speaker_cache=SpeakerEmbeddingCache(max_bytes=1024 * 1024),
+        _speaker_extractor=_FakeMingSpeakerExtractor(),
+        ming_config=SimpleNamespace(sample_rate=SAMPLE_RATE),
+    )
+
+
+def test_ming_stage0_speaker_cache_extracts_once_and_honors_created_at():
+    wrapper = _make_ming_speaker_wrapper()
+    waveform = torch.arange(12, dtype=torch.float32).reshape(1, -1)
+    info = {
+        KEY_SPEAKER_WAVEFORM: waveform,
+        KEY_SPEAKER_WAVEFORM_LENGTHS: torch.tensor([12]),
+        "voice_name": "UploadedVoice",
+        "voice_created_at": 123,
+    }
+
+    first = _resolve_speaker_embeddings(wrapper, info)
+    cached = _resolve_speaker_embeddings(
+        wrapper,
+        {
+            "voice_name": "uploadedvoice",
+            "voice_created_at": 123,
+        },
+    )
+    refreshed = _resolve_speaker_embeddings(
+        wrapper,
+        {
+            **info,
+            "voice_created_at": 124,
+        },
+    )
+
+    assert first is not None and cached is not None and refreshed is not None
+    torch.testing.assert_close(first[0], cached[0])
+    assert torch.all(first[0] == 1)
+    assert torch.all(refreshed[0] == 2)
+    assert len(wrapper._speaker_extractor.calls) == 2
+
+
+def test_ming_stage0_speaker_extraction_splits_multiple_waveforms():
+    wrapper = _make_ming_speaker_wrapper()
+
+    embeddings = _resolve_speaker_embeddings(
+        wrapper,
+        {
+            KEY_SPEAKER_WAVEFORM: torch.arange(10, dtype=torch.float32).reshape(1, -1),
+            KEY_SPEAKER_WAVEFORM_LENGTHS: torch.tensor([4, 6]),
+            KEY_SPEAKER_SAMPLE_RATES: torch.tensor([16000, 24000]),
+        },
+    )
+
+    assert embeddings is not None and len(embeddings) == 2
+    assert [call[0].shape[-1] for call in wrapper._speaker_extractor.calls] == [4, 6]
+    assert [call[1] for call in wrapper._speaker_extractor.calls] == [16000, 24000]
+
+
+def test_ming_stage0_speaker_extraction_rejects_mismatched_sample_rates():
+    wrapper = _make_ming_speaker_wrapper()
+
+    with pytest.raises(ValueError, match="Invalid Ming speaker sample rates"):
+        _resolve_speaker_embeddings(
+            wrapper,
+            {
+                KEY_SPEAKER_WAVEFORM: torch.arange(10, dtype=torch.float32).reshape(1, -1),
+                KEY_SPEAKER_WAVEFORM_LENGTHS: torch.tensor([4, 6]),
+                KEY_SPEAKER_SAMPLE_RATES: torch.tensor([16000]),
+            },
+        )
+
+
+def test_ming_stage0_prompt_waveform_does_not_trigger_speaker_extraction():
+    wrapper = _make_ming_speaker_wrapper()
+
+    embeddings = _resolve_speaker_embeddings(
+        wrapper,
+        {
+            "prompt_waveform": torch.ones((1, 10)),
+            "prompt_text": "Reference transcript.",
+        },
+    )
+
+    assert embeddings is None
+    assert not wrapper._speaker_extractor.calls
+
+
+def test_ming_stage0_direct_speaker_embedding_bypasses_extractor():
+    wrapper = _make_ming_speaker_wrapper()
+    direct = torch.full((SPEAKER_EMBEDDING_DIM,), 0.25)
+
+    embeddings = _resolve_speaker_embeddings(
+        wrapper,
+        {
+            KEY_SPEAKER_EMBEDDING: direct,
+            KEY_SPEAKER_WAVEFORM: torch.ones((1, 10)),
+            "voice_name": "direct",
+            "voice_created_at": 1,
+        },
+    )
+
+    assert embeddings is not None
+    torch.testing.assert_close(embeddings[0], direct)
+    assert not wrapper._speaker_extractor.calls
+
+
 def _make_ming_logits_model(vocab_size=8):
     model = object.__new__(MingLLMModel)
     model.ming_config = SimpleNamespace(

vllm_omni/entrypoints/openai/serving_speech.py

@@ -2124,37 +2124,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             waveform = resample_audio(waveform, int(sample_rate), SAMPLE_RATE)
         return waveform
 
-    def _build_ming_prompt_waveform(
+    def _build_ming_reference_waveforms(
         self,
         ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None,
-    ):
-        if isinstance(ref_audio_data, list):
-            return torch.cat(
-                [self._coerce_ming_prompt_waveform(item[0], item[1]) for item in ref_audio_data],
-                dim=-1,
-            )
-        if ref_audio_data is not None:
-            return self._coerce_ming_prompt_waveform(ref_audio_data[0], ref_audio_data[1])
-        return None
-
-    def _extract_ming_speaker_embeddings_from_ref_audio(
-        self,
-        ref_audio_data_list: list[tuple[list[float], int]],
-    ) -> list[list[float]]:
-        from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
-
-        extractor = MingSpeakerEmbeddingExtractor(self.engine_client.model_config.model, target_sr=16000)
-        embeddings = []
-        for wav_samples, sr in ref_audio_data_list:
-            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()) != SPEAKER_EMBEDDING_DIM:
-                raise ValueError(
-                    f"Ming speaker extractor returned {int(flat.numel())} dims; expected {SPEAKER_EMBEDDING_DIM}"
-                )
-            embeddings.append(flat.tolist())
-        return embeddings
+    ) -> tuple[list[torch.Tensor], list[int]]:
+        if ref_audio_data is None:
+            return [], []
+        items = ref_audio_data if isinstance(ref_audio_data, list) else [ref_audio_data]
+        waveforms = [torch.as_tensor(samples, dtype=torch.float32).reshape(1, -1) for samples, _ in items]
+        sample_rates = [int(sample_rate) for _, sample_rate in items]
+        if any(sample_rate <= 0 for sample_rate in sample_rates):
+            raise ValueError(f"Ming reference audio sample rates must be positive, got {sample_rates}")
+        return waveforms, sample_rates
 
     def _parse_ming_instruction_fields(
         self,
@@ -2197,11 +2178,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request: OpenAICreateSpeechRequest,
         *,
         ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None = None,
+        voice_name: str | None = None,
+        voice_created_at: int = 0,
     ) -> dict[str, Any]:
         """Build a Ming dense prompt directly from the OpenAI speech request."""
         from transformers import AutoTokenizer
 
-        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import KEY_MAX_DECODE_STEPS
+        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+            KEY_MAX_DECODE_STEPS,
+            KEY_SPEAKER_SAMPLE_RATES,
+            KEY_SPEAKER_WAVEFORM,
+            KEY_SPEAKER_WAVEFORM_LENGTHS,
+        )
         from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
 
         if self._tts_tokenizer is None:
@@ -2210,9 +2198,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             self._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code)
 
         ref_text = request.ref_text
-        prompt_waveform = self._build_ming_prompt_waveform(ref_audio_data) if ref_text is not None else None
+        reference_waveforms, reference_sample_rates = self._build_ming_reference_waveforms(ref_audio_data)
+        reference_waveform = torch.cat(reference_waveforms, dim=-1) if reference_waveforms else None
+        prompt_waveforms = (
+            [
+                self._coerce_ming_prompt_waveform(waveform, sample_rate)
+                for waveform, sample_rate in zip(reference_waveforms, reference_sample_rates, strict=True)
+            ]
+            if ref_text is not None
+            else []
+        )
+        prompt_waveform = torch.cat(prompt_waveforms, dim=-1) if prompt_waveforms else None
         speaker_embedding = request.speaker_embedding
-        use_zero_spk_emb = prompt_waveform is None and speaker_embedding is None
+        pending_speaker_extraction = speaker_embedding is None and reference_waveform is not None
+        use_zero_spk_emb = not pending_speaker_extraction and prompt_waveform is None and speaker_embedding is None
 
         runtime_controls = {}
         if request.max_new_tokens is not None:
@@ -2229,8 +2228,23 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             prompt_text=ref_text,
             prompt_waveform=prompt_waveform,
             speaker_embedding=speaker_embedding,
+            speaker_count=len(reference_waveforms) if pending_speaker_extraction else None,
             use_zero_spk_emb=use_zero_spk_emb,
         )
+        additional_information = prompt_dict["additional_information"]
+        if pending_speaker_extraction:
+            additional_information[KEY_SPEAKER_WAVEFORM] = reference_waveform
+            additional_information[KEY_SPEAKER_WAVEFORM_LENGTHS] = torch.tensor(
+                [int(waveform.shape[-1]) for waveform in reference_waveforms],
+                dtype=torch.int32,
+            )
+            additional_information[KEY_SPEAKER_SAMPLE_RATES] = torch.tensor(
+                reference_sample_rates,
+                dtype=torch.int32,
+            )
+            if voice_name:
+                additional_information["voice_name"] = voice_name
+                additional_information["voice_created_at"] = int(voice_created_at)
         prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"])
         prompt["prompt"] = prompt_dict["prompt"]
         prompt["text"] = prompt_dict["text"]

vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py

@@ -1,10 +1,8 @@
 # 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
@@ -36,7 +34,8 @@ class MingTTSAdapter(ARTTSAdapter):
     ) -> PreparedRequest:
         server = self.ctx.server
         ref_audio_source = request.ref_audio
-        speaker_cache_key: tuple[str, str, int] | None = None
+        uploaded_audio_voice = None
+        uploaded_audio_created_at = 0
         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]
@@ -46,51 +45,25 @@ 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")
                 if request.ref_text is None:
                     request.ref_text = speaker_info.get("ref_text")
+                uploaded_audio_voice = voice_lower
+                uploaded_audio_created_at = server._voice_created_at(voice_lower)
         ref_audio_data = None
         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 = 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:
-                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)
+        prompt = server._build_ming_dense_prompt(
+            request,
+            ref_audio_data=ref_audio_data,
+            voice_name=uploaded_audio_voice,
+            voice_created_at=uploaded_audio_created_at,
+        )
         tts_params = prompt.get("additional_information", {})
         # Ming stop-token / max_tokens sampling stays in the orchestrator tail.
         return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="ming_tts")

vllm_omni/model_executor/models/ming_tts/config_ming_tts.py

@@ -35,6 +35,9 @@ from .constants import (
     KEY_REQUEST_ID,
     KEY_SIGMA,
     KEY_SPEAKER_EMBEDDING,
+    KEY_SPEAKER_SAMPLE_RATES,
+    KEY_SPEAKER_WAVEFORM,
+    KEY_SPEAKER_WAVEFORM_LENGTHS,
     KEY_TEMPERATURE,
     KEY_TEXT_MODE,
     LATENT_CHUNK_SIZE,
@@ -326,6 +329,9 @@ __all__ = [
     "KEY_REQUEST_ID",
     "KEY_SIGMA",
     "KEY_SPEAKER_EMBEDDING",
+    "KEY_SPEAKER_SAMPLE_RATES",
+    "KEY_SPEAKER_WAVEFORM",
+    "KEY_SPEAKER_WAVEFORM_LENGTHS",
     "KEY_TEMPERATURE",
     "KEY_TEXT_MODE",
     "LATENT_CHUNK_SIZE",

vllm_omni/model_executor/models/ming_tts/constants.py

@@ -74,6 +74,9 @@ KEY_LAST_STOP_PROB = "ming_last_stop_prob"
 KEY_NEXT_EMBEDS = "ming_next_embeds"
 KEY_PROMPT_LATENTS = "ming_prompt_latents"
 KEY_SPEAKER_EMBEDDING = "ming_speaker_embedding"
+KEY_SPEAKER_WAVEFORM = "ming_speaker_waveform"
+KEY_SPEAKER_WAVEFORM_LENGTHS = "ming_speaker_waveform_lengths"
+KEY_SPEAKER_SAMPLE_RATES = "ming_speaker_sample_rates"
 KEY_REQUEST_ID = "ming_request_id"
 KEY_CHUNK_ID = "ming_chunk_id"
 KEY_CFG = "ming_cfg"

vllm_omni/model_executor/models/ming_tts/ming_tts.py

@@ -20,7 +20,6 @@ from .audio_prep import (
     _find_speaker_placeholder_positions,
     _initial_history,
     _take_scalar,
-    coerce_speaker_embeddings,
 )
 from .config_ming_tts import (
     AUDIO_START_TOKEN_ID,
@@ -33,13 +32,13 @@ from .config_ming_tts import (
     KEY_NEXT_EMBEDS,
     KEY_REQUEST_ID,
     KEY_SIGMA,
-    KEY_SPEAKER_EMBEDDING,
     KEY_TEMPERATURE,
     KEY_TEXT_MODE,
     MingTTSConfig,
 )
 from .patch_emission import MING_STOP_REASON_KEY
 from .prompt_encoder import _resolve_prompt_latents
+from .speaker_extractor import _resolve_speaker_embeddings
 
 
 class _ModelSampleAdapter(nn.Module):
@@ -64,8 +63,12 @@ class MingTTSForConditionalGeneration(nn.Module, SupportsPP, CustomProcessMixin)
         self.requires_raw_input_tokens = False
         self.model_stage = vllm_config.model_config.model_stage
         self._prompt_encoder = None
+        self._speaker_extractor = None
 
         if self.model_stage == "llm":
+            from vllm_omni.utils.speaker_cache import get_speaker_cache
+
+            self._speaker_cache = get_speaker_cache()
             self.model = init_vllm_registered_model(vllm_config=vllm_config, architectures=["MingLLMModel"])
             self.has_preprocess = True
             self.has_postprocess = True
@@ -174,13 +177,7 @@ class MingTTSForConditionalGeneration(nn.Module, SupportsPP, CustomProcessMixin)
         )
         update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous()
 
-        speaker_embedding = info_dict.get(KEY_SPEAKER_EMBEDDING, info_dict.get("speaker_embedding"))
-        speaker_embeddings = None
-        if speaker_embedding is not None:
-            speaker_embeddings = coerce_speaker_embeddings(
-                speaker_embedding,
-                use_zero_spk_emb=bool(info_dict.get("use_zero_spk_emb", False)),
-            )
+        speaker_embeddings = _resolve_speaker_embeddings(self, info_dict)
         if speaker_embeddings is not None and len(speaker_embeddings) > 0:
             speaker_slots = _find_speaker_placeholder_positions(input_ids, self.ming_config)
             if len(speaker_slots) < len(speaker_embeddings):

vllm_omni/model_executor/models/ming_tts/prompt_assembly.py

@@ -265,6 +265,7 @@ def build_ming_dense_prompt(
     prompt_waveform: Any = None,
     prompt_latents: Any = None,
     speaker_embedding: Any = None,
+    speaker_count: int | None = None,
     use_zero_spk_emb: bool = False,
     request_id: str | None = None,
     model_variant: str | None = None,
@@ -273,6 +274,17 @@ def build_ming_dense_prompt(
         model_variant = detect_model_variant(tokenizer)
     instruction_text = create_instruction(instruction)
     speaker_embeddings = coerce_speaker_embeddings(speaker_embedding, use_zero_spk_emb=use_zero_spk_emb)
+    if speaker_count is None:
+        effective_speaker_count = 0 if speaker_embeddings is None else len(speaker_embeddings)
+    else:
+        effective_speaker_count = int(speaker_count)
+        if effective_speaker_count < 0:
+            raise ValueError(f"speaker_count must be >= 0, got {effective_speaker_count}")
+        if speaker_embeddings is not None and effective_speaker_count != len(speaker_embeddings):
+            raise ValueError(
+                f"speaker_count={effective_speaker_count} does not match "
+                f"{len(speaker_embeddings)} provided speaker embeddings"
+            )
     effective_runtime_controls = resolve_effective_runtime_controls(text=text, runtime_controls=runtime_controls)
 
     prompt_waveform_tensor = None
@@ -305,7 +317,7 @@ def build_ming_dense_prompt(
         text=text,
         instruction=instruction_text,
         prompt_text=prompt_text if prompt_patch_count > 0 else None,
-        speaker_count=0 if speaker_embeddings is None else len(speaker_embeddings),
+        speaker_count=effective_speaker_count,
         prompt_patch_count=prompt_patch_count,
     )
 

vllm_omni/model_executor/models/ming_tts/speaker_extractor.py

@@ -1,14 +1,33 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/spkemb_extractor.py
+from __future__ import annotations
+
 import os
+import threading
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torchaudio
+from vllm.logger import init_logger
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.model_executor.models.ming_flash_omni.spk_embedding import SpkembExtractor
 
+from .audio_prep import coerce_prompt_waveform, coerce_speaker_embeddings
+from .constants import (
+    KEY_SPEAKER_EMBEDDING,
+    KEY_SPEAKER_SAMPLE_RATES,
+    KEY_SPEAKER_WAVEFORM,
+    KEY_SPEAKER_WAVEFORM_LENGTHS,
+)
+
+if TYPE_CHECKING:
+    from .ming_tts import MingTTSForConditionalGeneration
+
+logger = init_logger(__name__)
+_EXTRACTOR_LOAD_LOCK = threading.Lock()
+
 
 def resolve_model_to_local_path(model):
     if os.path.isdir(model):
@@ -49,3 +68,114 @@ class MingSpeakerEmbeddingExtractor:
 
     def extract_many(self, audio_paths):
         return [self.extract_from_file(path) for path in audio_paths]
+
+
+def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> MingSpeakerEmbeddingExtractor:
+    extractor = getattr(wrapper, "_speaker_extractor", None)
+    if extractor is not None:
+        return extractor
+
+    with _EXTRACTOR_LOAD_LOCK:
+        extractor = getattr(wrapper, "_speaker_extractor", None)
+        if extractor is None:
+            model_path = wrapper.vllm_config.model_config.model
+            extractor = MingSpeakerEmbeddingExtractor(model_path, target_sr=16000)
+            wrapper._speaker_extractor = extractor
+    return extractor
+
+
+def _speaker_waveforms_from_info(
+    info_dict: dict[str, Any],
+    default_sample_rate: int,
+) -> list[tuple[torch.Tensor, int]]:
+    raw_waveform = info_dict.get(KEY_SPEAKER_WAVEFORM)
+    if raw_waveform is None:
+        return []
+
+    waveform = coerce_prompt_waveform(raw_waveform)
+    raw_lengths = info_dict.get(KEY_SPEAKER_WAVEFORM_LENGTHS)
+    if raw_lengths is None:
+        parts = [waveform]
+    else:
+        lengths = torch.as_tensor(raw_lengths).detach().reshape(-1).to(torch.int64).cpu().tolist()
+        if not lengths or any(int(length) <= 0 for length in lengths):
+            raise ValueError(f"Invalid Ming speaker waveform lengths: {lengths}")
+        if sum(int(length) for length in lengths) > int(waveform.shape[-1]):
+            raise ValueError(
+                "Ming speaker waveform lengths exceed the waveform size: "
+                f"lengths={lengths}, waveform_samples={int(waveform.shape[-1])}"
+            )
+
+        parts = []
+        offset = 0
+        for length in lengths:
+            end = offset + int(length)
+            parts.append(waveform[:, offset:end])
+            offset = end
+
+    raw_sample_rates = info_dict.get(KEY_SPEAKER_SAMPLE_RATES)
+    if raw_sample_rates is None:
+        sample_rates = [int(default_sample_rate)] * len(parts)
+    else:
+        sample_rates = torch.as_tensor(raw_sample_rates).detach().reshape(-1).to(torch.int64).cpu().tolist()
+        if len(sample_rates) != len(parts) or any(int(sample_rate) <= 0 for sample_rate in sample_rates):
+            raise ValueError(
+                f"Invalid Ming speaker sample rates: sample_rates={sample_rates}, waveform_count={len(parts)}"
+            )
+    return [(waveform_part, int(sample_rate)) for waveform_part, sample_rate in zip(parts, sample_rates, strict=True)]
+
+
+def _resolve_speaker_embeddings(
+    wrapper: MingTTSForConditionalGeneration,
+    info_dict: dict[str, Any],
+) -> list[torch.Tensor] | None:
+    direct_embedding = info_dict.get(KEY_SPEAKER_EMBEDDING, info_dict.get("speaker_embedding"))
+    use_zero_spk_emb = bool(info_dict.get("use_zero_spk_emb", False))
+    if direct_embedding is not None or use_zero_spk_emb:
+        return coerce_speaker_embeddings(direct_embedding, use_zero_spk_emb=use_zero_spk_emb)
+
+    voice_name = info_dict.get("voice_name")
+    if isinstance(voice_name, (list, tuple)):
+        voice_name = voice_name[0] if voice_name else None
+    created_at = info_dict.get("voice_created_at", 0)
+    if isinstance(created_at, (list, tuple)):
+        created_at = created_at[0] if created_at else 0
+    if isinstance(created_at, torch.Tensor):
+        created_at = created_at.detach().reshape(-1)[0].item() if created_at.numel() else 0
+
+    cache = wrapper._speaker_cache
+    cache_key = None
+    if isinstance(voice_name, str) and voice_name.strip():
+        cache_key = cache.make_cache_key(voice_name, model_type="ming_tts", created_at=int(created_at or 0))
+        cached = cache.get(cache_key)
+        if cached is not None:
+            cached_embeddings = coerce_speaker_embeddings(cached.get("speaker_embedding"))
+            if cached_embeddings is not None:
+                logger.debug("Speaker cache HIT for Ming-TTS speaker '%s'", voice_name)
+                return cached_embeddings
+
+    waveform_inputs = _speaker_waveforms_from_info(info_dict, int(wrapper.ming_config.sample_rate))
+    if not waveform_inputs:
+        if cache_key is not None:
+            raise ValueError(f"Ming-TTS speaker '{voice_name}' was requested without reference audio on cache miss")
+        return None
+
+    extractor = _load_speaker_extractor(wrapper)
+    embeddings = [
+        extractor.extract_from_waveform(waveform, sample_rate).detach().reshape(-1).to(torch.float32).cpu()
+        for waveform, sample_rate in waveform_inputs
+    ]
+    resolved = coerce_speaker_embeddings(embeddings)
+    if cache_key is not None and resolved is not None:
+        cached_embedding = resolved[0] if len(resolved) == 1 else torch.stack(resolved, dim=0)
+        cache.put(cache_key, {"speaker_embedding": cached_embedding})
+        logger.debug("Speaker cache STORE for Ming-TTS speaker '%s'", voice_name)
+    return resolved
+
+
+__all__ = [
+    "MingSpeakerEmbeddingExtractor",
+    "_load_speaker_extractor",
+    "_resolve_speaker_embeddings",
+    "resolve_model_to_local_path",
+]