Update app.py

app.py

@@ -2,73 +2,92 @@ import gradio as gr
 import torch
 import numpy as np
 import scipy.io.wavfile
-import re
 from transformers import VitsModel, AutoTokenizer
+import re
 
-# 1. Load fine-tuned multi-speaker VITS model from HF Hub
+# Load fine-tuned model from Hugging Face Hub or local path
 model = VitsModel.from_pretrained("Somali-tts/somali_tts_model")
-# 2. Load tokenizer for Somali TTS
 tokenizer = AutoTokenizer.from_pretrained("saleolow/somali-mms-tts")
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 model.eval()
 
-# 3. Number-to-words helper
-number_map = {
-    0: "eber", 1: "koow", 2: "labo", 3: "seddex", 4: "afar",
-    5: "shan", 6: "lix", 7: "todobo", 8: "sideed", 9: "sagaal",
-    10: "toban"
+number_words = {
+    0: "eber", 1: "koow", 2: "labo", 3: "seddex", 4: "afar", 5: "shan",
+    6: "lix", 7: "todobo", 8: "sideed", 9: "sagaal", 10: "toban",
+    11: "toban iyo koow", 12: "toban iyo labo", 13: "toban iyo seddex",
+    14: "toban iyo afar", 15: "toban iyo shan", 16: "toban iyo lix",
+    17: "toban iyo todobo", 18: "toban iyo sideed", 19: "toban iyo sagaal",
+    20: "labaatan", 30: "sodon", 40: "afartan", 50: "konton",
+    60: "lixdan", 70: "todobaatan", 80: "sideetan", 90: "sagaashan",
+    100: "boqol", 1000: "kun"
 }
 
-def number_to_words(n: int) -> str:
-    if n <= 10:
-        return number_map[n]
-    if n < 20:
-        return "toban iyo " + number_map[n - 10]
-    tens, unit = divmod(n, 10)
-    s = number_map[tens * 10] if tens > 1 else "toban"
-    if unit:
-        s += " iyo " + number_map[unit]
-    return s
+def number_to_words(number):
+    number = int(number)
+    if number < 20:
+        return number_words[number]
+    elif number < 100:
+        tens, unit = divmod(number, 10)
+        return number_words[tens * 10] + (" iyo " + number_words[unit] if unit else "")
+    elif number < 1000:
+        hundreds, remainder = divmod(number, 100)
+        part = (number_words[hundreds] + " boqol") if hundreds > 1 else "boqol"
+        if remainder:
+            part += " iyo " + number_to_words(remainder)
+        return part
+    elif number < 1000000:
+        thousands, remainder = divmod(number, 1000)
+        words = []
+        if thousands == 1:
+            words.append("kun")
+        else:
+            words.append(number_to_words(thousands) + " kun")
+        if remainder >= 100:
+            hundreds, rem2 = divmod(remainder, 100)
+            if hundreds:
+                boqol_text = (number_words[hundreds] + " boqol") if hundreds > 1 else "boqol"
+                words.append(boqol_text)
+            if rem2:
+                words.append("iyo " + number_to_words(rem2))
+        elif remainder:
+            words.append("iyo " + number_to_words(remainder))
+        return " ".join(words)
+    elif number < 1000000000:
+        millions, remainder = divmod(number, 1000000)
+        words = []
+        if millions == 1:
+            words.append("milyan")
+        else:
+            words.append(number_to_words(millions) + " milyan")
+        if remainder:
+            words.append(number_to_words(remainder))
+        return " ".join(words)
+    else:
+        return str(number)
 
-# 4. Text normalization
-def normalize_text(text: str) -> str:
-    text = re.sub(r"\d+", lambda m: number_to_words(int(m.group())), text)
-    # basic orthographic fixes
-    text = text.replace("KH", "qa").replace("DH", "Dha").replace("SH", "Sha")
+def normalize_text(text):
+    numbers = re.findall(r'\d+', text)
+    for num in numbers:
+        text = text.replace(num, number_to_words(num))
+    text = text.replace("KH", "qa").replace("Z", "S")
+    text = text.replace("SH", "SHa'a").replace("DH", "Dha'a")
+    text = text.replace("ZamZam", "SamSam")
     return text
 
-# 5. TTS function with speaker selection
-# speaker: "Male" or "Female"
-def tts(text: str, speaker: str) -> str:
+def tts(text):
     text = normalize_text(text)
     inputs = tokenizer(text, return_tensors="pt").to(device)
-    spk_id = 0 if speaker == "Male" else 1
-    speaker_ids = torch.tensor([spk_id], device=device)
-
     with torch.no_grad():
-        out = model(**inputs, speaker_ids=speaker_ids)
-        wav = out.waveform.squeeze().cpu().numpy()
-
-    # write to WAV
+        waveform = model(**inputs).waveform.squeeze().cpu().numpy()
     filename = "output.wav"
-    scipy.io.wavfile.write(filename, model.config.sampling_rate,
-                             (wav * 32767).astype(np.int16))
+    scipy.io.wavfile.write(filename, rate=model.config.sampling_rate, data=(waveform * 32767).astype(np.int16))
     return filename
 
-# 6. Gradio interface
-demo = gr.Interface(
+gr.Interface(
     fn=tts,
-    inputs=[
-        gr.Textbox(lines=3, label="Geli qoraal Soomaali ah", placeholder="Qor qoraalka halkan…"),
-        gr.Radio(["Male", "Female"], label="Dooro codka")
-    ],
+    inputs=gr.Textbox(label="Geli qoraal Soomaali ah"),
     outputs=gr.Audio(label="Codka TTS"),
-    title="Somali TTS (Male & Female)",
-    description="Ku qor qoraal Soomaali ah, dooro cod Rag ama Dumar, kadib riix Submit.",
-    allow_flagging="never"
-)
-
-if __name__ == "__main__":
-    demo.launch()
+    title="Somali TTS",
+    description="Ku qor qoraal Soomaaliyeed si aad u maqasho cod dabiici ah.",
+).launch()