[Model] Reduce MOSS TTS reference audio preparation overhead (#7725)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/model_executor/models/moss_tts/test_reference_encoder.py

@@ -12,7 +12,9 @@ import torch
 
 from vllm_omni.model_executor.models.moss_tts.reference_encoder import (
     MossReferenceEncoder,
+    _prep_wav_sync,
     _RefEncodeBatcher,
+    _reference_resampler,
     build_reference_encoder,
     encode_request_references,
 )
@@ -581,3 +583,18 @@ async def test_idle_batcher_releases_all_completed_waveforms():
         assert batcher._drainer is not None and not batcher._drainer.done()
     finally:
         await batcher.aclose()
+
+
+def test_reference_resampler_is_cached_and_matches_functional_resample():
+    import torchaudio
+
+    _reference_resampler.cache_clear()
+    try:
+        waveform = torch.rand((1, 4800), generator=torch.Generator().manual_seed(123))
+        expected = torchaudio.functional.resample(waveform, 48000, _SR)
+        torch.testing.assert_close(_prep_wav_sync(waveform.numpy(), 48000, _SR), expected, rtol=0, atol=0)
+        torch.testing.assert_close(_prep_wav_sync(waveform.numpy(), 48000, _SR), expected, rtol=0, atol=0)
+        info = _reference_resampler.cache_info()
+        assert info.misses == 1 and info.hits == 1
+    finally:
+        _reference_resampler.cache_clear()

vllm_omni/model_executor/models/moss_tts/reference_encoder.py

@@ -27,6 +27,7 @@ from __future__ import annotations
 import asyncio
 import hashlib
 from collections.abc import Awaitable, Callable
+from functools import lru_cache
 from typing import Any
 
 import numpy as np
@@ -50,15 +51,24 @@ def _sha1(s: str) -> str:
     return hashlib.sha1((s or "").encode("utf-8")).hexdigest()
 
 
+@lru_cache(maxsize=16)
+def _reference_resampler(sr: int, sr_target: int) -> torch.nn.Module:
+    import torchaudio
+
+    # Match functional.resample's float32 kernel construction, rather than
+    # Resample's default float64 construction followed by a float32 cast.
+    return torchaudio.transforms.Resample(sr, sr_target, dtype=torch.float32)
+
+
 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(waveform, dtype=torch.float32)
+    # Keep the tensor independent from the server-side waveform/cache buffer.
+    wav_np = np.array(waveform, dtype=np.float32, order="C", copy=True)
+    wav = torch.from_numpy(wav_np)
     if wav.dim() == 1:
         wav = wav.unsqueeze(0)
     if sr != sr_target:
-        import torchaudio
-
-        wav = torchaudio.functional.resample(wav, sr, sr_target)
+        wav = _reference_resampler(sr, sr_target)(wav)
     return wav