diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index b4b2e7f55..c0b8fe5e2 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -521,6 +521,15 @@ are cached in-process with a shared LRU so repeated requests with the same all TTS model types; deleting a voice invalidates every model-type slot at once. +Decoded reference waveforms use a separate LRU cache of owned, contiguous +float32 arrays at the source sampling rate. Its byte budget counts numeric +buffers (four bytes per mono sample), excluding cache metadata. Numeric storage +avoids retaining a Python float object for every sample during garbage +collection. MOSS reference encoding consumes these arrays directly; other +list-based interfaces, including MOSS Nano, receive temporary lists that are +not retained in the resolve cache. The MOSS reference encoder also releases +completed batch inputs before waiting for more work. + ### Precomputed Custom Voices Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup. diff --git a/tests/entrypoints/openai_api/test_reference_waveform_cache.py b/tests/entrypoints/openai_api/test_reference_waveform_cache.py new file mode 100644 index 000000000..9d87d560b --- /dev/null +++ b/tests/entrypoints/openai_api/test_reference_waveform_cache.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import numpy as np +import pytest + +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech + +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) + + +@pytest.mark.parametrize("stereo", [False, True]) +def test_finalized_waveform_owns_contiguous_numeric_storage(server, stereo): + import gc + + source = np.linspace(-1, 1, 96000, dtype=np.float64).reshape(48000, 2) + decoded = source if stereo else source[:, 0] + expected = np.asarray(decoded, dtype=np.float32) + if stereo: + expected = expected.mean(axis=-1) + waveform, sr, artifact, duration = server._finalize_fetched_ref_audio(decoded, 24000) + assert waveform.dtype == np.float32 + assert waveform.flags.c_contiguous and waveform.flags.owndata + assert not np.shares_memory(waveform, decoded) + assert not gc.is_tracked(waveform) + np.testing.assert_array_equal(waveform, expected) + assert artifact == server._make_ref_audio_artifact_cache_key(expected, sr) + assert duration == 2.0 + + +@pytest.mark.asyncio +async def test_array_resolver_reuses_buffer_and_legacy_lists_are_temporary(server, mocker): + server._init_speaker_storage() + server._diffusion_mode = True + server._allowed_local_media_path = "" + server._ref_audio_cache_key = mocker.AsyncMock(return_value="locator") + server._media_connector = mocker.Mock() + server._media_connector.fetch_audio_async = mocker.AsyncMock( + return_value=(np.linspace(-1, 1, 24000, dtype=np.float32), 24000) + ) + first, sr, key = await server._resolve_ref_audio_array("reference") + second, _, _ = await server._resolve_ref_audio_array("reference") + legacy, _, _ = await server._resolve_ref_audio("reference") + assert first is second + assert isinstance(legacy, list) + np.testing.assert_array_equal(first, legacy) + legacy[0] = 100 + assert first[0] == -1 + assert server._ref_audio_resolve_cache[key][0] is first + assert server._ref_audio_resolve_cache_bytes == first.nbytes == 96000 + server._media_connector.fetch_audio_async.assert_awaited_once() + + +@pytest.mark.parametrize("variant", [None, "local", "delay", "realtime"]) +@pytest.mark.asyncio +async def test_moss_array_resolution_keeps_nano_list_transport(mocker, variant): + from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext + from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _MossTTSAdapterBase + + server = mocker.Mock() + server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0], 24000, "key")) + waveform = np.zeros(1, dtype=np.float32) + server._resolve_ref_audio_array = mocker.AsyncMock(return_value=(waveform, 24000, "key")) + adapter = _MossTTSAdapterBase.__new__(_MossTTSAdapterBase) + adapter.ctx = SpeechServingContext(server=server) + adapter._moss_variant = variant + result, _, _ = await adapter._resolve_ref_audio("reference") + if variant is None: + assert isinstance(result, list) + server._resolve_ref_audio_array.assert_not_awaited() + else: + assert result is waveform + server._resolve_ref_audio.assert_not_awaited() + + +def test_waveform_cache_eviction_releases_owned_buffer(server): + import weakref + + server._init_speaker_storage() + server._ref_audio_resolve_cache_max_entries = 1 + waveform = np.zeros(24000, dtype=np.float32) + ref = weakref.ref(waveform) + server._put_resolved_ref_audio("first", waveform, 24000, "artifact-first") + del waveform + assert ref() is not None + server._put_resolved_ref_audio("second", np.zeros(24000, dtype=np.float32), 24000, "artifact-second") + assert ref() is None + assert server._ref_audio_resolve_cache_bytes == 96000 + + +@pytest.mark.parametrize("entries, budget", [(1, 1000), (10, 8)]) +def test_numeric_waveform_cache_limits(server, entries, budget): + server._init_speaker_storage() + server._ref_audio_resolve_cache_max_entries = entries + server._ref_audio_resolve_cache_max_bytes = budget + server._put_resolved_ref_audio("a", np.zeros(2, dtype=np.float32), 24000, "artifact-a") + server._put_resolved_ref_audio("b", np.ones(2, dtype=np.float32), 24000, "artifact-b") + assert list(server._ref_audio_resolve_cache) == ["b"] + assert server._ref_audio_resolve_cache_bytes == 8 diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 36967e706..146488aa2 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1592,7 +1592,8 @@ class TestTTSMethods: assert first[1] == 24000 assert second[1] == 24000 - assert first[0] is second[0] + assert first[0] == second[0] + assert first[0] is not second[0] assert first[0][0] == pytest.approx(float(wav[0]), abs=1e-4) cache_key = first[2] assert speech_server._get_resolved_ref_audio_artifact_key( @@ -5430,7 +5431,7 @@ class TestTTSAsyncOffloading: ref_audio = "data:audio/wav;base64,same" qwen3_tts_server._put_resolved_ref_audio( hashlib.sha1(ref_audio.encode("utf-8")).hexdigest(), - wav_list, + np.asarray(wav_list, dtype=np.float32), 24000, artifact_key, ) @@ -5466,9 +5467,9 @@ class TestTTSAsyncOffloading: 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._put_resolved_ref_audio("ref-a", np.zeros(8, dtype=np.float32), 24000, "artifact-a") qwen3_tts_server._ref_audio_model_artifact_ready.add(("artifact-a", False)) - qwen3_tts_server._put_resolved_ref_audio("ref-b", [0.0] * 8, 24000, "artifact-b") + qwen3_tts_server._put_resolved_ref_audio("ref-b", np.zeros(8, dtype=np.float32), 24000, "artifact-b") assert ("artifact-a", False) 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()} @@ -5715,7 +5716,7 @@ class TestTTSAsyncOffloading: def test_qwen3_xvector_ready_artifact_does_not_enable_icl_artifact_only(self, qwen3_tts_server): # An x-vector-only artifact (speaker embedding, no ref_code) must not enable # the artifact-only path for a later ICL request with the same ref_audio (#5049). - qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a") + qwen3_tts_server._put_resolved_ref_audio("ref-a", np.zeros(8, dtype=np.float32), 24000, "artifact-a") qwen3_tts_server._track_ref_audio_artifact_warmup("req-xvec", "artifact-a", x_vector_only=True) qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-xvec") @@ -5723,7 +5724,7 @@ class TestTTSAsyncOffloading: assert qwen3_tts_server._adapter._qwen3_tts_can_use_ref_audio_artifact_only(icl_params, "artifact-a") is False def test_qwen3_xvector_ready_artifact_still_reusable_by_xvector_request(self, qwen3_tts_server): - qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a") + qwen3_tts_server._put_resolved_ref_audio("ref-a", np.zeros(8, dtype=np.float32), 24000, "artifact-a") qwen3_tts_server._track_ref_audio_artifact_warmup("req-xvec", "artifact-a", x_vector_only=True) qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-xvec") @@ -5731,7 +5732,7 @@ class TestTTSAsyncOffloading: assert qwen3_tts_server._adapter._qwen3_tts_can_use_ref_audio_artifact_only(xvec_params, "artifact-a") is True def test_qwen3_icl_ready_artifact_enables_icl_artifact_only(self, qwen3_tts_server): - qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a") + qwen3_tts_server._put_resolved_ref_audio("ref-a", np.zeros(8, dtype=np.float32), 24000, "artifact-a") qwen3_tts_server._track_ref_audio_artifact_warmup("req-icl", "artifact-a", x_vector_only=False) qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-icl") diff --git a/tests/model_executor/models/moss_tts/test_reference_encoder.py b/tests/model_executor/models/moss_tts/test_reference_encoder.py index bb03d556b..8cab6009e 100644 --- a/tests/model_executor/models/moss_tts/test_reference_encoder.py +++ b/tests/model_executor/models/moss_tts/test_reference_encoder.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for MossReferenceEncoder: content-addressed caching, single-flight, and micro-batched encoding.""" @@ -534,3 +534,50 @@ async def test_encode_request_references_keys_second_speaker_by_slot(make_encode assert proc.total_items == 2 # Concurrent, so both clips shared one batch window. assert proc.attempt_sizes == [2] + + +@pytest.mark.parametrize("sr", [24000, 48000]) +def test_numeric_waveform_prep_matches_list_and_does_not_mutate_cache(sr): + import numpy as np + + from vllm_omni.model_executor.models.moss_tts.reference_encoder import _prep_wav_sync + + waveform = np.linspace(-1, 1, sr, dtype=np.float32) + expected = waveform.copy() + tensor = _prep_wav_sync(waveform, sr, 24000) + assert torch.equal(tensor, _prep_wav_sync(waveform.tolist(), sr, 24000)) + tensor.zero_() + np.testing.assert_array_equal(waveform, expected) + + +@pytest.mark.asyncio +async def test_idle_batcher_releases_all_completed_waveforms(): + import weakref + + import numpy as np + + batcher = _RefEncodeBatcher( + lambda payload: [torch.zeros((1, 4), dtype=torch.long) for _ in payload], + window_ms=0, + max_batch=8, + ) + refs = [] + + async def submit_one(): + waveform = np.zeros(24000, dtype=np.float32) + refs.append(weakref.ref(waveform)) + await batcher.submit(waveform, 24000) + + try: + await asyncio.gather(submit_one(), submit_one()) + + # Future completion can race the executor releasing its work item. + # Wait without forcing GC or overwriting first/jobs with a new batch. + async def wait_for_release(): + while any(ref() is not None for ref in refs): + await asyncio.sleep(0.001) + + await asyncio.wait_for(wait_for_release(), timeout=2.0) + assert batcher._drainer is not None and not batcher._drainer.done() + finally: + await batcher.aclose() diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 8aa32899a..a348fdd09 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -248,7 +248,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._max_uploaded_speakers = 1000 self.uploaded_speakers: dict[str, dict[str, Any]] = {} self._ref_audio_data_url_cache: dict[str, str] = {} - self._ref_audio_resolve_cache: OrderedDict[str, tuple[list[float], int, int, str]] = OrderedDict() + 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 @@ -1222,7 +1222,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): allowed_local_media_path, ) - def _finalize_fetched_ref_audio(self, wav_np: np.ndarray, sr: int) -> tuple[list[float], int, str, float]: + def _finalize_fetched_ref_audio(self, wav_np: np.ndarray, sr: int) -> tuple[np.ndarray, int, str, float]: wav_np = np.asarray(wav_np, dtype=np.float32) if wav_np.ndim > 1: wav_np = np.mean(wav_np, axis=-1) @@ -1238,12 +1238,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): f"Reference audio too long ({duration:.1f}s). " f"Maximum {_REF_AUDIO_MAX_DURATION:.0f}s supported — use a shorter clip." ) - artifact_key = self._make_ref_audio_artifact_cache_key(wav_np, sr) - return wav_np.tolist(), sr, artifact_key, duration + # Own the buffer: a view could retain a much larger decoded allocation. + waveform = np.array(wav_np, dtype=np.float32, order="C", copy=True) + artifact_key = self._make_ref_audio_artifact_cache_key(waveform, sr) + return waveform, sr, artifact_key, duration async def _resolve_ref_audio(self, ref_audio_str: str) -> tuple[list[float], int, str]: + """Keep list-based consumers and their engine transport compatible. + + Only the numeric array is cached; this temporary list belongs to the + caller. Server-side reference encoders should use the array resolver. + """ + waveform, sr, cache_key = await self._resolve_ref_audio_array(ref_audio_str) + return waveform.tolist(), sr, cache_key + + async def _resolve_ref_audio_array(self, ref_audio_str: str) -> tuple[np.ndarray, int, str]: """Resolve ref_audio to (wav_samples, sample_rate, cache_key). + The float32 array can be shared with the cache. Callers must not mutate + it; copy before any in-place processing. + Delegates to upstream vLLM's MediaConnector which handles http(s) URLs, ``data:`` base64 URIs, and ``file:`` local paths (the latter gated by ``--allowed-local-media-path``). @@ -1266,7 +1280,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): else: allowed_path = getattr(self.model_config, "allowed_local_media_path", None) - wav_list: list[float] | None = None + waveform: np.ndarray | None = None sr = 0 artifact_key = "" post_key = "" @@ -1275,14 +1289,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): cached = self._ref_audio_resolve_cache.get(cache_key) if cached is not None: self._ref_audio_resolve_cache.move_to_end(cache_key) - wav_list, sr, _, _ = cached + waveform, sr, _, _ = cached logger.debug( "Resolved ref_audio from cache: samples=%d sr=%d duration_s=%.3f", - len(wav_list), + len(waveform), sr, - len(wav_list) / sr if sr > 0 else 0.0, + len(waveform) / sr if sr > 0 else 0.0, ) - return wav_list, sr, cache_key + return waveform, sr, cache_key if self._media_connector is None: model_config = self.model_config @@ -1294,21 +1308,21 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): fetch_start_s = time.perf_counter() 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) - tolist_ms = (time.perf_counter() - tolist_start_s) * 1000.0 + finalize_start_s = time.perf_counter() + waveform, sr, artifact_key, duration = self._finalize_fetched_ref_audio(wav_np, fetched_sr) + finalize_ms = (time.perf_counter() - finalize_start_s) * 1000.0 logger.debug( - "Resolved ref_audio: fetch_decode_ms=%.3f tolist_ms=%.3f samples=%d sr=%d duration_s=%.3f", + "Resolved ref_audio: fetch_decode_ms=%.3f finalize_ms=%.3f samples=%d sr=%d duration_s=%.3f", fetch_decode_ms, - tolist_ms, - len(wav_list), + finalize_ms, + len(waveform), sr, duration, ) post_key = await self._ref_audio_cache_key(ref_audio_str, allowed_path) if post_key == cache_key: - self._put_resolved_ref_audio(cache_key, wav_list, sr, artifact_key) - return wav_list, sr, cache_key + self._put_resolved_ref_audio(cache_key, waveform, sr, artifact_key) + return waveform, sr, cache_key logger.debug( "ref_audio metadata changed during fetch (attempt %d/%d); retrying", attempt + 1, @@ -1319,8 +1333,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "ref_audio file changed during fetch after %d attempts; skipping resolve cache", _REF_AUDIO_METADATA_FETCH_ATTEMPTS, ) - assert wav_list is not None and post_key - return wav_list, sr, post_key + assert waveform is not None and post_key + return waveform, sr, post_key @staticmethod def _make_ref_audio_artifact_cache_key(wav: np.ndarray, sr: int) -> str: @@ -1345,12 +1359,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._ref_audio_resolve_cache.move_to_end(cache_key) return cached[3] - def _put_resolved_ref_audio(self, cache_key: str, wav_list: list[float], sr: int, artifact_key: str) -> None: + def _put_resolved_ref_audio(self, cache_key: str, waveform: np.ndarray, sr: int, artifact_key: str) -> None: if self._ref_audio_resolve_cache_max_entries <= 0 or self._ref_audio_resolve_cache_max_bytes <= 0: return - # Approximate list[float] storage. CPython float objects add per-element - # overhead, so max_entries remains the hard cache cap. - size = len(wav_list) * 40 + size = waveform.nbytes if size > self._ref_audio_resolve_cache_max_bytes: return previous = self._ref_audio_resolve_cache.pop(cache_key, None) @@ -1358,7 +1370,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._ref_audio_resolve_cache_bytes -= previous[2] if previous[3] != artifact_key: self._discard_ref_audio_artifact_ready_if_unreferenced(previous[3]) - self._ref_audio_resolve_cache[cache_key] = (wav_list, int(sr), size, artifact_key) + self._ref_audio_resolve_cache[cache_key] = (waveform, int(sr), size, artifact_key) self._ref_audio_resolve_cache_bytes += size while len(self._ref_audio_resolve_cache) > self._ref_audio_resolve_cache_max_entries: _, (_, _, old_size, old_artifact_key) = self._ref_audio_resolve_cache.popitem(last=False) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index 5f5d8ecb9..ea7ff29ec 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -53,7 +53,9 @@ class _MossTTSAdapterBase(ARTTSAdapter): return self.ctx.server._speaker_cache async def _resolve_ref_audio(self, ref_audio: str): - return await self.ctx.server._resolve_ref_audio(ref_audio) + if self._moss_variant is None: # Nano sends lists through engine IPC. + return await self.ctx.server._resolve_ref_audio(ref_audio) + return await self.ctx.server._resolve_ref_audio_array(ref_audio) def _voice_created_at(self, voice: str) -> int: return self.ctx.server._voice_created_at(voice) diff --git a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py index a828d59e0..384336ebe 100644 --- a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py @@ -18,7 +18,7 @@ arriving via different locators shares one entry), single-flight (concurrent requests for one uncached clip join a single encode), and micro-batched encoding (cold encodes arriving close together share one processor forward). -Kept import-light (only ``asyncio`` / ``hashlib`` / ``torch`` plus the logger) +Kept import-light (``asyncio`` / ``hashlib`` / ``numpy`` / ``torch`` plus the logger) so importing it from the API-server process does not pull the talker/codec. """ @@ -29,7 +29,9 @@ import hashlib from collections.abc import Awaitable, Callable from typing import Any +import numpy as np import torch +from numpy.typing import NDArray from vllm.logger import init_logger logger = init_logger(__name__) @@ -41,14 +43,16 @@ _REF_ENCODE_MAX_BATCH = 8 _INT32_MAX = 2**31 +_Waveform = NDArray[np.float32] | list[float] + def _sha1(s: str) -> str: return hashlib.sha1((s or "").encode("utf-8")).hexdigest() -def _prep_wav_sync(wav_list: list, sr: int, sr_target: int) -> torch.Tensor: +def _prep_wav_sync(waveform: _Waveform, sr: int, sr_target: int) -> torch.Tensor: """Tensor-ise + resample one clip to ``sr_target`` (the blocking prep).""" - wav = torch.tensor(wav_list, dtype=torch.float32) + wav = torch.tensor(waveform, dtype=torch.float32) if wav.dim() == 1: wav = wav.unsqueeze(0) if sr != sr_target: @@ -101,12 +105,12 @@ class _RefEncodeBatcher: def __init__( self, - encode_batch_fn: Callable[[list[tuple[list, int]]], list], + encode_batch_fn: Callable[[list[tuple[_Waveform, int]]], list], *, window_ms: float, max_batch: int, ): - # encode_batch_fn: sync, takes [(wav_list, sr), ...], returns a list of + # encode_batch_fn: sync, takes [(waveform, sr), ...], returns a list of # (codes_tensor | Exception) aligned to the input order. self._encode_batch_fn = encode_batch_fn self._window_s = max(0.0, float(window_ms) / 1000.0) @@ -120,18 +124,23 @@ class _RefEncodeBatcher: if self._drainer is None or self._drainer.done(): self._drainer = asyncio.create_task(self._drain_loop()) - async def submit(self, wav_list: list, sr: int) -> torch.Tensor: + async def submit(self, waveform: _Waveform, sr: int) -> torch.Tensor: self._ensure_started() fut: asyncio.Future = asyncio.get_running_loop().create_future() - self._queue.put_nowait((wav_list, sr, fut)) # type: ignore[union-attr] + self._queue.put_nowait((waveform, sr, fut)) # type: ignore[union-attr] return await fut async def _drain_loop(self) -> None: assert self._queue is not None while True: first = await self._queue.get() - jobs = await self._coalesce(first) - await self._run_batch(jobs) + jobs = None + try: + jobs = await self._coalesce(first) + await self._run_batch(jobs) + finally: + # Do not retain the completed batch while waiting for new work. + del first, jobs async def _coalesce(self, first: tuple) -> list[tuple]: """Group ``first`` with jobs arriving within the batch window.""" @@ -155,8 +164,8 @@ class _RefEncodeBatcher: break return jobs - async def _run_batch(self, jobs: list[tuple[list, int, asyncio.Future]]) -> None: - payload = [(wav_list, sr) for wav_list, sr, _ in jobs] + async def _run_batch(self, jobs: list[tuple[_Waveform, int, asyncio.Future]]) -> None: + payload = [(waveform, sr) for waveform, sr, _ in jobs] futs = [fut for _, _, fut in jobs] try: results = await asyncio.to_thread(self._encode_batch_fn, payload) @@ -215,14 +224,14 @@ class MossReferenceEncoder: self, ref_str: str, *, - resolve_ref_audio: Callable[[str], Awaitable[tuple[list, int, str]]], + resolve_ref_audio: Callable[[str], Awaitable[tuple[_Waveform, int, str]]], get_artifact_key: Callable[[str], str | None], voice_name: str | None = None, voice_created_at: int = 0, ) -> tuple[torch.Tensor, str | None]: """Encode one reference clip into MOSS RVQ codes, reusing the cache. - ``resolve_ref_audio`` maps ``ref_str`` to ``(wav_list, sr, cache_key)`` + ``resolve_ref_audio`` maps ``ref_str`` to ``(waveform, sr, cache_key)`` where *cache_key* is the content-aware resolve key (it folds mtime/size for local files). ``get_artifact_key`` maps that resolve key to the waveform-content artifact key, or ``None`` when unknown. @@ -291,13 +300,13 @@ class MossReferenceEncoder: async def _resolve_and_encode( self, ref_str: str, - resolve_ref_audio: Callable[[str], Awaitable[tuple[list, int, str]]], + resolve_ref_audio: Callable[[str], Awaitable[tuple[_Waveform, int, str]]], get_artifact_key: Callable[[str], str | None], voice_name: str | None, created_at: int, ) -> tuple[torch.Tensor, str]: """Flight body: resolve → re-check cache by content hash → batch-encode.""" - wav_list, sr, resolve_key = await resolve_ref_audio(ref_str) + waveform, sr, resolve_key = await resolve_ref_audio(ref_str) # The content hash is available now that the clip is resolved; this also # catches the case where another flight populated the cache in between. @@ -312,7 +321,7 @@ class MossReferenceEncoder: if cached is not None: return cached["codes"], resolve_key - codes = await self._batcher.submit(wav_list, sr) + codes = await self._batcher.submit(waveform, sr) compact = _to_compact_codes(codes) self._speaker_cache.put(key, {"codes": compact}) logger.debug( @@ -323,15 +332,15 @@ class MossReferenceEncoder: ) return compact, resolve_key - def _encode_batch_sync(self, payload: list[tuple[list, int]]) -> list: + def _encode_batch_sync(self, payload: list[tuple[_Waveform, int]]) -> list: """Worker-thread body: prep each clip, then one batched forward.""" n = len(payload) results: list = [None] * n prepared: list[torch.Tensor] = [] prepared_idx: list[int] = [] - for i, (wav_list, sr) in enumerate(payload): + for i, (waveform, sr) in enumerate(payload): try: - prepared.append(_prep_wav_sync(wav_list, sr, self._sr_target)) + prepared.append(_prep_wav_sync(waveform, sr, self._sr_target)) prepared_idx.append(i) except Exception as exc: # noqa: BLE001 — isolate this clip's failure results[i] = exc @@ -416,7 +425,7 @@ async def encode_request_references( ref_audio: str, ref_audio_2: str | None = None, *, - resolve_ref_audio: Callable[[str], Awaitable[tuple[list, int, str]]], + resolve_ref_audio: Callable[[str], Awaitable[tuple[_Waveform, int, str]]], get_artifact_key: Callable[[str], str | None], voice_name: str | None = None, voice_created_at: int = 0,