From 89fd13cccf44a798d4a909d6cc54d111c3ccfe9c Mon Sep 17 00:00:00 2001 From: Jay D Date: Mon, 7 Sep 2026 20:07:41 -0400 Subject: [PATCH] Bound the voice conditioning cache to stop CUDA OOM on varied voices The voice conditioning cache added with the bf16 work was unbounded and kept GPU tensors alive for every distinct voice ever used. Each new voice pinned roughly 200-300 MiB of VRAM, so a server that cycled through a few dozen predefined voices and clones exhausted a 12 GiB GPU and every request failed with "TTS engine failed to synthesize audio for chunk 1". - engine.py: make the cache an LRU bounded by tts_engine.voice_cache_size (default 8, 0 disables), store entries on CPU and copy to the device on a hit, evict the oldest beyond the limit, and release CUDA allocator blocks after each newly encoded voice. Add clear_voice_cache(). - config.py: add tts_engine.voice_cache_size default. - docker-compose.yml: set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. - docs: document the new key. Verified on an RTX-class 12 GiB card: VRAM climbs while the cache fills, then holds flat at ~6.4 GiB across 31 distinct voices and 3 clones. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01AGsQBKMnHzENiqR5oRQrSy --- README.md | 2 +- config.py | 1 + docker-compose.yml | 2 ++ documentation.md | 1 + engine.py | 64 ++++++++++++++++++++++++++++++++++++++++++---- 5 files changed, 64 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index f31e77b..26a3d00 100644 --- a/README.md +++ b/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`. diff --git a/config.py b/config.py index 4244a3d..06051d5 100644 --- a/config.py +++ b/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( diff --git a/docker-compose.yml b/docker-compose.yml index 2a575bd..47a7ebb 100644 --- a/docker-compose.yml +++ b/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 diff --git a/documentation.md b/documentation.md index e3506a3..da4e014 100644 --- a/documentation.md +++ b/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` | diff --git a/engine.py b/engine.py index a2be1e9..197ed55 100644 --- a/engine.py +++ b/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