assistant/tts.py

from __future__ import annotations

import re
import wave
from functools import lru_cache
from typing import Iterator

import numpy as np
from piper.config import SynthesisConfig
from piper.voice import PiperVoice

from config import CONFIG

_SENTENCE_END = re.compile(r"(?<=[.!?…])\s+")
# Markdown and list symbols Piper would otherwise read aloud.
_SPEAK_STRIP = re.compile(r"[*_`#>]+|^\s*[-•]\s+", re.MULTILINE)


def clean_text(text: str) -> str:
    """Strip markdown so the voice speaks words, not symbols."""
    return re.sub(r"\s+", " ", _SPEAK_STRIP.sub("", text)).strip()


@lru_cache(maxsize=1)
def _voice():
    """Load the Piper voice once and cache it."""
    return PiperVoice.load(CONFIG.piper_path())


@lru_cache(maxsize=1)
def _syn_config():
    """Prosody knobs: slower and slightly varied reads warmer, less robotic."""
    return SynthesisConfig(
        length_scale=CONFIG.tts_length_scale,
        noise_scale=CONFIG.tts_noise_scale,
        noise_w_scale=CONFIG.tts_noise_w,
    )


def synth(text: str) -> tuple[int, np.ndarray]:
    """Synthesize text to (sample_rate, int16 mono samples) for gr.Audio."""
    voice = _voice()
    chunks = [
        np.frombuffer(audio.audio_int16_bytes, dtype=np.int16)
        for audio in voice.synthesize(text, syn_config=_syn_config())
    ]
    samples = (
        np.concatenate(chunks) if chunks else np.zeros(0, dtype=np.int16)
    )
    return voice.config.sample_rate, samples


def split_sentences(text: str) -> list[str]:
    """Split text into trimmed, non-empty sentences."""
    return [s.strip() for s in _SENTENCE_END.split(text) if s.strip()]


def synth_stream(text_chunks: Iterator[str]) -> Iterator[tuple[int, np.ndarray]]:
    """Buffer streamed text and yield audio for each completed sentence."""
    buffer = ""
    for chunk in text_chunks:
        buffer += chunk
        sentences = split_sentences(buffer)
        if len(sentences) > 1:
            for sentence in sentences[:-1]:
                yield synth(sentence)
            buffer = sentences[-1]
    if buffer.strip():
        yield synth(buffer)


def save_wav(path: str, sample_rate: int, samples: np.ndarray) -> None:
    """Write int16 mono samples to a WAV file."""
    with wave.open(path, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(samples.tobytes())