From 1e74807c356be5cba4239faae7c38763fe1a44be Mon Sep 17 00:00:00 2001 From: BruceLoveDecimal <39156883+BruceLoveDecimal@users.noreply.github.com> Date: Sat, 5 Sep 2026 07:36:26 +0800 Subject: [PATCH] [Feature][Moss-TTS] Reference Audio Preprocessing (#4982) Signed-off-by: BruceLoveDecimal <39156883+BruceLoveDecimal@users.noreply.github.com> Signed-off-by: liuqihao Co-authored-by: Claude Opus 5 --- .../openai_api/test_serving_speech.py | 121 ++-- .../models/moss_tts/test_reference_encoder.py | 536 ++++++++++++++++++ .../entrypoints/openai/serving_speech.py | 130 +++-- .../models/moss_tts/reference_encoder.py | 484 +++++++++++++--- 4 files changed, 1072 insertions(+), 199 deletions(-) create mode 100644 tests/model_executor/models/moss_tts/test_reference_encoder.py diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 028d84f08..47c11d740 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1711,16 +1711,16 @@ class TestTTSMethods: assert prompt_a["additional_information"]["ref_audio_cache_key"] == "key_aaa" assert prompt_b["additional_information"]["ref_audio_cache_key"] == "key_bbb" - # ── encode_reference_codes: speaker cache invalidation ── + # ── MossReferenceEncoder: speaker cache invalidation ── @pytest.mark.asyncio - async def test_encode_reference_codes_reencodes_on_key_change(self): + async def test_moss_ref_encoder_reencodes_on_key_change(self): """Changing the resolve_cache_key must trigger a fresh encode, not serve the stale cached tensor. This is the user-visible fix for delay-family voice clones that were reproducing the old speaker.""" from unittest.mock import MagicMock - from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes + from vllm_omni.model_executor.models.moss_tts.reference_encoder import MossReferenceEncoder from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache speaker_cache = SpeakerEmbeddingCache(max_bytes=64 * 1024 * 1024) @@ -1745,38 +1745,42 @@ class TestTTSMethods: async def mock_resolve(ref_str): return ([0.1, 0.2, 0.3], 24000, next(resolve_keys)) - # First call: cold miss, encodes and stores - result1 = await encode_reference_codes( - "data:audio/wav;base64,fake", - processor=processor, - resolve_ref_audio=mock_resolve, - speaker_cache=speaker_cache, + encoder = MossReferenceEncoder( + processor, variant="tts", n_vq=32, sr_target=24000, + speaker_cache=speaker_cache, ) - assert call_count == 1 - assert torch.equal(result1, codes_v1) + try: + # First call: cold miss, encodes and stores + result1, key1 = await encoder.encode( + "data:audio/wav;base64,fake", + resolve_ref_audio=mock_resolve, + get_artifact_key=lambda cache_key: None, + ) + assert call_count == 1 + assert torch.equal(result1, codes_v1) + assert key1 == "key_aaa" - # Second call: different resolve key → must re-encode, not serve cached - result2 = await encode_reference_codes( - "data:audio/wav;base64,fake", - processor=processor, - resolve_ref_audio=mock_resolve, - speaker_cache=speaker_cache, - variant="tts", - n_vq=32, - sr_target=24000, - ) - assert call_count == 2, "Different resolve key should trigger re-encode" - assert torch.equal(result2, codes_v2) + # Second call: different resolve key → must re-encode, not serve cached + result2, key2 = await encoder.encode( + "data:audio/wav;base64,fake", + resolve_ref_audio=mock_resolve, + get_artifact_key=lambda cache_key: None, + ) + assert call_count == 2, "Different resolve key should trigger re-encode" + assert torch.equal(result2, codes_v2) + assert key2 == "key_bbb" + finally: + await encoder.aclose() @pytest.mark.asyncio - async def test_encode_reference_codes_named_voice_skips_resolve_on_hit(self): + async def test_moss_ref_encoder_named_voice_skips_resolve_on_hit(self): """Named-voice cache hit must NOT call resolve_ref_audio.""" from unittest.mock import MagicMock - from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes + from vllm_omni.model_executor.models.moss_tts.reference_encoder import MossReferenceEncoder from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache speaker_cache = SpeakerEmbeddingCache(max_bytes=64 * 1024 * 1024) @@ -1791,33 +1795,37 @@ class TestTTSMethods: resolve_called += 1 return ([0.1], 24000, "some_key") - # First call: cold miss, resolves + encodes - await encode_reference_codes( - "data:audio/wav;base64,fake", - processor=processor, - resolve_ref_audio=mock_resolve, - speaker_cache=speaker_cache, + encoder = MossReferenceEncoder( + processor, variant="tts", n_vq=32, sr_target=24000, - voice_name="alice", - voice_created_at=100, + speaker_cache=speaker_cache, ) - assert resolve_called == 1 + try: + # First call: cold miss, resolves + encodes + _, key1 = await encoder.encode( + "data:audio/wav;base64,fake", + resolve_ref_audio=mock_resolve, + get_artifact_key=lambda cache_key: None, + voice_name="alice", + voice_created_at=100, + ) + assert resolve_called == 1 + assert key1 == "some_key" - # Second call: warm hit, must NOT resolve - await encode_reference_codes( - "data:audio/wav;base64,fake", - processor=processor, - resolve_ref_audio=mock_resolve, - speaker_cache=speaker_cache, - variant="tts", - n_vq=32, - sr_target=24000, - voice_name="alice", - voice_created_at=100, - ) - assert resolve_called == 1, "Named-voice cache hit should skip resolve_ref_audio" + # Second call: warm hit, must NOT resolve + _, key2 = await encoder.encode( + "data:audio/wav;base64,fake", + resolve_ref_audio=mock_resolve, + get_artifact_key=lambda cache_key: None, + voice_name="alice", + voice_created_at=100, + ) + assert resolve_called == 1, "Named-voice cache hit should skip resolve_ref_audio" + assert key2 is None # no resolve happened; salted by voice_created_at + finally: + await encoder.aclose() @pytest.mark.asyncio async def test_ttsd_second_reference_does_not_use_named_voice(self, speech_server, mocker): @@ -1835,14 +1843,12 @@ class TestTTSMethods: seen: list[tuple[str, str | None]] = [] - async def fake_encode(ref_str, **kwargs): - seen.append((ref_str, kwargs.get("voice_name"))) - return torch.zeros(3, dtype=torch.int64) + class _FakeEncoder: + async def encode(self, ref_str, **kwargs): + seen.append((ref_str, kwargs.get("voice_name"))) + return torch.zeros(3, dtype=torch.int64), f"rk:{ref_str}" - mocker.patch( - "vllm_omni.model_executor.models.moss_tts.reference_encoder.encode_reference_codes", - fake_encode, - ) + mocker.patch.object(speech_server, "_get_moss_ref_encoder", return_value=_FakeEncoder()) req = OpenAICreateSpeechRequest( input="hello", @@ -1850,11 +1856,14 @@ class TestTTSMethods: ref_audio="data:audio/wav;base64,aaa", ref_audio_2="data:audio/wav;base64,bbb", ) - await speech_server._build_moss_tts_params(req, has_inline_ref_audio=False) - assert seen == [ + params = await speech_server._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), ] + # Per-slot salt keys survive the concurrent (gather) encode order. + assert params["ref_audio_cache_key"] == "rk:data:audio/wav;base64,aaa" + 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): """Precomputed Qwen3 voices are reusable by name without per-request ref_audio.""" diff --git a/tests/model_executor/models/moss_tts/test_reference_encoder.py b/tests/model_executor/models/moss_tts/test_reference_encoder.py new file mode 100644 index 000000000..bb03d556b --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_reference_encoder.py @@ -0,0 +1,536 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for MossReferenceEncoder: content-addressed caching, +single-flight, and micro-batched encoding.""" + +import asyncio +import threading +import time + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.reference_encoder import ( + MossReferenceEncoder, + _RefEncodeBatcher, + build_reference_encoder, + encode_request_references, +) +from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts] + +_SR = 24000 +_N_VQ = 4 + + +class _FakeProcessor: + """Deterministic stand-in for the MOSS AutoProcessor.""" + + def __init__(self, *, n_vq=_N_VQ, sampling_rate=_SR, delay_s=0.0, fail_if_batched=False): + self.model_config = type("_Cfg", (), {"n_vq": n_vq, "sampling_rate": sampling_rate})() + self.attempt_sizes: list[int] = [] + self.total_items = 0 + self._delay_s = delay_s + self._fail_if_batched = fail_if_batched + self._fail_values: set[int] = set() + self._lock = threading.Lock() + + def set_fail_values(self, values) -> None: + self._fail_values = set(values) + + def encode_audios_from_wav(self, wav_list, sampling_rate, n_vq=None): + if self._delay_s: + time.sleep(self._delay_s) + with self._lock: + self.attempt_sizes.append(len(wav_list)) + if self._fail_if_batched and len(wav_list) > 1: + raise RuntimeError("batched forward not supported") + nq = n_vq or self.model_config.n_vq + out = [] + for wav in wav_list: + cid = round(float(wav.reshape(-1)[0].item())) + if cid in self._fail_values: + raise RuntimeError(f"encode failed for content {cid}") + out.append(torch.full((3, nq), cid, dtype=torch.long)) + with self._lock: + self.total_items += len(wav_list) + return out + + +class _FakeAudio: + """Fake resolve + content-hash side, mirroring serving_speech's contract: + ``resolve`` returns ``(wav, sr, cache_key)`` and ``artifact_key`` maps that + resolve cache key to the waveform-content key (or None when unknown).""" + + def __init__(self, *, resolve_delay_s=0.0): + self._wav: dict[str, tuple] = {} + self._artifact: dict[str, str | None] = {} + self.resolve_calls: list[str] = [] + self._resolve_delay_s = resolve_delay_s + self._lock = threading.Lock() + + def register(self, ref_str, content_id, *, artifact=..., sr=_SR, length=100, wav=...): + wav_list = [float(content_id)] * length if wav is ... else wav + self._wav[ref_str] = (wav_list, sr) + # Default: content hash keyed by content_id (same content -> same key). + self._artifact["rk:" + ref_str] = f"c{content_id}" if artifact is ... else artifact + + async def resolve(self, ref_str): + with self._lock: + self.resolve_calls.append(ref_str) + if self._resolve_delay_s: + await asyncio.sleep(self._resolve_delay_s) + wav_list, sr = self._wav[ref_str] + return wav_list, sr, "rk:" + ref_str + + def artifact_key(self, cache_key): + return self._artifact.get(cache_key) + + +@pytest.fixture +async def make_encoder(): + """Factory yielding encoders; tears each one's drainer down after the test.""" + created: list[MossReferenceEncoder] = [] + + def _make(processor, *, cache=None, window_ms=10.0, max_batch=8, variant="local"): + enc = MossReferenceEncoder( + processor, + variant=variant, + n_vq=_N_VQ, + sr_target=_SR, + speaker_cache=cache or SpeakerEmbeddingCache(max_bytes=8 * 1024**2), + window_ms=window_ms, + max_batch=max_batch, + ) + created.append(enc) + return enc + + yield _make + for enc in created: + await enc.aclose() + + +async def _encode(enc, audio, ref, **kw): + codes, _ = await enc.encode( + ref, + resolve_ref_audio=audio.resolve, + get_artifact_key=audio.artifact_key, + **kw, + ) + return codes + + +# --------------------------------------------------------------------------- # +# Group A — content-addressed cache (R1) # +# --------------------------------------------------------------------------- # + + +async def test_same_reference_encodes_once(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("a", 7) + enc = make_encoder(proc) + + r1 = await _encode(enc, audio, "a") + r2 = await _encode(enc, audio, "a") + + assert proc.total_items == 1 # second call is a cache hit + assert r1.dtype == torch.int64 and r2.dtype == torch.int64 + assert int(r1[0, 0]) == 7 and int(r2[0, 0]) == 7 + + +async def test_returned_tensor_does_not_alias_cache(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("a", 7) + enc = make_encoder(proc) + + r1 = await _encode(enc, audio, "a") + r1[:] = 999 # caller mutates its copy + r2 = await _encode(enc, audio, "a") + + assert int(r2[0, 0]) == 7 # cache untouched + assert proc.total_items == 1 + + +async def test_cached_dtype_is_int32(make_encoder): + cache = SpeakerEmbeddingCache(max_bytes=8 * 1024**2) + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("a", 7) + enc = make_encoder(proc, cache=cache) + + out = await _encode(enc, audio, "a") + + key = cache.make_cache_key("ref:c7", model_type=f"moss_tts_local_nq{_N_VQ}") + stored = cache.get(key) + assert stored is not None + assert stored["codes"].dtype == torch.int32 # compact on-disk + assert out.dtype == torch.int64 # int64 to the caller + + +async def test_same_content_via_different_refs_shares_cache(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + # base64 vs URL for the *same* clip resolve to the same content hash. + audio.register("data:base64,xxx", 7, artifact="c7") + audio.register("http://x/clip.wav", 7, artifact="c7") + enc = make_encoder(proc) + + await _encode(enc, audio, "data:base64,xxx") + await _encode(enc, audio, "http://x/clip.wav") + + assert proc.total_items == 1 # encoded once, shared by content hash + + +async def test_named_voice_created_at_isolation(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("ref", 7) + enc = make_encoder(proc) + + await _encode(enc, audio, "ref", voice_name="alice", voice_created_at=1) + await _encode(enc, audio, "ref", voice_name="alice", voice_created_at=2) + assert proc.total_items == 2 # re-upload (new created_at) re-encodes + + await _encode(enc, audio, "ref", voice_name="alice", voice_created_at=1) + assert proc.total_items == 2 # original generation still cached + + +async def test_named_voice_keys_are_case_insensitive(make_encoder): + """Flight key and cache key must agree on one spelling of the voice name.""" + proc, audio = _FakeProcessor(delay_s=0.02), _FakeAudio() + audio.register("ref", 7) + enc = make_encoder(proc) + + # Concurrent callers differing only by case share one flight... + r1, r2 = await asyncio.gather( + _encode(enc, audio, "ref", voice_name="Alice", voice_created_at=1), + _encode(enc, audio, "ref", voice_name="ALICE", voice_created_at=1), + ) + assert proc.total_items == 1 + assert int(r1[0, 0]) == 7 and int(r2[0, 0]) == 7 + + # ...and populate the slot every later spelling hits. + out = await _encode(enc, audio, "ref", voice_name="alice", voice_created_at=1) + assert proc.total_items == 1 + assert int(out[0, 0]) == 7 + + +async def test_unregistered_default_voice_uses_ref_content_key(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("first", 7, artifact="c7") + audio.register("second", 8, artifact="c8") + enc = make_encoder(proc) + + first = await _encode(enc, audio, "first", voice_name="default", voice_created_at=0) + second = await _encode(enc, audio, "second", voice_name="default", voice_created_at=0) + first_again = await _encode(enc, audio, "first", voice_name="default", voice_created_at=0) + + assert int(first[0, 0]) == 7 + assert int(second[0, 0]) == 8 + assert int(first_again[0, 0]) == 7 + assert proc.total_items == 2 # no collision on the placeholder voice name + + +async def test_artifact_key_unavailable_falls_back_to_ref_hash(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + # artifact=None simulates a cold/disabled resolve cache. + audio.register("ref", 7, artifact=None) + enc = make_encoder(proc) + + await _encode(enc, audio, "ref") + await _encode(enc, audio, "ref") + + # Falls back to the stable resolve cache key, so the second call still hits. + assert proc.total_items == 1 + + +async def test_encode_returns_resolve_key_for_salting(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("a", 7) + audio.register("named", 8) + enc = make_encoder(proc) + + _, k1 = await enc.encode("a", resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key) + _, k2 = await enc.encode("a", resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key) + # Anonymous refs always resolve, so warm hits still return a fresh key + # (the KV prefix-cache salt must track on-disk edits). + assert k1 == "rk:a" and k2 == "rk:a" + + _, k3 = await enc.encode( + "named", + resolve_ref_audio=audio.resolve, + get_artifact_key=audio.artifact_key, + voice_name="alice", + voice_created_at=1, + ) + _, k4 = await enc.encode( + "named", + resolve_ref_audio=audio.resolve, + get_artifact_key=audio.artifact_key, + voice_name="alice", + voice_created_at=1, + ) + assert k3 == "rk:named" + assert k4 is None # named-voice hot hit skips the resolve; salted by created_at + + +# --------------------------------------------------------------------------- # +# Group B — single-flight (R2) # +# --------------------------------------------------------------------------- # + + +async def test_concurrent_same_ref_encodes_once(make_encoder): + proc, audio = _FakeProcessor(delay_s=0.05), _FakeAudio() + audio.register("a", 7) + enc = make_encoder(proc) + + results = await asyncio.gather(*[_encode(enc, audio, "a") for _ in range(10)]) + + assert proc.total_items == 1 # single-flight collapsed 10 -> 1 encode + assert audio.resolve_calls == ["a"] # and one resolve/download + assert all(int(r[0, 0]) == 7 for r in results) + + +async def test_concurrent_results_are_isolated(make_encoder): + proc, audio = _FakeProcessor(delay_s=0.02), _FakeAudio() + audio.register("a", 7) + enc = make_encoder(proc) + + r1, r2 = await asyncio.gather(_encode(enc, audio, "a"), _encode(enc, audio, "a")) + r1[:] = 111 + assert int(r2[0, 0]) == 7 # separate storage per waiter + + +async def test_flight_exception_propagates_and_is_not_cached(make_encoder): + proc, audio = _FakeProcessor(delay_s=0.02), _FakeAudio() + audio.register("a", 9) + proc.set_fail_values({9}) + enc = make_encoder(proc) + + results = await asyncio.gather(*[_encode(enc, audio, "a") for _ in range(6)], return_exceptions=True) + assert all(isinstance(r, Exception) for r in results) # every waiter sees it + + proc.set_fail_values(set()) # failure was not cached: a retry now succeeds + ok = await _encode(enc, audio, "a") + assert int(ok[0, 0]) == 9 + + +async def test_initiator_cancel_does_not_break_other_waiter(make_encoder): + proc, audio = _FakeProcessor(delay_s=0.1), _FakeAudio() + audio.register("a", 7) + enc = make_encoder(proc) + + a = asyncio.create_task(_encode(enc, audio, "a")) + await asyncio.sleep(0.02) # let A create the flight + b = asyncio.create_task(_encode(enc, audio, "a")) + await asyncio.sleep(0.02) # let B join the same flight + a.cancel() + + res_b = await b + assert int(res_b[0, 0]) == 7 # B still completes off the shared flight + assert proc.total_items == 1 + with pytest.raises(asyncio.CancelledError): + await a + + +# --------------------------------------------------------------------------- # +# Group C — micro-batched encoding (R3) # +# --------------------------------------------------------------------------- # + + +async def test_distinct_refs_coalesce_into_one_batch(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + for i, ref in enumerate(("a", "b", "c")): + audio.register(ref, 10 + i) + enc = make_encoder(proc, window_ms=100.0, max_batch=8) + + results = await asyncio.gather(*[_encode(enc, audio, r) for r in ("a", "b", "c")]) + + assert proc.attempt_sizes == [3] # one batched forward of all three + assert sorted(int(r[0, 0]) for r in results) == [10, 11, 12] + + +async def test_max_batch_splits_work(): + # Drive the batcher directly: with everything already queued, max_batch=2 + # must split 4 submissions into batches of at most 2. + seen: list[int] = [] + + def encode_batch(payload): + seen.append(len(payload)) + return [torch.zeros(1) for _ in payload] + + batcher = _RefEncodeBatcher(encode_batch, window_ms=5.0, max_batch=2) + try: + await asyncio.gather(*[batcher.submit([1.0], _SR) for _ in range(4)]) + finally: + await batcher.aclose() + + assert sum(seen) == 4 + assert max(seen) == 2 and all(s <= 2 for s in seen) + + +async def test_single_submission_still_runs(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("a", 7) + enc = make_encoder(proc, window_ms=20.0) + + out = await _encode(enc, audio, "a") # not starved waiting for a full batch + assert int(out[0, 0]) == 7 + assert proc.attempt_sizes == [1] + + +async def test_window_zero_runs_immediately(): + seen: list[int] = [] + + def encode_batch(payload): + seen.append(len(payload)) + return [torch.zeros(1) for _ in payload] + + batcher = _RefEncodeBatcher(encode_batch, window_ms=0.0, max_batch=8) + try: + out = await batcher.submit([1.0], _SR) + finally: + await batcher.aclose() + assert out is not None + assert seen == [1] + + +def test_prep_failure_is_isolated(make_encoder_sync): + enc = make_encoder_sync() + good = [1.0] * 100 + bad = None # torch.tensor(None) raises inside prep + + results = enc._encode_batch_sync([(good, _SR), (bad, _SR), (good, _SR)]) + + assert isinstance(results[0], torch.Tensor) + assert isinstance(results[1], Exception) # only the bad clip fails + assert isinstance(results[2], torch.Tensor) + + +def test_batch_forward_failure_falls_back_per_item(make_encoder_sync): + proc = _FakeProcessor(fail_if_batched=True) + enc = make_encoder_sync(proc) + + results = enc._encode_batch_sync([([float(i)] * 100, _SR) for i in (1, 2, 3)]) + + assert all(isinstance(r, torch.Tensor) for r in results) + assert [int(r[0, 0]) for r in results] == [1, 2, 3] + assert any(s > 1 for s in proc.attempt_sizes) # a batched attempt happened + assert proc.attempt_sizes.count(1) == 3 # then per-item fallback + + +@pytest.fixture +def make_encoder_sync(): + """Encoder factory for the synchronous ``_encode_batch_sync`` tests.""" + + def _make(processor=None): + return MossReferenceEncoder( + processor or _FakeProcessor(), + variant="local", + n_vq=_N_VQ, + sr_target=_SR, + speaker_cache=SpeakerEmbeddingCache(max_bytes=8 * 1024**2), + ) + + return _make + + +# --------------------------------------------------------------------------- # +# Group D — integration with the real SpeakerEmbeddingCache # +# --------------------------------------------------------------------------- # + + +async def test_model_type_namespaces_cache(make_encoder): + # Two variants sharing one cache must not collide on the same ref. + cache = SpeakerEmbeddingCache(max_bytes=8 * 1024**2) + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("a", 7) + enc_tts = make_encoder(proc, cache=cache, variant="tts") + enc_ttsd = make_encoder(proc, cache=cache, variant="ttsd") + + await _encode(enc_tts, audio, "a") + await _encode(enc_ttsd, audio, "a") + + assert proc.total_items == 2 # distinct model_type keys -> no cross-hit + assert cache.stats()["entries"] == 2 + + +# --------------------------------------------------------------------------- # +# Group E — serving-layer entry points moved into the model package # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("variant", "expected_sr"), + [ + # Local-v1.5 encodes references at a fixed 24 kHz working rate even + # though its output codec is 48 kHz stereo. + ("local", 24000), + ("tts", 48000), + ("ttsd", 48000), + ], +) +def test_build_reference_encoder_derives_geometry(variant, expected_sr): + proc = _FakeProcessor(n_vq=24, sampling_rate=48000) + enc = build_reference_encoder( + proc, + variant=variant, + speaker_cache=SpeakerEmbeddingCache(max_bytes=8 * 1024**2), + ) + assert enc._n_vq == 24 + assert enc._sr_target == expected_sr + + +def test_build_reference_encoder_falls_back_without_model_config_fields(): + proc = _FakeProcessor() + proc.model_config = type("_Empty", (), {})() + enc = build_reference_encoder( + proc, + variant="tts", + speaker_cache=SpeakerEmbeddingCache(max_bytes=8 * 1024**2), + ) + assert enc._n_vq == 32 + assert enc._sr_target == 24000 + + +async def test_encode_request_references_single_speaker(make_encoder): + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("a", 1) + enc = make_encoder(proc) + + refs, resolve_keys = await encode_request_references( + enc, + "a", + None, + resolve_ref_audio=audio.resolve, + get_artifact_key=audio.artifact_key, + ) + + assert len(refs) == 1 + assert set(resolve_keys) == {0} + + +async def test_encode_request_references_keys_second_speaker_by_slot(make_encoder): + """Two-speaker resolve keys must land in slot order, not completion order, + and speaker 2 must stay content-addressed even under a named voice.""" + proc, audio = _FakeProcessor(), _FakeAudio() + audio.register("a", 1) + audio.register("b", 2) + enc = make_encoder(proc, variant="ttsd", window_ms=100.0) + + refs, resolve_keys = await encode_request_references( + enc, + "a", + "b", + resolve_ref_audio=audio.resolve, + get_artifact_key=audio.artifact_key, + voice_name="alice", + voice_created_at=7, + ) + + assert len(refs) == 2 + assert resolve_keys[0] == "rk:a" + assert resolve_keys[1] == "rk:b" + # Both speakers encoded; speaker 2 did not reuse the named-voice entry. + assert not torch.equal(refs[0], refs[1]) + assert proc.total_items == 2 + # Concurrent, so both clips shared one batch window. + assert proc.attempt_sizes == [2] diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 2eabe14c0..813d916fe 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1403,6 +1403,24 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._moss_processor_cache = proc return proc + def _get_moss_ref_encoder(self): + """Lazily build the per-server MOSS reference-audio encoder once.""" + cached = getattr(self, "_moss_ref_encoder", None) + if cached is not None: + return cached + from vllm_omni.model_executor.models.moss_tts.reference_encoder import build_reference_encoder + + # The variant's encode geometry (n_vq, working sample rate) is derived + # inside the model package; this layer only supplies the processor and + # the process-wide speaker cache. + encoder = build_reference_encoder( + self._get_moss_processor(), + variant=self._moss_variant, + speaker_cache=self._speaker_cache, + ) + self._moss_ref_encoder = encoder + return encoder + async def _build_moss_tts_params( self, request: OpenAICreateSpeechRequest, @@ -1474,63 +1492,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # consumes, so skipping this path silently drops all voice-clone # conditioning and produces unconditioned/garbage audio online). ---- proc = self._get_moss_processor() - n_vq = int(getattr(proc.model_config, "n_vq", 32)) - # Local-v1.5 encodes reference audio at a fixed 24 kHz working rate - # regardless of its 48 kHz stereo *output* codec -- mirrors the - # offline example's hardcoded encode_audios_from_wav(sampling_rate=24000) - # for this variant; proc.model_config.sampling_rate there is the - # output rate (48000), the wrong value to resample the reference into. - sr_target = 24000 if v == "local" else int(getattr(proc.model_config, "sampling_rate", 24000)) - - # Reference-audio encoding + speaker caching lives in the model package - # (moss_tts.reference_encoder), mirroring Fish Speech / CosyVoice3 / - # Qwen3-TTS which keep reference handling with the model rather than in - # this shared serving file. Imported lazily so the API-server process - # only pulls it on the delay-family path (alongside the upstream proc). - from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes - - # Named-voice speaker cache is only valid for uploaded speakers - # without an inline ref_audio. ``request.voice`` plus a file/URL - # otherwise keys on (name, created_at=0) and skips the content-aware - # resolve — the adapter already applies this guard when tagging - # ``voice_name`` on tts_params. - _raw_voice = getattr(request, "voice", None) - _raw_voice = _raw_voice.strip() if isinstance(_raw_voice, str) else "" - _voice_lower = _raw_voice.lower() if _raw_voice else "" - _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 - _resolve_keys: list[str] = [] - - async def _tracked_resolve(ref_str: str) -> tuple[list, int, str]: - wav_list, sr, cache_key = await self._resolve_ref_audio(ref_str) - _resolve_keys.append(cache_key) - return wav_list, sr, cache_key - - async def _encode_ref(ref_str: str, *, named_voice: bool) -> torch.Tensor: - # Named-voice cache is (voice_name, created_at). TTSD speaker 2 is a - # different clip, so it must stay anonymous or it reuses speaker 1. - use_named = named_voice and bool(_voice) - return await encode_reference_codes( - ref_str, - processor=proc, - resolve_ref_audio=_tracked_resolve, - speaker_cache=self._speaker_cache, - variant=v, - n_vq=n_vq, - sr_target=sr_target, - voice_name=_voice if use_named else None, - voice_created_at=_voice_created if use_named else 0, - ) user_kwargs: dict[str, Any] = {"text": request.input or ""} - if v in ("tts", "local"): - user_kwargs["reference"] = [await _encode_ref(request.ref_audio, named_voice=True)] - elif v == "ttsd": - refs = [await _encode_ref(request.ref_audio, named_voice=True)] - if request.ref_audio_2: - refs.append(await _encode_ref(request.ref_audio_2, named_voice=False)) - user_kwargs["reference"] = refs + resolve_keys: dict[int, str] = {} + if v in ("tts", "local", "ttsd"): + user_kwargs["reference"], resolve_keys = await self._encode_moss_references( + request, + has_inline_ref_audio=has_inline_ref_audio, + two_speaker=(v == "ttsd"), + ) elif v == "sound_effect": user_kwargs["text"] = request.input or "" # may be empty user_kwargs["ambient_sound"] = request.ambient_sound or "" @@ -1563,12 +1533,54 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): } if request.max_new_tokens is not None: params["max_new_frames"] = [request.max_new_tokens] - if _resolve_keys: - params["ref_audio_cache_key"] = _resolve_keys[0] - if len(_resolve_keys) > 1: - params["ref_audio_2_cache_key"] = _resolve_keys[1] + if 0 in resolve_keys: + params["ref_audio_cache_key"] = resolve_keys[0] + if 1 in resolve_keys: + params["ref_audio_2_cache_key"] = resolve_keys[1] return params + async def _encode_moss_references( + self, + request: OpenAICreateSpeechRequest, + *, + has_inline_ref_audio: bool, + two_speaker: bool, + ) -> tuple[list, dict[int, str]]: + """Encode the request's reference clip(s) into MOSS RVQ code tensors. + + Reference encoding + speaker caching lives in the model package + (moss_tts.reference_encoder), mirroring Fish Speech / CosyVoice3 / + Qwen3-TTS which keep reference handling with the model rather than in + this shared serving file. This method only resolves the serving-side + inputs the model helper needs: the generic audio resolver/artifact-key + lookups and whether the request may use the named-voice cache. + + Returns ``(codes_per_speaker, resolve_keys)`` where ``resolve_keys`` + maps the reference slot (0 = ref_audio, 1 = ref_audio_2) to its + content-aware resolve key, for salting the KV prefix cache. + """ + 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 + + return await encode_request_references( + self._get_moss_ref_encoder(), + request.ref_audio, + request.ref_audio_2 if two_speaker else None, + resolve_ref_audio=self._resolve_ref_audio, + get_artifact_key=self._get_resolved_ref_audio_artifact_key, + voice_name=voice or None, + voice_created_at=voice_created, + ) + async def _build_higgs_audio_v2_params(self, request: OpenAICreateSpeechRequest): """Build prompt_token_ids for higgs_audio_v2 via the upstream processor. 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 ef1abfa86..36b2616e7 100644 --- a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py @@ -4,17 +4,25 @@ This lives in the model package (not the shared serving layer) so all MOSS-specific reference handling stays with the model — mirroring how Fish Speech (``dac_encoder.encode_reference_audio_codes``), CosyVoice3, and Qwen3-TTS keep their reference/speaker extraction next to the model rather than -in ``serving_speech.py``. The serving layer just calls -:func:`encode_reference_codes` with its generic helpers (the audio resolver and -the process-wide speaker cache). +in ``serving_speech.py``. The serving layer constructs one +:class:`MossReferenceEncoder` per server (lazily, alongside the upstream MOSS +processor) and calls :meth:`MossReferenceEncoder.encode` with its generic +helpers (the audio resolver, the artifact-key lookup, and the process-wide +speaker cache). -Kept import-light (only ``asyncio`` / ``torch`` plus the logger) +On top of that cache the encoder adds content-addressed keys (the same clip +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) so importing it from the API-server process does not pull the talker/codec. """ from __future__ import annotations import asyncio +import hashlib from collections.abc import Awaitable, Callable from typing import Any @@ -23,9 +31,20 @@ from vllm.logger import init_logger logger = init_logger(__name__) +# Coalescing defaults. Kept as module constants (not env/CLI) to match the +# ``_REF_AUDIO_RESOLVE_CACHE_MAX_*`` convention in serving_speech.py. +_REF_ENCODE_BATCH_WINDOW_MS = 10.0 +_REF_ENCODE_MAX_BATCH = 8 -def _encode_wav_sync(processor: Any, wav_list: list, sr: int, sr_target: int, n_vq: int) -> torch.Tensor: - """Blocking resample + CPU codec encode (the expensive bit).""" +_INT32_MAX = 2**31 + + +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: + """Tensor-ise + resample one clip to ``sr_target`` (the blocking prep).""" wav = torch.tensor(wav_list, dtype=torch.float32) if wav.dim() == 1: wav = wav.unsqueeze(0) @@ -33,93 +52,390 @@ def _encode_wav_sync(processor: Any, wav_list: list, sr: int, sr_target: int, n_ import torchaudio wav = torchaudio.functional.resample(wav, sr, sr_target) - with torch.no_grad(): - codes_list = processor.encode_audios_from_wav([wav], sampling_rate=sr_target, n_vq=n_vq) - return codes_list[0] + return wav -async def encode_reference_codes( - ref_str: str, - *, +def _to_compact_codes(codes: torch.Tensor) -> torch.Tensor: + """Downcast RVQ codes to int32 for compact caching.""" + codes = codes.detach().cpu().contiguous() + if codes.numel() == 0: + return codes.to(torch.int32) + hi = int(codes.max().item()) + lo = int(codes.min().item()) + if -_INT32_MAX <= lo and hi < _INT32_MAX: + return codes.to(torch.int32) + logger.warning("MOSS ref codes out of int32 range (min=%d max=%d); caching as int64", lo, hi) + return codes.to(torch.int64) + + +def _clone_out(codes: torch.Tensor) -> torch.Tensor: + """Return an independent int64 copy for the caller.""" + return codes.to(torch.int64, copy=True) + + +def _registered_voice(voice_name: str | None, voice_created_at: int) -> tuple[str | None, int]: + """Return ``(name, created_at)`` for a registered uploaded voice, else ``(None, 0)``. + + The OpenAI speech API requires a ``voice`` field, and callers often send + placeholders such as "default" for ref-audio voice cloning. Only registered + uploaded voices have a positive created_at timestamp; other names must not + key the cache because the timbre comes from ref_audio. + + The name is lowercased here so every derived key agrees on one spelling: + if the cache key kept the caller's casing while the flight key normalized + it, two concurrent callers differing only by case would share a flight but + populate different cache slots. + """ + name = voice_name.strip().lower() if isinstance(voice_name, str) else "" + created_at = int(voice_created_at) + if name and created_at > 0: + return name, created_at + return None, 0 + + +class _RefEncodeBatcher: + """Coalesce cold reference encodes into batched processor forwards.""" + + def __init__( + self, + encode_batch_fn: Callable[[list[tuple[list, int]]], list], + *, + window_ms: float, + max_batch: int, + ): + # encode_batch_fn: sync, takes [(wav_list, 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) + self._max_batch = max(1, int(max_batch)) + self._queue: asyncio.Queue | None = None + self._drainer: asyncio.Task | None = None + + def _ensure_started(self) -> None: + if self._queue is None: + self._queue = asyncio.Queue() + 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: + self._ensure_started() + fut: asyncio.Future = asyncio.get_running_loop().create_future() + self._queue.put_nowait((wav_list, 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) + + async def _coalesce(self, first: tuple) -> list[tuple]: + """Group ``first`` with jobs arriving within the batch window.""" + jobs = [first] + if self._window_s > 0: + deadline = asyncio.get_running_loop().time() + self._window_s + while len(jobs) < self._max_batch: + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + break + try: + jobs.append(await asyncio.wait_for(self._queue.get(), remaining)) + except asyncio.TimeoutError: + break + else: + # window=0: coalesce only what is already queued, never wait. + while len(jobs) < self._max_batch: + try: + jobs.append(self._queue.get_nowait()) + except asyncio.QueueEmpty: + 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] + futs = [fut for _, _, fut in jobs] + try: + results = await asyncio.to_thread(self._encode_batch_fn, payload) + except Exception as exc: # noqa: BLE001 — propagate the batch failure to every waiter + for fut in futs: + if not fut.done(): + fut.set_exception(exc) + return + for fut, res in zip(futs, results): + if fut.done(): + continue + if isinstance(res, BaseException): + fut.set_exception(res) + else: + fut.set_result(res) + + async def aclose(self) -> None: + if self._drainer is not None: + self._drainer.cancel() + try: + await self._drainer + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + self._drainer = None + + +class MossReferenceEncoder: + """Content-addressed, single-flight, micro-batched reference-audio encoder.""" + + def __init__( + self, + processor: Any, + *, + variant: str, + n_vq: int, + sr_target: int, + speaker_cache: Any, + window_ms: float = _REF_ENCODE_BATCH_WINDOW_MS, + max_batch: int = _REF_ENCODE_MAX_BATCH, + ): + self._processor = processor + self._n_vq = int(n_vq) + self._sr_target = int(sr_target) + self._speaker_cache = speaker_cache + # ``created_at`` and the audio-content name vary per request; the + # model_type namespaces the whole family so a moss_tts server never + # collides with another model's speaker-cache entries. + self._model_type = f"moss_tts_{variant}_nq{int(n_vq)}" + self._inflight: dict[str, asyncio.Task] = {} + self._batcher = _RefEncodeBatcher(self._encode_batch_sync, window_ms=window_ms, max_batch=max_batch) + + def _make_cache_key(self, name: str, created_at: int) -> tuple: + return self._speaker_cache.make_cache_key(name, model_type=self._model_type, created_at=int(created_at)) + + async def encode( + self, + ref_str: str, + *, + resolve_ref_audio: Callable[[str], Awaitable[tuple[list, 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)`` + 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. + + Returns ``(codes, resolve_cache_key)``; the key is ``None`` when the + clip was served from the named-voice cache without resolving (the + caller salts those requests with ``voice_created_at`` instead). + """ + voice_name, created_at = _registered_voice(voice_name, voice_created_at) + + if voice_name: + # A named voice has a stable key that does not depend on the + # resolved audio, so the cache can be checked before resolving. + flight_key = f"voice:{voice_name}:{created_at}" + cached = self._speaker_cache.get(self._make_cache_key(voice_name, created_at)) + if cached is not None: + return _clone_out(cached["codes"]), None + else: + # Anonymous refs have no pre-resolve hot path: the content key must + # come from the resolve itself (mtime/size from a single stat) so an + # on-disk edit invalidates the cached codes; the flight body + # re-checks the speaker cache right after resolving, which is cheap + # when the resolve cache is warm. The flight key is the request-side + # reference (not the content hash), so concurrent requests for the + # same ref_str also share the resolve/download, not just the encode. + flight_key = "ref:" + _sha1(ref_str) + + codes, resolve_key = await self._single_flight( + flight_key, + lambda: self._resolve_and_encode(ref_str, resolve_ref_audio, get_artifact_key, voice_name, created_at), + ) + return _clone_out(codes), resolve_key + + async def _single_flight( + self, + flight_key: str, + start_flight: Callable[[], Awaitable[tuple[torch.Tensor, str]]], + ) -> tuple[torch.Tensor, str]: + """Join the in-flight encode for ``flight_key``, starting one if absent. + + The shared task is awaited through ``shield`` so a caller cancelling + its own request does not cancel the flight (asyncio would otherwise + propagate the cancel into the awaited task and take down every other + waiter with it). + """ + task = self._inflight.get(flight_key) + if task is not None: + return await asyncio.shield(task) + + task = asyncio.create_task(start_flight()) + self._inflight[flight_key] = task + + # Retire the slot when the flight *completes*, not when the creating + # caller returns: if the creator is cancelled the shielded task keeps + # running, and popping the slot early would let the next arrival start + # a duplicate resolve/encode instead of joining this one. Identity + # guard: only drop the slot if it still holds *our* task (a later + # request may have replaced it after ours completed). + def _retire(t: asyncio.Task, key: str = flight_key) -> None: + if self._inflight.get(key) is t: + self._inflight.pop(key, None) + + task.add_done_callback(_retire) + return await asyncio.shield(task) + + async def _resolve_and_encode( + self, + ref_str: str, + resolve_ref_audio: Callable[[str], Awaitable[tuple[list, 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) + + # The content hash is available now that the clip is resolved; this also + # catches the case where another flight populated the cache in between. + if voice_name: + key = self._make_cache_key(voice_name, created_at) + else: + artifact_key = get_artifact_key(resolve_key) + key_name = ("ref:" + artifact_key) if artifact_key else ("ref:" + resolve_key) + key = self._make_cache_key(key_name, 0) + + cached = self._speaker_cache.get(key) + if cached is not None: + return cached["codes"], resolve_key + + codes = await self._batcher.submit(wav_list, sr) + compact = _to_compact_codes(codes) + self._speaker_cache.put(key, {"codes": compact}) + logger.debug( + "MOSS ref encode STORE key=%s shape=%s dtype=%s", + key, + tuple(compact.shape), + compact.dtype, + ) + return compact, resolve_key + + def _encode_batch_sync(self, payload: list[tuple[list, 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): + try: + prepared.append(_prep_wav_sync(wav_list, sr, self._sr_target)) + prepared_idx.append(i) + except Exception as exc: # noqa: BLE001 — isolate this clip's failure + results[i] = exc + if not prepared: + return results + + try: + codes_list = self._encode_prepared(prepared) + for local_i, orig_i in enumerate(prepared_idx): + results[orig_i] = codes_list[local_i] + except Exception: # noqa: BLE001 — batch forward failed; retry item-by-item + logger.warning("MOSS ref batch encode (n=%d) failed; falling back to per-item", len(prepared)) + for local_i, orig_i in enumerate(prepared_idx): + try: + results[orig_i] = self._encode_prepared([prepared[local_i]])[0] + except Exception as exc: # noqa: BLE001 — isolate this clip's failure + results[orig_i] = exc + return results + + def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]: + with torch.no_grad(): + return self._processor.encode_audios_from_wav(prepared, sampling_rate=self._sr_target, n_vq=self._n_vq) + + async def aclose(self) -> None: + """Release the batcher drainer + any in-flight encodes (tests/shutdown).""" + for task in list(self._inflight.values()): + task.cancel() + self._inflight.clear() + await self._batcher.aclose() + + +def build_reference_encoder( processor: Any, - resolve_ref_audio: Callable[[str], Awaitable[tuple[list, int, str]]], - speaker_cache: Any, + *, variant: str, - n_vq: int, - sr_target: int, + speaker_cache: Any, +) -> MossReferenceEncoder: + """Build the per-server encoder for a MOSS-TTS ``variant``. + + Derives the encode geometry (``n_vq`` and the working sample rate) from the + upstream processor's ``model_config`` so the variant knowledge stays in the + model package rather than in ``serving_speech.py``. + """ + n_vq = int(getattr(processor.model_config, "n_vq", 32)) + # Local-v1.5 encodes reference audio at a fixed 24 kHz working rate + # regardless of its 48 kHz stereo *output* codec -- mirrors the offline + # example's hardcoded encode_audios_from_wav(sampling_rate=24000) for this + # variant; proc.model_config.sampling_rate there is the output rate + # (48000), the wrong value to resample the reference into. + sr_target = 24000 if variant == "local" else int(getattr(processor.model_config, "sampling_rate", 24000)) + return MossReferenceEncoder( + processor, + variant=variant, + n_vq=n_vq, + sr_target=sr_target, + speaker_cache=speaker_cache, + ) + + +async def encode_request_references( + encoder: MossReferenceEncoder, + ref_audio: str, + ref_audio_2: str | None = None, + *, + resolve_ref_audio: Callable[[str], Awaitable[tuple[list, int, str]]], + get_artifact_key: Callable[[str], str | None], voice_name: str | None = None, voice_created_at: int = 0, -) -> torch.Tensor: - """Encode one reference clip into MOSS RVQ codes, reusing the speaker cache. +) -> tuple[list[torch.Tensor], dict[int, str]]: + """Encode a request's reference clip(s) into MOSS RVQ code tensors. - The MOSS audio tokenizer sits on CPU (to spare ~6.7 GiB next to the 8B - talker), so re-encoding the same reference is a fixed per-request cost that - otherwise dominates the 8B voice-clone variants and serializes under - concurrency. Mirror CosyVoice3 / Qwen3-TTS: cache by named voice when one is - supplied (``voice_created_at`` invalidates on re-upload), else by a content - hash of the reference. The blocking encode runs in a worker thread via - ``asyncio.to_thread`` so cold/anonymous encodes from concurrent requests - overlap instead of serializing on the event loop. + ``ref_audio_2`` is the TTSD second speaker; pass ``None`` for the + single-speaker variants. The named-voice cache key is + ``(voice_name, created_at)`` and ignores the clip content, so only slot 0 + (the reference that actually belongs to the uploaded voice) may use it — + speaker 2 is a different clip and stays content-addressed, or it would + silently reuse speaker 1's codes. - For named voices the speaker cache is checked *before* resolving the audio, - avoiding the decode cost when the cache is warm. For anonymous references - the audio is resolved first so the cache key incorporates mtime/size from - a single stat (no TOCTOU window). - - Args: - ref_str: The raw reference audio (URL / path / data URL) as received. - processor: Upstream MOSS processor exposing ``encode_audios_from_wav``. - resolve_ref_audio: Async callable mapping ``ref_str`` to ``(wav_list, sr, cache_key)``. - speaker_cache: Process-wide ``SpeakerEmbeddingCache``. - variant: MOSS sub-variant (``tts`` / ``ttsd`` / ...), namespaces the cache key. - n_vq: Number of RVQ codebooks (also namespaces the cache key). - sr_target: Target sample rate for the codec. - voice_name: Named/uploaded voice, if any (enables stable cross-request caching). - voice_created_at: Upload timestamp; bumps the cache slot on re-upload. - - Returns: - The reference RVQ codes tensor (CPU), ready to pass to the processor. + Returns ``(codes_per_speaker, resolve_keys)`` where ``resolve_keys`` maps + the reference slot (0 = ``ref_audio``, 1 = ``ref_audio_2``) to its + content-aware resolve key, for salting the KV prefix cache. A dict rather + than an append list because the two-speaker encodes run concurrently, so + completion order is not slot order. """ - model_type = f"moss_tts_{variant}_nq{n_vq}" + resolve_keys: dict[int, str] = {} - # ---- Named-voice branch: check cache *before* resolving ---- - # When voice_name is set, the cache key is (voice_name, created_at) which - # does not depend on the resolved audio content. The serving layer only - # passes voice_name for uploaded speakers without an inline ref_audio; - # re-upload bumps created_at and clears the speaker-cache slot. - if voice_name: - speaker_name = voice_name - created_at = int(voice_created_at) - cache_key = speaker_cache.make_cache_key( - speaker_name, - model_type=model_type, - created_at=created_at, + async def encode_one(ref_str: str, *, named_voice: bool, slot: int) -> torch.Tensor: + codes, resolve_key = await encoder.encode( + ref_str, + resolve_ref_audio=resolve_ref_audio, + get_artifact_key=get_artifact_key, + voice_name=voice_name if named_voice else None, + voice_created_at=voice_created_at if named_voice else 0, + ) + if resolve_key is not None: + resolve_keys[slot] = resolve_key + return codes + + if ref_audio_2: + # Encode both speakers concurrently so they land in the same batch + # window / share single-flight instead of serializing. + refs = list( + await asyncio.gather( + encode_one(ref_audio, named_voice=True, slot=0), + encode_one(ref_audio_2, named_voice=False, slot=1), + ) ) - cached = speaker_cache.get(cache_key) - if cached is not None: - logger.debug("Speaker cache HIT for MOSS-TTS reference '%s'", speaker_name) - return cached["codes"].clone() - # Cache miss — resolve the audio for encoding. - wav_list, sr, _ = await resolve_ref_audio(ref_str) else: - # ---- Anonymous branch: resolve first for content-addressed key ---- - # The resolve incorporates mtime/size from a single stat so the key - # and the waveform always come from the same filesystem snapshot. - wav_list, sr, resolve_cache_key = await resolve_ref_audio(ref_str) - speaker_name = "ref:" + resolve_cache_key - created_at = 0 - cache_key = speaker_cache.make_cache_key( - speaker_name, - model_type=model_type, - created_at=created_at, - ) - cached = speaker_cache.get(cache_key) - if cached is not None: - logger.debug("Speaker cache HIT for MOSS-TTS reference '%s'", speaker_name) - return cached["codes"].clone() - - # ---- Cold miss: encode in a worker thread and store ---- - codes = await asyncio.to_thread(_encode_wav_sync, processor, wav_list, sr, sr_target, n_vq) - speaker_cache.put(cache_key, {"codes": codes.detach().cpu()}) - logger.debug("Speaker cache STORE for MOSS-TTS reference '%s'", speaker_name) - return codes + refs = [await encode_one(ref_audio, named_voice=True, slot=0)] + return refs, resolve_keys