diff --git a/plugins/_kokoro_tts/api/resolve_hf_repo.py b/plugins/_kokoro_tts/api/resolve_hf_repo.py new file mode 100644 index 0000000000..2f41233139 --- /dev/null +++ b/plugins/_kokoro_tts/api/resolve_hf_repo.py @@ -0,0 +1,61 @@ +import json +import urllib.error +import urllib.request + +from helpers.api import ApiHandler, Request, Response + + +class ResolveHfRepo(ApiHandler): + async def process(self, input: dict, request: Request) -> dict | Response: + repo = str(input.get("repo", "") or "").strip() + if not repo: + return {"success": False, "error": "No repo ID provided."} + + try: + url = f"https://huggingface.co/api/models/{repo}" + with urllib.request.urlopen(url, timeout=15) as resp: + data = json.loads(resp.read()) + except urllib.error.HTTPError as e: + if e.code == 404: + return {"success": False, "error": f"Repo '{repo}' not found."} + return {"success": False, "error": f"HuggingFace API error: {e}"} + except Exception as e: + return {"success": False, "error": f"Network error: {e}"} + + files = [s["rfilename"] for s in data.get("siblings", [])] + + # Find model: prefer full-precision .onnx + onnx_files = [f for f in files if f.endswith(".onnx")] + if not onnx_files: + return {"success": False, "error": f"No .onnx file found in {repo}."} + model_file = next( + ( + f + for f in onnx_files + if not any(q in f.lower() for q in ("int8", "quantized", "fp16", "q8")) + ), + onnx_files[0], + ) + + # Find voices: .npz or .bin (not .onnx) + voices_files = [ + f + for f in files + if f.endswith(".npz") or (f.endswith(".bin") and not f.endswith(".onnx")) + ] + if not voices_files: + return { + "success": False, + "error": f"No voices file (.npz or .bin) found in {repo}.", + } + # Prefer .npz over .bin + npz_files = [f for f in voices_files if f.endswith(".npz")] + voices_file = npz_files[0] if npz_files else voices_files[0] + + return { + "success": True, + "repo": repo, + "model_file": model_file, + "voices_file": voices_file, + "all_files": files, + } diff --git a/plugins/_kokoro_tts/api/status.py b/plugins/_kokoro_tts/api/status.py index 3b1321972e..50acc9f811 100644 --- a/plugins/_kokoro_tts/api/status.py +++ b/plugins/_kokoro_tts/api/status.py @@ -4,28 +4,37 @@ from plugins._kokoro_tts.helpers import migration, runtime +def _pkg_version(name: str) -> tuple[str, str]: + try: + return importlib.metadata.version(name), "" + except Exception as e: + return "", str(e) + + class Status(ApiHandler): async def process(self, input: dict, request: Request) -> dict | Response: migration.ensure_migrated() - package_version = "" - package_error = "" - try: - package_version = importlib.metadata.version("kokoro") - except Exception as e: - package_error = str(e) + cfg = runtime.get_config() + py_version, py_error = _pkg_version("kokoro") + onnx_version, onnx_error = _pkg_version("kokoro_onnx") return { "plugin": "_kokoro_tts", "enabled": runtime.is_globally_enabled(), - "config": runtime.get_config(), + "config": cfg, + "engine": cfg.get("engine", "kokoro_py"), "model": { "ready": await runtime.is_downloaded(), "loading": await runtime.is_downloading(), }, "package": { - "version": package_version, - "error": package_error, + "version": py_version, + "error": py_error, + }, + "onnx_package": { + "version": onnx_version, + "error": onnx_error, }, "fallback": "Browser-native speechSynthesis remains the fallback when Kokoro is disabled.", } diff --git a/plugins/_kokoro_tts/default_config.yaml b/plugins/_kokoro_tts/default_config.yaml index bd1c4ec635..2cb7963523 100644 --- a/plugins/_kokoro_tts/default_config.yaml +++ b/plugins/_kokoro_tts/default_config.yaml @@ -1,3 +1,10 @@ voice: am_puck,am_onyx voice_weights: {} speed: 1.1 +engine: kokoro_py # kokoro_py | kokoro_onnx +lang: en-us # espeak-ng language code: en-us, de, fr, es, it, etc. +onnx_hf_repo: '' # HuggingFace repo ID, e.g. Godelaune/Kokoro-82M-ONNX-German-Martin +onnx_model_file: '' # auto-detected from repo if empty; e.g. kokoro-martin.onnx +onnx_voices_file: '' # auto-detected from repo if empty; e.g. voices-martin.npz +onnx_mixed_lang: en-us # secondary espeak-ng language for matched terms +onnx_mixed_lang_terms: '' # comma/newline-separated terms to phonemize with secondary language diff --git a/plugins/_kokoro_tts/helpers/runtime.py b/plugins/_kokoro_tts/helpers/runtime.py index c8237c6aec..1fe3d6f6ae 100644 --- a/plugins/_kokoro_tts/helpers/runtime.py +++ b/plugins/_kokoro_tts/helpers/runtime.py @@ -4,13 +4,15 @@ import base64 import io import math +import os import re +import urllib.request import warnings from typing import Any import soundfile as sf -from helpers import plugins +from helpers import files, plugins from helpers.notification import ( NotificationManager, NotificationPriority, @@ -29,10 +31,28 @@ "voice": "am_puck,am_onyx", "voice_weights": {}, "speed": 1.1, + "engine": "kokoro_py", # kokoro_py | kokoro_onnx + "lang": "en-us", # espeak-ng language code: en-us, de, fr, es, it, etc. + "onnx_hf_repo": "", # HuggingFace repo ID for ONNX model download + "onnx_model_file": "", # filename in repo, e.g. kokoro-martin.onnx + "onnx_voices_file": "", # filename in repo, e.g. voices-martin.npz + "onnx_mixed_lang": "en-us", # secondary espeak-ng language for matched terms + "onnx_mixed_lang_terms": "", # comma/newline-separated terms to phonemize with secondary language } VOICE_ID_PATTERN = re.compile(r"^[a-z]{2}_[a-z0-9_]+$") +# Map espeak-ng language codes to kokoro-py KPipeline lang_code +_ESPEAK_TO_KOKORO_LANG = { + "en-us": "a", + "en-gb": "b", + "ja": "j", + "zh": "z", + "ko": "k", +} + _pipeline = None +_pipeline_lang_code: str | None = None +_onnx_pipeline = None is_updating_model = False @@ -68,6 +88,28 @@ def normalize_config(config: dict[str, Any] | None) -> dict[str, Any]: except (TypeError, ValueError): pass + # Engine selection + engine = str(config.get("engine", normalized["engine"]) or "").strip().lower() + if engine in ("kokoro_py", "kokoro_onnx"): + normalized["engine"] = engine + + # Language code (espeak-ng format) + lang = str(config.get("lang", normalized["lang"]) or "").strip().lower() + if lang: + normalized["lang"] = lang + + # ONNX HuggingFace repo and filenames + for key in ("onnx_hf_repo", "onnx_model_file", "onnx_voices_file"): + val = str(config.get(key, normalized[key]) or "").strip() + normalized[key] = val + + mixed_lang = str(config.get("onnx_mixed_lang", normalized["onnx_mixed_lang"]) or "").strip().lower() + if mixed_lang: + normalized["onnx_mixed_lang"] = mixed_lang + normalized["onnx_mixed_lang_terms"] = str( + config.get("onnx_mixed_lang_terms", normalized["onnx_mixed_lang_terms"]) or "" + ).strip() + return normalized @@ -84,18 +126,26 @@ def is_globally_enabled() -> bool: async def preload(config: dict[str, Any] | None = None): - return await _preload() + cfg = normalize_config(config or get_config()) + if cfg["engine"] == "kokoro_onnx": + return await _preload_onnx(cfg) + return await _preload(cfg) + +async def _preload(config: dict[str, Any] | None = None): + global _pipeline, _pipeline_lang_code, is_updating_model -async def _preload(): - global _pipeline, is_updating_model + cfg = normalize_config(config or get_config()) + lang_code = _ESPEAK_TO_KOKORO_LANG.get( + cfg["lang"], cfg["lang"][0] if cfg["lang"] else "a" + ) while is_updating_model: await asyncio.sleep(0.1) try: is_updating_model = True - if not _pipeline: + if _pipeline is None or _pipeline_lang_code != lang_code: NotificationManager.send_notification( NotificationType.INFO, NotificationPriority.NORMAL, @@ -106,7 +156,8 @@ async def _preload(): PrintStyle.standard("Loading Kokoro TTS model...") from kokoro import KPipeline - _pipeline = KPipeline(lang_code="a", repo_id="hexgrad/Kokoro-82M") + _pipeline = KPipeline(lang_code=lang_code, repo_id="hexgrad/Kokoro-82M") + _pipeline_lang_code = lang_code NotificationManager.send_notification( NotificationType.INFO, NotificationPriority.NORMAL, @@ -118,11 +169,159 @@ async def _preload(): is_updating_model = False +def _detect_hf_files(hf_repo: str) -> tuple[str, str]: + """Auto-detect ONNX model and voices filenames from a HuggingFace repo. + + Returns ``(model_file, voices_file)``. Raises :class:`ValueError` if + either file cannot be found. + """ + import json + + url = f"https://huggingface.co/api/models/{hf_repo}" + with urllib.request.urlopen(url, timeout=15) as resp: + data = json.loads(resp.read()) + + files = [s["rfilename"] for s in data.get("siblings", [])] + + # Find model: prefer full-precision .onnx + onnx_files = [f for f in files if f.endswith(".onnx")] + if not onnx_files: + raise ValueError(f"No .onnx file found in {hf_repo}") + model_file = next( + ( + f + for f in onnx_files + if not any(q in f.lower() for q in ("int8", "quantized", "fp16", "q8")) + ), + onnx_files[0], + ) + + # Find voices: .npz or .bin (but not .onnx) + voices_files = [ + f + for f in files + if f.endswith(".npz") or (f.endswith(".bin") and not f.endswith(".onnx")) + ] + if not voices_files: + raise ValueError(f"No voices file (.npz or .bin) found in {hf_repo}") + # Prefer .npz over .bin + npz = [f for f in voices_files if f.endswith(".npz")] + voices_file = npz[0] if npz else voices_files[0] + + return model_file, voices_file + + +async def _ensure_onnx_model(cfg: dict[str, Any]) -> tuple[str, str]: + """Ensure ONNX model and voices files are available locally. + + Downloads from HuggingFace if ``onnx_hf_repo`` is set and files are not yet cached. + Returns ``(model_path, voices_path)``. Returns empty strings if ``onnx_hf_repo`` + is empty (files must exist locally in that case). + """ + hf_repo = cfg.get("onnx_hf_repo", "") + model_file = cfg.get("onnx_model_file", "") + voices_file = cfg.get("onnx_voices_file", "") + + if not hf_repo: + return "", "" + + if hf_repo and (not model_file or not voices_file): + detected_model, detected_voices = _detect_hf_files(hf_repo) + if not model_file: + model_file = detected_model + if not voices_file: + voices_file = detected_voices + + sanitized = hf_repo.replace("/", "_") + cache_dir = files.get_abs_path("usr/models", sanitized) + model_path = os.path.join(cache_dir, model_file) if model_file else "" + voices_path = os.path.join(cache_dir, voices_file) if voices_file else "" + + model_ok = bool(model_path) and os.path.isfile(model_path) + voices_ok = bool(voices_path) and os.path.isfile(voices_path) + + if model_ok and voices_ok: + return model_path, voices_path + + os.makedirs(cache_dir, exist_ok=True) + + if not model_ok and model_file: + url = f"https://huggingface.co/{hf_repo}/resolve/main/{model_file}" + PrintStyle.standard(f"Downloading ONNX model: {url}") + NotificationManager.send_notification( + NotificationType.INFO, + NotificationPriority.NORMAL, + "Downloading ONNX model from HuggingFace...", + display_time=99, + group="kokoro-onnx-download", + ) + urllib.request.urlretrieve(url, model_path) + + if not voices_ok and voices_file: + url = f"https://huggingface.co/{hf_repo}/resolve/main/{voices_file}" + PrintStyle.standard(f"Downloading ONNX voices: {url}") + urllib.request.urlretrieve(url, voices_path) + + NotificationManager.send_notification( + NotificationType.INFO, + NotificationPriority.NORMAL, + "ONNX model download complete.", + display_time=2, + group="kokoro-onnx-download", + ) + + return model_path, voices_path + + +async def _preload_onnx(config: dict[str, Any]): + global _onnx_pipeline, is_updating_model + + while is_updating_model: + await asyncio.sleep(0.1) + + try: + is_updating_model = True + if not _onnx_pipeline: + NotificationManager.send_notification( + NotificationType.INFO, + NotificationPriority.NORMAL, + "Loading Kokoro ONNX TTS model...", + display_time=99, + group="kokoro-onnx-preload", + ) + PrintStyle.standard("Loading Kokoro ONNX TTS model...") + + model_path, voices_path = await _ensure_onnx_model(config) + + if not model_path or not voices_path: + raise ValueError( + "ONNX engine requires onnx_hf_repo, onnx_model_file, and " + "onnx_voices_file to be configured, or model files to exist locally." + ) + + from kokoro_onnx import Kokoro + + _onnx_pipeline = Kokoro(model_path, voices_path) + + NotificationManager.send_notification( + NotificationType.INFO, + NotificationPriority.NORMAL, + "Kokoro ONNX TTS model loaded.", + display_time=2, + group="kokoro-onnx-preload", + ) + finally: + is_updating_model = False + + async def is_downloading() -> bool: return is_updating_model async def is_downloaded() -> bool: + cfg = get_config() + if cfg.get("engine") == "kokoro_onnx": + return _onnx_pipeline is not None return _pipeline is not None @@ -130,12 +329,7 @@ async def synthesize_sentences( sentences: list[str], config: dict[str, Any] | None = None ) -> str: cfg = normalize_config(config or get_config()) - return await _synthesize_sentences( - sentences, - voice=str(cfg["voice"]), - voice_weights=dict(cfg["voice_weights"]), - speed=float(cfg["speed"]), - ) + return await _synthesize_sentences(sentences, cfg=cfg) def _resolve_voice( @@ -155,9 +349,21 @@ def _resolve_voice( async def _synthesize_sentences( - sentences: list[str], *, voice: str, voice_weights: dict[str, float], speed: float + sentences: list[str], *, cfg: dict[str, Any] +) -> str: + if cfg["engine"] == "kokoro_onnx": + return await _synthesize_onnx(sentences, cfg=cfg) + return await _synthesize_kokoro_py(sentences, cfg=cfg) + + +async def _synthesize_kokoro_py( + sentences: list[str], *, cfg: dict[str, Any] ) -> str: - await _preload() + await _preload(cfg) + + voice = str(cfg["voice"]) + voice_weights = dict(cfg["voice_weights"]) + speed = float(cfg["speed"]) combined_audio: list[float] = [] resolved_voice = _resolve_voice(_pipeline, voice, voice_weights) @@ -178,9 +384,87 @@ async def _synthesize_sentences( if not combined_audio: return "" - buffer = io.BytesIO() - sf.write(buffer, combined_audio, 24000, format="WAV") - return base64.b64encode(buffer.getvalue()).decode("utf-8") + return _encode_wav_base64(combined_audio) except Exception as e: PrintStyle.error(f"Error in Kokoro TTS synthesis: {e}") raise + + +def _encode_wav_base64(samples: list[float]) -> str: + buffer = io.BytesIO() + sf.write(buffer, samples, 24000, format="WAV") + return base64.b64encode(buffer.getvalue()).decode("utf-8") + + +def _parse_mixed_lang_terms(raw_terms: str) -> list[str]: + terms = [term.strip() for term in re.split(r"[,\n]", raw_terms) if term.strip()] + return sorted(dict.fromkeys(terms), key=len, reverse=True) + + +def _mixed_lang_phonemes( + pipeline: Any, text: str, *, primary_lang: str, mixed_lang: str, terms: list[str] +) -> str: + if not terms: + return pipeline.tokenizer.phonemize(text, primary_lang) + + pattern = re.compile( + r"(? last_end: + phoneme_parts.append( + pipeline.tokenizer.phonemize(text[last_end : match.start()], primary_lang) + ) + phoneme_parts.append(pipeline.tokenizer.phonemize(match.group(0), mixed_lang)) + last_end = match.end() + if last_end < len(text): + phoneme_parts.append(pipeline.tokenizer.phonemize(text[last_end:], primary_lang)) + return " ".join(part for part in phoneme_parts if part) + + +async def _synthesize_onnx( + sentences: list[str], *, cfg: dict[str, Any] +) -> str: + await _preload_onnx(cfg) + + voice = str(cfg["voice"]) + speed = float(cfg["speed"]) + lang = str(cfg["lang"]) + mixed_lang = str(cfg["onnx_mixed_lang"]) + mixed_terms = _parse_mixed_lang_terms(str(cfg["onnx_mixed_lang_terms"])) + + combined_audio: list[float] = [] + + try: + for sentence in sentences: + text = sentence.strip() + if not text: + continue + + if mixed_terms: + text = _mixed_lang_phonemes( + _onnx_pipeline, + text, + primary_lang=lang, + mixed_lang=mixed_lang, + terms=mixed_terms, + ) + samples, sample_rate = _onnx_pipeline.create( + text, voice=voice, speed=speed, lang=lang, is_phonemes=True + ) + else: + samples, sample_rate = _onnx_pipeline.create( + text, voice=voice, speed=speed, lang=lang + ) + combined_audio.extend(samples.tolist()) + + if not combined_audio: + return "" + + return _encode_wav_base64(combined_audio) + except Exception as e: + PrintStyle.error(f"Error in Kokoro ONNX TTS synthesis: {e}") + raise diff --git a/plugins/_kokoro_tts/webui/config.html b/plugins/_kokoro_tts/webui/config.html index 010baeffaa..13699203c0 100644 --- a/plugins/_kokoro_tts/webui/config.html +++ b/plugins/_kokoro_tts/webui/config.html @@ -43,6 +43,19 @@ if (Object.keys(sanitized).length) this.syncVoice(Object.keys(sanitized)); const speed = Number(config.speed); config.speed = Number.isFinite(speed) && speed > 0 ? speed : 1.1; + + // Engine selection (kokoro_py | kokoro_onnx) + if (!['kokoro_py', 'kokoro_onnx'].includes(config.engine)) config.engine = 'kokoro_py'; + + // espeak-ng language code + if (!config.lang) config.lang = 'en-us'; + + // ONNX fields + if (!config.onnx_hf_repo) config.onnx_hf_repo = ''; + if (!config.onnx_model_file) config.onnx_model_file = ''; + if (!config.onnx_voices_file) config.onnx_voices_file = ''; + if (!config.onnx_mixed_lang) config.onnx_mixed_lang = 'en-us'; + if (!config.onnx_mixed_lang_terms) config.onnx_mixed_lang_terms = ''; }, voiceIds() { const weighted = Object.keys(config.voice_weights || {}); @@ -110,6 +123,31 @@ syncVoice(ids) { config.voice = ids.join(','); }, + resolving: false, + resolveError: '', + async resolveHfRepo() { + if (!config.onnx_hf_repo) return; + this.resolving = true; + this.resolveError = ''; + try { + const resp = await fetch('/api/plugins/_kokoro_tts/resolve_hf_repo', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ repo: config.onnx_hf_repo }), + }); + const data = await resp.json(); + if (data.success) { + config.onnx_model_file = data.model_file; + config.onnx_voices_file = data.voices_file; + } else { + this.resolveError = data.error || 'Could not resolve repo files'; + } + } catch (e) { + this.resolveError = String(e); + } finally { + this.resolving = false; + } + }, }">