From 90b68520ac260bcd37abf3521389da37f07239d3 Mon Sep 17 00:00:00 2001 From: Tikkanaditya Siddartha Jyothi Date: Tue, 28 Jul 2026 09:30:08 -0700 Subject: [PATCH] fix(tts): bound the Piper/KittenTTS model caches with a small LRU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Salvaged from PR #62977 (@Vissirexa) — TTS model-cache half only (the hindsight turn-buffer half is a different subsystem and was dropped). _piper_voice_cache and _kittentts_model_cache were keyed by voice/model with no eviction, and each entry is a whole loaded model (tens of MB). A surface that sweeps voices pinned one model per voice for the process lifetime. New _tts_cache_get_or_load() get-or-loads through a small LRU (_TTS_MODEL_CACHE_MAX=3), refreshing recency on a hit and evicting the least-recently-used model on a cold miss. --- tests/tools/test_tts_model_cache_lru.py | 44 +++++++++++++++++++++++++ tools/tts_tool.py | 44 ++++++++++++++++++++----- 2 files changed, 80 insertions(+), 8 deletions(-) create mode 100644 tests/tools/test_tts_model_cache_lru.py diff --git a/tests/tools/test_tts_model_cache_lru.py b/tests/tools/test_tts_model_cache_lru.py new file mode 100644 index 000000000000..18b249f2e7d9 --- /dev/null +++ b/tests/tools/test_tts_model_cache_lru.py @@ -0,0 +1,44 @@ +"""LRU bound on the Piper/KittenTTS model caches. + +Each cached entry is a whole loaded TTS model (tens of MB). Without a cap the +caches pinned one model per distinct voice/model for the process lifetime, so a +surface that sweeps voices grew memory with no ceiling. `_tts_cache_get_or_load` +loads on a miss and evicts the least-recently-used entry beyond the cap. +""" +from __future__ import annotations + +import tools.tts_tool as tts + + +def test_loads_on_miss_and_serves_from_cache_on_hit(): + cache: dict = {} + calls = [] + + def load(): + calls.append(1) + return "model" + + assert tts._tts_cache_get_or_load(cache, "a", load) == "model" + assert tts._tts_cache_get_or_load(cache, "a", load) == "model" + assert len(calls) == 1 # loaded once, second call served from cache + + +def test_evicts_least_recently_used_beyond_cap(monkeypatch): + monkeypatch.setattr(tts, "_TTS_MODEL_CACHE_MAX", 2) + cache: dict = {} + for k in ("a", "b", "c"): + tts._tts_cache_get_or_load(cache, k, lambda k=k: k) + assert set(cache) == {"b", "c"} # "a", the oldest, was evicted + assert len(cache) == 2 + + +def test_hit_refreshes_recency_so_eviction_is_lru_not_fifo(monkeypatch): + monkeypatch.setattr(tts, "_TTS_MODEL_CACHE_MAX", 2) + cache: dict = {} + tts._tts_cache_get_or_load(cache, "a", lambda: "a") + tts._tts_cache_get_or_load(cache, "b", lambda: "b") + # Touch "a": "b" is now the least recently used. + tts._tts_cache_get_or_load(cache, "a", lambda: "a") + tts._tts_cache_get_or_load(cache, "c", lambda: "c") + assert "b" not in cache + assert set(cache) == {"a", "c"} diff --git a/tools/tts_tool.py b/tools/tts_tool.py index 42e446cf9239..9353e72a42db 100644 --- a/tools/tts_tool.py +++ b/tools/tts_tool.py @@ -2292,6 +2292,32 @@ def _generate_neutts(text: str, output_path: str, tts_config: Dict[str, Any]) -> # Provider: Piper (local, neural VITS, 44 languages) # =========================================================================== +# Each cached entry below is a whole loaded TTS model (tens of MB). An +# unbounded dict pins one model per distinct voice/model for the process +# lifetime, so a surface that sweeps voices grows memory with no ceiling. Cap +# each cache with a small LRU — most sessions use one or two voices, and a +# reload on a cold miss is cheap next to keeping every model resident. +_TTS_MODEL_CACHE_MAX = 3 + + +def _tts_cache_get_or_load(cache: Dict[str, Any], key: str, load: Callable[[], Any]) -> Any: + """Get ``key`` from ``cache`` or load it, keeping the cache LRU-bounded. + + Refreshes recency on a hit (insertion-ordered dict: pop + reinsert), loads + on a miss, then evicts least-recently-used entries beyond the cap. An entry + evicted while a caller still holds its returned reference stays alive for + that caller; only the cache slot is released. + """ + if key in cache: + cache[key] = cache.pop(key) + return cache[key] + value = load() + cache[key] = value + while len(cache) > _TTS_MODEL_CACHE_MAX: + cache.pop(next(iter(cache)), None) + return value + + # Module-level cache for Piper voice instances. Voices are keyed on their # absolute .onnx model path so switching voices doesn't invalidate older # cached voices. @@ -2403,12 +2429,14 @@ def _generate_piper_tts(text: str, output_path: str, tts_config: Dict[str, Any]) # PiperVoice instance serves all speakers, so it stays out of the cache # key. Multi-speaker workflows share one model load. cache_key = f"{model_path}::cuda={use_cuda}" - global _piper_voice_cache - if cache_key not in _piper_voice_cache: + + def _load_piper_voice(): logger.info("[Piper] Loading voice: %s", model_path) - _piper_voice_cache[cache_key] = PiperVoice.load(model_path, use_cuda=use_cuda) + v = PiperVoice.load(model_path, use_cuda=use_cuda) logger.info("[Piper] Voice loaded") - voice = _piper_voice_cache[cache_key] + return v + + voice = _tts_cache_get_or_load(_piper_voice_cache, cache_key, _load_piper_voice) # Optional synthesis knobs — only pass a SynthesisConfig when at least # one advanced knob is configured, so we don't depend on a newer Piper @@ -2500,13 +2528,13 @@ def _generate_kittentts(text: str, output_path: str, tts_config: Dict[str, Any]) clean_text = kt_config.get("clean_text", True) # Use cached model instance if available - global _kittentts_model_cache - if model_name not in _kittentts_model_cache: + def _load_kittentts_model(): logger.info("[KittenTTS] Loading model: %s", model_name) - _kittentts_model_cache[model_name] = KittenTTS(model_name) + m = KittenTTS(model_name) logger.info("[KittenTTS] Model loaded successfully") + return m - model = _kittentts_model_cache[model_name] + model = _tts_cache_get_or_load(_kittentts_model_cache, model_name, _load_kittentts_model) # Generate audio (returns numpy array at 24kHz) audio = model.generate(text, voice=voice, speed=speed, clean_text=clean_text)