From fb89ab43ead813afa8a5580f8aa986f2b111771c Mon Sep 17 00:00:00 2001 From: LHXuuu Date: Mon, 27 Jul 2026 22:55:47 +0800 Subject: [PATCH] [Refactor][Ming-TTS] Move speaker extraction and caching to Stage-0 (#5351) Signed-off-by: LHXuuu Co-authored-by: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> --- .../openai_api/test_serving_speech.py | 147 ++++++++++++------ .../ming_tts/test_ming_shared_modules.py | 123 +++++++++++++++ .../entrypoints/openai/serving_speech.py | 78 ++++++---- .../openai/tts_adapters/ming_tts.py | 47 ++---- .../models/ming_tts/config_ming_tts.py | 6 + .../models/ming_tts/constants.py | 3 + .../models/ming_tts/ming_tts.py | 15 +- .../models/ming_tts/prompt_assembly.py | 14 +- .../models/ming_tts/speaker_extractor.py | 130 ++++++++++++++++ 9 files changed, 440 insertions(+), 123 deletions(-) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index aab1acfbc..b6c7c7c87 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/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 == "" 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 == "" 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 ── diff --git a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py index c5966dd60..7a0ad68de 100644 --- a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py +++ b/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( diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 2bb402caf..782b1c884 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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"] diff --git a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py index dff392c28..adce42867 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py +++ b/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") diff --git a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py index fdac2a30e..a6dc270c3 100644 --- a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py +++ b/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", diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py index 3db8ff3b8..397de3288 100644 --- a/vllm_omni/model_executor/models/ming_tts/constants.py +++ b/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" diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts.py b/vllm_omni/model_executor/models/ming_tts/ming_tts.py index 70a414186..d6e2ac92d 100644 --- a/vllm_omni/model_executor/models/ming_tts/ming_tts.py +++ b/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): diff --git a/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py b/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py index 36817022a..5960af742 100644 --- a/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py +++ b/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, ) diff --git a/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py b/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py index 00523c26f..8e9f69068 100644 --- a/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py +++ b/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", +]