diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index b9f2e5c71..99e58fce2 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -739,7 +739,11 @@ async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None: ) state = State() - await api_server.omni_init_app_state(engine, state, _minimal_args()) + await api_server.omni_init_app_state( + engine, + state, + _minimal_args(), + ) _assert_app_state_snapshot( state, @@ -750,6 +754,55 @@ async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None: assert state.diffusion_engine is engine +@pytest.mark.asyncio +async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch) -> None: + stage = SimpleNamespace(engine_args={}) + engine = _FakeEngineClient(stage_configs=[stage]) + speech_kwargs = {} + + def _for_diffusion_factory(label: str): + @classmethod + def _factory(cls, *args, **kwargs): + return _marker(label) + + return _factory + + @classmethod + def _speech_factory(cls, *args, **kwargs): + speech_kwargs.update(kwargs) + return _marker("speech") + + monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion") + monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat")) + monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch")) + monkeypatch.setattr( + api_server.OmniOpenAIServingAudioGenerate, + "for_diffusion", + _for_diffusion_factory("audio_generate"), + ) + monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video")) + monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None) + monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _speech_factory) + monkeypatch.setattr( + api_server.ServingRealtimeRobotOpenPI, + "create_policy_server", + classmethod(lambda cls, *a, **k: _marker("openpi")), + ) + + state = State() + await api_server.omni_init_app_state( + engine, + state, + _minimal_args( + allowed_local_media_path="/allowed/media", + allowed_media_domains=["media.example.com"], + ), + ) + + assert speech_kwargs["allowed_local_media_path"] == "/allowed/media" + assert speech_kwargs["allowed_media_domains"] == ["media.example.com"] + + @pytest.mark.asyncio async def test_multistage_app_state_key_snapshot(monkeypatch) -> None: """Lock multi-stage ``app.state`` keys after init, including live vs None. diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index a9a87d1d7..b42de54f2 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1245,6 +1245,46 @@ class TestTTSMethods: cache_key ) == speech_server._make_ref_audio_artifact_cache_key(np.asarray(first[0], dtype=np.float32), 24000) + @pytest.mark.asyncio + async def test_diffusion_ref_audio_uses_media_access_config( + self, + mocker: MockerFixture, + monkeypatch, + tmp_path, + ): + monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path / "speakers")) + connector = mocker.MagicMock() + connector.fetch_audio_async = mocker.AsyncMock(return_value=(np.zeros(24000, dtype=np.float32), 24000)) + connector_cls = mocker.patch.object( + serving_speech_module, + "MediaConnector", + return_value=connector, + ) + media_dir = tmp_path / "media" + media_dir.mkdir() + ref_audio = media_dir / "reference.wav" + ref_audio.write_bytes(b"v1") + + server = OmniOpenAIServingSpeech.for_diffusion( + diffusion_engine=mocker.MagicMock(), + model_name="test-model", + allowed_local_media_path=str(media_dir), + allowed_media_domains=["media.example.com"], + ) + try: + first = await server._resolve_ref_audio(ref_audio.as_uri()) + ref_audio.write_bytes(b"v2-with-different-size") + second = await server._resolve_ref_audio(ref_audio.as_uri()) + finally: + server.shutdown() + + connector_cls.assert_called_once_with( + allowed_local_media_path=str(media_dir), + allowed_media_domains=["media.example.com"], + ) + assert first[2] != second[2] + assert connector.fetch_audio_async.await_count == 2 + # ── ref-audio cache key tests (static helper) ── def test_local_file_cache_key_invalidation(self): diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 6838e31ae..cdc36e588 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -825,6 +825,8 @@ async def omni_init_app_state( diffusion_engine=engine_client, model_name=model_name, stage_configs=diffusion_stage_configs, + allowed_local_media_path=getattr(args, "allowed_local_media_path", ""), + allowed_media_domains=getattr(args, "allowed_media_domains", None), ) state.openai_serving_duplex = None state.openai_streaming_speech = None diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index c2c200039..61b2200cf 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import base64 import hashlib @@ -275,6 +278,8 @@ def _conditioning_cache_salt(request, tts_params: dict | None = None) -> str: class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): _diffusion_mode: bool = False + _media_connector: MediaConnector | None = None + _allowed_local_media_path: str = "" _tts_executor: ThreadPoolExecutor | None = None def _init_speaker_storage(self) -> None: @@ -391,6 +396,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): diffusion_engine: "Any", model_name: str, stage_configs: "list[Any] | None" = None, + allowed_local_media_path: str = "", + allowed_media_domains: list[str] | None = None, ) -> "OmniOpenAIServingSpeech": """Create a speech serving instance for pure diffusion TTS models. @@ -402,6 +409,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): instance._diffusion_engine = diffusion_engine instance._diffusion_model_name = model_name instance._diffusion_stage_configs = stage_configs + instance._allowed_local_media_path = allowed_local_media_path + instance._media_connector = MediaConnector( + allowed_local_media_path=allowed_local_media_path, + allowed_media_domains=allowed_media_domains, + ) instance._tts_model_type = "omnivoice" instance._is_tts = False # Diffusion-only instances don't have a TTS stage; set None so any @@ -412,6 +424,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return instance def __init__(self, *args, **kwargs): + self._media_connector = None + self._allowed_local_media_path = "" self.model_name = kwargs.pop("model_name", None) # True when the server was launched with --forced-aligner (a pooling # aligner stage is appended to the pipeline). Gates word_timestamps. @@ -1860,15 +1874,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """ # Pass the allowed-local-media-path so the stat is restricted to # paths the server operator has explicitly permitted. - allowed_path = None - if not self._diffusion_mode: + allowed_path: str | None + if self._diffusion_mode: + allowed_path = self._allowed_local_media_path + else: allowed_path = getattr(self.model_config, "allowed_local_media_path", None) wav_list: list[float] | None = None sr = 0 artifact_key = "" post_key = "" - connector: MediaConnector | None = None for attempt in range(_REF_AUDIO_METADATA_FETCH_ATTEMPTS): cache_key = await self._ref_audio_cache_key(ref_audio_str, allowed_path) cached = self._ref_audio_resolve_cache.get(cache_key) @@ -1883,19 +1898,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) return wav_list, sr, cache_key - if connector is None: - # In diffusion mode, model_config may not be available - if self._diffusion_mode: - connector = MediaConnector() - else: - model_config = self.model_config - connector = MediaConnector( - allowed_local_media_path=model_config.allowed_local_media_path, - allowed_media_domains=model_config.allowed_media_domains, - ) + if self._media_connector is None: + model_config = self.model_config + self._media_connector = MediaConnector( + allowed_local_media_path=model_config.allowed_local_media_path, + allowed_media_domains=model_config.allowed_media_domains, + ) fetch_start_s = time.perf_counter() - wav_np, fetched_sr = await connector.fetch_audio_async(ref_audio_str) + wav_np, fetched_sr = await self._media_connector.fetch_audio_async(ref_audio_str) fetch_decode_ms = (time.perf_counter() - fetch_start_s) * 1000.0 tolist_start_s = time.perf_counter() wav_list, sr, artifact_key, duration = self._finalize_fetched_ref_audio(wav_np, fetched_sr)