Update app.py

app.py

@@ -1,30 +1,130 @@
 import gradio as gr
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
 import torch
-import soundfile as sf
-import tempfile
+import numpy as np
+import scipy.io.wavfile
+from transformers import VitsModel, AutoTokenizer
+import re
 
-model_id = "jellecali8/somali_tts_model"
+# Load model and tokenizer
+model = VitsModel.from_pretrained("Somali-tts/somali_tts_model")
+tokenizer = AutoTokenizer.from_pretrained("saleolow/somali-mms-tts")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+model.eval()
 
-processor = AutoProcessor.from_pretrained(model_id)
-model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id)
+# Numbers in Somali
+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(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)
+
+def normalize_text(text):
+    text = re.sub(r'(\d{1,3})(,\d{3})+', lambda m: m.group(0).replace(",", ""), text)
+    text = re.sub(r'\.\d+', '', text)
+    def replace_num(match):
+        return number_to_words(match.group())
+    text = re.sub(r'\d+', replace_num, text)
+    symbol_map = {
+        '$': 'doolar',
+        '=': 'egwal',
+        '+': 'balaas',
+        '#': 'haash'
+    }
+    for sym, word in symbol_map.items():
+        text = text.replace(sym, ' ' + word + ' ')
+    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
 
 def tts(text):
-    inputs = processor(text, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(**inputs)
-    audio = outputs[0].cpu().numpy()
+    paragraphs = text.strip().split("\n")
+    audio_list = []
+    max_chars = 500  # Qiyaasta ugu badan 2 daqiiqo
+    warn_msg = ""
 
-    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-        sf.write(f.name, audio, samplerate=16000)
-        return f.name
+    for i, para in enumerate(paragraphs):
+        para = para.strip()
+        if not para:
+            continue
 
-iface = gr.Interface(
-    fn=tts,
-    inputs=gr.Textbox(lines=2, placeholder="Ku qor qoraalka Somali halkan...", label="Qoraalka Somali"),
-    outputs=gr.Audio(label="Codka la soo saaray"),
-    title="Somali TTS Demo",
-    description="Qoraal ku qor af-Soomaali kadib dhageyso codka."
-)
+        if len(para) > max_chars:
+            warn_msg += f"❗ Qaybta {i+1} aad ayaa ka badan 2 daqiiqo. Waan kala jaray.\n"
+            sub_parts = [para[j:j+max_chars] for j in range(0, len(para), max_chars)]
+        else:
+            sub_parts = [para]
+
+        for part in sub_parts:
+            norm_para = normalize_text(part)
+            inputs = tokenizer(norm_para, return_tensors="pt").to(device)
+            with torch.no_grad():
+                waveform = model(**inputs).waveform.squeeze().cpu().numpy()
 
-iface.launch()
+            pause = np.zeros(int(model.config.sampling_rate * 0.8))  # 0.8s pause
+            audio_list.append(np.concatenate((waveform, pause)))
+
+    final_audio = np.concatenate(audio_list)
+    filename = "output.wav"
+    scipy.io.wavfile.write(filename, rate=model.config.sampling_rate, data=(final_audio * 32767).astype(np.int16))
+
+    if warn_msg:
+        print(warn_msg)
+    return filename
+
+# Gradio interface
+gr.Interface(
+    fn=tts,
+    inputs=gr.Textbox(label="Geli qoraal Soomaali ah", lines=10, placeholder="Ku qor 1 ama in ka badan paragraph..."),
+    outputs=gr.Audio(label="Codka TTS"),
+    title="Somali TTS",
+    description="Ku qor qoraal Soomaaliyeed si aad u maqasho cod dabiici ah. Qoraalka ha ka badnaan 2 daqiiqo per jumlad."
+).launch()