src/tts/mms_tts.py

"""
Facebook MMS-TTS engine for Bambara, Fula, French, and English.

Usage:
    engine = MMSTTSEngine()
    wav_np, sample_rate = engine.synthesize("Foro fɛ ji.", "bam", device="cuda")
    wav_bytes = engine.text_to_audio_bytes("Foro fɛ ji.", "bam", device="cuda")
"""
from __future__ import annotations

import io
import re
from typing import Dict, Tuple

import numpy as np
import soundfile as sf

MODEL_IDS: Dict[str, str] = {
    "bam": "facebook/mms-tts-bam",
    "ful": "facebook/mms-tts-ful",
    "fr":  "facebook/mms-tts-fra",
    "en":  "facebook/mms-tts-eng",
}

# Fallback for unknown languages — use French
_DEFAULT_LANG = "fr"

# MMS-TTS quality degrades beyond ~15 words; split longer text at sentence boundaries
_MAX_WORDS_PER_CHUNK = 15

# Sentence-boundary split pattern (period, exclamation, question mark followed by space or end)
_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+")


class MMSTTSEngine:
    """Lazy-loading MMS-TTS engine. Models are loaded on first use and cached in CPU RAM."""

    def __init__(self) -> None:
        # {language_code: (VitsModel, VitsTokenizer)}
        self._cache: Dict[str, tuple] = {}

    # ── private helpers ──────────────────────────────────────────────────────

    def _get_model(self, language: str):
        """Return (model, tokenizer) for the requested language, loading if needed."""
        lang = language if language in MODEL_IDS else _DEFAULT_LANG
        if lang not in self._cache:
            from transformers import VitsModel, VitsTokenizer  # type: ignore
            model_id = MODEL_IDS[lang]
            tokenizer = VitsTokenizer.from_pretrained(model_id)
            model = VitsModel.from_pretrained(model_id)
            model.eval()
            # Keep on CPU until synthesize() moves it to the target device
            self._cache[lang] = (model, tokenizer)
        return self._cache[lang]

    @staticmethod
    def _split_sentences(text: str) -> list[str]:
        """Split text into chunks of ≤ _MAX_WORDS_PER_CHUNK words."""
        sentences = _SENTENCE_RE.split(text.strip())
        chunks: list[str] = []
        current: list[str] = []
        current_words = 0

        for sent in sentences:
            words = sent.split()
            if current_words + len(words) > _MAX_WORDS_PER_CHUNK and current:
                chunks.append(" ".join(current))
                current = words
                current_words = len(words)
            else:
                current.extend(words)
                current_words += len(words)

        if current:
            chunks.append(" ".join(current))

        return chunks or [text]

    def _synthesize_chunk(
        self, text: str, model, tokenizer, device: str
    ) -> np.ndarray:
        """Synthesize a single short text chunk. Returns 1-D float32 numpy array."""
        import torch

        model.to(device)
        inputs = tokenizer(text, return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            output = model(**inputs)

        waveform = output.waveform[0].cpu().numpy()  # shape: (samples,)
        return waveform

    # ── public API ───────────────────────────────────────────────────────────

    def synthesize(
        self, text: str, language: str, device: str = "cuda"
    ) -> Tuple[np.ndarray, int]:
        """
        Convert text to speech waveform.

        Args:
            text:     Text to synthesize (any length — long text is split automatically).
            language: Language code: "bam", "ful", "fr", or "en".
            device:   "cuda" or "cpu".

        Returns:
            (waveform_np, sample_rate) — float32 numpy array, sample rate in Hz.
        """
        lang = language if language in MODEL_IDS else _DEFAULT_LANG
        model, tokenizer = self._get_model(lang)

        chunks = self._split_sentences(text)
        waveforms: list[np.ndarray] = []

        for chunk in chunks:
            if not chunk.strip():
                continue
            waveforms.append(self._synthesize_chunk(chunk, model, tokenizer, device))

        # Free device memory before returning
        model.to("cpu")

        if not waveforms:
            return np.zeros(1, dtype=np.float32), model.config.sampling_rate

        combined = np.concatenate(waveforms)
        return combined, model.config.sampling_rate

    def text_to_audio_bytes(
        self, text: str, language: str, device: str = "cuda"
    ) -> bytes:
        """
        Convert text to WAV bytes suitable for gr.Audio or HTTP response.

        Returns raw WAV file bytes (16-bit PCM).
        """
        waveform, sample_rate = self.synthesize(text, language, device=device)

        buf = io.BytesIO()
        # soundfile expects float32 in [-1, 1]; MMS output is already normalised
        sf.write(buf, waveform, sample_rate, format="WAV", subtype="PCM_16")
        return buf.getvalue()