From 235c43032646bce93950b6746dd7fe6b6d8e06ba Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Tue, 29 Sep 2026 12:10:17 +0800 Subject: [PATCH] [Frontend] Tune reference cache limits and registered MOSS voice requests (#7883) Signed-off-by: Canlin Guo --- docs/serving/speech_api.md | 61 ++++++ tests/config/test_config_merge.py | 37 ++++ tests/config/test_speech_cache.py | 72 +++++++ .../openai_api/test_api_server_guards.py | 26 ++- .../test_moss_registered_reference.py | 200 ++++++++++++++++++ .../openai_api/test_omni_sleep_wakeup.py | 1 + .../openai_api/test_ref_audio_cache.py | 16 +- .../test_reference_cache_settings.py | 58 +++++ .../test_reference_waveform_cache.py | 5 +- .../openai_api/test_serving_speech.py | 10 +- .../openai_api/test_tts_adapter.py | 46 +++- tests/utils/test_speaker_cache.py | 20 ++ vllm_omni/config/speech_cache.py | 15 ++ vllm_omni/config/stage_config.py | 35 +-- vllm_omni/deploy/moss_tts_local.yaml | 4 + vllm_omni/deploy/moss_tts_local_h200.yaml | 4 + vllm_omni/entrypoints/openai/api_server.py | 10 + .../entrypoints/openai/protocol/audio.py | 15 +- .../entrypoints/openai/serving_speech.py | 42 +++- .../entrypoints/openai/tts_adapters/base.py | 13 +- .../openai/tts_adapters/moss_tts.py | 101 +++++++-- vllm_omni/utils/speaker_cache.py | 31 ++- 22 files changed, 752 insertions(+), 70 deletions(-) create mode 100644 tests/config/test_config_merge.py create mode 100644 tests/config/test_speech_cache.py create mode 100644 tests/entrypoints/openai_api/test_moss_registered_reference.py create mode 100644 tests/entrypoints/openai_api/test_reference_cache_settings.py create mode 100644 vllm_omni/config/speech_cache.py diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 2759fffc0..a9750fd68 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -1029,3 +1029,64 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ --enforce-eager \ --uvicorn-log-level debug ``` + +### Reusing reference audio + +The voice upload API and cache settings below are shared serving features. +For repeated voices, upload the reference once and reuse its name. The example +uses MOSS-TTS Local: uploading stores the audio, and the first speech request +encodes it. MOSS-TTS Local can then reuse cached reference codes without resolving +the waveform while the speaker cache entry remains resident. After a server +restart or speaker-cache eviction, another speech request warms the codes again. +Reference processing and cache-hit behavior for other models depend on their +serving adapters. + +```bash +curl --fail http://localhost:8123/v1/audio/voices \ + -F 'name=my-speaker' \ + -F 'consent=user_consent_id' \ + -F 'audio_sample=@/path/to/reference.wav' + +# Run once to warm the reference codes, then reuse the same voice. +curl --fail http://localhost:8123/v1/audio/speech \ + -H 'Content-Type: application/json' \ + -d '{"model":"OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5","voice":"my-speaker","input":"Hello, world.","response_format":"wav"}' \ + --output speech.wav +``` + +If you have the reference transcript, upload it with `-F 'ref_text=...'`. +MOSS-TTS Local uses a reference transcript for continuation conditioning; omitting +it uses reference-audio conditioning instead. Keep this choice consistent when +comparing benchmarks. Supplying inline `ref_audio` overrides the uploaded voice +and does not use the named-voice cache shortcut. + +The separate decoded-waveform resolve cache defaults to **4 GiB and 2048 entries +per speech server instance**. These are LRU limits, not preallocated memory or a +limit on total server RSS. Waveforms retain their source sampling rate and are +stored as owned, contiguous float32 arrays in **CPU memory**. The byte limit +counts numeric buffers (four bytes per mono sample), not the encoded file size +or GPU memory. +Either limit can cause eviction. Configure the API-process caches in the deploy +YAML passed to `--deploy-config` (alongside `stages`, not inside a stage): + +```yaml +speech_cache: + resolve_max_bytes: 4294967296 + resolve_max_entries: 2048 + speaker_max_bytes: 536870912 +``` + +All values must be non-negative integers. Either resolve limit set to `0` +disables waveform storage; `speaker_max_bytes: 0` disables speaker artifact +storage. Invalid values fail configuration loading. `base_config` inheritance +merges this section by field, preserving unspecified values from the base. +Restart the service to apply changes. The effective settings are logged at startup. + +The speaker budget defaults to **512 MiB** and counts cached tensor bytes, +not Python metadata or total process RSS. It configures the API process's +speaker cache. MOSS reference codes are cached as **CPU int32 tensors**; this +budget does not reserve GPU memory or include the reference encoder's GPU +weights and working memory. The shared cache itself does not move tensors +between devices; other adapters determine the device of their cached artifacts. +Separate model-worker caches are not controlled by this section. The process-wide singleton rejects a +conflicting explicit budget instead of silently ignoring it. diff --git a/tests/config/test_config_merge.py b/tests/config/test_config_merge.py new file mode 100644 index 000000000..0bcc15b0b --- /dev/null +++ b/tests/config/test_config_merge.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from copy import deepcopy + +import pytest + +from vllm_omni.config.stage_config import _merge_config_fields, resolve_deploy_yaml + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_selected_fields_merge_recursively_without_mutating_inputs(): + base = {"section": {"nested": {"a": 1, "b": 2}}, "replace": {"old": 1}} + overlay = {"section": {"nested": {"b": 0}}, "replace": {"new": 2}} + before = deepcopy((base, overlay)) + merged = _merge_config_fields(base, overlay, deep_merge_keys=frozenset({"section"})) + assert merged == {"section": {"nested": {"a": 1, "b": 0}}, "replace": {"new": 2}} + assert (base, overlay) == before + + +@pytest.mark.parametrize("value", [None, False, 0, [1, 2]]) +def test_non_mapping_overlay_replaces_mapping(value): + assert _merge_config_fields({"section": {"a": 1}}, {"section": value}, deep_merge_keys=frozenset({"section"})) == { + "section": value + } + + +def test_recursive_deploy_inheritance(tmp_path): + # Test raw YAML merging independently of the current flat speech schema. + (tmp_path / "base.yaml").write_text("speech_cache:\n nested:\n a: 1\n b: 2\n") + (tmp_path / "middle.yaml").write_text("base_config: base.yaml\nspeech_cache:\n nested:\n b: 0\n") + leaf = tmp_path / "leaf.yaml" + leaf.write_text("base_config: middle.yaml\nasync_chunk: false\n") + merged = resolve_deploy_yaml(leaf) + assert merged["speech_cache"] == {"nested": {"a": 1, "b": 0}} + assert merged["async_chunk"] is False diff --git a/tests/config/test_speech_cache.py b/tests/config/test_speech_cache.py new file mode 100644 index 000000000..0aa6a656a --- /dev/null +++ b/tests/config/test_speech_cache.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from dataclasses import FrozenInstanceError + +import pytest +from pydantic import ValidationError + +from vllm_omni.config.omni_config import _get_deploy_config +from vllm_omni.config.speech_cache import SpeechCacheConfig +from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig, load_deploy_config + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("field", ["resolve_max_bytes", "resolve_max_entries", "speaker_max_bytes"]) +@pytest.mark.parametrize("value", [-1, True, 1.0, 1.5, "1024", None]) +def test_invalid_limits(field, value): + with pytest.raises(ValueError, match=field): + SpeechCacheConfig(**{field: value}) + + +def test_unknown_limit_rejected(): + with pytest.raises(ValidationError, match="resolve_max_byte"): + SpeechCacheConfig(resolve_max_byte=1024) + + +def test_limits_are_immutable(): + config = SpeechCacheConfig() + with pytest.raises(FrozenInstanceError): + config.resolve_max_entries = 0 + + +def test_yaml_inheritance(tmp_path): + (tmp_path / "base.yaml").write_text( + "cuda_mps: true\nspeech_cache:\n resolve_max_bytes: 1234\n resolve_max_entries: 8\n speaker_max_bytes: 16\n" + ) + overlay = tmp_path / "overlay.yaml" + overlay.write_text("base_config: base.yaml\nspeech_cache:\n resolve_max_entries: 0\n") + deploy = load_deploy_config(overlay) + expected = SpeechCacheConfig(resolve_max_bytes=1234, resolve_max_entries=0, speaker_max_bytes=16) + assert deploy.speech_cache == expected + assert deploy.cuda_mps is True + + +@pytest.mark.parametrize("value", ["null", "[]", "false", "12", "{resolve_max_bytes: -1}", "{unknown: 10}"]) +@pytest.mark.parametrize("inherited", [False, True]) +def test_invalid_yaml(tmp_path, value, inherited): + path = tmp_path / "deploy.yaml" + (tmp_path / "base.yaml").write_text("speech_cache:\n resolve_max_entries: 8\n") + prefix = "base_config: base.yaml\n" if inherited else "" + path.write_text(f"{prefix}speech_cache: {value}\n") + with pytest.raises((ValueError, TypeError)): + load_deploy_config(path) + + +@pytest.mark.parametrize("explicit", [True, False]) +def test_selected_deploy_path_preserves_speech_cache(tmp_path, monkeypatch, explicit): + import vllm_omni.config.omni_config as module + + path = tmp_path / "test.yaml" + path.write_text("speech_cache:\n resolve_max_entries: 17\n") + monkeypatch.setattr(module, "_DEPLOY_DIR", tmp_path) + pipeline = PipelineConfig( + model_type="test", + default_deploy_config_name="test.yaml", + stages=(StagePipelineConfig(stage_id=0, model_stage="test", final_output=True),), + ) + deploy, selected_path = _get_deploy_config(pipeline, None, str(path) if explicit else None) + assert selected_path == str(path) + assert load_deploy_config(selected_path).speech_cache == deploy.speech_cache + assert deploy.speech_cache.resolve_max_entries == 17 diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index dbbc03cdf..b7e2d449a 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -366,6 +366,8 @@ class _FakeSocket: class _FakeEngineClient: + config_path: str | None = None + def __init__(self, *, stage_configs=None, endpoint_restrictions=None, vllm_config=None) -> None: self.stage_configs = stage_configs if stage_configs is not None else [] self.endpoint_restrictions = endpoint_restrictions if endpoint_restrictions is not None else {} @@ -999,9 +1001,14 @@ async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None: @pytest.mark.asyncio -async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch) -> None: +async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch, tmp_path) -> None: stage = SimpleNamespace(stage_type="diffusion", engine_args={}) engine = _FakeEngineClient(stage_configs=[stage]) + base = tmp_path / "base.yaml" + base.write_text("speech_cache:\n resolve_max_bytes: 1234\n speaker_max_bytes: 0\n") + deploy = tmp_path / "deploy.yaml" + deploy.write_text("base_config: base.yaml\nspeech_cache:\n resolve_max_entries: 17\n") + engine.config_path = str(deploy) speech_kwargs = {} def _for_diffusion_factory(label: str): @@ -1042,10 +1049,14 @@ async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch) -> assert speech_kwargs["allowed_local_media_path"] == "/allowed/media" assert speech_kwargs["allowed_media_domains"] == ["media.example.com"] + config = speech_kwargs["speech_cache_config"] + assert config.resolve_max_bytes == 1234 + assert config.resolve_max_entries == 17 + assert config.speaker_max_bytes == 0 @pytest.mark.asyncio -async def test_multistage_app_state_key_snapshot(monkeypatch) -> None: +async def test_multistage_app_state_key_snapshot(monkeypatch, tmp_path) -> None: """Lock multi-stage ``app.state`` keys after init, including live vs None. Fails if chat/speech/video/realtime/tokenization keys disappear or are @@ -1075,7 +1086,15 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None: def warmup(self): return None + deploy = tmp_path / "deploy.yaml" + deploy.write_text("speech_cache:\n resolve_max_bytes: 1024\n speaker_max_bytes: 8\n") + engine.config_path = str(deploy) + speech_kwargs = {} + class _FakeSpeech(_FakeCtor): + def __init__(self, *args, **kwargs): + speech_kwargs.update(kwargs) + async def warmup(self): return None @@ -1104,6 +1123,9 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None: state = State() await api_server.omni_init_app_state(engine, state, _minimal_args()) + assert speech_kwargs["speech_cache_config"].resolve_max_bytes == 1024 + assert speech_kwargs["speech_cache_config"].resolve_max_entries == 2048 + assert speech_kwargs["speech_cache_config"].speaker_max_bytes == 8 _assert_app_state_snapshot( state, diff --git a/tests/entrypoints/openai_api/test_moss_registered_reference.py b/tests/entrypoints/openai_api/test_moss_registered_reference.py new file mode 100644 index 000000000..d622ba2a6 --- /dev/null +++ b/tests/entrypoints/openai_api/test_moss_registered_reference.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from safetensors.torch import save_file + +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext, conditioning_cache_salt +from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import MossTTSAdapter +from vllm_omni.model_executor.models.moss_tts.reference_encoder import MossReferenceEncoder +from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def setup(tmp_path): + server = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + server.uploaded_speakers_dir = tmp_path + server.uploaded_speakers = { + "alice": dict( + embedding_source="audio", created_at=123, file_path=str(tmp_path / "old.safetensors"), ref_text="Reference." + ) + } + server._speaker_cache = SpeakerEmbeddingCache() + server._ref_audio_data_url_cache = {} + server._ref_audio_resolve_cache = {} + server._voice_created_at = Mock(side_effect=[123, 124]) + adapter = MossTTSAdapter(SpeechServingContext(server=server)) + adapter._moss_variant = "local" + encoder = MossReferenceEncoder(None, variant="local", n_vq=12, sr_target=24000, speaker_cache=server._speaker_cache) + adapter._get_moss_ref_encoder = lambda: encoder + + class Processor: + def build_user_message(self, **kwargs): + return kwargs + + def build_assistant_message(self, **kwargs): + return kwargs + + def __call__(self, *, conversations, mode): + message = conversations[0][-1] + codes = message.get("audio_codes_list", message.get("reference"))[0] + text = torch.zeros((len(codes), 1), dtype=torch.int64) + return {"input_ids": [torch.cat([text, codes], dim=1)]} + + adapter._get_moss_processor = lambda: Processor() + return server, adapter, encoder + + +@pytest.mark.asyncio +async def test_hot_registered_voice_skips_uri_and_resolve(setup): + server, adapter, encoder = setup + codes = torch.arange(36, dtype=torch.int32).reshape(3, 12) + server._speaker_cache.put(encoder._make_cache_key("alice", 123), {"codes": codes}) + for name in [ + "_get_uploaded_audio_data", + "_validate_ref_audio_format", + "_load_registered_reference", + "_resolve_ref_audio_array", + ]: + setattr(server, name, Mock(side_effect=AssertionError(name))) + request = OpenAICreateSpeechRequest(input="Hello", voice="alice") + assert adapter.validate(request) is None + result = await adapter.build(request, [], False) + assert torch.equal(result.tts_params["codes"]["ref"], codes) + assert request.ref_audio is None + assert not server._ref_audio_data_url_cache + server._voice_created_at.assert_not_called() + + +@pytest.mark.asyncio +async def test_reupload_during_encoding_keeps_codes_and_salt_generation(setup): + server, adapter, encoder = setup + entered, release = asyncio.Event(), asyncio.Event() + server._load_registered_reference = Mock(return_value=(np.zeros(24000, dtype=np.float32), 24000)) + + async def submit(waveform, sr): + entered.set() + await release.wait() + return torch.full((3, 12), 7, dtype=torch.int64) + + encoder._batcher.submit = submit + request = OpenAICreateSpeechRequest(input="Hello", voice="alice") + assert adapter.validate(request) is None + task = asyncio.create_task(adapter.build(request, [], False)) + await entered.wait() + server.uploaded_speakers["alice"] = dict(server.uploaded_speakers["alice"], created_at=124, ref_text="New.") + release.set() + result = await task + assert torch.equal(result.tts_params["codes"]["ref"], torch.full((3, 12), 7)) + assert result.tts_params["voice_created_at"] == [123] + assert request.ref_text == "Reference." + assert result.prompt["cache_salt"] == conditioning_cache_salt( + request, result.tts_params, registered_voice=("alice", 123) + ) + assert server._speaker_cache.get(encoder._make_cache_key("alice", 124)) is None + server._voice_created_at.assert_not_called() + + +def test_inline_override_is_not_bound(setup): + server, adapter, _ = setup + request = OpenAICreateSpeechRequest(input="Hello", voice="alice", ref_audio="file:///reference.wav") + assert adapter.validate(request) is None + assert request._registered_voice_reference is None + + +def test_private_snapshot_cannot_be_supplied_as_json(setup): + request = OpenAICreateSpeechRequest.model_validate( + dict(input="Hello", voice="alice", _registered_voice_reference={"name": "alice", "created_at": 999}) + ) + assert request._registered_voice_reference is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("channels", [1, 2]) +async def test_loaded_waveform_matches_legacy_data_uri(setup, channels): + server, adapter, _ = setup + samples = np.random.default_rng(1).uniform(-1.1, 1.1, (24000, channels)).astype(np.float32).squeeze() + info = server.uploaded_speakers["alice"] + save_file( + {"audio": torch.from_numpy(samples)}, info["file_path"], metadata={"created_at": "123", "sample_rate": "24000"} + ) + request = OpenAICreateSpeechRequest(input="Hello", voice="alice") + assert adapter.validate(request) is None + waveform, sr = server._load_registered_reference(request._registered_voice_reference) + from vllm.multimodal.media import MediaConnector + + legacy, legacy_sr = await MediaConnector().fetch_audio_async(server._get_uploaded_audio_data("alice")) + expected, _, _, _ = server._finalize_fetched_ref_audio(legacy, legacy_sr) + np.testing.assert_array_equal(waveform, expected) + assert sr == legacy_sr + # A deleted captured generation must not fall back to the new upload. + from pathlib import Path + + Path(info["file_path"]).unlink() + with pytest.raises((FileNotFoundError, OSError)): + server._load_registered_reference(request._registered_voice_reference) + + +@pytest.mark.asyncio +async def test_registered_cold_and_hot_salt_match(setup): + server, adapter, encoder = setup + server._load_registered_reference = Mock(return_value=(np.zeros(24000, dtype=np.float32), 24000)) + + async def submit(waveform, sr): + return torch.full((3, 12), 7, dtype=torch.int64) + + encoder._batcher.submit = submit + results = [] + for _ in range(2): + request = OpenAICreateSpeechRequest(input="Hello", voice="alice", ref_text="Override.") + assert adapter.validate(request) is None + assert request.ref_text == "Override." + results.append(await adapter.build(request, [], False)) + assert results[0].prompt["cache_salt"] == results[1].prompt["cache_salt"] + server._load_registered_reference.assert_called_once() + + +@pytest.mark.parametrize("generation", [0, -1]) +def test_invalid_registered_generation_is_rejected(setup, generation): + server, adapter, _ = setup + server.uploaded_speakers["alice"]["created_at"] = generation + request = OpenAICreateSpeechRequest(input="Hello", voice="alice") + assert "no valid generation" in adapter.validate(request) + assert request._registered_voice_reference is None + + +@pytest.mark.asyncio +async def test_new_request_uses_reuploaded_generation(setup): + server, adapter, encoder = setup + results = [] + for generation in [123, 124]: + server.uploaded_speakers["alice"]["created_at"] = generation + codes = torch.full((3, 12), generation, dtype=torch.int32) + server._speaker_cache.put(encoder._make_cache_key("alice", generation), {"codes": codes}) + request = OpenAICreateSpeechRequest(input="Hello", voice="alice") + assert adapter.validate(request) is None + result = await adapter.build(request, [], False) + assert torch.equal(result.tts_params["codes"]["ref"], codes) + results.append(result) + assert results[0].prompt["cache_salt"] != results[1].prompt["cache_salt"] + server.uploaded_speakers.clear() + assert "requires 'ref_audio'" in adapter.validate(OpenAICreateSpeechRequest(input="Hello", voice="alice")) + + +def test_cold_load_rejects_mismatched_generation(setup): + server, adapter, _ = setup + info = server.uploaded_speakers["alice"] + save_file({"audio": torch.zeros(24000)}, info["file_path"], metadata={"created_at": "124", "sample_rate": "24000"}) + request = OpenAICreateSpeechRequest(input="Hello", voice="alice") + assert adapter.validate(request) is None + with pytest.raises(ValueError, match="generation changed"): + server._load_registered_reference(request._registered_voice_reference) diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py index 7374483a4..2da64d274 100644 --- a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py +++ b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py @@ -275,6 +275,7 @@ def test_sleep_then_wakeup_roundtrip(sleep_capable_engine): @pytest.fixture def pure_diffusion_engine(mocker): engine = mocker.MagicMock() + engine.config_path = None engine.stage_configs = [{"stage_type": "diffusion"}] engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)]) engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)]) diff --git a/tests/entrypoints/openai_api/test_ref_audio_cache.py b/tests/entrypoints/openai_api/test_ref_audio_cache.py index 8026864b0..50adf9ef2 100644 --- a/tests/entrypoints/openai_api/test_ref_audio_cache.py +++ b/tests/entrypoints/openai_api/test_ref_audio_cache.py @@ -8,11 +8,8 @@ from collections import OrderedDict import numpy as np import pytest -from vllm_omni.entrypoints.openai.serving_speech import ( - _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES, - _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES, - OmniOpenAIServingSpeech, -) +from vllm_omni.config.speech_cache import SpeechCacheConfig +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -24,8 +21,9 @@ def cache(): server.model_config = None server._ref_audio_resolve_cache = OrderedDict() server._ref_audio_resolve_cache_bytes = 0 - server._ref_audio_resolve_cache_max_entries = _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES - server._ref_audio_resolve_cache_max_bytes = _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES + config = SpeechCacheConfig() + server._ref_audio_resolve_cache_max_entries = config.resolve_max_entries + server._ref_audio_resolve_cache_max_bytes = config.resolve_max_bytes server._ref_audio_model_artifact_ready = set() server._request_ref_audio_artifact_keys = {} return server @@ -36,8 +34,8 @@ def _put(cache, key, samples, artifact=None): def test_compact_storage_and_default_capacity(cache): - assert cache._ref_audio_resolve_cache_max_entries == 1024 - assert cache._ref_audio_resolve_cache_max_bytes == 512 * 1024 * 1024 + assert cache._ref_audio_resolve_cache_max_entries == 2048 + assert cache._ref_audio_resolve_cache_max_bytes == 4 * 1024**3 _put(cache, "a", [0.0, 0.25, -0.5, 1.0]) value = cache._ref_audio_resolve_cache["a"] assert value[0].dtype == np.float32 diff --git a/tests/entrypoints/openai_api/test_reference_cache_settings.py b/tests/entrypoints/openai_api/test_reference_cache_settings.py new file mode 100644 index 000000000..3a2f20fd5 --- /dev/null +++ b/tests/entrypoints/openai_api/test_reference_cache_settings.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import numpy as np +import pytest + +from vllm_omni.config.speech_cache import SpeechCacheConfig +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def server(monkeypatch, tmp_path): + import vllm_omni.utils.speaker_cache as module + + monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path)) + monkeypatch.setattr(module, "_SINGLETON", None) + instance = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + instance.speech_cache_config = SpeechCacheConfig() + return instance + + +def test_defaults(server): + server._init_speaker_storage() + assert server._ref_audio_resolve_cache_max_bytes == 4 * 1024**3 + assert server._ref_audio_resolve_cache_max_entries == 2048 + assert server._speaker_cache.stats()["max_bytes"] == 512 * 1024**2 + + +@pytest.mark.parametrize("entries,budget", [(1, 1000), (10, 8)]) +def test_reference_cache_enforces_configured_limits(server, entries, budget): + server.speech_cache_config = SpeechCacheConfig( + resolve_max_entries=entries, resolve_max_bytes=budget, speaker_max_bytes=16 + ) + server._init_speaker_storage() + server._put_resolved_ref_audio("a", np.asarray([0.0, 0.1], dtype=np.float32), 24000, "artifact-a") + server._put_resolved_ref_audio("b", np.asarray([0.2, 0.3], dtype=np.float32), 24000, "artifact-b") + assert list(server._ref_audio_resolve_cache) == ["b"] + assert server._ref_audio_resolve_cache_bytes == 8 + assert server._speaker_cache.stats()["max_bytes"] == 16 + + +@pytest.mark.parametrize("field", ["resolve_max_entries", "resolve_max_bytes"]) +def test_zero_disables_storage(server, field): + server.speech_cache_config = SpeechCacheConfig(**{field: 0}) + server._init_speaker_storage() + server._put_resolved_ref_audio("a", np.asarray([0.0], dtype=np.float32), 24000, "artifact-a") + assert not server._ref_audio_resolve_cache + assert server._ref_audio_resolve_cache_bytes == 0 + + +def test_diffusion_receives_config(server): + config = SpeechCacheConfig(resolve_max_bytes=80, resolve_max_entries=1, speaker_max_bytes=0) + speech = OmniOpenAIServingSpeech.for_diffusion(object(), "test", speech_cache_config=config) + assert speech.speech_cache_config is config + assert speech._ref_audio_resolve_cache_max_bytes == 80 + assert speech._speaker_cache.stats()["max_bytes"] == 0 diff --git a/tests/entrypoints/openai_api/test_reference_waveform_cache.py b/tests/entrypoints/openai_api/test_reference_waveform_cache.py index 9d87d560b..963e782bc 100644 --- a/tests/entrypoints/openai_api/test_reference_waveform_cache.py +++ b/tests/entrypoints/openai_api/test_reference_waveform_cache.py @@ -4,6 +4,7 @@ import numpy as np import pytest +from vllm_omni.config.speech_cache import SpeechCacheConfig from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -12,7 +13,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @pytest.fixture def server(monkeypatch, tmp_path): monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path)) - return OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + instance = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + instance.speech_cache_config = SpeechCacheConfig() + return instance @pytest.mark.parametrize("stereo", [False, True]) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index fb2e9839d..1ec32a95e 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -2144,7 +2144,9 @@ class TestTTSMethods: speech_server._tts_model_type = "moss_tts" speech_server._adapter = speech_server._get_tts_adapter() speech_server._adapter._moss_variant = "ttsd" - speech_server.uploaded_speakers = {"alice": {}} + speech_server.uploaded_speakers = { + "alice": {"embedding_source": "audio", "created_at": 42, "file_path": "/test.safetensors"} + } mocker.patch.object(speech_server, "_voice_created_at", return_value=42) proc = mocker.MagicMock() @@ -2166,16 +2168,16 @@ class TestTTSMethods: req = OpenAICreateSpeechRequest( input="hello", voice="alice", - ref_audio="data:audio/wav;base64,aaa", ref_audio_2="data:audio/wav;base64,bbb", ) + assert speech_server._adapter._bind_registered_reference(req) is None params = await speech_server._adapter._build_moss_tts_params(req, has_inline_ref_audio=False) assert sorted(seen) == [ - ("data:audio/wav;base64,aaa", "alice"), ("data:audio/wav;base64,bbb", None), + ("registered:alice:42", "alice"), ] # Per-slot salt keys survive the concurrent (gather) encode order. - assert params["ref_audio_cache_key"] == "rk:data:audio/wav;base64,aaa" + assert "ref_audio_cache_key" not in params assert params["ref_audio_2_cache_key"] == "rk:data:audio/wav;base64,bbb" def test_precomputed_qwen3_voice_infers_base_without_ref_audio(self, speech_server): diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index ad89f4954..babf8a2b9 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -219,7 +219,12 @@ def test_moss_reference_transcript_mode(variant, ref_text, mode, mocker): reference = [torch.ones((3, 12), dtype=torch.int64)] unified = torch.arange(52, dtype=torch.int64).reshape(1, 4, 13) processor = mocker.Mock(return_value={"input_ids": unified}) - server = mocker.Mock(_moss_variant="local", uploaded_speakers={"speaker": {}}) + server = mocker.Mock( + _moss_variant="local", + uploaded_speakers={ + "speaker": {"embedding_source": "audio", "created_at": 123, "file_path": "/test.safetensors"} + }, + ) server._voice_created_at.return_value = 123 request = OpenAICreateSpeechRequest( input="Target.", ref_text=ref_text, language="English", voice="Speaker", seed=0, max_new_tokens=2048 @@ -780,3 +785,42 @@ def test_qwen3_validate_rejects_no_voice_no_default_then_accepts(): if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +def test_registered_voice_salt_does_not_materialize_audio(): + from vllm_omni.entrypoints.openai.tts_adapters.base import conditioning_cache_salt + + class NoRepr(str): + def __repr__(self): + raise AssertionError("Uploaded audio must not be serialized on the hot path") + + request = OpenAICreateSpeechRequest(input="Target.", voice="speaker").model_copy( + update={"ref_audio": NoRepr("data:audio/wav;base64,AAAA")} + ) + salt = conditioning_cache_salt(request, registered_voice=("speaker", 123)) + assert salt + with pytest.raises(AssertionError, match="must not be serialized"): + conditioning_cache_salt(request) + + +@pytest.mark.parametrize( + "field, value", + [("input", "Different target."), ("ref_text", "Different reference."), ("language", "Chinese")], +) +def test_registered_voice_salt_preserves_request_conditioning(field, value): + from vllm_omni.entrypoints.openai.tts_adapters.base import conditioning_cache_salt + + request = OpenAICreateSpeechRequest(input="Target.", voice="speaker") + salt = conditioning_cache_salt(request, registered_voice=("speaker", 123)) + changed = request.model_copy(update={field: value}) + assert conditioning_cache_salt(changed, registered_voice=("speaker", 123)) != salt + + +@pytest.mark.parametrize("key", ["ref_audio_2_cache_key", "task_type", "ref_text"]) +def test_registered_voice_salt_preserves_resolved_conditioning(key): + from vllm_omni.entrypoints.openai.tts_adapters.base import conditioning_cache_salt + + request = OpenAICreateSpeechRequest(input="Target.", voice="speaker") + first = conditioning_cache_salt(request, {key: "a"}, registered_voice=("speaker", 123)) + second = conditioning_cache_salt(request, {key: "b"}, registered_voice=("speaker", 123)) + assert first != second diff --git a/tests/utils/test_speaker_cache.py b/tests/utils/test_speaker_cache.py index ea50c5881..e2e4b86e5 100644 --- a/tests/utils/test_speaker_cache.py +++ b/tests/utils/test_speaker_cache.py @@ -151,3 +151,23 @@ class TestSingleton: a = get_speaker_cache() b = get_speaker_cache() assert a is b + + +class TestSpeakerCacheConfiguration: + def test_conflicting_singleton_budget(self, monkeypatch): + import vllm_omni.utils.speaker_cache as module + + monkeypatch.setattr(module, "_SINGLETON", None) + cache = get_speaker_cache(max_bytes=8) + assert get_speaker_cache() is cache + assert get_speaker_cache(max_bytes=8) is cache + with pytest.raises( + ValueError, + match=r"speech_cache\.speaker_max_bytes: .*initialized with 8 bytes, .*requested budget is 16 bytes", + ): + get_speaker_cache(max_bytes=16) + + def test_zero_budget_disables_metadata_storage(self): + cache = SpeakerEmbeddingCache(max_bytes=0) + cache.put(_k("m", "a"), {"name": "a"}) + assert cache.get(_k("m", "a")) is None diff --git a/vllm_omni/config/speech_cache.py b/vllm_omni/config/speech_cache.py new file mode 100644 index 000000000..4fa0b42f1 --- /dev/null +++ b/vllm_omni/config/speech_cache.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""API-process speech cache budgets, independent of stage engine caches.""" + +from pydantic import ConfigDict, Field +from vllm.config.utils import config + + +@config(config=ConfigDict(strict=True, extra="forbid", frozen=True)) +class SpeechCacheConfig: + """LRU limits, not preallocated memory. Zero disables the respective cache.""" + + resolve_max_bytes: int = Field(default=4 * 1024**3, ge=0) + resolve_max_entries: int = Field(default=2048, ge=0) + speaker_max_bytes: int = Field(default=512 * 1024**2, ge=0) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 0fe7efe4f..1721a338c 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -18,6 +18,7 @@ from vllm.logger import init_logger from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm_omni.config.endpoint_policy import EndpointRestriction +from vllm_omni.config.speech_cache import SpeechCacheConfig from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler @@ -585,6 +586,7 @@ class DeployConfig: # Experimental local NVIDIA MPS; disabled unless a deploy explicitly opts in. cuda_mps: bool = False duplex_session: DuplexSessionRuntimeConfig = field(default_factory=DuplexSessionRuntimeConfig) + speech_cache: SpeechCacheConfig = field(default_factory=SpeechCacheConfig) connectors: dict[str, Any] | None = None edges: list[dict[str, Any]] | None = None stages: list[StageDeployConfig] = field(default_factory=list) @@ -685,16 +687,20 @@ _DEEP_MERGE_KEYS = frozenset( ) +_DEPLOY_DEEP_MERGE_KEYS = frozenset({"speech_cache"}) + + +def _merge_config_fields(base: dict, overlay: dict, *, deep_merge_keys: frozenset[str]) -> dict: + """Recursively merge selected fields; overlay replaces all other fields.""" + base_nested = {k: v for k, v in base.items() if k in deep_merge_keys} + overlay_nested = {k: v for k, v in overlay.items() if k in deep_merge_keys} + merged_nested = _get_recursively_merged_dict(original=base_nested, update=overlay_nested) + return {**base, **overlay, **merged_nested} + + def _deep_merge_stage(base: dict, overlay: dict) -> dict: """Deep-merge ``_DEEP_MERGE_KEYS`` so thin overlays don't drop base keys.""" - # Deep merge _DEEP_MERGE_KEYS recursively - base_merge_dict = {k: v for k, v in base.items() if k in _DEEP_MERGE_KEYS} - overlay_merge_dict = {k: v for k, v in overlay.items() if k in _DEEP_MERGE_KEYS} - - # Get the merge dict; priority is base < overlay < merged sub - merged_subdict = _get_recursively_merged_dict(original=base_merge_dict, update=overlay_merge_dict) - merged_dict = {**base, **overlay, **merged_subdict} - return merged_dict + return _merge_config_fields(base, overlay, deep_merge_keys=_DEEP_MERGE_KEYS) def _get_recursively_merged_dict(original: dict, update: dict) -> dict: @@ -776,10 +782,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: # Merge top-level scalars: overlay wins. Structured sections are merged # below so a thin overlay does not discard inherited runtime contracts. - merged = { - **base_dict, - **{k: v for k, v in raw_dict.items() if k not in ("connectors", "stages", "platforms")}, - } + merged = _merge_config_fields( + base_dict, + {k: v for k, v in raw_dict.items() if k not in ("connectors", "stages", "platforms")}, + deep_merge_keys=_DEPLOY_DEEP_MERGE_KEYS, + ) merged_connectors = _merge_connectors(base_dict.get("connectors"), raw_dict.get("connectors")) if merged_connectors is not None: merged["connectors"] = merged_connectors @@ -800,6 +807,9 @@ def load_deploy_config(path: str | Path) -> DeployConfig: "define topology in PipelineConfig and deployment overrides under `stages`." ) + speech_cache = raw_dict.get("speech_cache", {}) + if not isinstance(speech_cache, dict): + raise ValueError("speech_cache must be a mapping") stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])] model_runner = raw_dict.get("model_runner", "v1") @@ -810,6 +820,7 @@ def load_deploy_config(path: str | Path) -> DeployConfig: raise ValueError("cuda_mps must be a boolean") kwargs: dict[str, Any] = { + "speech_cache": SpeechCacheConfig(**speech_cache), "cuda_mps": raw_dict.get("cuda_mps", False), "async_chunk": raw_dict.get("async_chunk", True), "session_mode": raw_dict.get("session_mode", "turn"), diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index 41dca59f3..77c2a42b6 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -5,6 +5,10 @@ # frame — same per-frame timing (80 ms/frame, 12.5 fps) as MOSS-TTS-Realtime. async_chunk: true +speech_cache: + resolve_max_bytes: 2147483648 # 2 GiB + resolve_max_entries: 2048 + connectors: shm: name: SharedMemoryConnector diff --git a/vllm_omni/deploy/moss_tts_local_h200.yaml b/vllm_omni/deploy/moss_tts_local_h200.yaml index 6fb5e010d..260d44a0b 100644 --- a/vllm_omni/deploy/moss_tts_local_h200.yaml +++ b/vllm_omni/deploy/moss_tts_local_h200.yaml @@ -3,6 +3,10 @@ # Capacity and graph capture require more than 80 GiB of GPU memory. async_chunk: true +speech_cache: + resolve_max_bytes: 2147483648 # 2 GiB + resolve_max_entries: 2048 + connectors: shm: name: SharedMemoryConnector diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index a821c0a26..5d474306b 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -91,6 +91,8 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.endpoint_policy import ( shutdown_unsupported_routes, ) +from vllm_omni.config.speech_cache import SpeechCacheConfig +from vllm_omni.config.stage_config import load_deploy_config from vllm_omni.engine.stage_init_utils import set_death_signal from vllm_omni.engine.stage_runtime import OmniClientConfig from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni @@ -764,6 +766,12 @@ async def omni_init_app_state( state: FastAPI application state object to initialize args: Parsed command-line arguments """ + # AsyncOmni already exposes the selected deploy path, including the model's + # default YAML. Reuse the loader for base_config inheritance; API-only cache + # settings need not be carried through the engine or stage configurations. + deploy_path = getattr(engine_client, "config_path", None) + speech_cache_config = load_deploy_config(deploy_path).speech_cache if deploy_path else SpeechCacheConfig() + # Get vllm_config from engine_client (following 0.14.0 pattern) vllm_config = await openai_app_state._get_vllm_config(engine_client) @@ -846,6 +854,7 @@ async def omni_init_app_state( diffusion_engine=engine_client, model_name=model_name, stage_configs=diffusion_stage_configs, + speech_cache_config=speech_cache_config, allowed_local_media_path=getattr(args, "allowed_local_media_path", ""), allowed_media_domains=getattr(args, "allowed_media_domains", None), ) @@ -1155,6 +1164,7 @@ async def omni_init_app_state( state.openai_serving_models, request_logger=request_logger, model_name=model_name, + speech_cache_config=speech_cache_config, forced_aligner_enabled=build_forced_aligner_config( getattr(args, "forced_aligner", None), getattr(args, "forced_aligner_config", None), diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 8a1b48304..3b8df1c98 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math +from dataclasses import dataclass from typing import Annotated, Any, Literal import numpy as np -from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator +from pydantic import AliasChoices, BaseModel, Field, PrivateAttr, field_validator, model_validator # Bound int request fields to avoid overflow issues. _INT64_MIN = -(2**63) @@ -55,7 +56,19 @@ def _normalize_speaker_embedding_value(value): return [float(x) for x in value] +@dataclass(frozen=True) +class RegisteredVoiceReference: + """Server-only snapshot of an uploaded audio generation.""" + + name: str + created_at: int + file_path: str + ref_text: str | None + + class OpenAICreateSpeechRequest(BaseModel): + _registered_voice_reference: RegisteredVoiceReference | None = PrivateAttr(default=None) + input: str model: str | None = None # Accept both "voice" (OpenAI convention) and "speaker" (model/internal diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index bb27e6fd6..02a6c8cf6 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -36,6 +36,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.utils import random_uuid from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError +from vllm_omni.config.speech_cache import SpeechCacheConfig from vllm_omni.config.stage_config import StagePipelineConfig from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin, StreamingAudioResampler from vllm_omni.entrypoints.openai.errors import InvalidPresetVoiceReferenceError, InvalidVoiceReferenceError @@ -45,6 +46,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( BatchSpeechResponse, CreateAudio, OpenAICreateSpeechRequest, + RegisteredVoiceReference, SpeechBatchItem, SpeechBatchItemResult, SpeechInputTokenDetails, @@ -111,8 +113,6 @@ _REF_AUDIO_MIN_DURATION = 1.0 # seconds _REF_AUDIO_MAX_DURATION = 30.0 # seconds _REF_AUDIO_METADATA_FETCH_ATTEMPTS = 3 _REMOTE_REF_AUDIO_SCHEMES = frozenset({"http", "https", "data"}) -_REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 1024 -_REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 512 * 1024 * 1024 _TTS_MAX_INSTRUCTIONS_LENGTH = 500 _DEFAULT_VOICE_NAME = "default" @@ -257,14 +257,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._ref_audio_data_url_cache: dict[str, str] = {} self._ref_audio_resolve_cache: OrderedDict[str, tuple[np.ndarray, int, int, str]] = OrderedDict() self._ref_audio_resolve_cache_bytes = 0 - self._ref_audio_resolve_cache_max_entries = _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES - self._ref_audio_resolve_cache_max_bytes = _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES + config = self.speech_cache_config + self._ref_audio_resolve_cache_max_entries = config.resolve_max_entries + self._ref_audio_resolve_cache_max_bytes = config.resolve_max_bytes + logger.info("Speech cache configuration: %s", config) # Readiness is keyed by (artifact_key, x_vector_only). An x-vector-only # request caches a speaker embedding but no ref_code, so its artifact # must not satisfy a later ICL request that needs ref_code (#5049). self._ref_audio_model_artifact_ready: set[tuple[str, bool]] = set() self._request_ref_audio_artifact_keys: dict[str, tuple[str, bool]] = {} - self._speaker_cache = get_speaker_cache() + self._speaker_cache = get_speaker_cache(max_bytes=config.speaker_max_bytes) self._last_upload_ts = 0 self._upload_lock = asyncio.Lock() self._restore_uploaded_speakers() @@ -355,6 +357,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): stage_configs: "list[Any] | None" = None, allowed_local_media_path: str = "", allowed_media_domains: list[str] | None = None, + speech_cache_config: SpeechCacheConfig | None = None, ) -> "OmniOpenAIServingSpeech": """Create a speech serving instance for pure diffusion TTS models. @@ -362,6 +365,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): engine client that pure diffusion engines don't provide. """ instance = cls.__new__(cls) + instance.speech_cache_config = speech_cache_config or SpeechCacheConfig() instance._diffusion_mode = True instance._diffusion_engine = diffusion_engine instance._diffusion_model_name = model_name @@ -384,6 +388,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._media_connector = None self._allowed_local_media_path = "" self.model_name = kwargs.pop("model_name", None) + self.speech_cache_config = kwargs.pop("speech_cache_config", None) or SpeechCacheConfig() # True when the server was launched with --forced-aligner (a pooling # aligner stage is appended to the pipeline). Gates word_timestamps. self.forced_aligner_enabled: bool = bool(kwargs.pop("forced_aligner_enabled", False)) @@ -692,6 +697,33 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return None return samples, sr + def _load_registered_reference(self, reference: RegisteredVoiceReference) -> tuple[np.ndarray, int]: + """Load one captured generation without retaining a data URI or waveform. + + Keep the legacy WAV conversion on cache misses to preserve PCM16 + quantization. A missing/deleted generation fails instead of substituting + a newer upload with the same name. Called off the event loop. + """ + from safetensors import safe_open + + path = Path(reference.file_path) + if not _validate_path_within_directory(path, self.uploaded_speakers_dir): + raise ValueError("Invalid registered reference path") + with safe_open(str(path), framework="pt") as f: + metadata = f.metadata() or {} + if int(metadata.get("created_at", 0)) != reference.created_at: + raise ValueError("Registered reference generation changed") + samples = f.get_tensor("audio").numpy() + sr = int(metadata["sample_rate"]) + # Match _get_uploaded_audio_data's WAV subtype and the resolver's float32 + # channel mixing, without base64 serialization or any persistent copy. + buf = io.BytesIO() + sf.write(buf, samples, sr, format="WAV") + buf.seek(0) + waveform, sr = sf.read(buf, dtype="float32") + waveform, sr, _, _ = self._finalize_fetched_ref_audio(waveform, sr) + return waveform, sr + def _get_uploaded_audio_data(self, voice_name: str) -> str | None: """Return a base64-encoded WAV data URL for an uploaded voice. diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index 8daaca0ea..b2297d826 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -58,7 +58,12 @@ def resolve_stage_model_path(engine_client: Any) -> str | None: return str(model_path) if model_path else None -def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: dict | None = None) -> str: +def conditioning_cache_salt( + request: "OpenAICreateSpeechRequest", + tts_params: dict | None = None, + *, + registered_voice: tuple[str, int] | None = None, +) -> str: """Stable hash of the real Stage 0 conditioning for the prefix cache. The talker's vLLM prompt is placeholder token ids; the real inputs are @@ -71,7 +76,9 @@ def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: di conditioning such as ``voice_created_at`` and content-aware ref-audio cache keys must also be folded in. This distinguishes delete/re-upload of the same voice and same-path local audio rewrites without hashing decoded - waveform arrays. + waveform arrays. An adapter may supply ``registered_voice=(name, created_at)`` + only after verifying the uploaded voice and excluding inline audio overrides. + This replaces the uploaded data URI in the hash without changing the request. """ h = hashlib.sha256() for part in ( @@ -80,7 +87,7 @@ def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: di request.language, request.voice, request.ref_text, - request.ref_audio, + registered_voice if registered_voice is not None else request.ref_audio, request.instructions, request.x_vector_only_mode, request.speaker_embedding, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index ea7ff29ec..db2833f42 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -15,6 +15,7 @@ from transformers.dynamic_module_utils import get_class_from_dynamic_module from vllm.inputs import tokens_input from vllm.logger import init_logger +from vllm_omni.entrypoints.openai.protocol.audio import RegisteredVoiceReference from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ( ARTTSAdapter, @@ -103,25 +104,60 @@ class _MossTTSAdapterBase(ARTTSAdapter): """ from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_request_references - # Named-voice caching is only valid for uploaded speakers without an - # inline ref_audio: ``request.voice`` plus a file/URL would otherwise - # key on (name, created_at=0) and skip the content-aware resolve. - raw_voice = getattr(request, "voice", None) - raw_voice = raw_voice.strip() if isinstance(raw_voice, str) else "" - voice_lower = raw_voice.lower() - use_named_voice = bool(voice_lower) and voice_lower in self.uploaded_speakers and not has_inline_ref_audio - voice = voice_lower if use_named_voice else "" - voice_created = self._voice_created_at(voice) if voice else 0 + snapshot = request._registered_voice_reference if not has_inline_ref_audio else None + ref_str = cast(str, request.ref_audio) + resolver = self._resolve_ref_audio + if snapshot is not None: + # This key never enters URI validation. Only a cache miss loads the + # immutable generation's file; speaker 2 keeps the ordinary resolver. + ref_str = f"registered:{snapshot.name}:{snapshot.created_at}" - return await encode_request_references( + async def resolve_registered(ref_audio: str): + if ref_audio == ref_str: + waveform, sr = await asyncio.to_thread(self.ctx.server._load_registered_reference, snapshot) + return waveform, sr, ref_str + return await self._resolve_ref_audio(ref_audio) + + resolver = resolve_registered + + references, resolve_keys = await encode_request_references( self._get_moss_ref_encoder(), - cast(str, request.ref_audio), + ref_str, request.ref_audio_2 if two_speaker else None, - resolve_ref_audio=self._resolve_ref_audio, + resolve_ref_audio=resolver, get_artifact_key=self._get_resolved_ref_audio_artifact_key, - voice_name=voice or None, - voice_created_at=voice_created, + voice_name=snapshot.name if snapshot is not None else None, + voice_created_at=snapshot.created_at if snapshot is not None else 0, ) + if snapshot is not None: + # The generation identifies slot 0 on both cold and hot requests. + # Do not let a cold-only resolve key change its prefix-cache salt. + resolve_keys.pop(0, None) + return references, resolve_keys + + def _bind_registered_reference(self, request: "OpenAICreateSpeechRequest") -> str | None: + if request._registered_voice_reference is not None: + return None + voice = (request.voice or "").strip().lower() + if not voice: + return None + info = self.uploaded_speakers.get(voice) + if info is None: + return None + if info.get("embedding_source") != "audio": + return f"Uploaded voice '{voice}' must contain reference audio for MOSS-TTS." + created_at = int(info.get("created_at", 0)) + if created_at <= 0: + return f"Uploaded voice '{voice}' has no valid generation. Re-upload the voice." + request._registered_voice_reference = RegisteredVoiceReference( + name=voice, + created_at=created_at, + file_path=info["file_path"], + ref_text=info.get("ref_text"), + ) + if not request.ref_text or not request.ref_text.strip(): + request.ref_text = request._registered_voice_reference.ref_text + return None def _detect_moss_variant(self) -> str: """Sub-classify a ``moss_tts``-stage server into the actual MOSS-TTS @@ -407,7 +443,12 @@ class _MossTTSAdapterBase(ARTTSAdapter): we fall through to the original nano contract (ref_audio only). """ server = self.ctx.server - err = server._apply_uploaded_speaker(request) + if self._moss_variant is None: + err = server._apply_uploaded_speaker(request) + elif request.ref_audio is None: + err = self._bind_registered_reference(request) + else: + err = None if err: return err @@ -418,7 +459,7 @@ class _MossTTSAdapterBase(ARTTSAdapter): v = self._moss_variant if v in (None, "tts", "realtime", "local"): - if request.ref_audio is None: + if request.ref_audio is None and request._registered_voice_reference is None: label = ( "MOSS-TTS-Nano" if v is None @@ -429,12 +470,12 @@ class _MossTTSAdapterBase(ARTTSAdapter): ) ) return f"{label} requires 'ref_audio' (reference audio for voice cloning)." - return server._validate_ref_audio_format(request.ref_audio) + return server._validate_ref_audio_format(request.ref_audio) if request.ref_audio is not None else None if v == "ttsd": - if request.ref_audio is None: + if request.ref_audio is None and request._registered_voice_reference is None: return "MOSS-TTSD requires 'ref_audio' (speaker 1 reference)." - fmt_err = server._validate_ref_audio_format(request.ref_audio) + fmt_err = server._validate_ref_audio_format(request.ref_audio) if request.ref_audio is not None else None if fmt_err: return fmt_err if request.ref_audio_2 is not None: @@ -463,12 +504,26 @@ class _MossTTSAdapterBase(ARTTSAdapter): self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: server = self.ctx.server + if self._moss_variant is not None and not has_inline_ref_audio: + error = self._bind_registered_reference(request) + if error: + raise ValueError(error) + snapshot = request._registered_voice_reference if not has_inline_ref_audio else None tts_params = await self._build_moss_tts_params(request, has_inline_ref_audio=has_inline_ref_audio) - if request.voice: + registered_voice = None + if snapshot is not None: + tts_params["voice_name"] = [snapshot.name] + tts_params["voice_created_at"] = [snapshot.created_at] + registered_voice = (snapshot.name, snapshot.created_at) + elif self._moss_variant is None and request.voice and not has_inline_ref_audio: + # Nano retains its existing waveform-based uploaded-voice path. voice_lower = request.voice.lower() - if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: + if voice_lower in server.uploaded_speakers: + created_at = server._voice_created_at(voice_lower) tts_params["voice_name"] = [voice_lower] - tts_params["voice_created_at"] = [server._voice_created_at(voice_lower)] + tts_params["voice_created_at"] = [created_at] + if created_at > 0: + registered_voice = (voice_lower, created_at) # MOSS samples internally from additional_information. build() runs # before the shared path applies request.seed to SamplingParams. seed = request.seed @@ -482,7 +537,7 @@ class _MossTTSAdapterBase(ARTTSAdapter): else: prompt = tokens_input(prompt_token_ids=[1]) prompt["additional_information"] = tts_params - prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) + prompt["cache_salt"] = conditioning_cache_salt(request, tts_params, registered_voice=registered_voice) return PreparedRequest( prompt=prompt, tts_params=tts_params, diff --git a/vllm_omni/utils/speaker_cache.py b/vllm_omni/utils/speaker_cache.py index d571a36d4..e81b5ad56 100644 --- a/vllm_omni/utils/speaker_cache.py +++ b/vllm_omni/utils/speaker_cache.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Process-wide thread-safe LRU cache for speaker extraction artifacts. Keyed by ``(model_type, speaker_name, created_at)`` so each upload generation @@ -17,9 +20,10 @@ from typing import Any import torch from vllm.logger import init_logger +from vllm_omni.config.speech_cache import SpeechCacheConfig + logger = init_logger(__name__) -_MAX_BYTES = 512 * 1024**2 # 512 MiB _CUSTOM_VOICE_MANIFEST = "custom_voice_manifest.json" _CUSTOM_VOICE_SCHEMA_VERSION = 1 @@ -228,7 +232,9 @@ def load_validated_profile_tensors( class SpeakerEmbeddingCache: """Thread-safe in-memory LRU cache for speaker extraction artifacts.""" - def __init__(self, *, max_bytes: int = _MAX_BYTES): + def __init__(self, *, max_bytes: int | None = None): + if max_bytes is None: + max_bytes = SpeechCacheConfig().speaker_max_bytes self._cache: OrderedDict[tuple[str, str, int], dict[str, Any]] = OrderedDict() self._sizes: dict[tuple[str, str, int], int] = {} self._total_bytes = 0 @@ -265,6 +271,8 @@ class SpeakerEmbeddingCache: self._insert_locked(key, artifacts) def _insert_locked(self, key: tuple[str, str, int], artifacts: dict[str, Any]) -> None: + if self._max_bytes <= 0: + return size = _estimate_tensor_bytes(artifacts) if size > self._max_bytes: logger.warning("Speaker cache skip: entry %s size=%dB exceeds max_bytes=%dB", key, size, self._max_bytes) @@ -320,11 +328,16 @@ class SpeakerEmbeddingCache: } -def get_speaker_cache() -> SpeakerEmbeddingCache: - """Return the process-wide speaker cache singleton.""" +def get_speaker_cache(*, max_bytes: int | None = None) -> SpeakerEmbeddingCache: + """Return the process-wide cache; reject conflicting explicit budgets.""" global _SINGLETON - if _SINGLETON is None: - with _SINGLETON_LOCK: - if _SINGLETON is None: - _SINGLETON = SpeakerEmbeddingCache() - return _SINGLETON + with _SINGLETON_LOCK: + if _SINGLETON is None: + _SINGLETON = SpeakerEmbeddingCache(max_bytes=max_bytes) + elif max_bytes is not None and _SINGLETON.stats()["max_bytes"] != max_bytes: + raise ValueError( + "Conflicting speech_cache.speaker_max_bytes: " + f"the process-wide speaker cache is already initialized with {_SINGLETON.stats()['max_bytes']} bytes, " + f"but the requested budget is {max_bytes} bytes" + ) + return _SINGLETON