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