[Feature] Add precomputed custom voices and Qwen3-TTS ref-context cache (#3492)

Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Dan250124 <416947747@qq.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1,9 +1,13 @@
 # tests/entrypoints/openai/test_serving_speech.py
 import asyncio
+import base64
+import hashlib
+import io
 import json
 import logging
 import os
 import struct
+import wave
 from inspect import Signature, signature
 from pathlib import Path
 from types import SimpleNamespace
@@ -170,6 +174,27 @@ def create_mock_audio_output_for_test(
     )
 
 
+def _write_custom_voice_manifest(root: Path, *, model_type: str, voices: dict) -> None:
+    payload = {
+        "schema_version": 1,
+        "model_type": model_type,
+        "voices": voices,
+    }
+    (root / "custom_voice_manifest.json").write_text(json.dumps(payload), encoding="utf-8")
+
+
+def _wav_data_url(samples: np.ndarray, sample_rate: int) -> str:
+    pcm = (np.clip(samples, -1.0, 1.0) * 32767).astype("<i2")
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wav:
+        wav.setnchannels(1)
+        wav.setsampwidth(2)
+        wav.setframerate(sample_rate)
+        wav.writeframes(pcm.tobytes())
+    encoded = base64.b64encode(buf.getvalue()).decode("ascii")
+    return f"data:audio/wav;base64,{encoded}"
+
+
 @pytest.fixture
 def test_app(mocker: MockerFixture):
     # Mock the engine client
@@ -972,6 +997,233 @@ class TestTTSMethods:
         params = speech_server._build_tts_params(req)
         assert "voice_clone_prompt" not in params
 
+    @pytest.mark.asyncio
+    async def test_resolve_ref_audio_reuses_decoded_audio_for_same_source(self, speech_server):
+        wav = np.linspace(-0.5, 0.5, 48000, dtype=np.float32)
+        ref_audio = _wav_data_url(wav, 24000)
+        speech_server.model_config.allowed_local_media_path = ""
+        speech_server.model_config.allowed_media_domains = None
+
+        first = await speech_server._resolve_ref_audio(ref_audio)
+        second = await speech_server._resolve_ref_audio(ref_audio)
+
+        assert first[1] == 24000
+        assert second[1] == 24000
+        assert first[0] is second[0]
+        assert first[0][0] == pytest.approx(float(wav[0]), abs=1e-4)
+        assert speech_server._get_resolved_ref_audio_artifact_key(
+            ref_audio
+        ) == speech_server._make_ref_audio_artifact_cache_key(np.asarray(first[0], dtype=np.float32), 24000)
+
+    def test_precomputed_qwen3_voice_infers_base_without_ref_audio(self, speech_server):
+        """Precomputed Qwen3 voices are reusable by name without per-request ref_audio."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.precomputed_speakers = {
+            "alice": {
+                "name": "Alice",
+                "model_type": "qwen3_tts",
+                "mode": "icl",
+                "ref_text": "reference transcript",
+                "ref_code_length": 3,
+            }
+        }
+        speech_server.supported_speakers = {"alice"}
+
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
+        assert speech_server._validate_tts_request(req) is None
+        assert req.task_type == "Base"
+
+        params = speech_server._build_tts_params(req)
+        assert params["task_type"] == ["Base"]
+        assert params["speaker"] == ["alice"]
+        assert params["x_vector_only_mode"] == [False]
+        assert params["ref_text"] == ["reference transcript"]
+        assert params["ref_code_length"] == [3]
+        assert "ref_audio" not in params
+
+    def test_uploaded_qwen3_voice_wins_over_same_named_precomputed_voice(self, speech_server, tmp_path):
+        """Qwen3 uploaded voices should take precedence over same-name precomputed voices."""
+        from safetensors.torch import save_file
+
+        uploaded_path = tmp_path / "alice.safetensors"
+        save_file({"speaker_embedding": torch.tensor([0.1] * 4)}, str(uploaded_path))
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.uploaded_speakers_dir = tmp_path
+        speech_server.uploaded_speakers = {
+            "alice": {
+                "name": "alice",
+                "created_at": 123,
+                "file_path": str(uploaded_path),
+                "embedding_source": "direct",
+            }
+        }
+        speech_server.precomputed_speakers = {
+            "alice": {
+                "name": "Alice",
+                "model_type": "qwen3_tts",
+                "mode": "icl",
+                "ref_text": "precomputed transcript",
+                "ref_code_length": 3,
+            }
+        }
+
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
+        assert speech_server._validate_tts_request(req) is None
+
+        params = speech_server._build_tts_params(req)
+        assert params["speaker"] == ["alice"]
+        assert params["voice_created_at"] == [123]
+        assert params["task_type"] == ["Base"]
+        assert "voice_clone_prompt" in params
+        assert "ref_code_length" not in params
+        assert params.get("ref_text") != ["precomputed transcript"]
+
+    def test_precomputed_qwen3_missing_safetensors_is_not_registered(self, speech_server, tmp_path):
+        """Manifest entries are not supported speakers unless their safetensors load."""
+        _write_custom_voice_manifest(
+            tmp_path,
+            model_type="qwen3_tts",
+            voices={"Alice": {"file": "missing.safetensors", "mode": "xvec", "embedding_dim": 4}},
+        )
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(
+                custom_voice_dir=str(tmp_path),
+                talker_config=SimpleNamespace(hidden_size=4),
+            )
+        )
+
+        profiles = speech_server._load_precomputed_speakers()
+        assert profiles == {}
+
+        speech_server.precomputed_speakers = profiles
+        speech_server.supported_speakers = set(profiles)
+        assert "alice" not in speech_server.supported_speakers
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
+        assert speech_server._validate_tts_request(req) is not None
+
+    def test_precomputed_qwen3_icl_without_ref_code_is_not_registered(self, speech_server, tmp_path):
+        """Qwen3 ICL profiles without ref_code cannot be exposed by the API layer."""
+        from safetensors.torch import save_file
+
+        save_file({"speaker_embedding": torch.arange(4, dtype=torch.float32)}, str(tmp_path / "alice.safetensors"))
+        _write_custom_voice_manifest(
+            tmp_path,
+            model_type="qwen3_tts",
+            voices={
+                "Alice": {
+                    "file": "alice.safetensors",
+                    "mode": "icl",
+                    "ref_text": "reference transcript",
+                    "embedding_dim": 4,
+                }
+            },
+        )
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(
+                custom_voice_dir=str(tmp_path),
+                talker_config=SimpleNamespace(hidden_size=4),
+            )
+        )
+
+        profiles = speech_server._load_precomputed_speakers()
+        assert profiles == {}
+
+        speech_server.precomputed_speakers = profiles
+        speech_server.supported_speakers = set(profiles)
+        assert "alice" not in speech_server.supported_speakers
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
+        assert speech_server._validate_tts_request(req) is not None
+
+    def test_precomputed_voxcpm2_missing_safetensors_is_not_registered(self, speech_server, tmp_path):
+        """VoxCPM2 must not advertise a manifest-only voice that cannot hit prompt cache."""
+        _write_custom_voice_manifest(
+            tmp_path,
+            model_type="voxcpm2",
+            voices={"Bob": {"file": "missing.safetensors", "mode": "reference", "ref_audio_feat_len": 2}},
+        )
+        speech_server._tts_model_type = "voxcpm2"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(custom_voice_dir=str(tmp_path))
+        )
+
+        profiles = speech_server._load_precomputed_speakers()
+        assert profiles == {}
+
+        speech_server.precomputed_speakers = profiles
+        speech_server.supported_speakers = set(profiles)
+        assert "bob" not in speech_server.supported_speakers
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Bob")
+        assert speech_server._validate_tts_request(req) is not None
+        assert speech_server._validate_tts_request(OpenAICreateSpeechRequest(input="Hello")) is None
+
+    def test_prepare_voxcpm2_rejects_supported_speaker_without_voice_profile(self, speech_server, mocker):
+        """VoxCPM2 named voices must be uploaded or precomputed on the real request path."""
+        speech_server._tts_model_type = "voxcpm2"
+        speech_server.supported_speakers = {"bob", "default"}
+        speech_server.uploaded_speakers = {}
+        speech_server.precomputed_speakers = {}
+        speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        speech_server.engine_client.generate = mocker.MagicMock(return_value="generator")
+        speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
+            return_value={"prompt_token_ids": [1], "additional_information": {}}
+        )
+
+        with pytest.raises(ValueError, match="Invalid voice 'bob'"):
+            asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="Bob")))
+
+        speech_server._build_voxcpm2_prompt.assert_not_awaited()
+        speech_server.engine_client.generate.assert_not_called()
+
+    def test_prepare_voxcpm2_accepts_default_voice(self, speech_server, mocker):
+        """VoxCPM2 default voice preserves the built-in zero-shot request path."""
+        speech_server._tts_model_type = "voxcpm2"
+        speech_server.supported_speakers = {"default"}
+        speech_server.uploaded_speakers = {}
+        speech_server.precomputed_speakers = {}
+        speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(()))
+        speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
+            return_value={"prompt_token_ids": [1], "additional_information": {}}
+        )
+
+        asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="default")))
+
+        speech_server._build_voxcpm2_prompt.assert_awaited_once()
+        speech_server.engine_client.generate.assert_called_once()
+
+    def test_prepare_voxcpm2_precomputed_voice_sets_model_cache_key(self, speech_server, mocker):
+        """VoxCPM2 precomputed voices must carry voice metadata to the model cache lookup."""
+        speech_server._tts_model_type = "voxcpm2"
+        speech_server.supported_speakers = {"alice"}
+        speech_server.uploaded_speakers = {}
+        speech_server.precomputed_speakers = {
+            "alice": {
+                "name": "Alice",
+                "model_type": "voxcpm2",
+                "mode": "reference",
+                "ref_audio_feat_len": 2,
+            }
+        }
+        speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(()))
+        speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
+            return_value={
+                "prompt_token_ids": [1],
+                "additional_information": {
+                    "voice_profile": speech_server.precomputed_speakers["alice"],
+                },
+            }
+        )
+
+        asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="Alice")))
+
+        prompt = speech_server.engine_client.generate.call_args.kwargs["prompt"]
+        additional = prompt["additional_information"]
+        assert additional["voice_name"] == "alice"
+        assert additional["voice_created_at"] == 0
+
     def test_build_tts_params(self, speech_server):
         """Test TTS parameter building."""
         req = OpenAICreateSpeechRequest(input="Hello", voice="Ryan", language="English")
@@ -2497,6 +2749,95 @@ class TestTTSAsyncOffloading:
         qwen3_tts_server._build_tts_params.assert_called_once()
         qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
 
+    def test_qwen3_repeated_ref_audio_hot_path_sends_cache_key_without_waveform(self, qwen3_tts_server):
+        """After a ref artifact is marked ready, repeated requests avoid ref_audio payload IPC."""
+        wav_list = [0.0] * 48000
+        artifact_key = "a" * 40
+        ref_audio = "data:audio/wav;base64,same"
+        qwen3_tts_server._put_resolved_ref_audio(
+            hashlib.sha1(ref_audio.encode("utf-8")).hexdigest(),
+            wav_list,
+            24000,
+            artifact_key,
+        )
+        qwen3_tts_server._ref_audio_model_artifact_ready.add(artifact_key)
+        qwen3_tts_server._codec_frame_rate = 25.0
+        qwen3_tts_server._tts_tokenizer = lambda _text, padding=False: {"input_ids": list(range(10))}
+        qwen3_tts_server.engine_client.model_config.hf_config.talker_config = SimpleNamespace(
+            codec_language_id={},
+            spk_is_dialect={},
+        )
+
+        request = OpenAICreateSpeechRequest(
+            input="hello",
+            task_type="Base",
+            ref_audio=ref_audio,
+            ref_text="reference",
+        )
+        request_id, _generator, tts_params = asyncio.run(
+            qwen3_tts_server._prepare_speech_generation(request, request_id="req-hot")
+        )
+
+        assert request_id == "req-hot"
+        assert "ref_audio" not in tts_params
+        assert tts_params["_qwen3_tts_ref_audio_cache_key"] == [artifact_key]
+        assert tts_params["ref_code_length"] == [50]
+        prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
+        assert prompt["additional_information"] is tts_params
+
+    def test_qwen3_ref_audio_artifact_ready_is_evicted_with_resolve_cache(self, qwen3_tts_server):
+        qwen3_tts_server._ref_audio_resolve_cache_max_entries = 1
+        qwen3_tts_server._ref_audio_resolve_cache_max_bytes = 1_000_000
+
+        qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a")
+        qwen3_tts_server._ref_audio_model_artifact_ready.add("artifact-a")
+        qwen3_tts_server._put_resolved_ref_audio("ref-b", [0.0] * 8, 24000, "artifact-b")
+
+        assert "artifact-a" not in qwen3_tts_server._ref_audio_model_artifact_ready
+        assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()}
+
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server):
+        async def failing_generator():
+            raise ValueError("boom")
+            yield  # pragma: no cover
+
+        qwen3_tts_server._request_ref_audio_artifact_keys["req-fail"] = "artifact-fail"
+
+        with pytest.raises(ValueError, match="boom"):
+            await anext(qwen3_tts_server._generate_audio_chunks(failing_generator(), "req-fail"))
+
+        assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys
+        assert "artifact-fail" not in qwen3_tts_server._ref_audio_model_artifact_ready
+
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
+        async def pcm_generator():
+            yield SimpleNamespace(
+                multimodal_output={
+                    "audio": torch.zeros(16, dtype=torch.float32),
+                    "sr": 24000,
+                }
+            )
+            await asyncio.sleep(0)
+
+        qwen3_tts_server._request_ref_audio_artifact_keys["req-close"] = "artifact-close"
+
+        stream = qwen3_tts_server._generate_audio_chunks(pcm_generator(), "req-close")
+        assert await anext(stream)
+        await stream.aclose()
+
+        assert "req-close" not in qwen3_tts_server._request_ref_audio_artifact_keys
+        assert "artifact-close" not in qwen3_tts_server._ref_audio_model_artifact_ready
+
+    def test_qwen3_ref_audio_artifact_ready_requires_live_resolve_cache_entry(self, qwen3_tts_server):
+        qwen3_tts_server._request_ref_audio_artifact_keys["req-evicted"] = "artifact-evicted"
+
+        qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-evicted")
+
+        assert "req-evicted" not in qwen3_tts_server._request_ref_audio_artifact_keys
+        assert "artifact-evicted" not in qwen3_tts_server._ref_audio_model_artifact_ready
+
     def test_shutdown_is_idempotent(self, mocker: MockerFixture):
         """Calling shutdown() twice should not raise."""
         mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())