Upload app.py with huggingface_hub

app.py

@@ -4,188 +4,132 @@ import torch, yaml, os, sys, glob, re
 import librosa
 import soundfile as sf
 import torchaudio
-import numpy as np
 from huggingface_hub import snapshot_download
 from munch import Munch
 from nltk.tokenize import word_tokenize
 import nltk
 nltk.download("punkt_tab", quiet=True)
 
-# ══════════════════════════════════════════════════════════
-# KIKUYU — load at startup (small, 145MB)
-# ══════════════════════════════════════════════════════════
-from transformers import VitsModel, AutoTokenizer
-
-print("Loading Kikuyu TTS...")
-kikuyu_model     = VitsModel.from_pretrained("gateremark/kikuyu-tts-v1")
-kikuyu_tokenizer = AutoTokenizer.from_pretrained("gateremark/kikuyu-tts-v1")
-kikuyu_model.eval()
-print("✅ Kikuyu model loaded")
-
-def synthesize_kikuyu(text):
-    inputs = kikuyu_tokenizer(text=text.strip(), return_tensors="pt")
-    with torch.no_grad():
-        output = kikuyu_model(**inputs)
-    waveform = output.waveform.squeeze().cpu().numpy()
-    sr = kikuyu_model.config.sampling_rate
-    sf.write("/tmp/kikuyu_output.wav", waveform, sr)
-    return "/tmp/kikuyu_output.wav"
-
-# ══════════════════════════════════════════════════════════
-# LUHYA — lazy load on first request (heavy, 2.18GB)
-# ══════════════════════════════════════════════════════════
-_luhya_loaded   = False
-_luhya_model    = None
-_luhya_sampler  = None
-_default_style  = None
-_model_params   = None
-_textcleaner    = None
-_model_dir      = None
-
-def load_luhya():
-    global _luhya_loaded, _luhya_model, _luhya_sampler
-    global _default_style, _model_params, _textcleaner, _model_dir
-
-    if _luhya_loaded:
-        return
-
-    print("Loading Luhya TTS (first request)...")
-
-    # Clone + patch StyleTTS2
-    os.system("git clone https://github.com/yl4579/StyleTTS2 /app/StyleTTS2 2>/dev/null || true")
-    for fpath in ["/app/StyleTTS2/models.py", "/app/StyleTTS2/utils.py"]:
-        with open(fpath) as f:
-            code = f.read()
-        patched = re.sub(
-            r'torch\.load\(([^)]+)\)',
-            lambda m: m.group(0) if 'weights_only' in m.group(1)
-                      else f'torch.load({m.group(1)}, weights_only=False)',
-            code
-        )
-        with open(fpath, "w") as f:
-            f.write(patched)
-
-    if "/app/StyleTTS2" not in sys.path:
-        sys.path.insert(0, "/app/StyleTTS2")
-    os.chdir("/app/StyleTTS2")
-
-    _model_dir = snapshot_download("crazydev919/luhya-tts")
-
-    from models import build_model, load_ASR_models, load_F0_models
-    from utils import recursive_munch
-    from text_utils import TextCleaner
-    from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule
-
-    _textcleaner = TextCleaner()
-    device = "cpu"
-
-    config = yaml.safe_load(open(f"{_model_dir}/config.yml"))
-    config["ASR_path"]   = f"{_model_dir}/Utils/ASR/epoch_00080.pth"
-    config["ASR_config"] = f"{_model_dir}/Utils/ASR/config.yml"
-    config["F0_path"]    = f"{_model_dir}/Utils/JDC/bst.t7"
-    config["PLBERT_dir"] = f"{_model_dir}/Utils/PLBERT/"
-
-    text_aligner    = load_ASR_models(config["ASR_path"], config["ASR_config"])
-    pitch_extractor = load_F0_models(config["F0_path"])
-
-    from Utils.PLBERT.util import load_plbert
-    plbert = load_plbert(config["PLBERT_dir"])
-
-    _model_params = recursive_munch(config["model_params"])
-    _luhya_model  = build_model(_model_params, text_aligner, pitch_extractor, plbert)
-    _ = [_luhya_model[k].eval() for k in _luhya_model]
-    _ = [_luhya_model[k].to(device) for k in _luhya_model]
-
-    params = torch.load(
-        f"{_model_dir}/model.pth", map_location="cpu", weights_only=False
-    )["net"]
-    for key in _luhya_model:
-        if key in params:
-            try:
-                _luhya_model[key].load_state_dict(params[key])
-            except:
-                from collections import OrderedDict
-                sd = OrderedDict()
-                for k, v in params[key].items():
-                    sd[k[7:] if k.startswith("module.") else k] = v
-                _luhya_model[key].load_state_dict(sd, strict=False)
-    _ = [_luhya_model[k].eval() for k in _luhya_model]
-
-    _luhya_sampler = DiffusionSampler(
-        _luhya_model.diffusion.diffusion,
-        sampler=ADPM2Sampler(),
-        sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0),
-        clamp=False
+# ── Clone + patch StyleTTS2 ────────────────────────────────
+os.system("git clone https://github.com/yl4579/StyleTTS2 /app/StyleTTS2 2>/dev/null || true")
+
+for fpath in ["/app/StyleTTS2/models.py", "/app/StyleTTS2/utils.py"]:
+    with open(fpath) as f:
+        code = f.read()
+    patched = re.sub(
+        r'torch\.load\(([^)]+)\)',
+        lambda m: m.group(0) if 'weights_only' in m.group(1)
+                  else f'torch.load({m.group(1)}, weights_only=False)',
+        code
     )
+    with open(fpath, "w") as f:
+        f.write(patched)
+print("Patched torch.load")
+
+sys.path.insert(0, "/app/StyleTTS2")
+os.chdir("/app/StyleTTS2")
+
+model_dir = snapshot_download("crazydev919/luhya-tts")
+
+from models import *
+from utils import *
+from text_utils import TextCleaner
+from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule
+
+device      = "cpu"
+textcleaner = TextCleaner()
+
+config = yaml.safe_load(open(f"{model_dir}/config.yml"))
+config["ASR_path"]   = f"{model_dir}/Utils/ASR/epoch_00080.pth"
+config["ASR_config"] = f"{model_dir}/Utils/ASR/config.yml"
+config["F0_path"]    = f"{model_dir}/Utils/JDC/bst.t7"
+config["PLBERT_dir"] = f"{model_dir}/Utils/PLBERT/"
+
+text_aligner    = load_ASR_models(config["ASR_path"], config["ASR_config"])
+pitch_extractor = load_F0_models(config["F0_path"])
+from Utils.PLBERT.util import load_plbert
+plbert = load_plbert(config["PLBERT_dir"])
+
+model_params = recursive_munch(config["model_params"])
+model = build_model(model_params, text_aligner, pitch_extractor, plbert)
+_ = [model[key].eval() for key in model]
+_ = [model[key].to(device) for key in model]
+
+params = torch.load(
+    f"{model_dir}/model.pth", map_location="cpu", weights_only=False
+)["net"]
+for key in model:
+    if key in params:
+        try:
+            model[key].load_state_dict(params[key])
+        except:
+            from collections import OrderedDict
+            sd = OrderedDict()
+            for k, v in params[key].items():
+                sd[k[7:] if k.startswith("module.") else k] = v
+            model[key].load_state_dict(sd, strict=False)
+_ = [model[key].eval() for key in model]
+print("Model loaded")
+
+sampler = DiffusionSampler(
+    model.diffusion.diffusion,
+    sampler=ADPM2Sampler(),
+    sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0),
+    clamp=False
+)
 
-    # Precompute default style
-    ref_candidates = (
-        glob.glob(f"{_model_dir}/ref_wavs/*.wav") +
-        glob.glob(f"{_model_dir}/*.wav")
-    )
-    DEFAULT_REF = sorted(ref_candidates)[0]
+to_mel = torchaudio.transforms.MelSpectrogram(
+    n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
+mean, std = -4, 4
 
-    to_mel = torchaudio.transforms.MelSpectrogram(
-        n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
-    mean, std = -4, 4
+def length_to_mask(lengths):
+    mask = torch.arange(lengths.max()).unsqueeze(0).expand(
+        lengths.shape[0], -1).type_as(lengths)
+    mask = torch.gt(mask + 1, lengths.unsqueeze(1))
+    return mask
 
-    wave, sr = librosa.load(DEFAULT_REF, sr=24000)
-    audio, _ = librosa.effects.trim(wave, top_db=30)
-    wave_t   = torch.from_numpy(audio).float()
-    mel      = (torch.log(1e-5 + to_mel(wave_t).unsqueeze(0)) - mean) / std
-    mel      = mel.to(device)
+def preprocess(wave):
+    wave_tensor = torch.from_numpy(wave).float()
+    mel_tensor  = to_mel(wave_tensor)
+    mel_tensor  = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
+    return mel_tensor
 
+def compute_style(path):
+    wave, sr = librosa.load(path, sr=24000)
+    audio, _ = librosa.effects.trim(wave, top_db=30)
+    mel      = preprocess(audio).to(device)
     with torch.no_grad():
-        ref_s = _luhya_model.style_encoder(mel.unsqueeze(1))
-        ref_p = _luhya_model.predictor_encoder(mel.unsqueeze(1))
-    _default_style = torch.cat([ref_s, ref_p], dim=1)
-
-    _luhya_loaded = True
-    print("✅ Luhya model loaded and cached")
+        ref_s = model.style_encoder(mel.unsqueeze(1))
+        ref_p = model.predictor_encoder(mel.unsqueeze(1))
+    return torch.cat([ref_s, ref_p], dim=1)
 
-def synthesize_luhya(text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
-    load_luhya()  # lazy load
+ref_candidates = (
+    glob.glob(f"{model_dir}/ref_wavs/*.wav") +
+    glob.glob(f"{model_dir}/*.wav")
+)
+DEFAULT_REF   = sorted(ref_candidates)[0]
+DEFAULT_STYLE = compute_style(DEFAULT_REF)
+print(f"Reference: {DEFAULT_REF}")
 
+def synthesize(text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
     import phonemizer
-    device = "cpu"
-
-    to_mel = torchaudio.transforms.MelSpectrogram(
-        n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
-    mean, std = -4, 4
-
-    def length_to_mask(lengths):
-        mask = torch.arange(lengths.max()).unsqueeze(0).expand(
-            lengths.shape[0], -1).type_as(lengths)
-        mask = torch.gt(mask + 1, lengths.unsqueeze(1))
-        return mask
-
-    def compute_style(path):
-        wave, sr = librosa.load(path, sr=24000)
-        audio, _ = librosa.effects.trim(wave, top_db=30)
-        wave_t   = torch.from_numpy(audio).float()
-        mel      = (torch.log(1e-5 + to_mel(wave_t).unsqueeze(0)) - mean) / std
-        with torch.no_grad():
-            ref_s = _luhya_model.style_encoder(mel.unsqueeze(1))
-            ref_p = _luhya_model.predictor_encoder(mel.unsqueeze(1))
-        return torch.cat([ref_s, ref_p], dim=1)
-
     pb    = phonemizer.backend.EspeakBackend(
         language="sw", preserve_punctuation=True, with_stress=True)
-    ref_s = compute_style(ref_audio) if ref_audio else _default_style
+    ref_s = compute_style(ref_audio) if ref_audio else DEFAULT_STYLE
 
     ps     = " ".join(word_tokenize(pb.phonemize([text.strip()])[0]))
-    tokens = _textcleaner(ps)
+    tokens = textcleaner(ps)
     tokens.insert(0, 0)
     tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
 
     with torch.no_grad():
         il   = torch.LongTensor([tokens.shape[-1]]).to(device)
         tm   = length_to_mask(il).to(device)
-        t_en = _luhya_model.text_encoder(tokens, il, tm)
-        bd   = _luhya_model.bert(tokens, attention_mask=(~tm).int())
-        d_en = _luhya_model.bert_encoder(bd).transpose(-1, -2)
+        t_en = model.text_encoder(tokens, il, tm)
+        bd   = model.bert(tokens, attention_mask=(~tm).int())
+        d_en = model.bert_encoder(bd).transpose(-1, -2)
 
-        sp = _luhya_sampler(
+        sp = sampler(
             noise=torch.randn((1, 256)).unsqueeze(1).to(device),
             embedding=bd, embedding_scale=1,
             features=ref_s, num_steps=int(steps)
@@ -194,9 +138,9 @@ def synthesize_luhya(text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
         s   = beta  * sp[:, 128:] + (1 - beta)  * ref_s[:, 128:]
         ref = alpha * sp[:, :128] + (1 - alpha) * ref_s[:, :128]
 
-        d    = _luhya_model.predictor.text_encoder(d_en, s, il, tm)
-        x, _ = _luhya_model.predictor.lstm(d)
-        dur  = torch.sigmoid(_luhya_model.predictor.duration_proj(x)).sum(axis=-1)
+        d    = model.predictor.text_encoder(d_en, s, il, tm)
+        x, _ = model.predictor.lstm(d)
+        dur  = torch.sigmoid(model.predictor.duration_proj(x)).sum(axis=-1)
         pd   = torch.round(dur.squeeze()).clamp(min=1)
 
         at = torch.zeros(il, int(pd.sum().data))
@@ -208,7 +152,7 @@ def synthesize_luhya(text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
         en  = d.transpose(-1, -2) @ at.unsqueeze(0).to(device)
         asr = t_en @ at.unsqueeze(0).to(device)
 
-        if _model_params.decoder.type == "hifigan":
+        if model_params.decoder.type == "hifigan":
             en_new  = torch.zeros_like(en)
             asr_new = torch.zeros_like(asr)
             en_new[:, :, 0]   = en[:, :, 0]
@@ -217,75 +161,24 @@ def synthesize_luhya(text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
             asr_new[:, :, 1:] = asr[:, :, :-1]
             en, asr = en_new, asr_new
 
-        F0, N = _luhya_model.predictor.F0Ntrain(en, s)
-        out   = _luhya_model.decoder(asr, F0, N, ref.squeeze().unsqueeze(0))
+        F0, N = model.predictor.F0Ntrain(en, s)
+        out   = model.decoder(asr, F0, N, ref.squeeze().unsqueeze(0))
 
     wav = out.squeeze().cpu().numpy()[..., :-50]
-    sf.write("/tmp/luhya_output.wav", wav, 24000)
-    return "/tmp/luhya_output.wav"
-
-# ══════════════════════════════════════════════════════════
-# UNIFIED ENTRY POINT
-# ══════════════════════════════════════════════════════════
-def synthesize(language, text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
-    if language == "Kikuyu":
-        return synthesize_kikuyu(text)
-    else:
-        return synthesize_luhya(text, ref_audio, alpha, beta, steps)
-
-# ══════════════════════════════════════════════════════════
-# GRADIO UI
-# ══════════════════════════════════════════════════════════
-with gr.Blocks(title="Kenyan Languages TTS") as demo:
-    gr.Markdown("# 🗣️ Kenyan Languages TTS\nLuhya (Lunyore) and Kikuyu text-to-speech.")
-
-    language = gr.Radio(
-        choices=["Kikuyu", "Luhya"],
-        value="Kikuyu",
-        label="Language"
-    )
-    text = gr.Textbox(
-        label="Text",
-        value="Mũtũũrĩre wa ndũire nĩ kĩheo kĩa mwanya.",
-        lines=4
-    )
-
-    with gr.Group(visible=False) as luhya_controls:
-        gr.Markdown("**Luhya voice controls**")
-        ref_audio = gr.Audio(label="Reference Voice (optional)", type="filepath")
-        with gr.Row():
-            alpha = gr.Slider(0.0, 1.0, value=0.3, step=0.1, label="Alpha")
-            beta  = gr.Slider(0.0, 1.0, value=0.7, step=0.1, label="Beta")
-            steps = gr.Slider(1, 10,   value=5,   step=1,   label="Steps")
-
-    output_audio = gr.Audio(label="Generated Speech", type="filepath")
-    btn          = gr.Button("Generate", variant="primary")
-
-    def toggle(lang):
-        texts = {
-            "Kikuyu": "Mũtũũrĩre wa ndũire nĩ kĩheo kĩa mwanya.",
-            "Luhya" : "mirembe. obulani lwa bwana nyasaye.",
-        }
-        return gr.update(visible=(lang=="Luhya")), gr.update(value=texts[lang])
-
-    language.change(fn=toggle, inputs=language, outputs=[luhya_controls, text])
-
-    btn.click(
-        fn      = synthesize,
-        inputs  = [language, text, ref_audio, alpha, beta, steps],
-        outputs = output_audio
-    )
-
-    gr.Markdown("""
-    **API:**
-    POST /api/predict
-{"data": ["Kikuyu", "your kikuyu text", null, 0.3, 0.7, 5]}
-{"data": ["Luhya",  "your luhya text",  null, 0.3, 0.7, 5]}
-Note: First Luhya request takes ~30s to load the model.
-    """)
-
-demo.queue(max_size=3).launch(
-    prevent_thread_lock = True,
-    max_threads         = 1,
-    show_error          = True,
+    sf.write("/tmp/output.wav", wav, 24000)
+    return "/tmp/output.wav"
+
+demo = gr.Interface(
+    fn      = synthesize,
+    inputs  = [
+        gr.Textbox(label="Luhya text", value="mirembe. obulani lwa bwana nyasaye.", lines=3),
+        gr.Audio(label="Reference voice (optional)", type="filepath", value=None),
+        gr.Slider(0.0, 1.0, value=0.3, step=0.1, label="Alpha"),
+        gr.Slider(0.0, 1.0, value=0.7, step=0.1, label="Beta"),
+        gr.Slider(1, 10, value=5, step=1, label="Diffusion steps"),
+    ],
+    outputs = gr.Audio(label="Generated speech", type="filepath"),
+    title   = "Luhya (Lunyore) TTS",
 )
+
+demo.launch(server_name="0.0.0.0", server_port=7860)