kokoro_benchmark.py

#!/usr/bin/env python3
"""
Kokoro ONNX v1.0 Performance Benchmark
Comprehensive stress testing for the new voice system.
"""

import asyncio
import time
import statistics
import os
import sys
from pathlib import Path
from typing import List, Dict, Any
import gc
import psutil
from concurrent.futures import ThreadPoolExecutor

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

from core.kokoro_tts import KokoroTTSEngine

class KokoroBenchmark:
    """Comprehensive benchmark suite for Kokoro ONNX v1.0."""
    
    def __init__(self):
        self.engine = None
        self.results = {
            'initialization': {},
            'single_generation': {},
            'batch_generation': {},
            'stress_test': {},
            'memory_usage': {},
            'voice_variety': {},
            'long_text': {}
        }
        
        # Test data
        self.test_texts = [
            "Hello, this is a simple test.",
            "The quick brown fox jumps over the lazy dog.",
            "Artificial intelligence is revolutionizing how we interact with technology.",
            "In the heart of autumn, when leaves dance in golden spirals, Autumn AI brings warmth to digital conversations.",
            "Performance testing requires systematic evaluation of speed, accuracy, and resource utilization under various load conditions."
        ]
        
        self.long_text = """
        Welcome to the comprehensive performance evaluation of the Kokoro ONNX text-to-speech system.
        This extended passage serves multiple purposes in our benchmark suite. First, it tests the system's
        ability to handle longer text inputs without degradation in quality or significant performance impact.
        Second, it evaluates memory management and resource utilization during extended synthesis operations.
        Third, it provides insight into the consistency of voice quality across longer passages.
        The Kokoro neural voice synthesis technology represents a significant advancement in artificial
        intelligence-driven speech generation, offering natural-sounding voices with emotional expression
        capabilities. Through rigorous testing, we can validate its performance characteristics and ensure
        it meets the demanding requirements of modern voice applications.
        """
        
        self.emotions = ["neutral", "happy", "excited", "calm", "serious", "friendly"]
        
    async def initialize_engine(self) -> Dict[str, Any]:
        """Initialize the Kokoro engine and measure startup time."""
        print("🚀 Initializing Kokoro ONNX v1.0 Engine...")
        
        start_time = time.time()
        self.engine = KokoroTTSEngine()
        
        init_start = time.time()
        success = await self.engine.initialize()
        init_time = time.time() - init_start
        
        total_time = time.time() - start_time
        
        result = {
            'success': success,
            'initialization_time': init_time,
            'total_startup_time': total_time,
            'available_voices': len(getattr(self.engine, 'available_voices', [])),
            'using_onnx': self.engine.use_onnx if self.engine else False
        }
        
        self.results['initialization'] = result
        
        print(f"✅ Initialization: {init_time:.3f}s")
        print(f"🎤 Available voices: {result['available_voices']}")
        print(f"🔥 Using ONNX: {result['using_onnx']}")
        
        return result
    
    async def benchmark_single_generation(self, iterations: int = 10) -> Dict[str, Any]:
        """Benchmark single text generation performance."""
        print(f"\n📊 Single Generation Benchmark ({iterations} iterations)...")
        
        if not self.engine or not self.engine.is_initialized:
            return {'error': 'Engine not initialized'}
        
        times = []
        file_sizes = []
        success_count = 0
        
        for i in range(iterations):
            text = self.test_texts[i % len(self.test_texts)]
            emotion = self.emotions[i % len(self.emotions)]
            
            print(f"  Test {i+1}/{iterations}: {emotion} - {text[:30]}...")
            
            start_time = time.time()
            
            try:
                output_file = await self.engine.generate_speech(text, emotion=emotion)
                generation_time = time.time() - start_time
                
                if output_file and os.path.exists(output_file):
                    file_size = os.path.getsize(output_file)
                    times.append(generation_time)
                    file_sizes.append(file_size)
                    success_count += 1
                    
                    # Clean up
                    os.remove(output_file)
                    
                    print(f"    ✅ {generation_time:.3f}s ({file_size:,} bytes)")
                else:
                    print(f"    ❌ Failed")
                    
            except Exception as e:
                print(f"    ❌ Error: {e}")
        
        if times:
            result = {
                'iterations': iterations,
                'success_count': success_count,
                'success_rate': success_count / iterations * 100,
                'avg_time': statistics.mean(times),
                'min_time': min(times),
                'max_time': max(times),
                'median_time': statistics.median(times),
                'std_dev': statistics.stdev(times) if len(times) > 1 else 0,
                'avg_file_size': statistics.mean(file_sizes),
                'total_time': sum(times)
            }
        else:
            result = {'error': 'No successful generations'}
        
        self.results['single_generation'] = result
        return result
    
    async def benchmark_batch_generation(self, batch_size: int = 5) -> Dict[str, Any]:
        """Benchmark concurrent batch generation."""
        print(f"\n🔄 Batch Generation Benchmark ({batch_size} concurrent)...")
        
        if not self.engine or not self.engine.is_initialized:
            return {'error': 'Engine not initialized'}
        
        async def generate_single(text: str, emotion: str, index: int):
            start_time = time.time()
            try:
                output_file = await self.engine.generate_speech(text, emotion=emotion)
                generation_time = time.time() - start_time
                
                if output_file and os.path.exists(output_file):
                    file_size = os.path.getsize(output_file)
                    os.remove(output_file)  # Clean up
                    return {
                        'index': index,
                        'success': True,
                        'time': generation_time,
                        'file_size': file_size
                    }
                else:
                    return {'index': index, 'success': False, 'time': generation_time}
            except Exception as e:
                return {'index': index, 'success': False, 'error': str(e), 'time': time.time() - start_time}
        
        # Create batch tasks
        tasks = []
        for i in range(batch_size):
            text = self.test_texts[i % len(self.test_texts)]
            emotion = self.emotions[i % len(self.emotions)]
            tasks.append(generate_single(text, emotion, i))
        
        # Execute batch
        batch_start = time.time()
        results = await asyncio.gather(*tasks, return_exceptions=True)
        batch_time = time.time() - batch_start
        
        # Analyze results
        successful = [r for r in results if isinstance(r, dict) and r.get('success')]
        failed = [r for r in results if not (isinstance(r, dict) and r.get('success'))]
        
        if successful:
            times = [r['time'] for r in successful]
            file_sizes = [r['file_size'] for r in successful]
            
            result = {
                'batch_size': batch_size,
                'success_count': len(successful),
                'failure_count': len(failed),
                'success_rate': len(successful) / batch_size * 100,
                'total_batch_time': batch_time,
                'avg_individual_time': statistics.mean(times),
                'min_time': min(times),
                'max_time': max(times),
                'avg_file_size': statistics.mean(file_sizes),
                'throughput_per_second': len(successful) / batch_time
            }
        else:
            result = {'error': 'No successful generations in batch'}
        
        self.results['batch_generation'] = result
        return result
    
    async def benchmark_stress_test(self, duration_seconds: int = 30) -> Dict[str, Any]:
        """Continuous stress test for specified duration."""
        print(f"\n⚡ Stress Test ({duration_seconds}s continuous generation)...")
        
        if not self.engine or not self.engine.is_initialized:
            return {'error': 'Engine not initialized'}
        
        start_time = time.time()
        end_time = start_time + duration_seconds
        
        generation_count = 0
        success_count = 0
        times = []
        errors = []
        
        while time.time() < end_time:
            text = self.test_texts[generation_count % len(self.test_texts)]
            emotion = self.emotions[generation_count % len(self.emotions)]
            
            gen_start = time.time()
            try:
                output_file = await self.engine.generate_speech(text, emotion=emotion)
                gen_time = time.time() - gen_start
                
                if output_file and os.path.exists(output_file):
                    times.append(gen_time)
                    success_count += 1
                    os.remove(output_file)  # Clean up immediately
                
                generation_count += 1
                
                # Brief pause to prevent overwhelming
                await asyncio.sleep(0.1)
                
            except Exception as e:
                errors.append(str(e))
                generation_count += 1
        
        actual_duration = time.time() - start_time
        
        result = {
            'duration_seconds': actual_duration,
            'total_attempts': generation_count,
            'successful_generations': success_count,
            'failed_generations': generation_count - success_count,
            'success_rate': success_count / generation_count * 100 if generation_count > 0 else 0,
            'avg_generation_time': statistics.mean(times) if times else 0,
            'generations_per_second': success_count / actual_duration,
            'total_errors': len(errors),
            'unique_errors': len(set(errors))
        }
        
        self.results['stress_test'] = result
        return result
    
    def benchmark_memory_usage(self) -> Dict[str, Any]:
        """Monitor memory usage."""
        print(f"\n🧠 Memory Usage Analysis...")
        
        process = psutil.Process()
        memory_info = process.memory_info()
        
        result = {
            'rss_mb': memory_info.rss / (1024 * 1024),  # Resident Set Size
            'vms_mb': memory_info.vms / (1024 * 1024),  # Virtual Memory Size
            'cpu_percent': process.cpu_percent(),
            'memory_percent': process.memory_percent(),
            'num_threads': process.num_threads()
        }
        
        self.results['memory_usage'] = result
        return result
    
    async def benchmark_voice_variety(self) -> Dict[str, Any]:
        """Test different voices and emotions."""
        print(f"\n🎭 Voice Variety Benchmark...")
        
        if not self.engine or not self.engine.is_initialized:
            return {'error': 'Engine not initialized'}
        
        if not hasattr(self.engine, 'available_voices') or not self.engine.available_voices:
            return {'error': 'No available voices'}
        
        # Test first 10 voices to keep benchmark reasonable
        test_voices = self.engine.available_voices[:10]
        test_text = "This is a voice variety test for performance benchmarking."
        
        voice_results = {}
        times = []
        
        for voice in test_voices:
            print(f"  Testing voice: {voice}")
            
            voice_times = []
            for emotion in self.emotions:
                start_time = time.time()
                try:
                    # Temporarily override emotion voice mapping
                    original_voice = self.engine.emotion_styles[emotion].get('voice')
                    self.engine.emotion_styles[emotion]['voice'] = voice
                    
                    output_file = await self.engine.generate_speech(test_text, emotion=emotion)
                    gen_time = time.time() - start_time
                    
                    if output_file and os.path.exists(output_file):
                        voice_times.append(gen_time)
                        os.remove(output_file)
                    
                    # Restore original voice
                    if original_voice:
                        self.engine.emotion_styles[emotion]['voice'] = original_voice
                        
                except Exception as e:
                    print(f"    Error with {voice}/{emotion}: {e}")
            
            if voice_times:
                voice_results[voice] = {
                    'avg_time': statistics.mean(voice_times),
                    'min_time': min(voice_times),
                    'max_time': max(voice_times),
                    'successful_emotions': len(voice_times)
                }
                times.extend(voice_times)
        
        result = {
            'voices_tested': len(test_voices),
            'successful_voices': len(voice_results),
            'avg_time_all_voices': statistics.mean(times) if times else 0,
            'voice_details': voice_results
        }
        
        self.results['voice_variety'] = result
        return result
    
    async def benchmark_long_text(self) -> Dict[str, Any]:
        """Test performance with longer text passages."""
        print(f"\n📖 Long Text Benchmark...")
        
        if not self.engine or not self.engine.is_initialized:
            return {'error': 'Engine not initialized'}
        
        # Test with different text lengths
        text_lengths = [
            ("Short", self.test_texts[0]),
            ("Medium", " ".join(self.test_texts)),
            ("Long", self.long_text),
            ("Very Long", self.long_text * 2)
        ]
        
        length_results = {}
        
        for length_name, text in text_lengths:
            print(f"  Testing {length_name} text ({len(text)} chars)...")
            
            start_time = time.time()
            try:
                output_file = await self.engine.generate_speech(text, emotion="neutral")
                gen_time = time.time() - start_time
                
                if output_file and os.path.exists(output_file):
                    file_size = os.path.getsize(output_file)
                    audio_duration = file_size / (24000 * 2)  # Approximate duration
                    
                    length_results[length_name] = {
                        'char_count': len(text),
                        'generation_time': gen_time,
                        'file_size': file_size,
                        'estimated_audio_duration': audio_duration,
                        'chars_per_second': len(text) / gen_time,
                        'realtime_factor': audio_duration / gen_time if gen_time > 0 else 0
                    }
                    
                    os.remove(output_file)
                    print(f"    ✅ {gen_time:.3f}s ({file_size:,} bytes)")
                else:
                    print(f"    ❌ Failed")
                    
            except Exception as e:
                print(f"    ❌ Error: {e}")
        
        result = {
            'length_tests': len(text_lengths),
            'successful_tests': len(length_results),
            'length_details': length_results
        }
        
        self.results['long_text'] = result
        return result
    
    def print_summary(self):
        """Print comprehensive benchmark summary."""
        print("\n" + "="*80)
        print("🏆 KOKORO ONNX v1.0 BENCHMARK RESULTS SUMMARY")
        print("="*80)
        
        # Initialization
        init = self.results.get('initialization', {})
        if init:
            print(f"\n🚀 INITIALIZATION:")
            print(f"   ⏱️  Startup Time: {init.get('total_startup_time', 0):.3f}s")
            print(f"   🎤 Available Voices: {init.get('available_voices', 0)}")
            print(f"   🔥 Using ONNX: {init.get('using_onnx', False)}")
        
        # Single Generation
        single = self.results.get('single_generation', {})
        if single and 'avg_time' in single:
            print(f"\n📊 SINGLE GENERATION:")
            print(f"   ✅ Success Rate: {single['success_rate']:.1f}%")
            print(f"   ⏱️  Average Time: {single['avg_time']:.3f}s")
            print(f"   ⚡ Range: {single['min_time']:.3f}s - {single['max_time']:.3f}s")
            print(f"   📁 Average File Size: {single['avg_file_size']/1024:.1f} KB")
        
        # Batch Generation
        batch = self.results.get('batch_generation', {})
        if batch and 'throughput_per_second' in batch:
            print(f"\n🔄 BATCH GENERATION:")
            print(f"   ✅ Success Rate: {batch['success_rate']:.1f}%")
            print(f"   🚀 Throughput: {batch['throughput_per_second']:.2f} generations/second")
            print(f"   ⏱️  Batch Time: {batch['total_batch_time']:.3f}s")
        
        # Stress Test
        stress = self.results.get('stress_test', {})
        if stress and 'generations_per_second' in stress:
            print(f"\n⚡ STRESS TEST:")
            print(f"   ✅ Success Rate: {stress['success_rate']:.1f}%")
            print(f"   🚀 Sustained Rate: {stress['generations_per_second']:.2f} gen/sec")
            print(f"   📊 Total Generations: {stress['successful_generations']}")
            print(f"   ❌ Total Errors: {stress['total_errors']}")
        
        # Memory Usage
        memory = self.results.get('memory_usage', {})
        if memory:
            print(f"\n🧠 MEMORY USAGE:")
            print(f"   💾 RAM Usage: {memory['rss_mb']:.1f} MB")
            print(f"   🧮 CPU Usage: {memory['cpu_percent']:.1f}%")
            print(f"   🧵 Threads: {memory['num_threads']}")
        
        # Voice Variety
        voices = self.results.get('voice_variety', {})
        if voices and 'voices_tested' in voices:
            print(f"\n🎭 VOICE VARIETY:")
            print(f"   🎤 Voices Tested: {voices['voices_tested']}")
            print(f"   ✅ Successful: {voices['successful_voices']}")
            print(f"   ⏱️  Average Time: {voices['avg_time_all_voices']:.3f}s")
        
        # Long Text
        long_text = self.results.get('long_text', {})
        if long_text and 'length_details' in long_text:
            print(f"\n📖 LONG TEXT PERFORMANCE:")
            for length_name, details in long_text['length_details'].items():
                rt_factor = details.get('realtime_factor', 0)
                print(f"   {length_name}: {details['chars_per_second']:.0f} chars/sec (RT factor: {rt_factor:.2f}x)")
        
        print("\n" + "="*80)
        print("🎉 BENCHMARK COMPLETE!")
        print("="*80)

async def main():
    """Run the complete benchmark suite."""
    print("🔥 KOKORO ONNX v1.0 PERFORMANCE BENCHMARK")
    print("=" * 50)
    
    benchmark = KokoroBenchmark()
    
    try:
        # Initialize
        await benchmark.initialize_engine()
        
        if not benchmark.engine or not benchmark.engine.is_initialized:
            print("❌ Failed to initialize engine. Benchmark aborted.")
            return
        
        # Run all benchmarks
        await benchmark.benchmark_single_generation(iterations=15)
        await benchmark.benchmark_batch_generation(batch_size=8)
        await benchmark.benchmark_stress_test(duration_seconds=20)
        benchmark.benchmark_memory_usage()
        await benchmark.benchmark_voice_variety()
        await benchmark.benchmark_long_text()
        
        # Print summary
        benchmark.print_summary()
        
    except KeyboardInterrupt:
        print("\n\n⏹️  Benchmark interrupted by user")
    except Exception as e:
        print(f"\n\n❌ Benchmark error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        # Cleanup
        if benchmark.engine:
            await benchmark.engine.shutdown()
        
        # Force garbage collection
        gc.collect()

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