optimization_benchmark.py

#!/usr/bin/env python3
"""
Kokoro TTS Optimization Benchmark
Comprehensive performance comparison between original and optimized Kokoro implementations.
"""

import asyncio
import time
import logging
import statistics
import os
import sys
from pathlib import Path
import psutil
import gc

# Add the project root to Python path
sys.path.append(str(Path(__file__).parent))

# Import both implementations
from core.kokoro_tts import KokoroTTSEngine
from core.kokoro_tts_optimized import OptimizedKokoroTTSEngine

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class PerformanceMonitor:
    """Monitor system performance during benchmarks."""
    
    def __init__(self):
        self.process = psutil.Process()
        self.initial_memory = self.process.memory_info().rss / 1024 / 1024  # MB
        self.peak_memory = self.initial_memory
        
    def update_peak_memory(self):
        """Update peak memory usage."""
        current_memory = self.process.memory_info().rss / 1024 / 1024
        self.peak_memory = max(self.peak_memory, current_memory)
        return current_memory
    
    def get_memory_delta(self):
        """Get memory increase from initial."""
        current = self.process.memory_info().rss / 1024 / 1024
        return current - self.initial_memory

class OptimizationBenchmark:
    """Comprehensive optimization benchmark suite."""
    
    def __init__(self):
        self.test_texts = [
            "Hello, this is a simple test.",
            "The quick brown fox jumps over the lazy dog in the beautiful moonlight.",
            "Artificial intelligence is revolutionizing the way we interact with technology, bringing new possibilities to voice synthesis.",
            "In a world where communication transcends boundaries, advanced text-to-speech systems enable seamless interaction between humans and machines.",
            "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation.",
        ]
        
        self.emotions = ["neutral", "happy", "excited", "calm", "serious", "friendly"]
        self.performance_monitor = PerformanceMonitor()
        self.results = {
            "original": {},
            "optimized": {}
        }
    
    async def benchmark_initialization(self, engine_class, engine_name):
        """Benchmark engine initialization time."""
        logger.info(f"🚀 Benchmarking {engine_name} initialization...")
        
        start_time = time.time()
        engine = engine_class()
        
        init_start = time.time()
        success = await engine.initialize()
        init_time = time.time() - init_start
        
        total_time = time.time() - start_time
        
        logger.info(f"✅ {engine_name} initialization: {init_time:.3f}s (total: {total_time:.3f}s)")
        
        return {
            "init_time": init_time,
            "total_time": total_time,
            "success": success,
            "engine": engine
        }
    
    async def benchmark_single_generation(self, engine, engine_name, iterations=10):
        """Benchmark single speech generation performance."""
        logger.info(f"🎵 Benchmarking {engine_name} single generation ({iterations} iterations)...")
        
        generation_times = []
        file_sizes = []
        memory_usage = []
        
        for i in range(iterations):
            text = self.test_texts[i % len(self.test_texts)]
            emotion = self.emotions[i % len(self.emotions)]
            
            # Clean up before each test
            gc.collect()
            
            start_memory = self.performance_monitor.update_peak_memory()
            start_time = time.time()
            
            # Generate speech
            if hasattr(engine, 'generate_speech_optimized'):
                output_file = await engine.generate_speech_optimized(text, emotion=emotion)
            else:
                output_file = await engine.generate_speech(text, emotion=emotion)
            
            generation_time = time.time() - start_time
            end_memory = self.performance_monitor.update_peak_memory()
            
            if output_file and os.path.exists(output_file):
                file_size = os.path.getsize(output_file) / 1024  # KB
                file_sizes.append(file_size)
                
                # Clean up test file
                try:
                    os.remove(output_file)
                except:
                    pass
            else:
                file_size = 0
            
            generation_times.append(generation_time)
            memory_usage.append(end_memory - start_memory)
            
            logger.debug(f"  Generation {i+1}: {generation_time:.3f}s, {file_size:.1f}KB")
        
        return {
            "avg_time": statistics.mean(generation_times),
            "min_time": min(generation_times),
            "max_time": max(generation_times),
            "std_time": statistics.stdev(generation_times) if len(generation_times) > 1 else 0,
            "avg_file_size": statistics.mean(file_sizes) if file_sizes else 0,
            "avg_memory_delta": statistics.mean(memory_usage),
            "times": generation_times
        }
    
    async def benchmark_batch_generation(self, engine, engine_name, batch_size=5):
        """Benchmark batch generation performance."""
        logger.info(f"🔄 Benchmarking {engine_name} batch generation ({batch_size} items)...")
        
        test_batch = self.test_texts[:batch_size]
        
        start_time = time.time()
        start_memory = self.performance_monitor.update_peak_memory()
        
        if hasattr(engine, 'batch_generate_speech'):
            # Use optimized batch method
            results = await engine.batch_generate_speech(test_batch, emotion="neutral")
            successful = len([r for r in results if r is not None])
        else:
            # Simulate batch with individual calls
            results = []
            for text in test_batch:
                result = await engine.generate_speech(text, emotion="neutral")
                results.append(result)
            successful = len([r for r in results if r is not None])
        
        batch_time = time.time() - start_time
        end_memory = self.performance_monitor.update_peak_memory()
        
        # Clean up test files
        for result in results:
            if result and os.path.exists(result):
                try:
                    os.remove(result)
                except:
                    pass
        
        throughput = successful / batch_time if batch_time > 0 else 0
        
        logger.info(f"  Batch results: {successful}/{batch_size} successful in {batch_time:.3f}s")
        
        return {
            "batch_time": batch_time,
            "successful": successful,
            "total": batch_size,
            "throughput": throughput,
            "memory_delta": end_memory - start_memory
        }
    
    async def benchmark_memory_efficiency(self, engine, engine_name):
        """Benchmark memory efficiency and cleanup."""
        logger.info(f"🧠 Benchmarking {engine_name} memory efficiency...")
        
        initial_memory = self.performance_monitor.update_peak_memory()
        
        # Generate multiple speeches to test memory accumulation
        for i in range(10):
            text = f"Memory test iteration {i+1}: " + self.test_texts[i % len(self.test_texts)]
            
            if hasattr(engine, 'generate_speech_optimized'):
                output_file = await engine.generate_speech_optimized(text)
            else:
                output_file = await engine.generate_speech(text)
            
            if output_file and os.path.exists(output_file):
                try:
                    os.remove(output_file)
                except:
                    pass
        
        memory_after_generation = self.performance_monitor.update_peak_memory()
        
        # Test cleanup
        if hasattr(engine, 'cleanup'):
            await engine.cleanup()
        
        gc.collect()
        memory_after_cleanup = self.performance_monitor.update_peak_memory()
        
        return {
            "initial_memory": initial_memory,
            "peak_memory": memory_after_generation,
            "final_memory": memory_after_cleanup,
            "memory_growth": memory_after_generation - initial_memory,
            "cleanup_effectiveness": memory_after_generation - memory_after_cleanup
        }
    
    async def benchmark_voice_switching(self, engine, engine_name, switches=10):
        """Benchmark voice switching performance."""
        logger.info(f"🎤 Benchmarking {engine_name} voice switching ({switches} switches)...")
        
        available_voices = ["af_heart", "af_sarah", "af_bella", "af_nova", "af_jessica"]
        switch_times = []
        
        text = "Voice switching performance test."
        
        for i in range(switches):
            voice = available_voices[i % len(available_voices)]
            
            start_time = time.time()
            
            if hasattr(engine, 'generate_speech_optimized'):
                output_file = await engine.generate_speech_optimized(text, voice=voice)
            else:
                output_file = await engine.generate_speech(text, emotion="neutral")
            
            switch_time = time.time() - start_time
            switch_times.append(switch_time)
            
            if output_file and os.path.exists(output_file):
                try:
                    os.remove(output_file)
                except:
                    pass
        
        return {
            "avg_switch_time": statistics.mean(switch_times),
            "min_switch_time": min(switch_times),
            "max_switch_time": max(switch_times),
            "switch_times": switch_times
        }
    
    def get_performance_stats(self, engine):
        """Get engine-specific performance statistics."""
        if hasattr(engine, 'get_performance_stats'):
            return engine.get_performance_stats()
        return {}
    
    async def run_comprehensive_benchmark(self):
        """Run the complete benchmark suite."""
        logger.info("🔥 Starting Comprehensive Kokoro TTS Optimization Benchmark")
        logger.info("=" * 80)
        
        # Test original implementation
        logger.info("📊 PHASE 1: Benchmarking Original Kokoro TTS")
        logger.info("-" * 50)
        
        original_init = await self.benchmark_initialization(KokoroTTSEngine, "Original Kokoro")
        if original_init["success"]:
            original_engine = original_init["engine"]
            
            self.results["original"] = {
                "initialization": original_init,
                "single_generation": await self.benchmark_single_generation(original_engine, "Original", 8),
                "batch_generation": await self.benchmark_batch_generation(original_engine, "Original", 5),
                "memory_efficiency": await self.benchmark_memory_efficiency(original_engine, "Original"),
                "voice_switching": await self.benchmark_voice_switching(original_engine, "Original", 6),
                "performance_stats": self.get_performance_stats(original_engine)
            }
        else:
            logger.error("❌ Original Kokoro TTS initialization failed")
            return False
        
        # Clear memory between tests
        del original_engine
        gc.collect()
        time.sleep(2)
        
        # Test optimized implementation
        logger.info("📊 PHASE 2: Benchmarking Optimized Kokoro TTS")
        logger.info("-" * 50)
        
        optimized_init = await self.benchmark_initialization(OptimizedKokoroTTSEngine, "Optimized Kokoro")
        if optimized_init["success"]:
            optimized_engine = optimized_init["engine"]
            
            self.results["optimized"] = {
                "initialization": optimized_init,
                "single_generation": await self.benchmark_single_generation(optimized_engine, "Optimized", 8),
                "batch_generation": await self.benchmark_batch_generation(optimized_engine, "Optimized", 5),
                "memory_efficiency": await self.benchmark_memory_efficiency(optimized_engine, "Optimized"),
                "voice_switching": await self.benchmark_voice_switching(optimized_engine, "Optimized", 6),
                "performance_stats": self.get_performance_stats(optimized_engine)
            }
            
            # Cleanup optimized engine
            await optimized_engine.cleanup()
        else:
            logger.error("❌ Optimized Kokoro TTS initialization failed")
            return False
        
        # Generate comparison report
        self.generate_comparison_report()
        return True
    
    def generate_comparison_report(self):
        """Generate detailed comparison report."""
        logger.info("📈 OPTIMIZATION COMPARISON REPORT")
        logger.info("=" * 80)
        
        # Initialization comparison
        orig_init = self.results["original"]["initialization"]["init_time"]
        opt_init = self.results["optimized"]["initialization"]["init_time"]
        init_improvement = ((orig_init - opt_init) / orig_init) * 100
        
        logger.info("🚀 INITIALIZATION PERFORMANCE:")
        logger.info(f"  Original:  {orig_init:.3f}s")
        logger.info(f"  Optimized: {opt_init:.3f}s")
        logger.info(f"  Improvement: {init_improvement:+.1f}%")
        logger.info("")
        
        # Generation time comparison
        orig_gen = self.results["original"]["single_generation"]["avg_time"]
        opt_gen = self.results["optimized"]["single_generation"]["avg_time"]
        gen_improvement = ((orig_gen - opt_gen) / orig_gen) * 100
        
        logger.info("🎵 GENERATION PERFORMANCE:")
        logger.info(f"  Original Average:  {orig_gen:.3f}s")
        logger.info(f"  Optimized Average: {opt_gen:.3f}s")
        logger.info(f"  Speed Improvement: {gen_improvement:+.1f}%")
        logger.info(f"  Original Range: {self.results['original']['single_generation']['min_time']:.3f}s - {self.results['original']['single_generation']['max_time']:.3f}s")
        logger.info(f"  Optimized Range: {self.results['optimized']['single_generation']['min_time']:.3f}s - {self.results['optimized']['single_generation']['max_time']:.3f}s")
        logger.info("")
        
        # Batch performance comparison
        orig_batch = self.results["original"]["batch_generation"]["throughput"]
        opt_batch = self.results["optimized"]["batch_generation"]["throughput"]
        batch_improvement = ((opt_batch - orig_batch) / orig_batch) * 100 if orig_batch > 0 else 0
        
        logger.info("🔄 BATCH PERFORMANCE:")
        logger.info(f"  Original Throughput:  {orig_batch:.3f} generations/sec")
        logger.info(f"  Optimized Throughput: {opt_batch:.3f} generations/sec")
        logger.info(f"  Throughput Improvement: {batch_improvement:+.1f}%")
        logger.info("")
        
        # Memory efficiency comparison
        orig_memory = self.results["original"]["memory_efficiency"]["memory_growth"]
        opt_memory = self.results["optimized"]["memory_efficiency"]["memory_growth"]
        memory_improvement = ((orig_memory - opt_memory) / orig_memory) * 100 if orig_memory > 0 else 0
        
        logger.info("🧠 MEMORY EFFICIENCY:")
        logger.info(f"  Original Memory Growth:  {orig_memory:.1f}MB")
        logger.info(f"  Optimized Memory Growth: {opt_memory:.1f}MB")
        logger.info(f"  Memory Improvement: {memory_improvement:+.1f}%")
        logger.info("")
        
        # Voice switching comparison
        orig_switch = self.results["original"]["voice_switching"]["avg_switch_time"]
        opt_switch = self.results["optimized"]["voice_switching"]["avg_switch_time"]
        switch_improvement = ((orig_switch - opt_switch) / orig_switch) * 100
        
        logger.info("🎤 VOICE SWITCHING PERFORMANCE:")
        logger.info(f"  Original Average:  {orig_switch:.3f}s")
        logger.info(f"  Optimized Average: {opt_switch:.3f}s")
        logger.info(f"  Switch Improvement: {switch_improvement:+.1f}%")
        logger.info("")
        
        # Overall assessment
        overall_improvements = [gen_improvement, batch_improvement, memory_improvement, switch_improvement]
        avg_improvement = statistics.mean([x for x in overall_improvements if x != 0])
        
        logger.info("🏆 OVERALL ASSESSMENT:")
        logger.info(f"  Average Performance Improvement: {avg_improvement:+.1f}%")
        
        if avg_improvement > 20:
            logger.info("  🌟 EXCELLENT optimization results!")
        elif avg_improvement > 10:
            logger.info("  ✅ GOOD optimization results!")
        elif avg_improvement > 0:
            logger.info("  👍 POSITIVE optimization results!")
        else:
            logger.info("  ⚠️ Mixed optimization results")
        
        # Performance stats
        if self.results["optimized"]["performance_stats"]:
            stats = self.results["optimized"]["performance_stats"]
            logger.info("")
            logger.info("📊 OPTIMIZED ENGINE STATISTICS:")
            logger.info(f"  Cache Hit Rate: {stats.get('cache_hit_rate', 0):.1%}")
            logger.info(f"  Cached Voices: {stats.get('cached_voices', 0)}")
            logger.info(f"  Total Generations: {stats.get('total_generations', 0)}")

async def main():
    """Main benchmark execution."""
    try:
        benchmark = OptimizationBenchmark()
        success = await benchmark.run_comprehensive_benchmark()
        
        if success:
            logger.info("✅ Optimization benchmark completed successfully!")
        else:
            logger.error("❌ Optimization benchmark failed!")
        
    except KeyboardInterrupt:
        logger.info("⏸️ Benchmark interrupted by user")
    except Exception as e:
        logger.error(f"❌ Benchmark error: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(main())