""" Model Optimization Module Advanced optimizations for ChatterboxTTS models including mixed precision, caching, and memory management """ import torch import logging import time import hashlib from typing import Dict, Any, Optional, Tuple from pathlib import Path import gc class MixedPrecisionOptimizer: """Mixed precision optimization for faster inference""" def __init__(self, enabled=True): self.enabled = enabled and torch.cuda.is_available() self.autocast_context = None self.logger = logging.getLogger(__name__) if self.enabled: self.logger.info("🚀 Mixed precision optimization enabled") else: self.logger.info("📋 Mixed precision optimization disabled") def __enter__(self): if self.enabled: self.autocast_context = torch.autocast(device_type='cuda', dtype=torch.float16) return self.autocast_context.__enter__() return None def __exit__(self, exc_type, exc_val, exc_tb): if self.enabled and self.autocast_context: return self.autocast_context.__exit__(exc_type, exc_val, exc_tb) return False class TextTokenizationCache: """Smart caching system for text tokenization""" def __init__(self, max_size=1000): self.cache = {} self.max_size = max_size self.hit_count = 0 self.miss_count = 0 self.logger = logging.getLogger(__name__) def _get_cache_key(self, text: str, params: Dict = None) -> str: """Generate cache key from text and parameters""" key_str = text if params: # Include relevant parameters that affect tokenization param_str = str(sorted(params.items())) key_str += f"|{param_str}" return hashlib.md5(key_str.encode()).hexdigest() def get(self, text: str, params: Dict = None) -> Optional[Any]: """Get cached tokenization result""" key = self._get_cache_key(text, params) if key in self.cache: self.hit_count += 1 self.logger.debug(f"📋 Tokenization cache HIT for: '{text[:30]}...'") return self.cache[key] self.miss_count += 1 self.logger.debug(f"🔍 Tokenization cache MISS for: '{text[:30]}...'") return None def put(self, text: str, result: Any, params: Dict = None): """Store tokenization result in cache""" if len(self.cache) >= self.max_size: # Simple LRU eviction - remove oldest entry oldest_key = next(iter(self.cache)) del self.cache[oldest_key] self.logger.debug("🗑️ Evicted oldest cache entry") key = self._get_cache_key(text, params) self.cache[key] = result self.logger.debug(f"💾 Cached tokenization for: '{text[:30]}...'") def get_stats(self) -> Dict[str, Any]: """Get cache performance statistics""" total_requests = self.hit_count + self.miss_count hit_rate = (self.hit_count / total_requests * 100) if total_requests > 0 else 0 return { 'hit_count': self.hit_count, 'miss_count': self.miss_count, 'hit_rate': hit_rate, 'cache_size': len(self.cache), 'max_size': self.max_size } def clear(self): """Clear the cache""" self.cache.clear() self.hit_count = 0 self.miss_count = 0 self.logger.info("🗑️ Tokenization cache cleared") class MemoryOptimizer: """Advanced memory optimization strategies""" def __init__(self): self.logger = logging.getLogger(__name__) self.memory_threshold_gb = 6.0 # Memory threshold for aggressive cleanup def optimize_memory_usage(self, aggressive=False): """Optimize memory usage with various strategies""" initial_memory = self.get_gpu_memory_usage() # Standard cleanup gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() if aggressive: # More aggressive cleanup if torch.cuda.is_available(): torch.cuda.ipc_collect() # Clear unused memory torch.cuda.reset_peak_memory_stats() final_memory = self.get_gpu_memory_usage() freed_memory = initial_memory - final_memory if freed_memory > 0.1: # More than 100MB freed self.logger.info(f"🗑️ Memory cleanup freed {freed_memory:.2f}GB") return freed_memory def get_gpu_memory_usage(self) -> float: """Get current GPU memory usage in GB""" if torch.cuda.is_available(): return torch.cuda.memory_allocated() / 1024**3 return 0.0 def should_optimize_memory(self) -> bool: """Check if memory optimization is needed""" current_memory = self.get_gpu_memory_usage() return current_memory > self.memory_threshold_gb def monitor_memory_usage(self, operation_name: str = ""): """Context manager for monitoring memory usage during operations""" return MemoryMonitor(self, operation_name) class MemoryMonitor: """Context manager for monitoring memory usage""" def __init__(self, optimizer: MemoryOptimizer, operation_name: str): self.optimizer = optimizer self.operation_name = operation_name self.start_memory = 0.0 def __enter__(self): self.start_memory = self.optimizer.get_gpu_memory_usage() return self def __exit__(self, exc_type, exc_val, exc_tb): end_memory = self.optimizer.get_gpu_memory_usage() memory_diff = end_memory - self.start_memory if memory_diff > 0.1: # Significant memory increase self.optimizer.logger.info( f"📈 {self.operation_name}: +{memory_diff:.2f}GB memory usage" ) elif memory_diff < -0.1: # Significant memory decrease self.optimizer.logger.info( f"📉 {self.operation_name}: {abs(memory_diff):.2f}GB memory freed" ) class InferenceOptimizer: """Optimizations specifically for inference performance""" def __init__(self): self.logger = logging.getLogger(__name__) self.setup_inference_optimizations() def setup_inference_optimizations(self): """Setup general inference optimizations""" # Disable gradient computation globally for inference torch.set_grad_enabled(False) # Enable cuDNN benchmarking for consistent input sizes if torch.backends.cudnn.is_available(): torch.backends.cudnn.benchmark = True torch.backends.cudnn.deterministic = False self.logger.info("🏃 cuDNN benchmark mode enabled") # Enable optimized attention if available if hasattr(torch.backends.cuda, 'enable_flash_sdp'): torch.backends.cuda.enable_flash_sdp(True) self.logger.info("⚡ Flash attention enabled") def optimize_model_for_inference(self, model): """Apply inference-specific optimizations to a model""" optimizations_applied = [] # Set model to eval mode if possible if hasattr(model, 'eval'): model.eval() optimizations_applied.append("eval_mode") # Apply eval mode to submodules if hasattr(model, 't3') and hasattr(model.t3, 'eval'): model.t3.eval() optimizations_applied.append("t3_eval") if hasattr(model, 's3gen') and hasattr(model.s3gen, 'eval'): model.s3gen.eval() optimizations_applied.append("s3gen_eval") # Enable inference mode context if hasattr(torch, 'inference_mode'): # Note: This would need to be applied as a context manager optimizations_applied.append("inference_mode_available") self.logger.info(f"🔧 Applied inference optimizations: {', '.join(optimizations_applied)}") return model class OptimizedTTSWrapper: """Wrapper that applies all optimizations to TTS operations""" def __init__(self, model, enable_mixed_precision=True, enable_caching=True): self.model = model self.mixed_precision = MixedPrecisionOptimizer(enable_mixed_precision) self.tokenization_cache = TextTokenizationCache() if enable_caching else None self.memory_optimizer = MemoryOptimizer() self.inference_optimizer = InferenceOptimizer() self.logger = logging.getLogger(__name__) # Apply inference optimizations to the model self.inference_optimizer.optimize_model_for_inference(self.model) # Performance tracking self.stats = { 'optimized_generations': 0, 'cache_hits': 0, 'memory_optimizations': 0, 'total_time_saved': 0.0 } def generate(self, text: str, **kwargs) -> torch.Tensor: """Optimized generation with all optimizations applied""" start_time = time.time() # Check if memory optimization is needed if self.memory_optimizer.should_optimize_memory(): freed = self.memory_optimizer.optimize_memory_usage(aggressive=True) if freed > 0: self.stats['memory_optimizations'] += 1 # Use mixed precision if enabled with self.mixed_precision: with self.memory_optimizer.monitor_memory_usage(f"TTS generation"): # Check cache for identical requests cache_key = None if self.tokenization_cache: cache_key = f"{text}|{str(sorted(kwargs.items()))}" cached_result = self.tokenization_cache.get(text, kwargs) if cached_result is not None: self.stats['cache_hits'] += 1 self.logger.debug("📋 Using cached generation result") return cached_result # Generate audio with optimizations with torch.inference_mode() if hasattr(torch, 'inference_mode') else torch.no_grad(): audio = self.model.generate(text, **kwargs) # Cache the result if caching is enabled if self.tokenization_cache and cache_key: self.tokenization_cache.put(text, audio, kwargs) generation_time = time.time() - start_time self.stats['optimized_generations'] += 1 # Estimate time saved (rough heuristic) estimated_baseline_time = generation_time * 1.15 # Assume 15% overhead without optimizations time_saved = estimated_baseline_time - generation_time self.stats['total_time_saved'] += time_saved return audio def prepare_conditionals(self, *args, **kwargs): """Optimized voice conditioning preparation""" with self.mixed_precision: with self.memory_optimizer.monitor_memory_usage("Voice conditioning"): return self.model.prepare_conditionals(*args, **kwargs) def get_optimization_stats(self) -> Dict[str, Any]: """Get comprehensive optimization statistics""" stats = self.stats.copy() if self.tokenization_cache: cache_stats = self.tokenization_cache.get_stats() stats.update({ 'cache_hit_rate': cache_stats['hit_rate'], 'cache_size': cache_stats['cache_size'] }) stats['current_memory_gb'] = self.memory_optimizer.get_gpu_memory_usage() return stats def log_performance_summary(self): """Log comprehensive performance summary""" stats = self.get_optimization_stats() self.logger.info("📊 OPTIMIZATION PERFORMANCE SUMMARY") self.logger.info("=" * 50) self.logger.info(f"Optimized generations: {stats['optimized_generations']}") self.logger.info(f"Memory optimizations: {stats['memory_optimizations']}") self.logger.info(f"Total time saved: {stats['total_time_saved']:.2f}s") if 'cache_hit_rate' in stats: self.logger.info(f"Cache hit rate: {stats['cache_hit_rate']:.1f}%") self.logger.info(f"Cache size: {stats['cache_size']} entries") self.logger.info(f"Current memory usage: {stats['current_memory_gb']:.2f}GB") # Calculate estimated performance improvement if stats['optimized_generations'] > 0: avg_time_saved = stats['total_time_saved'] / stats['optimized_generations'] self.logger.info(f"Average time saved per generation: {avg_time_saved:.2f}s") def create_optimized_model(model, enable_mixed_precision=True, enable_caching=True): """Create an optimized wrapper around a ChatterboxTTS model""" return OptimizedTTSWrapper( model, enable_mixed_precision=enable_mixed_precision, enable_caching=enable_caching )