Upload app.py with huggingface_hub

app.py

@@ -4,50 +4,56 @@ 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)
 
-# ── Clone StyleTTS2 ────────────────────────────────────────
-os.system("git clone https://github.com/yl4579/StyleTTS2 /app/StyleTTS2 2>/dev/null || true")
+# ══════════════════════════════════════════════════════════
+# KIKUYU — lightweight VITS model
+# ══════════════════════════════════════════════════════════
+from transformers import VitsModel, AutoTokenizer
 
-# ── Patch models.py for PyTorch 2.6 weights_only ──────────
-models_path = "/app/StyleTTS2/models.py"
-with open(models_path) as f:
-    code = f.read()
+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")
 
-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(models_path, "w") as f:
-    f.write(patched)
-
-# Also patch utils.py
-utils_path = "/app/StyleTTS2/utils.py"
-with open(utils_path) as f:
-    code2 = f.read()
-
-patched2 = 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)',
-    code2
-)
-with open(utils_path, "w") as f:
-    f.write(patched2)
+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 — StyleTTS2 fine-tuned model
+# ══════════════════════════════════════════════════════════
+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 calls")
 
 sys.path.insert(0, "/app/StyleTTS2")
 os.chdir("/app/StyleTTS2")
 
-MODEL_REPO = "crazydev919/luhya-tts"
-model_dir  = snapshot_download(MODEL_REPO)
+LUHYA_REPO = "crazydev919/luhya-tts"
+model_dir  = snapshot_download(LUHYA_REPO)
 
 from models import *
 from utils import *
@@ -58,49 +64,45 @@ device      = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Running on: {device}")
 textcleaner = TextCleaner()
 
-# ── Load config ────────────────────────────────────────────
 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/"
 
-# ── Load models ────────────────────────────────────────────
 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]
+luhya_model = build_model(model_params, text_aligner, pitch_extractor, plbert)
+_ = [luhya_model[key].eval() for key in luhya_model]
+_ = [luhya_model[key].to(device) for key in luhya_model]
 
 params = torch.load(
     f"{model_dir}/model.pth", map_location="cpu", weights_only=False
 )["net"]
-for key in model:
+for key in luhya_model:
     if key in params:
         try:
-            model[key].load_state_dict(params[key])
+            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
-            model[key].load_state_dict(sd, strict=False)
-_ = [model[key].eval() for key in model]
-print("✅ Model loaded")
+            luhya_model[key].load_state_dict(sd, strict=False)
+_ = [luhya_model[key].eval() for key in luhya_model]
+print("✅ Luhya model loaded")
 
-# ── Diffusion sampler ──────────────────────────────────────
 sampler = DiffusionSampler(
-    model.diffusion.diffusion,
+    luhya_model.diffusion.diffusion,
     sampler=ADPM2Sampler(),
     sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0),
     clamp=False
 )
 
-# ── Helpers ────────────────────────────────────────────────
 to_mel = torchaudio.transforms.MelSpectrogram(
     n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
 mean, std = -4, 4
@@ -122,16 +124,19 @@ def compute_style(path):
     audio, _ = librosa.effects.trim(wave, top_db=30)
     mel      = preprocess(audio).to(device)
     with torch.no_grad():
-        ref_s = model.style_encoder(mel.unsqueeze(1))
-        ref_p = model.predictor_encoder(mel.unsqueeze(1))
+        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)
 
-DEFAULT_REF   = sorted(glob.glob(f"{model_dir}/ref_wavs/*.wav"))[0]
+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"Default reference: {DEFAULT_REF}")
+print(f"✅ Default reference: {DEFAULT_REF}")
 
-# ── Inference ──────────────────────────────────────────────
-def synthesize(text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
+def synthesize_luhya(text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
     import phonemizer
     pb    = phonemizer.backend.EspeakBackend(
         language="sw", preserve_punctuation=True, with_stress=True)
@@ -145,9 +150,9 @@ def synthesize(text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
     with torch.no_grad():
         il   = torch.LongTensor([tokens.shape[-1]]).to(device)
         tm   = length_to_mask(il).to(device)
-        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)
+        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)
 
         sp = sampler(
             noise=torch.randn((1, 256)).unsqueeze(1).to(device),
@@ -158,9 +163,9 @@ def synthesize(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    = 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)
+        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)
         pd   = torch.round(dur.squeeze()).clamp(min=1)
 
         at = torch.zeros(il, int(pd.sum().data))
@@ -181,43 +186,86 @@ def synthesize(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 = model.predictor.F0Ntrain(en, s)
-        out   = model.decoder(asr, F0, N, ref.squeeze().unsqueeze(0))
+        F0, N = luhya_model.predictor.F0Ntrain(en, s)
+        out   = luhya_model.decoder(asr, F0, N, ref.squeeze().unsqueeze(0))
 
     wav = out.squeeze().cpu().numpy()[..., :-50]
-    sf.write("/tmp/output.wav", wav, 24000)
-    return "/tmp/output.wav"
-
-# ── Gradio UI ──────────────────────────────────────────────
-demo = gr.Interface(
-    fn          = synthesize,
-    inputs      = [
-        gr.Textbox(
-            label = "Luhya Text",
+    sf.write("/tmp/luhya_output.wav", wav, 24000)
+    return "/tmp/luhya_output.wav"
+
+# ══════════════════════════════════════════════════════════
+# UNIFIED FUNCTION — language param routes to correct model
+# ══════════════════════════════════════════════════════════
+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\nText-to-speech for Luhya (Lunyore) and Kikuyu.")
+
+    with gr.Row():
+        language = gr.Radio(
+            choices = ["Luhya", "Kikuyu"],
+            value   = "Luhya",
+            label   = "Language"
+        )
+
+    with gr.Row():
+        text = gr.Textbox(
+            label = "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 (style)"),
-        gr.Slider(0.0, 1.0, value=0.7, step=0.1, label="Beta (prosody)"),
-        gr.Slider(1,   10,  value=5,   step=1,   label="Diffusion steps"),
-    ],
-    outputs     = gr.Audio(label="Generated Speech", type="filepath"),
-    title       = "Luhya (Lunyore) TTS",
-    description = (
-        "Type Luhya text and get speech audio.\n"
-        "Fine-tuned StyleTTS2 on Luhya Lunyore data.\n\n"
-        "**API:** POST to `/api/predict` with "
-        "`{\"data\": [\"your luhya text\", null, 0.3, 0.7, 5]}`"
-    ),
-    examples    = [
-        ["mirembe. obulani lwa bwana nyasaye.", None, 0.3, 0.7, 5],
-        ["nyasaye nareba, adam olithena?",       None, 0.3, 0.7, 5],
-    ],
-    cache_examples = False,
-)
+            lines = 4
+        )
+
+    # Luhya-only controls
+    with gr.Group(visible=True) as luhya_controls:
+        gr.Markdown("**Luhya voice controls**")
+        with gr.Row():
+            ref_audio = gr.Audio(
+                label = "Reference Voice (optional)",
+                type  = "filepath",
+                value = None
+            )
+        with gr.Row():
+            alpha = gr.Slider(0.0, 1.0, value=0.3, step=0.1, label="Alpha (style)")
+            beta  = gr.Slider(0.0, 1.0, value=0.7, step=0.1, label="Beta (prosody)")
+            steps = gr.Slider(1, 10, value=5, step=1, label="Diffusion steps")
+
+    output_audio = gr.Audio(label="Generated Speech", type="filepath")
+    generate_btn = gr.Button("Generate", variant="primary")
+
+    # Show/hide luhya controls based on language
+    def toggle_controls(lang):
+        return gr.update(visible=(lang == "Luhya"))
+
+    language.change(fn=toggle_controls, inputs=language, outputs=luhya_controls)
+
+    # Update placeholder text based on language
+    def update_placeholder(lang):
+        examples = {
+            "Luhya"  : "mirembe. obulani lwa bwana nyasaye.",
+            "Kikuyu" : "Mũtũũrĩre wa ndũire nĩ kĩheo kĩa mwanya mũno.",
+        }
+        return gr.update(value=examples[lang])
+
+    language.change(fn=update_placeholder, inputs=language, outputs=text)
+
+    generate_btn.click(
+        fn      = synthesize,
+        inputs  = [language, text, ref_audio, alpha, beta, steps],
+        outputs = output_audio
+    )
+
+    gr.Markdown("""
+    **API usage:**
+    POST /api/predict
+{"data": ["Luhya", "mirembe.", null, 0.3, 0.7, 5]}
+{"data": ["Kikuyu", "Mũtũũrĩre wa ndũire nĩ kĩheo.", null, 0.3, 0.7, 5]}
+""")
+
 demo.launch(prevent_thread_lock=True)