app.py

"""
🍵 Ghana Speech Nano — multilingual Matcha-TTS demo

Serves checkpoints from https://huggingface.co/ghananlpcommunity/ghana-speech-nano
(checkpoints/*.ckpt). The model was finetuned on ghananlpcommunity/ghana-speech across
42 languages, with the Matcha-TTS speaker-conditioning slot repurposed as a language ID
(n_spks=42). Text is phonemized with the same espeak-ng "lfn" voice pipeline used at
training time (matcha.text.cleaners.twi_cleaners), then filtered down to the model's
symbol set exactly like scripts/prep_ghana_speech.py does.
"""

import functools
import os
from pathlib import Path

import gradio as gr
import spaces
import torch
from huggingface_hub import hf_hub_download, list_repo_files

# Importing matcha.cli applies its torch.load(weights_only=False) monkeypatch, which is
# required to unpickle these Lightning/OmegaConf checkpoints on modern PyTorch.
from matcha.cli import (
    VOCODER_URLS,
    assert_model_downloaded,
    load_hifigan,
    load_matcha,
)
from matcha.hifigan.denoiser import Denoiser
from matcha.text import cleaned_text_to_sequence
from matcha.text.cleaners import twi_cleaners
from matcha.text.symbols import symbols
from matcha.utils.utils import get_user_data_dir, intersperse, plot_tensor

MODEL_REPO = "ghananlpcommunity/ghana-speech-nano"
CHECKPOINT_DIR_IN_REPO = "checkpoints"
DATASET_REPO = "ghananlpcommunity/ghana-speech"  # used only as a fallback lookup
SAMPLE_RATE = 22050
N_SPKS = 42

# Verified directly from ghananlpcommunity/ghana-speech's dataset card (the `configs:` list),
# which is exactly the `sorted(set(configs))` order scripts/prep_ghana_speech.py used to
# assign lang_id = index. This is the source of truth; the dataset API call below is only a
# fallback in case the dataset is ever restructured.
GHANA_SPEECH_LANGUAGES = [
    "Akuapem_Twi_twi",
    "Anyin_any",
    "Asante_Twi_twi",
    "Avatime_avn",
    "Bassar_Ntcham_bud",
    "Bimoba_bim",
    "Birifor_Southern_biv",
    "Bissa_bib",
    "Buli_bwu",
    "Chumburung_ncu",
    "Dagaare_dga",
    "Dagbani_dag",
    "Dangme_ada",
    "Deg_mzw",
    "Ewe_ewe",
    "Fante_fat",
    "Fulfulde_Maasina_ffm",
    "Gikyode_acd",
    "Gonja_gjn",
    "Hausa_hau",
    "Kabiye_kbp",
    "Kasem_xsm",
    "Konkomba_xon",
    "Konni_kma",
    "Kusaal_kus",
    "Lelemi_lef",
    "Mampruli_maw",
    "Nawuri_naw",
    "Ninkare_gur",
    "Nkonya_nko",
    "Ntrubo_ntr",
    "Nzema_nzi",
    "Paasaal_sig",
    "Sehwi_sfw",
    "Sekpele_lip",
    "Selee_snw",
    "Sisaala_Tumulung_sil",
    "Siwu_akp",
    "Tampulma_tpm",
    "Tem_kdh",
    "Tuwuli_bov",
    "Vagla_vag",
]
assert len(GHANA_SPEECH_LANGUAGES) == N_SPKS

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

LOCATION = Path(get_user_data_dir())

# --------------------------------------------------------------------------------------
# Mel statistics (must match the training data for correct synthesis)
# --------------------------------------------------------------------------------------
MEL_MEAN = -5.670213434000389
MEL_STD  =  3.317951496927415


# --------------------------------------------------------------------------------------
# Language names: recompute the exact same lang_id -> language mapping that
# scripts/prep_ghana_speech.py used when it built the training filelists
# (lang_to_id = {lang: i for i, lang in enumerate(sorted(set(configs)))}).
# Falls back to generic "Language N" labels if the dataset can't be reached.
# --------------------------------------------------------------------------------------
@functools.lru_cache(maxsize=1)
def get_lang_map():
    # Primary: verified list extracted from the dataset's own README/card (see
    # GHANA_SPEECH_LANGUAGES above). This avoids a network call and is guaranteed correct
    # for the checkpoints currently in the model repo.
    return {i: name for i, name in enumerate(GHANA_SPEECH_LANGUAGES)}


@functools.lru_cache(maxsize=1)
def get_lang_map_live():
    """Fallback: recompute the mapping from the live dataset config list, in case the
    dataset is ever restructured and GHANA_SPEECH_LANGUAGES goes stale."""
    try:
        from datasets import get_dataset_config_names

        configs = sorted(set(get_dataset_config_names(DATASET_REPO)))
        if configs:
            return {i: name for i, name in enumerate(configs)}
    except Exception as e:  # pylint: disable=broad-except
        print(f"[!] Could not fetch language names from {DATASET_REPO}: {e}")
    return None


@functools.lru_cache(maxsize=1)
def list_checkpoints():
    """List available .ckpt files under checkpoints/ in the model repo, last.ckpt first."""
    try:
        files = list_repo_files(MODEL_REPO)
    except Exception as e:  # pylint: disable=broad-except
        print(f"[!] Could not list files in {MODEL_REPO}: {e}")
        return ["last.ckpt"]
    ckpts = sorted(
        os.path.basename(f) for f in files if f.startswith(f"{CHECKPOINT_DIR_IN_REPO}/") and f.endswith(".ckpt")
    )
    if "last.ckpt" in ckpts:
        ckpts.remove("last.ckpt")
        ckpts = ["last.ckpt"] + ckpts
    return ckpts or ["last.ckpt"]


# --------------------------------------------------------------------------------------
# Model / vocoder loading (cached so repeated calls with the same checkpoint are free)
# --------------------------------------------------------------------------------------
@functools.lru_cache(maxsize=2)
def load_model(checkpoint_name: str):
    print(f"[!] Downloading + loading checkpoint: {checkpoint_name}")
    local_path = hf_hub_download(repo_id=MODEL_REPO, filename=f"{CHECKPOINT_DIR_IN_REPO}/{checkpoint_name}")
    model = load_matcha(checkpoint_name, local_path, DEVICE)

    # Override mel normalisation stats with the correct training values.
    # The checkpoint may not contain them, leaving defaults (0/1) that cause garbled speech.
    if hasattr(model, 'mel_mean') and model.mel_mean is not None:
        current_mean = model.mel_mean.item()
        current_std  = model.mel_std.item()
        print(f"[!] Loaded mel_mean = {current_mean:.4f}, mel_std = {current_std:.4f}")
        if not torch.isclose(model.mel_mean, torch.tensor(MEL_MEAN), atol=1e-3) or \
           not torch.isclose(model.mel_std, torch.tensor(MEL_STD), atol=1e-3):
            print(f"[!] Overriding to training stats: mean={MEL_MEAN}, std={MEL_STD}")
            model.mel_mean.fill_(MEL_MEAN)
            model.mel_std.fill_(MEL_STD)
    else:
        print("[!] mel_mean/mel_std not found – registering as buffers.")
        model.register_buffer('mel_mean', torch.tensor(MEL_MEAN))
        model.register_buffer('mel_std',  torch.tensor(MEL_STD))

    return model


@functools.lru_cache(maxsize=1)
def load_vocoder_cached():
    vocoder_name = "hifigan_univ_v1"
    vocoder_path = LOCATION / vocoder_name
    assert_model_downloaded(vocoder_path, VOCODER_URLS[vocoder_name])
    vocoder = load_hifigan(vocoder_path, DEVICE)
    denoiser = Denoiser(vocoder, mode="zeros")
    return vocoder, denoiser


def to_waveform(mel, vocoder, denoiser, denoiser_strength=0.00025):
    audio = vocoder(mel).clamp(-1, 1)
    audio = denoiser(audio.squeeze(), strength=denoiser_strength).cpu().squeeze()
    return audio.cpu().squeeze().numpy()


# --------------------------------------------------------------------------------------
# Text -> phonemes -> symbol ids (mirrors the training-data preprocessing exactly)
# --------------------------------------------------------------------------------------
def phonemize(text: str, direct_phonemes: bool = False) -> str:
    """If direct_phonemes is True, skip twi_cleaners and only filter symbols."""
    if direct_phonemes:
        # Filter out any characters not in the model's symbol set
        phon = "".join(c for c in text if c in _SYMBOL_SET)
        phon = " ".join(phon.split())
        return phon
    else:
        phon = twi_cleaners(text)
        phon = "".join(c for c in phon if c in _SYMBOL_SET)
        phon = " ".join(phon.split())
        return phon


def process_text(text: str, device, direct_phonemes: bool = False):
    phon = phonemize(text, direct_phonemes)
    if not phon.strip():
        raise gr.Error(
            "Couldn't turn this text into any recognised phonemes. Try rephrasing, "
            "or check that the input isn't empty / only punctuation."
        )
    seq = cleaned_text_to_sequence(phon)
    x = torch.tensor(intersperse(seq, 0), dtype=torch.long, device=device)[None]
    x_lengths = torch.tensor([x.shape[-1]], dtype=torch.long, device=device)
    return phon, x, x_lengths


# --------------------------------------------------------------------------------------
# Synthesis (the GPU‑decorated function)
# --------------------------------------------------------------------------------------
@spaces.GPU
@torch.inference_mode()
def synthesise(text, language_label, checkpoint_name, n_timesteps, length_scale, temperature, direct_phonemes):
    if not text or not text.strip():
        raise gr.Error("Please enter some text to synthesise.")

    lang_id = LANG_LABEL_TO_ID.get(language_label, 0)

    model = load_model(checkpoint_name)  # mel stats are already fixed here
    vocoder, denoiser = load_vocoder_cached()

    phon, x, x_lengths = process_text(text, DEVICE, direct_phonemes)
    spk = torch.tensor([lang_id], device=DEVICE, dtype=torch.long) if N_SPKS > 1 else None

    output = model.synthesise(
        x,
        x_lengths,
        n_timesteps=int(n_timesteps),
        temperature=temperature,
        spks=spk,
        length_scale=length_scale,
    )
    waveform = to_waveform(output["mel"], vocoder, denoiser)

    mel_plot = plot_tensor(output["mel"].squeeze().cpu().numpy())

    return (SAMPLE_RATE, waveform), phon, mel_plot


def refresh_checkpoints():
    list_checkpoints.cache_clear()
    choices = list_checkpoints()
    return gr.update(choices=choices, value=choices[0])


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
LANG_MAP = get_lang_map()


def _display_label(config_name: str) -> str:
    # "Akuapem_Twi_twi" -> "Akuapem Twi (twi)" -- cosmetic only, lookup still keyed on the
    # exact config_name so language IDs stay unambiguous.
    parts = config_name.split("_")
    code, name_parts = parts[-1], parts[:-1]
    return f"{' '.join(name_parts)} ({code})"


LANG_LABEL_TO_ID = {_display_label(LANG_MAP[i]): i for i in sorted(LANG_MAP)}
LANG_CHOICES = list(LANG_LABEL_TO_ID.keys())
CHECKPOINT_CHOICES = list_checkpoints()

DESCRIPTION = f"""# 🍵 Ghana Speech Nano — Multilingual Matcha-TTS

A single [Matcha-TTS](https://github.com/shivammehta25/Matcha-TTS) acoustic model finetuned by
[Ghana NLP Community](https://huggingface.co/ghananlpcommunity) on **42 languages**
(`ghananlpcommunity/ghana-speech`), using the speaker-conditioning slot as a language ID.
Checkpoints: [ghananlpcommunity/ghana-speech-nano](https://huggingface.co/ghananlpcommunity/ghana-speech-nano/tree/main/checkpoints).

**Two input modes:**
- **Orthographic text** (default): Type normal words and the app will phonemise them using the same `lfn` voice used during training.
- **Direct phoneme input**: Check the box to paste pre‑phonemised strings (exact training‑time phonemes) – no further phonemisation is applied.

Detected **{len(LANG_CHOICES)}** languages, verified against the
[dataset card](https://huggingface.co/datasets/ghananlpcommunity/ghana-speech) for
`ghananlpcommunity/ghana-speech` (language ID = alphabetical index, matching how the
training filelists were built).
"""

with gr.Blocks(title="🍵 Ghana Speech Nano — Multilingual Matcha-TTS") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        checkpoint_dd = gr.Dropdown(
            choices=CHECKPOINT_CHOICES,
            value=CHECKPOINT_CHOICES[0],
            label="Checkpoint",
            scale=3,
        )
        refresh_btn = gr.Button("🔄 Refresh checkpoint list", scale=1)

    with gr.Row():
        text = gr.Textbox(
            label="Text to synthesise",
            placeholder="Type text in any of the supported languages (or paste phonemes)...",
            lines=3,
            scale=3,
        )
        language_dd = gr.Dropdown(
            choices=LANG_CHOICES,
            value=LANG_CHOICES[0] if LANG_CHOICES else None,
            label="Language / Speaker",
            scale=1,
        )

    direct_phoneme_checkbox = gr.Checkbox(label="Direct phoneme input (skip phonemisation)", value=False)

    with gr.Row():
        n_timesteps = gr.Slider(minimum=1, maximum=100, step=1, value=10, label="Number of ODE steps")
        length_scale = gr.Slider(minimum=0.5, maximum=1.5, step=0.05, value=0.85, label="Length scale (speaking rate)")
        temperature = gr.Slider(minimum=0.0, maximum=2.0, step=0.05, value=0.667, label="Sampling temperature")

    synth_btn = gr.Button("🍵 Synthesise", variant="primary")

    with gr.Row():
        audio_out = gr.Audio(label="Synthesised audio", type="numpy")
        mel_out = gr.Image(label="Mel spectrogram")

    phonemes_out = gr.Textbox(label="Phonemised text (what the model actually sees)", interactive=False)

    synth_btn.click(
        fn=synthesise,
        inputs=[text, language_dd, checkpoint_dd, n_timesteps, length_scale, temperature, direct_phoneme_checkbox],
        outputs=[audio_out, phonemes_out, mel_out],
        api_name="tts",
    )
    refresh_btn.click(fn=refresh_checkpoints, outputs=[checkpoint_dd])

    gr.Examples(
        examples=[
            ["fˈa ˌaduˈan bˈi fˈi fˈie bɾˈe mˈe", LANG_CHOICES[0] if LANG_CHOICES else None, True],  # direct phoneme input
        ],
        inputs=[text, language_dd, direct_phoneme_checkbox],
    )

if __name__ == "__main__":
    demo.queue().launch()