Update app.py

app.py

@@ -1,11 +1,4 @@
 import os
-
-# Set cache directories to /tmp to avoid permission issues in the container
-os.environ["HF_HOME"] = "/tmp"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp"
-os.environ["TORCH_HOME"] = "/tmp"
-os.environ["XDG_CACHE_HOME"] = "/tmp"
-
 import io
 import re
 import numpy as np
@@ -16,15 +9,23 @@ from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from transformers import VitsModel, AutoTokenizer
 
+# Set environment variables to avoid permission issues in container
+os.environ["HF_HOME"] = "/tmp"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp"
+os.environ["TORCH_HOME"] = "/tmp"
+os.environ["XDG_CACHE_HOME"] = "/tmp"
+
 app = FastAPI()
 
-# Load model and tokenizer once at startup
+# Load model and tokenizer at startup
 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()
 
+# Number-to-Somali words dictionary
 number_words = {
     0: "eber", 1: "koow", 2: "labo", 3: "seddex", 4: "afar", 5: "shan",
     6: "lix", 7: "todobo", 8: "sideed", 9: "sagaal", 10: "toban",
@@ -36,8 +37,7 @@ number_words = {
     100: "boqol", 1000: "kun"
 }
 
-def number_to_words(number):
-    number = int(number)
+def number_to_words(number: int) -> str:
     if number < 20:
         return number_words[number]
     elif number < 100:
@@ -56,14 +56,7 @@ def number_to_words(number):
             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:
+        if remainder:
             words.append("iyo " + number_to_words(remainder))
         return " ".join(words)
     elif number < 1000000000:
@@ -79,10 +72,12 @@ def number_to_words(number):
     else:
         return str(number)
 
-def normalize_text(text):
+def normalize_text(text: str) -> str:
+    # Replace numbers with words
     numbers = re.findall(r'\d+', text)
     for num in numbers:
-        text = text.replace(num, number_to_words(num))
+        text = text.replace(num, number_to_words(int(num)))
+    # Additional Somali text normalization rules
     text = text.replace("KH", "qa").replace("Z", "S")
     text = text.replace("SH", "SHa'a").replace("DH", "Dha'a")
     text = text.replace("ZamZam", "SamSam")
@@ -93,23 +88,38 @@ class TextIn(BaseModel):
 
 @app.post("/synthesize")
 async def synthesize(data: TextIn):
+    # Normalize and convert text input
     text = normalize_text(data.inputs)
+
+    # Tokenize and move to device
     inputs = tokenizer(text, return_tensors="pt").to(device)
+
+    # Generate waveform with no_grad
     with torch.no_grad():
-        waveform = model(**inputs).waveform.squeeze().cpu().numpy()
+        output = model(**inputs)
+        waveform = output.waveform.squeeze().cpu().numpy()
 
-    # If stereo or multi-channel, convert to mono by averaging channels
+    # If multi-channel audio, average channels to mono
     if waveform.ndim > 1:
         waveform = waveform.mean(axis=0)
 
-    print(f"Audio length: {waveform.shape[0]}, min: {waveform.min()}, max: {waveform.max()}")
+    # Convert to float32 if not already
+    waveform = waveform.astype(np.float32)
 
+    # Clip waveform to [-1.0, 1.0]
     waveform = np.clip(waveform, -1.0, 1.0)
 
-    buf = io.BytesIO()
+    # Convert to 16-bit PCM
+    pcm_waveform = (waveform * 32767).astype(np.int16)
 
+    # Prepare WAV file in memory buffer
+    buf = io.BytesIO()
     sample_rate = getattr(model.config, "sampling_rate", 22050)
-    scipy.io.wavfile.write(buf, rate=sample_rate, data=(waveform * 32767).astype(np.int16))
+    scipy.io.wavfile.write(buf, rate=sample_rate, data=pcm_waveform)
     buf.seek(0)
 
+    # Debug info
+    print(f"Generated audio length: {pcm_waveform.shape[0]} samples, Sample rate: {sample_rate}")
+
+    # Stream response as WAV audio
     return StreamingResponse(buf, media_type="audio/wav")