flozi00/Chatterbox-Multilingual-TTS

1"""2Caching system for generated audio.3Supports local and Hugging Face Hub storage.4"""5 6import hashlib7import os8from dataclasses import dataclass9from pathlib import Path10from typing import Optional11 12from loguru import logger13 14 15@dataclass16class CacheConfig:17    """Configuration for audio caching."""18 19    enabled: bool = True20    local_cache_dir: Optional[str] = None  # Local cache directory21    hf_repo_id: Optional[str] = None  # Hugging Face Hub repo for remote cache22    max_duration_seconds: float = 30.0  # Only cache audio shorter than this23 24 25class AudioCache:26    """27    Cache for generated TTS audio.28    Supports both local filesystem and Hugging Face Hub storage.29    """30 31    def __init__(self, config: Optional[CacheConfig] = None):32        self.config = config or CacheConfig()33        self._hf_fs = None34 35    def _get_cache_key(self, text: str, voice_id: str, backend: str) -> str:36        """Generate a unique cache key for the given parameters."""37        content = f"{backend}:{voice_id}:{text}"38        return hashlib.md5(content.encode()).hexdigest()39 40    def _get_hf_fs(self):41        """Get HuggingFace filesystem (lazy initialization)."""42        if self._hf_fs is None and self.config.hf_repo_id:43            try:44                from huggingface_hub import HfFileSystem45 46                self._hf_fs = HfFileSystem(token=os.environ.get("HF_TOKEN"))47            except Exception as e:48                logger.warning(f"Could not initialize HF filesystem: {e}")49        return self._hf_fs50 51    def get(self, text: str, voice_id: str, backend: str) -> Optional[bytes]:52        """53        Retrieve cached audio if it exists.54 55        Args:56            text: Original text that was synthesized57            voice_id: Voice identifier used58            backend: Backend name used for synthesis59 60        Returns:61            Cached audio bytes or None if not found62        """63        if not self.config.enabled:64            return None65 66        cache_key = self._get_cache_key(text, voice_id, backend)67 68        # Try local cache first69        if self.config.local_cache_dir:70            local_path = Path(self.config.local_cache_dir) / f"{cache_key}.mp3"71            if local_path.exists():72                logger.debug(f"Cache hit (local): {cache_key}")73                return local_path.read_bytes()74 75        # Try HF Hub cache76        if self.config.hf_repo_id:77            fs = self._get_hf_fs()78            if fs:79                hf_path = f"{self.config.hf_repo_id}/{voice_id}/{cache_key}.mp3"80                try:81                    if fs.exists(hf_path):82                        with fs.open(hf_path, "rb") as f:83                            logger.debug(f"Cache hit (HF Hub): {cache_key}")84                            return f.read()85                except Exception as e:86                    logger.debug(f"HF cache lookup failed: {e}")87 88        return None89 90    def set(91        self,92        text: str,93        voice_id: str,94        backend: str,95        audio_data: bytes,96        duration_seconds: Optional[float] = None,97    ) -> bool:98        """99        Store audio in cache.100 101        Args:102            text: Original text that was synthesized103            voice_id: Voice identifier used104            backend: Backend name used for synthesis105            audio_data: Audio bytes to cache106            duration_seconds: Duration of the audio (for max duration check)107 108        Returns:109            True if cached successfully, False otherwise110        """111        if not self.config.enabled:112            return False113 114        # Check duration limit115        if duration_seconds and duration_seconds > self.config.max_duration_seconds:116            logger.debug(117                f"Audio too long to cache: {duration_seconds}s > {self.config.max_duration_seconds}s"118            )119            return False120 121        cache_key = self._get_cache_key(text, voice_id, backend)122        success = False123 124        # Save to local cache125        if self.config.local_cache_dir:126            try:127                cache_dir = Path(self.config.local_cache_dir)128                cache_dir.mkdir(parents=True, exist_ok=True)129                local_path = cache_dir / f"{cache_key}.mp3"130                local_path.write_bytes(audio_data)131                logger.debug(f"Cached locally: {cache_key}")132                success = True133            except Exception as e:134                logger.warning(f"Failed to cache locally: {e}")135 136        # Save to HF Hub137        if self.config.hf_repo_id:138            fs = self._get_hf_fs()139            if fs:140                try:141                    voice_dir = f"{self.config.hf_repo_id}/{voice_id}"142                    if not fs.exists(voice_dir):143                        fs.makedirs(voice_dir, exist_ok=True)144 145                    hf_path = f"{voice_dir}/{cache_key}.mp3"146                    with fs.open(hf_path, "wb") as f:147                        f.write(audio_data)148 149                    logger.debug(f"Cached to HF Hub: {cache_key}")150                    success = True151                except Exception as e:152                    logger.warning(f"Failed to cache to HF Hub: {e}")153 154        return success155 156    def clear_local(self) -> int:157        """Clear local cache. Returns number of files deleted."""158        if not self.config.local_cache_dir:159            return 0160 161        cache_dir = Path(self.config.local_cache_dir)162        if not cache_dir.exists():163            return 0164 165        count = 0166        for file in cache_dir.glob("*.mp3"):167            file.unlink()168            count += 1169 170        logger.info(f"Cleared {count} files from local cache")171        return count172