app.py

"""ThendoLabs — ASR (10 languages) + TTS (isiZulu, Setswana, Yoruba).

Public demo of the african_speak fine-tuned models. Champions and inference
settings mirror the shipped configurations (see the About tab for lineage,
attribution, and limitations).
"""
import spaces  # noqa: F401  — MUST be imported before torch (ZeroGPU)

import os
import subprocess
import sys
import tempfile

from huggingface_hub import hf_hub_download, snapshot_download

APP_ROOT = os.path.dirname(os.path.abspath(__file__))
os.chdir(APP_ROOT)
sys.path.insert(0, os.path.join(APP_ROOT, "tts"))

# ---------------------------------------------------------------- downloads
REFS = {
    "zul": ["nchlt_zul_125f_0672", "nchlt_zul_171m_0695", "nchlt_zul_202f_0317"],
    "tn": ["tsn_3342_2507009650", "tsn_3342_5431483414", "tsn_3342_8497760639"],
}


def _bootstrap() -> None:
    os.makedirs("refs/zul", exist_ok=True)
    os.makedirs("refs/tn", exist_ok=True)
    for lang, uids in REFS.items():
        sub = "zu" if lang == "zul" else "tn"
        for uid in uids:
            p = hf_hub_download("Or4kool/setswana-tts-data", f"gateway_refs/{sub}/{uid}.wav",
                                repo_type="dataset")
            dst = f"refs/{lang}/{uid}.wav"
            if not os.path.exists(dst):
                os.symlink(p, dst)
    os.makedirs("tts/checkpoints/st2-zul-b", exist_ok=True)
    os.makedirs("tts/checkpoints/st2-tn-br3", exist_ok=True)
    for repo, fname, dst in [
        ("Or4kool/st2-zul-wip", "st2_zul_champion_stageB.pth", "tts/checkpoints/st2-zul-b/best.pth"),
        ("Or4kool/vits-tn-wip", "st2_tn_br3_best.pth", "tts/checkpoints/st2-tn-br3/best.pth"),
    ]:
        p = hf_hub_download(repo, fname)
        if not os.path.exists(dst):
            os.symlink(p, dst)
    globals()["YOR_CKPT"] = hf_hub_download("Or4kool/vits-yo-v1", "best_model.pth")
    globals()["YOR_CFG"] = hf_hub_download("Or4kool/vits-yo-v1", "config.json")
    snapshot_download("papercup-ai/multilingual-pl-bert",
                      local_dir="tts/styletts2/models/multilingual-pl-bert")
    snapshot_download("Or4kool/w2v-bert-african-lm", local_dir="lm_dir",
                      allow_patterns=["*/*.arpa", "*/unigrams.txt", "decode_params.json"])
    # Vendored StyleTTS2 repo: clone pinned SHA + apply inference patches.
    subprocess.run(["git", "lfs", "install"], check=True)
    if not os.path.isdir("tts/styletts2/repo"):
        subprocess.run([sys.executable, "tts/styletts2/setup_st2.py"], check=True)


_bootstrap()

# ------------------------------------------------------- torch-side imports
import gradio as gr  # noqa: E402
import librosa  # noqa: E402
import numpy as np  # noqa: E402
import pyloudnorm as pyln  # noqa: E402
import torch  # noqa: E402
import yaml  # noqa: E402
from transformers import AutoFeatureExtractor, AutoTokenizer, Wav2Vec2BertForCTC  # noqa: E402

from src.data.num_expand import load_digit_lexicon  # noqa: E402
from src.data.zulu_text import clean_zulu_text  # noqa: E402

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

ASR_LANGS = {
    "Dagbani": "african", "Igbo": "african", "Hausa": "african", "Ibibio": "african",
    "Yoruba": "african", "Nigerian Pidgin": "african", "Setswana": "african",
    "Adamawa Fulfulde": "african", "Twi/Akan": "african", "isiZulu": "zul",
}
LANG_CODES = {
    "Dagbani": "dag", "Igbo": "ig", "Hausa": "ha", "Ibibio": "ibb", "Yoruba": "yo",
    "Nigerian Pidgin": "pcm", "Setswana": "tn", "Adamawa Fulfulde": "fub",
    "Twi/Akan": "tw", "isiZulu": "zul",
}

# ASR models: standard transformers stacks — eager module-scope load (ZeroGPU packs them).
_ASR = {}
for key, repo, sub in [("african", "Or4kool/w2v-bert-african", "sp2v2"),
                       ("zul", "Or4kool/w2v-bert-zul", None)]:
    kw = {"subfolder": sub} if sub else {}
    model = Wav2Vec2BertForCTC.from_pretrained(repo, **kw).to(DEVICE).eval()
    fe = AutoFeatureExtractor.from_pretrained(repo)
    tok = AutoTokenizer.from_pretrained(repo)
    _ASR[key] = (model, fe, tok)

# Per-language KenLM beam-search decoders (production decode path; tuned
# alpha/beta from the phase-5 sweep). Languages without an LM (isiZulu)
# fall back to greedy.
import json as _json  # noqa: E402

from lm.decoder import ctc_labels, decode_logits, load_decoders  # noqa: E402

_LABELS = ctc_labels(_ASR["african"][2])
_DEC_PARAMS = _json.load(open("lm_dir/decode_params.json"))
DECODERS = {}
for _lang, _p in _DEC_PARAMS.items():
    DECODERS.update(load_decoders("lm_dir", [_lang], _LABELS, _p["alpha"], _p["beta"]))
BEAM_WIDTH = 100

# ------------------------------------------------------------- TTS configs
ZUL_LEX = load_digit_lexicon("tts/configs/zulu_numbers.yaml")

TTS_CFGS = {
    "isiZulu": {
        "st2_repo": "tts/styletts2/repo",
        "st2_config": "tts/styletts2/config_ft_zul_b.yml",
        "st2_checkpoint": "tts/checkpoints/st2-zul-b/best.pth",
        "g2p": "zul",
        "wavs_24k": "refs/zul",
        "ref_uids": REFS["zul"],
        "style_beta": 0.4,
    },
    "Setswana": {
        "st2_repo": "tts/styletts2/repo",
        "st2_config": "tts/styletts2/config_ft_tn_br3.yml",
        "st2_checkpoint": "tts/checkpoints/st2-tn-br3/best.pth",
        "frontend": "mixed",
        "g2p": "v1",
        "lexicon": "tts/configs/setswana_lexicon.yaml",
        "numbers": "tts/configs/setswana_numbers.yaml",
        "wavs_24k": "refs/tn",
        "ref_uids": REFS["tn"],
    },
}
for name, cfg in TTS_CFGS.items():
    path = os.path.join(tempfile.gettempdir(), f"cfg_{name}.yaml")
    with open(path, "w") as f:
        yaml.safe_dump(cfg, f)
    cfg["_path"] = path

_SYNTHS = {}  # worker-local cache (each ZeroGPU worker holds its own copy)


def _get_synth(language: str):
    if language in _SYNTHS:
        return _SYNTHS[language]
    if language == "Yoruba":
        from TTS.utils.synthesizer import Synthesizer
        synth = Synthesizer(tts_checkpoint=YOR_CKPT, tts_config_path=YOR_CFG,
                            use_cuda=torch.cuda.is_available())
    else:
        from src.tts_st2_synth import StyleTTS2Synthesizer
        synth = StyleTTS2Synthesizer(TTS_CFGS[language]["_path"])
    _SYNTHS[language] = synth
    return synth


def _postprocess(wav: np.ndarray, sr: int) -> np.ndarray:
    wav = np.asarray(wav, dtype=np.float32)
    wav, _ = librosa.effects.trim(wav, top_db=35)
    loud = pyln.Meter(sr).integrated_loudness(wav)
    if np.isfinite(loud):
        wav = pyln.normalize.loudness(wav, loud, -16.0)
    peak = np.abs(wav).max()
    if peak > 0.99:
        wav = wav * (0.99 / peak)
    return wav


def _tts_duration(text: str, language: str) -> int:
    return min(90, 35 + len(text or "") // 4)


@spaces.GPU(duration=_tts_duration)
def synthesize(text: str, language: str):
    """Synthesize speech from text in isiZulu, Setswana, or Yoruba.

    Args:
        text: Input sentence (max 300 characters).
        language: One of "isiZulu", "Setswana", "Yoruba".
    Returns:
        Audio waveform at the model's native sample rate.
    """
    text = (text or "").strip()
    if not text:
        raise gr.Error("Please enter some text.")
    if len(text) > 300:
        raise gr.Error("Please keep the text under 300 characters for the demo.")
    try:
        if language == "isiZulu":
            text = clean_zulu_text(text, ZUL_LEX)
            synth = _get_synth(language)
            wav, sr = synth.say(text)
        elif language == "Setswana":
            synth = _get_synth(language)
            wav, sr = synth.say(text)
        elif language == "Yoruba":
            synth = _get_synth(language)
            wav = np.asarray(synth.tts(text.lower()), dtype=np.float32)
            sr = synth.output_sample_rate
        else:
            raise gr.Error(f"Unknown language: {language}")
    except gr.Error:
        raise
    except Exception as exc:  # surface model errors as friendly messages
        raise gr.Error(f"Synthesis failed: {exc}") from exc
    return int(sr), _postprocess(np.asarray(wav, dtype=np.float32), int(sr))


@spaces.GPU(duration=30)
def transcribe(audio_path: str, language: str) -> str:
    """Transcribe speech (1-20 s clip) in one of 10 African languages.

    Args:
        audio_path: Path to an uploaded or recorded audio clip.
        language: The language spoken in the clip.
    Returns:
        The transcript (per-language KenLM beam decode; greedy where no LM exists).
    """
    if not audio_path:
        raise gr.Error("Please record or upload an audio clip.")
    model, fe, tok = _ASR[ASR_LANGS[language]]
    wav, _ = librosa.load(audio_path, sr=16000, mono=True)
    if len(wav) < 1600:
        raise gr.Error("Clip too short — please record at least half a second.")
    wav = wav[: 16000 * 30]
    inputs = fe(wav, sampling_rate=16000, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        logits = model(**inputs).logits
    lm_text = decode_logits(DECODERS, LANG_CODES[language],
                            logits[0].cpu().numpy(), BEAM_WIDTH)
    if lm_text is not None:
        return lm_text
    ids = torch.argmax(logits, dim=-1)[0]
    return tok.decode(ids, skip_special_tokens=True)


# ------------------------------------------------------------------ UI
TTS_EXAMPLES = [
    ["sawubona mngani wami, ngiyakwemukela namuhla", "isiZulu"],
    ["ingcosi igqoke kahle namuhla ekuseni", "isiZulu"],
    ["dumela tsala ya me, o amogetswe mo lefelong la rona", "Setswana"],
    ["bawo ni, ore mi, kaabo si ile wa", "Yoruba"],
]

ABOUT = """
## About these models

**ASR** — `w2v-bert-2.0` fine-tunes: a 9-language polyglot head (Dagbani, Igbo, Hausa,
Ibibio, Yoruba, Nigerian Pidgin, Setswana, Adamawa Fulfulde, Twi/Akan) and a dedicated
isiZulu head. Decoding uses the production per-language KenLM beam search (tuned per
language); isiZulu decodes greedily (no KenLM built for it yet).

**TTS** — StyleTTS 2 fine-tunes for **isiZulu** (NCHLT corpus, CC-BY 3.0; explicit
click phonemes ǀ ǃ ǁ via a rule-based G2P) and **Setswana** (SLR32; the BR3 champion,
with English code-switch routing), plus a **Yoruba** VITS voice (OpenSLR-129).
Voices derive from corpus speakers; native-listener evaluation gated every release.

**Limitations** — lexical tone is not modeled (no digital tone lexicons exist for these
languages); typing English into the isiZulu box will produce Zulu-rule pronunciations
(including clicks on the letter *c*); single-listener validation per language so far.

**Attribution** — NCHLT isiZulu (SADiLaR, CC-BY 3.0) · OpenSLR SLR32 / SLR129 ·
LibriTTS (CC-BY 4.0) · `facebook/w2v-bert-2.0` (MIT) · StyleTTS 2 (MIT) ·
Coqui VITS (MPL-2.0). The MMS-300m adversarial critic (CC-BY-NC) was used at
training time only and is not part of, nor distributed with, these models.
"""

with gr.Blocks(title="ThendoLabs") as demo:
    gr.Markdown("# 🌍 ThendoLabs\nASR for 10 African languages · TTS for isiZulu, Setswana, and Yoruba")
    with gr.Tab("Text → Speech"):
        tts_lang = gr.Dropdown(list(TTS_CFGS) + ["Yoruba"], value="isiZulu", label="Language")
        tts_text = gr.Textbox(label="Text", placeholder="sawubona mngani wami... (max 300 chars)")
        tts_btn = gr.Button("Speak", variant="primary")
        tts_out = gr.Audio(label="Synthesized speech")
        gr.Examples(TTS_EXAMPLES, inputs=[tts_text, tts_lang], outputs=tts_out,
                    fn=synthesize, cache_examples=True, cache_mode="lazy")
        tts_btn.click(synthesize, [tts_text, tts_lang], tts_out)
    with gr.Tab("Speech → Text"):
        asr_lang = gr.Dropdown(list(ASR_LANGS), value="isiZulu", label="Language")
        asr_audio = gr.Audio(sources=["microphone", "upload"], type="filepath",
                             label="Speak or upload a short clip (1-20 s)")
        asr_btn = gr.Button("Transcribe", variant="primary")
        asr_out = gr.Textbox(label="Transcript", lines=6, max_lines=24, show_copy_button=True)
        asr_btn.click(transcribe, [asr_audio, asr_lang], asr_out)
    with gr.Tab("About"):
        gr.Markdown(ABOUT)

demo.launch(mcp_server=True)