kokoro_engine.py

"""
Kokoro TTS Engine - Simple ONNX implementation
Works with Python 3.9+ without complex dependencies
"""
import os
import logging
import numpy as np
from huggingface_hub import hf_hub_download

log = logging.getLogger("kokoro_engine")


class KokoroEngine:
    """
    محرك Kokoro TTS بسيط باستخدام kokoro-onnx من PyPI
    """
    
    def __init__(self, voice: str = "af_alloy", sample_rate: int = 24000):
        """
        تهيئة محرك Kokoro
        
        Args:
            voice: اسم الصوت الافتراضي
            sample_rate: معدل العينة (دائماً 24000 لـ Kokoro)
        """
        self.sample_rate = sample_rate
        self.voice = voice
        self.kokoro = None
        
        try:
            # استخدام kokoro-onnx من PyPI مباشرة
            from kokoro_onnx import Kokoro
            
            # تحميل ملفات النموذج
            model_path = self._download_model()
            voices_path = self._download_voices()
            
            # تهيئة Kokoro
            self.kokoro = Kokoro(model_path, voices_path)
            
            log.info(f"✅ Kokoro-ONNX initialized successfully")
            
        except Exception as e:
            log.error(f"❌ Failed to initialize Kokoro-ONNX: {e}")
            import traceback
            traceback.print_exc()
            raise
    
    def _download_model(self) -> str:
        """تحميل ملف النموذج"""
        try:
            # جرب أولاً من onnx-community (الأكثر موثوقية)
            model_path = hf_hub_download(
                repo_id="fastrtc/kokoro-onnx",
                filename="kokoro-v1.0.onnx",
                cache_dir="./models"
                )
            log.info(f"Model downloaded from onnx-community: {model_path}")
            return model_path
        except Exception as e1:
            log.warning(f"Failed to download from onnx-community: {e1}")
            try:
                # جرب من NeuML كبديل
                model_path = hf_hub_download(
                repo_id="onnx-community/Kokoro-82M-ONNX",
                filename="onnx/model.onnx",
                cache_dir="./models"
                )
                log.info(f"Model downloaded from NeuML: {model_path}")
                return model_path
            except Exception as e2:
                log.error(f"Failed to download model from all sources")
                raise RuntimeError(f"Could not download model: {e2}")
    
    def _download_voices(self) -> str:
        """تحميل ملف الأصوات"""
        try:
            # جرب أولاً voices.bin من onnx-community
            voices_path = hf_hub_download(
                repo_id="fastrtc/kokoro-onnx",
                filename="voices-v1.0.bin",
                cache_dir="./models"
            )
            log.info(f"Voices downloaded from onnx-community: {voices_path}")
            return voices_path
        except Exception as e1:
            log.warning(f"Failed to download voices.bin: {e1}")
            try:
                # جرب voices.json من NeuML كبديل
                voices_path = hf_hub_download(
                repo_id="onnx-community/Kokoro-82M-ONNX",
                filename="voices/voices.bin",
                cache_dir="./models"
                )
                log.info(f"Voices downloaded from NeuML: {voices_path}")
                return voices_path
            except Exception as e2:
                log.error(f"Failed to download voices from all sources")
                raise RuntimeError(f"Could not download voices: {e2}")
    
    def set_voice(self, voice: str):
        """
        تغيير الصوت
        
        Args:
            voice: اسم الصوت الجديد
        """
        self.voice = voice
        log.debug(f"Voice changed to: {voice}")
    
    def synthesize(self, text: str, speed: float = 1.0) -> np.ndarray:
        """
        تحويل النص إلى صوت
        
        Args:
            text: النص المراد تحويله
            speed: سرعة التحدث (1.0 = عادي)
            
        Returns:
            np.ndarray: البيانات الصوتية كـ numpy array (float32)
        """
        if self.kokoro is None:
            raise RuntimeError("Kokoro engine not initialized")
        
        try:
            # استخدام kokoro-onnx للتوليد
            samples, sample_rate = self.kokoro.create(
                text,
                voice=self.voice,
                speed=speed,
                lang='en-us'
            )
            
            log.info(f"✅ Audio generated: {len(samples)} samples, duration: {len(samples)/sample_rate:.2f}s")
            
            # التأكد من أن النوع float32
            return samples.astype(np.float32)
            
        except Exception as e:
            log.error(f"❌ Synthesis failed: {e}")
            import traceback
            traceback.print_exc()
            raise
    
    def get_available_voices(self):
        """
        الحصول على قائمة الأصوات المتاحة
        """
        return [
            # British Female
            "bf_alice", "bf_emma", "bf_isabella", "bf_lily",
            # American Female
            "af_alloy", "af_aoede", "af_bella", "af_heart",
            "af_jessica", "af_kore", "af_nicole", "af_nova",
            "af_river", "af_sarah", "af_sky",
            # British Male
            "bm_daniel", "bm_fable", "bm_george", "bm_lewis",
            # American Male
            "am_adam", "am_echo", "am_eric", "am_fenrir",
            "am_liam", "am_michael", "am_onyx", "am_puck"
        ]