127 lines
4.5 KiB
Python
127 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
|
|
from config import kokoro_bundle_dir, settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _kokoro_root() -> Path:
|
|
p = kokoro_bundle_dir().resolve()
|
|
if not p.is_dir():
|
|
raise FileNotFoundError(f"Kokoro model directory missing: {p}")
|
|
cfg = p / "config.json"
|
|
w0 = p / "kokoro-v1_0.pth"
|
|
w1 = p / "kokoro-v1_1-zh.pth"
|
|
weight = w0 if w0.is_file() else (w1 if w1.is_file() else None)
|
|
if not cfg.is_file() or weight is None:
|
|
raise FileNotFoundError(f"Kokoro bundle incomplete (need config.json and weight .pth): {p}")
|
|
return p
|
|
|
|
|
|
class TTSService:
|
|
def __init__(self) -> None:
|
|
self._pipeline = None
|
|
self._kokoro_root: Optional[Path] = None
|
|
self._ready = False
|
|
self._attempted = False
|
|
self._init_error: Optional[str] = None
|
|
|
|
def _voice_pt_path(self) -> Path:
|
|
root = self._kokoro_root
|
|
if root is None:
|
|
raise RuntimeError("Kokoro pipeline not initialized")
|
|
name = (settings.tts_voice or "").strip()
|
|
if not name:
|
|
raise ValueError("TTS_VOICE must be set (voice name, e.g. zf_xiaoxiao)")
|
|
p = root / "voices" / f"{name}.pt"
|
|
if not p.is_file():
|
|
raise FileNotFoundError(f"Voice file missing: {p}")
|
|
return p
|
|
|
|
def _ensure_loaded(self) -> None:
|
|
if self._ready or self._attempted:
|
|
return
|
|
self._attempted = True
|
|
try:
|
|
import kokoro
|
|
|
|
root = _kokoro_root()
|
|
w0 = root / "kokoro-v1_0.pth"
|
|
w1 = root / "kokoro-v1_1-zh.pth"
|
|
weight = w0 if w0.is_file() else w1
|
|
hf_repo_id = (
|
|
"hexgrad/Kokoro-82M-v1.1-zh" if weight.name.startswith("kokoro-v1_1") else "hexgrad/Kokoro-82M"
|
|
)
|
|
kmodel = kokoro.KModel(
|
|
repo_id=hf_repo_id,
|
|
config=str(root / "config.json"),
|
|
model=str(weight),
|
|
)
|
|
self._pipeline = kokoro.KPipeline(lang_code="z", model=kmodel, repo_id=str(root))
|
|
self._kokoro_root = root
|
|
self._ready = True
|
|
except Exception as exc: # pragma: no cover
|
|
self._init_error = str(exc)
|
|
logger.error("TTS init failed: %s", exc)
|
|
|
|
@staticmethod
|
|
def _normalize_text(text: str) -> str:
|
|
cleaned = (text or "").strip()
|
|
if not cleaned:
|
|
return ""
|
|
cleaned = cleaned.replace("\r", "\n")
|
|
cleaned = re.sub(r"[\t\f\v]+", " ", cleaned)
|
|
cleaned = re.sub(r"\s*([，。！？；：、,.!?;:])\s*", r"\1", cleaned)
|
|
cleaned = re.sub(r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])", "", cleaned)
|
|
cleaned = re.sub(r"\n{2,}", "\n", cleaned)
|
|
cleaned = re.sub(r"\s{2,}", " ", cleaned)
|
|
return cleaned.strip()
|
|
|
|
@staticmethod
|
|
def _apply_edge_fade(audio: np.ndarray, sr: int) -> np.ndarray:
|
|
if audio.size == 0:
|
|
return audio
|
|
fade_samples = max(1, min(int(sr * max(settings.tts_fade_ms, 0) / 1000.0), audio.shape[0] // 8))
|
|
if fade_samples <= 1:
|
|
return audio
|
|
out = np.array(audio, copy=True)
|
|
ramp = np.linspace(0.0, 1.0, fade_samples, dtype=np.float32)
|
|
out[:fade_samples] *= ramp
|
|
out[-fade_samples:] *= ramp[::-1]
|
|
return out
|
|
|
|
async def synthesize(self, text: str) -> tuple[np.ndarray, int]:
|
|
text = self._normalize_text(text)
|
|
if not text:
|
|
return np.zeros(1, dtype=np.float32), 24000
|
|
self._ensure_loaded()
|
|
if not self._ready or self._pipeline is None:
|
|
return np.zeros(1, dtype=np.float32), 24000
|
|
try:
|
|
voice_path = str(self._voice_pt_path())
|
|
except (OSError, ValueError, RuntimeError) as exc:
|
|
logger.error("TTS voice: %s", exc)
|
|
return np.zeros(1, dtype=np.float32), 24000
|
|
chunks = []
|
|
for _, _, audio in self._pipeline(
|
|
text,
|
|
voice=voice_path,
|
|
speed=max(0.8, min(1.1, settings.tts_speed)),
|
|
split_pattern=r"\n+",
|
|
):
|
|
chunks.append(self._apply_edge_fade(np.asarray(audio, dtype=np.float32), 24000))
|
|
if chunks:
|
|
return np.concatenate(chunks), 24000
|
|
return np.zeros(1, dtype=np.float32), 24000
|
|
|
|
@property
|
|
def health(self) -> dict:
|
|
return {"ready": self._ready, "attempted": self._attempted, "error": self._init_error}
|