src/tts/waxal_tts.py

"""
WaxalTTSEngine — lightweight VITS-based TTS for Sahel-Voice-Lab.

Bambara : ynnov/ekodi-bambara-tts-female  (VitsModel + AutoTokenizer)
Fula    : placeholder — returns None until ous-sow/fula-tts is trained
"""
from __future__ import annotations

import logging
import os
import threading
from typing import Optional

import numpy as np

logger = logging.getLogger(__name__)

BAMBARA_TTS_REPO = os.environ.get("BAMBARA_TTS_REPO", "ynnov/ekodi-bambara-tts-female")
FULA_TTS_REPO    = os.environ.get("FULA_TTS_REPO",    "ous-sow/fula-tts")
HF_TOKEN         = os.environ.get("HF_TOKEN")


def generate_pular_tts(text: str) -> None:
    """
    Placeholder for Fula (Pulaar) TTS.
    Returns None until ous-sow/fula-tts is trained and pushed to the Hub.
    Run notebooks/train_fula_tts.ipynb on Kaggle T4 to produce the model.
    """
    logger.info("generate_pular_tts: model not yet trained — returning None")
    return None


class WaxalTTSEngine:
    """Unified TTS engine: Bambara (VITS) + Fula (placeholder)."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        # Bambara
        self._bam_model     = None
        self._bam_tokenizer = None
        self._bam_ready     = False
        self._bam_error: Optional[str] = None

    # ── Public API ────────────────────────────────────────────────────────────

    def synthesize(self, text: str, lang: str) -> Optional[tuple[np.ndarray, int]]:
        """
        Returns (audio_float32, sample_rate) or None if TTS unavailable.
        Never raises — all errors are logged.
        """
        text = text.strip()
        if not text:
            return None
        try:
            if lang == "bam":
                return self._synthesize_bambara(text)
            elif lang == "ful":
                return generate_pular_tts(text)
            else:
                return None
        except Exception as exc:
            logger.error("WaxalTTS.synthesize(%s) unexpected error: %s", lang, exc)
            return None

    def get_status(self) -> dict:
        bam = "ready" if self._bam_ready else (
            f"error: {self._bam_error}" if self._bam_error else "loading…"
        )
        return {"bam": bam, "ful": "not trained yet"}

    def preload(self) -> None:
        """Start background thread to load the Bambara model."""
        threading.Thread(target=self._load_bambara, daemon=True).start()

    # ── Bambara (ynnov/ekodi-bambara-tts-female, VITS) ───────────────────────

    def _load_bambara(self) -> None:
        try:
            from transformers import VitsModel, AutoTokenizer
            logger.info("WaxalTTS: loading Bambara TTS from %s …", BAMBARA_TTS_REPO)
            tok = AutoTokenizer.from_pretrained(BAMBARA_TTS_REPO, token=HF_TOKEN)
            mdl = VitsModel.from_pretrained(BAMBARA_TTS_REPO, token=HF_TOKEN)
            mdl.eval()
            with self._lock:
                self._bam_tokenizer = tok
                self._bam_model     = mdl
                self._bam_ready     = True
            logger.info("WaxalTTS: Bambara TTS ready")
        except Exception as exc:
            self._bam_error = str(exc)
            logger.error("WaxalTTS: Bambara TTS load failed: %s", exc)

    def _synthesize_bambara(self, text: str) -> Optional[tuple[np.ndarray, int]]:
        if not self._bam_ready:
            self._load_bambara()
        if not self._bam_ready:
            logger.warning("WaxalTTS: Bambara TTS not ready — %s", self._bam_error)
            return None
        try:
            import torch
            with self._lock:
                inputs = self._bam_tokenizer(text, return_tensors="pt")
                with torch.no_grad():
                    output = self._bam_model(**inputs)
                audio = output.waveform[0].cpu().numpy().astype(np.float32)
                sr    = self._bam_model.config.sampling_rate
            return audio, sr
        except Exception as exc:
            logger.error("WaxalTTS: Bambara synthesis failed: %s", exc)
            self._bam_error = str(exc)
            self._bam_ready = False
            return None

    # ── Utility ───────────────────────────────────────────────────────────────

    @staticmethod
    def audio_to_gradio(audio: np.ndarray, sr: int) -> tuple[int, np.ndarray]:
        """Convert float32 → int16 tuple that gr.Audio expects."""
        pcm = (audio * 32767).clip(-32768, 32767).astype(np.int16)
        return sr, pcm