javis/src/javis/tts/piper_backend.py

"""Piper TTS inference backend.

Wraps ``piper-tts`` (``pip install piper-tts``) to provide a fast,
CPU-capable ONNX TTS backend.

A trained Piper model consists of two files:

    piper_model.onnx        — ONNX weights
    piper_model.onnx.json   — model config (phonemes, sample rate, …)

The model path is resolved in this order:
  1. ``config.tts.piper_model`` (explicit path in config.yaml)
  2. ``{voices_dir}/{tts.default_voice}/piper_model.onnx`` (active voice dir)

Inference is synchronous (ONNX runtime) so it runs in a thread-pool executor
to avoid blocking the asyncio event loop.

CUDA support: install ``onnxruntime-gpu`` instead of ``onnxruntime`` and Piper
will use CUDA automatically.
"""
from __future__ import annotations

import asyncio
import re
import threading
from pathlib import Path
from typing import Optional

import numpy as np

from ..core.config import Config
from ..core.logger import get_logger
from .base import SynthesisResult, TTSBackend

log = get_logger("tts.piper")

# Markdown patterns that Piper would speak literally.
_MD_PATTERNS: list[tuple] = [
    (re.compile(r"```[\s\S]*?```"), ""),
    (re.compile(r"`([^`]+)`"), r"\1"),
    (re.compile(r"\*{1,3}([^*]*)\*{1,3}"), r"\1"),
    (re.compile(r"_{1,2}([^_]*)_{1,2}"), r"\1"),
    (re.compile(r"^#{1,6}\s+", re.MULTILINE), ""),
    (re.compile(r"\[([^\]]+)\]\([^)]+\)"), r"\1"),
    (re.compile(r"^>\s+", re.MULTILINE), ""),
    (re.compile(r"^\d+\.\s+", re.MULTILINE), ""),
    (re.compile(r"^[-*+]\s+", re.MULTILINE), ""),
]
_PUNCT_ONLY = re.compile(r"^[^\w\s]+$")


def _clean(text: str) -> str:
    for pat, repl in _MD_PATTERNS:
        text = pat.sub(repl, text)
    text = re.sub(r"  +", " ", text).strip()
    return "" if _PUNCT_ONLY.match(text) else text


class PiperBackend(TTSBackend):
    """Piper ONNX TTS backend."""

    def __init__(self, config: Config) -> None:
        self.config = config
        self._voice: Optional[object] = None   # piper.voice.PiperVoice
        self._sample_rate: int = 22050
        self._lock = threading.Lock()

    # ── TTSBackend interface ──────────────────────────────────────────────────

    def load(self) -> None:
        """Load the ONNX model.  Called once at startup."""
        try:
            from piper.voice import PiperVoice  # type: ignore[import]
        except ImportError as exc:
            raise RuntimeError(
                "piper-tts is not installed. Run: pip install piper-tts"
            ) from exc

        model_path = self._resolve_model_path()
        config_path = Path(str(model_path) + ".json")

        log.info(f"Loading Piper model: {model_path}")
        self._voice = PiperVoice.load(
            str(model_path),
            config_path=str(config_path) if config_path.exists() else None,
            use_cuda=False,  # ONNX runtime picks CUDA automatically when available
        )
        self._sample_rate = (
            self._voice.config.sample_rate  # type: ignore[union-attr]
            if hasattr(self._voice, "config")
            else 22050
        )
        log.info(
            f"Piper model loaded — sample_rate={self._sample_rate} Hz, "
            f"language={getattr(getattr(self._voice, 'config', None), 'espeak_voice', '?')}"
        )

    async def synthesize(
        self,
        text: str,
        voice_profile: object,
        language: Optional[str] = None,
    ) -> SynthesisResult:
        text = _clean(text)
        if not text:
            return SynthesisResult(
                audio=np.zeros(0, dtype=np.float32),
                sample_rate=self._sample_rate,
            )

        audio = await asyncio.to_thread(self._synthesize_sync, text)
        return SynthesisResult(audio=audio, sample_rate=self._sample_rate)

    # ── Internal ──────────────────────────────────────────────────────────────

    def _synthesize_sync(self, text: str) -> np.ndarray:
        """Run Piper inference synchronously (called from thread-pool)."""
        if self._voice is None:
            raise RuntimeError("PiperBackend.load() has not been called.")

        from piper.config import SynthesisConfig  # type: ignore[import]

        syn_config = SynthesisConfig(
            length_scale=self.config.tts.piper_length_scale,
            noise_scale=self.config.tts.piper_noise_scale,
            noise_w_scale=self.config.tts.piper_noise_w,
        )

        chunks: list[np.ndarray] = []
        with self._lock:
            for chunk in self._voice.synthesize(text, syn_config):
                arr = chunk.audio_float_array
                if arr is not None and len(arr) > 0:
                    chunks.append(np.asarray(arr, dtype=np.float32))

        if not chunks:
            return np.zeros(0, dtype=np.float32)
        return np.concatenate(chunks)

    def _resolve_model_path(self) -> Path:
        """Return the ONNX model path from config or voice profile dir."""
        explicit = self.config.tts.piper_model.strip()
        if explicit:
            p = Path(explicit)
            if p.exists():
                return p
            log.warning(f"piper_model path not found: {p}")

        # Fall back to active voice profile directory.
        voice_name = self.config.tts.default_voice
        voices_dir = self.config.paths.voices_dir
        candidate = voices_dir / voice_name / "piper_model.onnx"
        if candidate.exists():
            return candidate

        from ..voices.piper_train import has_checkpoints
        if has_checkpoints(voices_dir / voice_name):
            hint = (f"Training checkpoints exist — export them with: python -m "
                    f"javis.voices.piper_train --voice {voice_name} --export-snapshot")
        else:
            hint = "Run: python -m javis.voices.piper_train --voice <name>"
        raise FileNotFoundError(
            f"Piper ONNX model not found. Expected at {candidate} "
            f"or set tts.piper_model in config.yaml. {hint}"
        )