| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | import base64 |
| 5 | import io |
| 6 | import math |
| 7 | import re |
| 8 | import warnings |
| 9 | from typing import Any |
| 10 | |
| 11 | import soundfile as sf |
| 12 | |
| 13 | from helpers import plugins |
| 14 | from helpers.notification import ( |
| 15 | NotificationManager, |
| 16 | NotificationPriority, |
| 17 | NotificationType, |
| 18 | ) |
| 19 | from helpers.print_style import PrintStyle |
| 20 | from plugins._kokoro_tts.helpers import migration |
| 21 | |
| 22 | |
| 23 | warnings.filterwarnings("ignore", category=FutureWarning) |
| 24 | warnings.filterwarnings("ignore", category=UserWarning) |
| 25 | |
| 26 | |
| 27 | PLUGIN_NAME = "_kokoro_tts" |
| 28 | DEFAULT_CONFIG = { |
| 29 | "voice": "am_puck,am_onyx", |
| 30 | "voice_weights": {}, |
| 31 | "speed": 1.1, |
| 32 | } |
| 33 | VOICE_ID_PATTERN = re.compile(r"^[a-z]{2}_[a-z0-9_]+$") |
| 34 | |
| 35 | _pipeline = None |
| 36 | is_updating_model = False |
| 37 | |
| 38 | |
| 39 | def normalize_config(config: dict[str, Any] | None) -> dict[str, Any]: |
| 40 | normalized = {**DEFAULT_CONFIG, "voice_weights": {}} |
| 41 | if not isinstance(config, dict): |
| 42 | return normalized |
| 43 | |
| 44 | voice = str(config.get("voice", normalized["voice"]) or "").strip() |
| 45 | if voice: |
| 46 | normalized["voice"] = voice |
| 47 | |
| 48 | weights = config.get("voice_weights") |
| 49 | if isinstance(weights, dict): |
| 50 | for raw_voice, raw_weight in weights.items(): |
| 51 | voice_id = str(raw_voice or "").strip() |
| 52 | if not VOICE_ID_PATTERN.fullmatch(voice_id): |
| 53 | continue |
| 54 | try: |
| 55 | weight = float(raw_weight) |
| 56 | except (TypeError, ValueError): |
| 57 | continue |
| 58 | if math.isfinite(weight) and weight > 0: |
| 59 | normalized["voice_weights"][voice_id] = weight |
| 60 | |
| 61 | if normalized["voice_weights"]: |
| 62 | normalized["voice"] = ",".join(normalized["voice_weights"]) |
| 63 | |
| 64 | try: |
| 65 | speed = float(config.get("speed", normalized["speed"])) |
| 66 | if math.isfinite(speed) and speed > 0: |
| 67 | normalized["speed"] = speed |
| 68 | except (TypeError, ValueError): |
| 69 | pass |
| 70 | |
| 71 | return normalized |
| 72 | |
| 73 | |
| 74 | def get_config() -> dict[str, Any]: |
| 75 | config = plugins.get_plugin_config(PLUGIN_NAME) or {} |
| 76 | return normalize_config(config) |
| 77 | |
| 78 | |
| 79 | def is_globally_enabled() -> bool: |
| 80 | migration.ensure_migrated() |
| 81 | return plugins.determined_toggle_from_paths( |
| 82 | True, reversed(plugins.get_plugin_roots(PLUGIN_NAME)) |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | async def preload(config: dict[str, Any] | None = None): |
| 87 | return await _preload() |
| 88 | |
| 89 | |
| 90 | async def _preload(): |
| 91 | global _pipeline, is_updating_model |
| 92 | |
| 93 | while is_updating_model: |
| 94 | await asyncio.sleep(0.1) |
| 95 | |
| 96 | try: |
| 97 | is_updating_model = True |
| 98 | if not _pipeline: |
| 99 | NotificationManager.send_notification( |
| 100 | NotificationType.INFO, |
| 101 | NotificationPriority.NORMAL, |
| 102 | "Loading Kokoro TTS model...", |
| 103 | display_time=99, |
| 104 | group="kokoro-preload", |
| 105 | ) |
| 106 | PrintStyle.standard("Loading Kokoro TTS model...") |
| 107 | from kokoro import KPipeline |
| 108 | |
| 109 | _pipeline = KPipeline(lang_code="a", repo_id="hexgrad/Kokoro-82M") |
| 110 | NotificationManager.send_notification( |
| 111 | NotificationType.INFO, |
| 112 | NotificationPriority.NORMAL, |
| 113 | "Kokoro TTS model loaded.", |
| 114 | display_time=2, |
| 115 | group="kokoro-preload", |
| 116 | ) |
| 117 | finally: |
| 118 | is_updating_model = False |
| 119 | |
| 120 | |
| 121 | async def is_downloading() -> bool: |
| 122 | return is_updating_model |
| 123 | |
| 124 | |
| 125 | async def is_downloaded() -> bool: |
| 126 | return _pipeline is not None |
| 127 | |
| 128 | |
| 129 | async def synthesize_sentences( |
| 130 | sentences: list[str], config: dict[str, Any] | None = None |
| 131 | ) -> str: |
| 132 | cfg = normalize_config(config or get_config()) |
| 133 | return await _synthesize_sentences( |
| 134 | sentences, |
| 135 | voice=str(cfg["voice"]), |
| 136 | voice_weights=dict(cfg["voice_weights"]), |
| 137 | speed=float(cfg["speed"]), |
| 138 | ) |
| 139 | |
| 140 | |
| 141 | def _resolve_voice( |
| 142 | pipeline: Any, voice: str, voice_weights: dict[str, float] |
| 143 | ) -> Any: |
| 144 | if not voice_weights: |
| 145 | return voice |
| 146 | |
| 147 | total = sum(voice_weights.values()) |
| 148 | if not math.isfinite(total) or total <= 0: |
| 149 | return voice |
| 150 | blend = None |
| 151 | for voice_id, weight in voice_weights.items(): |
| 152 | weighted_pack = pipeline.load_single_voice(voice_id) * (weight / total) |
| 153 | blend = weighted_pack if blend is None else blend + weighted_pack |
| 154 | return blend |
| 155 | |
| 156 | |
| 157 | async def _synthesize_sentences( |
| 158 | sentences: list[str], *, voice: str, voice_weights: dict[str, float], speed: float |
| 159 | ) -> str: |
| 160 | await _preload() |
| 161 | |
| 162 | combined_audio: list[float] = [] |
| 163 | resolved_voice = _resolve_voice(_pipeline, voice, voice_weights) |
| 164 | |
| 165 | try: |
| 166 | for sentence in sentences: |
| 167 | if not sentence.strip(): |
| 168 | continue |
| 169 | |
| 170 | segments = _pipeline( # type: ignore[misc] |
| 171 | sentence.strip(), voice=resolved_voice, speed=speed |
| 172 | ) |
| 173 | for segment in list(segments): |
| 174 | audio_tensor = segment.audio |
| 175 | audio_numpy = audio_tensor.detach().cpu().numpy() # type: ignore[union-attr] |
| 176 | combined_audio.extend(audio_numpy.tolist()) |
| 177 | |
| 178 | if not combined_audio: |
| 179 | return "" |
| 180 | |
| 181 | buffer = io.BytesIO() |
| 182 | sf.write(buffer, combined_audio, 24000, format="WAV") |
| 183 | return base64.b64encode(buffer.getvalue()).decode("utf-8") |
| 184 | except Exception as e: |
| 185 | PrintStyle.error(f"Error in Kokoro TTS synthesis: {e}") |
| 186 | raise |