Bound the voice conditioning cache to stop CUDA OOM after switching voices

README.md

@@ -787,7 +787,7 @@ The server relies exclusively on `config.yaml` for runtime configuration.
 
 *   `server`: `host`, `port`, logging settings.
 *   `model`: `repo_id` (e.g., "ResembleAI/chatterbox").
-*   `tts_engine`: `device` ('auto', 'cuda', 'mps', 'cpu'), `predefined_voices_path`, `reference_audio_path`, `default_voice_id`.
+*   `tts_engine`: `device` ('auto', 'cuda', 'mps', 'cpu'), `predefined_voices_path`, `reference_audio_path`, `default_voice_id`, `voice_cache_size` (LRU size for cached voice conditionals, default 8).
 *   `paths`: `model_cache` (for `download_model.py`), `output`.
 *   `generation_defaults`: Default UI values for `temperature`, `exaggeration`, `cfg_weight`, `seed`, `speed_factor`, `language`.
 *   `audio_output`: `format`, `sample_rate`, `max_reference_duration_sec`.

config.py

@@ -63,6 +63,7 @@
             DEFAULT_REFERENCE_AUDIO_PATH
         ),  # Directory for reference audio files for cloning.
         "default_voice_id": "default_sample.wav",  # Default voice file to use if none is specified.
+        "voice_cache_size": 8,  # Max distinct voices whose conditionals are cached (LRU, CPU-resident). 0 disables.
     },
     "paths": {  # General configurable paths for the application.
         "model_cache": str(

docker-compose.yml

@@ -41,6 +41,8 @@ services:
       - HF_TOKEN=YOUR_TOKEN_HERE
       # Enable faster Hugging Face downloads inside the container
       - HF_HUB_ENABLE_HF_TRANSFER=1
+      # Reduce CUDA allocator fragmentation on long-running servers
+      - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
       # Make NVIDIA GPUs visible and specify capabilities for PyTorch
       - NVIDIA_VISIBLE_DEVICES=all
       - NVIDIA_DRIVER_CAPABILITIES=compute,utility

documentation.md

@@ -468,6 +468,7 @@ The following table describes the main sections and some key parameters you migh
 |                       | `predefined_voices_path`      | string        | Directory for predefined voice audio files.                                                                   | `voices`                 |
 |                       | `reference_audio_path`        | string        | Directory for user-uploaded reference audio files for voice cloning.                                          | `reference_audio`        |
 |                       | `default_voice_id`            | string        | Filename of the default predefined voice to use if none selected (primarily for UI).                            | `default_sample.wav`     |
+|                       | `voice_cache_size`            | integer       | Max distinct voices whose encoded conditionals are cached (LRU, stored on CPU). `0` disables caching.           | `8`                      |
 | **`paths`**           | `model_cache`                 | string        | Directory for caching models downloaded by `download_model.py`. **Note:** Runtime uses global HF cache.         | `./model_cache`          |
 |                       | `output`                      | string        | Default directory for audio files saved from the UI or API.                                                   | `./outputs`              |
 | **`generation_defaults`**| `temperature`                | float         | Controls randomness (0.0-1.5). Lower is more deterministic.                                                   | `0.8`                    |

engine.py

@@ -1,12 +1,14 @@
 # File: engine.py
 # Core TTS model loading and speech generation logic.
 
+import copy
 import gc
 import logging
 import os
 import random
 import numpy as np
 import torch
+from collections import OrderedDict
 from typing import Optional, Tuple
 from pathlib import Path
 
@@ -114,7 +116,14 @@ def _resolve_bf16_setting() -> bool:
 
 # Voice conditioning cache: avoids re-encoding the same voice file on every request.
 # Key: (resolved_path, file_mtime, exaggeration) — mtime invalidates if file changes.
-_conds_cache: dict = {}
+#
+# The cache is a bounded LRU and entries are stored on the CPU. An unbounded
+# GPU-resident cache leaked roughly 200-300 MiB of VRAM per distinct voice
+# (allocator blocks pinned by the cached tensors), which exhausted a 12 GiB GPU
+# after a few dozen voices. Size is controlled by `tts_engine.voice_cache_size`
+# (0 disables caching).
+_conds_cache: "OrderedDict[tuple, object]" = OrderedDict()
+_VOICE_CACHE_SIZE_DEFAULT = 8
 
 
 def _conds_cache_key(path: str, exaggeration: float) -> tuple:
@@ -125,6 +134,40 @@ def _conds_cache_key(path: str, exaggeration: float) -> tuple:
     return (path, mtime, exaggeration)
 
 
+def _voice_cache_size() -> int:
+    try:
+        return max(0, int(config_manager.get_int("tts_engine.voice_cache_size", _VOICE_CACHE_SIZE_DEFAULT)))
+    except Exception:
+        return _VOICE_CACHE_SIZE_DEFAULT
+
+
+def _conds_copy_to(conds, device):
+    """
+    Return a copy of a Conditionals object with its tensors on `device`.
+    Both Conditionals.to() and T3Cond.to() mutate in place, so shallow-copy the
+    containers first to keep the cached CPU copy untouched.
+    """
+    new = copy.copy(conds)
+    new.t3 = copy.copy(conds.t3)
+    new.gen = dict(conds.gen)
+    return new.to(device)
+
+
+def _release_cuda_memory():
+    """Return freed allocator blocks to the driver so VRAM does not creep upward."""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def clear_voice_cache() -> int:
+    """Drop all cached voice conditionals and release GPU memory. Returns entries removed."""
+    n = len(_conds_cache)
+    _conds_cache.clear()
+    _release_cuda_memory()
+    return n
+
+
 def set_seed(seed_value: int):
     """
     Sets the seed for torch, random, and numpy for reproducibility.
@@ -480,7 +523,10 @@ def synthesize(
             ex_for_key = 0.0 if loaded_model_type == "turbo" else exaggeration
             conds_key = _conds_cache_key(audio_prompt_path, ex_for_key)
             if conds_key in _conds_cache:
-                chatterbox_model.conds = _conds_cache[conds_key]
+                _conds_cache.move_to_end(conds_key)
+                chatterbox_model.conds = _conds_copy_to(
+                    _conds_cache[conds_key], chatterbox_model.device
+                )
                 effective_prompt = None  # conds already set, skip prepare_conditionals
                 logger.debug(f"Voice cache hit: {audio_prompt_path}")
 
@@ -506,11 +552,19 @@ def synthesize(
                     cfg_weight=cfg_weight,
                 )
 
-        # Store conds in cache after first compute for this voice.
+        # Store a CPU copy of the conds in the LRU cache after first compute for
+        # this voice, then hand the allocator blocks used while encoding the
+        # reference clip back to the driver.
         if conds_key is not None and effective_prompt is not None:
-            if chatterbox_model.conds is not None:
-                _conds_cache[conds_key] = chatterbox_model.conds
+            max_size = _voice_cache_size()
+            if max_size > 0 and chatterbox_model.conds is not None:
+                _conds_cache[conds_key] = _conds_copy_to(chatterbox_model.conds, "cpu")
+                _conds_cache.move_to_end(conds_key)
+                while len(_conds_cache) > max_size:
+                    evicted_key, _ = _conds_cache.popitem(last=False)
+                    logger.debug(f"Evicted voice conditionals for: {evicted_key[0]}")
                 logger.debug(f"Cached voice conditionals for: {audio_prompt_path}")
+            _release_cuda_memory()
 
         # The ChatterboxTTS.generate method already returns a CPU tensor.
         return wav_tensor, chatterbox_model.sr