| from pathlib import Path |
| from kokoro_onnx import Kokoro |
| from misaki import espeak, ja, en, zh |
| from misaki.espeak import EspeakG2P |
| import re |
| from functools import lru_cache |
| from loguru import logger |
| import onnxruntime |
| import os |
| from lib.utils import Timer, write_audio |
|
|
| providers = onnxruntime.get_available_providers() |
| print(f"Available onnx runtime providers: {providers}") |
| MODEL_DIR = Path(r"D:\yujuan\yoyo-translator-win\models\kokoro") |
|
|
| def create_session(model_path): |
| |
| providers = onnxruntime.get_available_providers() |
| print(f"Available onnx runtime providers: {providers}") |
|
|
| |
| sess_options = onnxruntime.SessionOptions() |
| cpu_count = os.cpu_count() // 2 |
| print(f"Setting threads to CPU cores count: {cpu_count}") |
| sess_options.intra_op_num_threads = cpu_count |
| session = onnxruntime.InferenceSession( |
| model_path, providers=providers, sess_options=sess_options |
| ) |
| return session |
|
|
|
|
| class KokoroTTS: |
| language_voice_mapping = { |
| "JP": "jf_alpha", |
| "JA": "jf_alpha", |
| "ZH": "zf_xiaoyi", |
| "EN": "af_heart", |
| "FR": "ff_siwis", |
| "IT": "im_nicola", |
| "HI": "hf_alpha", |
| "PT": "im_nicola", |
| "ES": "im_nicola" |
| } |
|
|
| def __init__(self, model_path: str, voice_model_path: str, vocab_config=None, gcp=None, voice=None): |
| self._session = create_session(model_path) |
| self.model = Kokoro.from_session(self._session, voice_model_path, vocab_config=vocab_config) |
| self.g2p = gcp |
| self.voice = voice |
|
|
| @classmethod |
| def from_language(cls, language: str, model_dir: Path=MODEL_DIR): |
| model_path: str = str(model_dir/"kokoro-quant.onnx") |
| voice_model_path: str = str(model_dir/"voices-v1.0.bin") |
| voice = cls.language_voice_mapping.get(language.upper()) |
| logger.info(f"[TTS] language: {language}") |
| if not voice: |
| raise ValueError(f"Unsupported language: {language}, voice: {voice}") |
| if language.upper() == "ZH": |
| tts = cls(model_path, voice_model_path, vocab_config=model_dir / "zh_config.json", gcp=zh.ZHG2P(), |
| voice=voice) |
| tts.generate("你好") |
| elif language.upper() in ['JP', 'JA']: |
| tts = cls(model_path, voice_model_path, vocab_config=model_dir / "ja_config.json", gcp=ja.JAG2P(), |
| voice=voice) |
| elif language.upper() == 'EN': |
| fallback = espeak.EspeakFallback(british=False) |
| tts = cls(model_path, voice_model_path, gcp=en.G2P(trf=False, british=False, fallback=fallback), |
| voice=voice) |
| tts.generate("hello") |
| elif language.upper() == "HI": |
| g2p = EspeakG2P(language="hi") |
| tts = cls(model_path, voice_model_path, gcp=g2p, voice=voice) |
| tts.generate("हेलो") |
| elif language.upper() == "IT": |
| g2p = EspeakG2P(language="it") |
| tts = cls(model_path, voice_model_path, gcp=g2p, voice=voice) |
| tts.generate("Ciao") |
| elif language.upper() == "PT": |
| g2p = EspeakG2P(language="pt-br") |
| tts = cls(model_path, voice_model_path, gcp=g2p, voice=voice) |
| tts.generate("Olá") |
| elif language.upper() == "ES": |
| g2p = EspeakG2P(language="es") |
| tts = cls(model_path, voice_model_path, gcp=g2p, voice=voice) |
| tts.generate("Hola") |
| elif language.upper() == "FR": |
| g2p = EspeakG2P(language="fr-fr") |
| tts = cls(model_path, voice_model_path, gcp=g2p, voice=voice) |
| tts.generate("Bonjour") |
| else: |
| tts = cls(model_path, voice_model_path, gcp=EspeakG2P(language.lower()), voice=voice) |
| return tts |
|
|
| def generate(self, text, speed=1.2): |
| with Timer("tts inference") as t: |
| phonemes, _ = self.g2p(text) |
| samples, sample_rate = self.model.create(phonemes, self.voice, is_phonemes=True, speed=speed) |
|
|
| return samples, sample_rate, t.duration |
|
|
| async def stream(self, text, speed=1.2): |
| phonemes, _ = self.g2p(text) |
| stream = self.model.create_stream(phonemes, self.voice, is_phonemes=True, speed=speed) |
| async for samples, sample_rate in stream: |
| yield samples, sample_rate |
|
|
|
|
| @lru_cache |
| def get_model(language): |
| return KokoroTTS.from_language(language=language, model_dir_path=resource_path('models/kokoro')) |
|
|
