Upload app.py with huggingface_hub

app.py

@@ -12,7 +12,7 @@ import nltk
 nltk.download("punkt_tab", quiet=True)
 
 # ══════════════════════════════════════════════════════════
-# KIKUYU — lightweight VITS model
+# KIKUYU — load at startup (small, 145MB)
 # ══════════════════════════════════════════════════════════
 from transformers import VitsModel, AutoTokenizer
 
@@ -32,129 +32,160 @@ def synthesize_kikuyu(text):
     return "/tmp/kikuyu_output.wav"
 
 # ══════════════════════════════════════════════════════════
-# LUHYA — StyleTTS2 fine-tuned model
+# LUHYA — lazy load on first request (heavy, 2.18GB)
 # ══════════════════════════════════════════════════════════
-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
+_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
     )
-    with open(fpath, "w") as f:
-        f.write(patched)
-print("✅ Patched torch.load calls")
-
-sys.path.insert(0, "/app/StyleTTS2")
-os.chdir("/app/StyleTTS2")
-
-LUHYA_REPO = "crazydev919/luhya-tts"
-model_dir  = snapshot_download(LUHYA_REPO)
-
-from models import *
-from utils import *
-from text_utils import TextCleaner
-from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule
-
-device      = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Running on: {device}")
-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"])
-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 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[key].eval() for key in luhya_model]
-print("✅ Luhya model loaded")
-
-sampler = DiffusionSampler(
-    luhya_model.diffusion.diffusion,
-    sampler=ADPM2Sampler(),
-    sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0),
-    clamp=False
-)
-
-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
+    # 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]
 
-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
+    to_mel = torchaudio.transforms.MelSpectrogram(
+        n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
+    mean, std = -4, 4
 
-def compute_style(path):
-    wave, sr = librosa.load(path, sr=24000)
+    wave, sr = librosa.load(DEFAULT_REF, sr=24000)
     audio, _ = librosa.effects.trim(wave, top_db=30)
-    mel      = preprocess(audio).to(device)
+    wave_t   = torch.from_numpy(audio).float()
+    mel      = (torch.log(1e-5 + to_mel(wave_t).unsqueeze(0)) - mean) / std
+    mel      = mel.to(device)
+
     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)
+        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)
 
-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}")
+    _luhya_loaded = True
+    print("✅ Luhya model loaded and cached")
 
 def synthesize_luhya(text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
+    load_luhya()  # lazy load
+
     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 = _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(
+        sp = _luhya_sampler(
             noise=torch.randn((1, 256)).unsqueeze(1).to(device),
             embedding=bd, embedding_scale=1,
             features=ref_s, num_steps=int(steps)
@@ -163,9 +194,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    = _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))
@@ -177,7 +208,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]
@@ -186,15 +217,15 @@ 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 = _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/luhya_output.wav", wav, 24000)
     return "/tmp/luhya_output.wav"
 
 # ══════════════════════════════════════════════════════════
-# UNIFIED FUNCTION — language param routes to correct model
+# UNIFIED ENTRY POINT
 # ══════════════════════════════════════════════════════════
 def synthesize(language, text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
     if language == "Kikuyu":
@@ -206,66 +237,55 @@ def synthesize(language, text, ref_audio=None, alpha=0.3, beta=0.7, steps=5):
 # 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"
-        )
+    gr.Markdown("# 🗣️ Kenyan Languages TTS\nLuhya (Lunyore) and Kikuyu text-to-speech.")
 
-    with gr.Row():
-        text = gr.Textbox(
-            label = "Text",
-            value = "mirembe. obulani lwa bwana nyasaye.",
-            lines = 4
-        )
+    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
+    )
 
-    # Luhya-only controls
-    with gr.Group(visible=True) as luhya_controls:
+    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():
-            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")
+            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")
-    generate_btn = gr.Button("Generate", variant="primary")
+    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.",
+    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(value=examples[lang])
+        return gr.update(visible=(lang=="Luhya")), gr.update(value=texts[lang])
 
-    language.change(fn=update_placeholder, inputs=language, outputs=text)
+    language.change(fn=toggle, inputs=language, outputs=[luhya_controls, text])
 
-    generate_btn.click(
+    btn.click(
         fn      = synthesize,
         inputs  = [language, text, ref_audio, alpha, beta, steps],
         outputs = output_audio
     )
 
     gr.Markdown("""
-    **API usage:**
+    **API:**
     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)
+{"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,
+)