Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 
-# === IMPORTANT ===
-# Set cache directories BEFORE any imports that use Hugging Face or PyTorch caching
+# Set cache dirs before imports to fix permission errors
 os.environ["HF_HOME"] = "/tmp"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 os.environ["TORCH_HOME"] = "/tmp"
@@ -9,17 +8,18 @@ os.environ["XDG_CACHE_HOME"] = "/tmp"
 
 import io
 import re
+import math
 import numpy as np
 import scipy.io.wavfile
 import torch
-from fastapi import FastAPI
+from fastapi import FastAPI, Query
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from transformers import VitsModel, AutoTokenizer
 
 app = FastAPI()
 
-# Load model and tokenizer once at startup
+# Load model/tokenizer once at startup
 model = VitsModel.from_pretrained("Somali-tts/somali_tts_model")
 tokenizer = AutoTokenizer.from_pretrained("saleolow/somali-mms-tts")
 
@@ -27,7 +27,6 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 model.eval()
 
-# Somali number words dictionary for normalization
 number_words = {
     0: "eber", 1: "koow", 2: "labo", 3: "seddex", 4: "afar", 5: "shan",
     6: "lix", 7: "todobo", 8: "sideed", 9: "sagaal", 10: "toban",
@@ -75,11 +74,9 @@ def number_to_words(number: int) -> str:
         return str(number)
 
 def normalize_text(text: str) -> str:
-    # Replace digits with Somali words
     numbers = re.findall(r'\d+', text)
     for num in numbers:
         text = text.replace(num, number_to_words(int(num)))
-    # Additional Somali text normalizations
     text = text.replace("KH", "qa").replace("Z", "S")
     text = text.replace("SH", "SHa'a").replace("DH", "Dha'a")
     text = text.replace("ZamZam", "SamSam")
@@ -89,28 +86,60 @@ class TextIn(BaseModel):
     inputs: str
 
 @app.post("/synthesize")
-async def synthesize(data: TextIn):
+async def synthesize(data: TextIn, test: bool = Query(False, description="Set true to generate test tone instead of TTS")):
+    if test:
+        # Generate 2-second 440Hz sine wave for testing playback
+        duration_s = 2.0
+        sample_rate = 22050
+        t = np.linspace(0, duration_s, int(sample_rate*duration_s), endpoint=False)
+        freq = 440
+        waveform = 0.5 * np.sin(2 * math.pi * freq * t).astype(np.float32)
+        pcm_waveform = (waveform * 32767).astype(np.int16)
+
+        buf = io.BytesIO()
+        scipy.io.wavfile.write(buf, rate=sample_rate, data=pcm_waveform)
+        buf.seek(0)
+
+        print(f"[TEST MODE] Generated test tone: {pcm_waveform.shape[0]} samples, Sample rate: {sample_rate}")
+
+        return StreamingResponse(buf, media_type="audio/wav")
+
+    # Normalize input text
     text = normalize_text(data.inputs)
+
+    # Tokenize and move to device
     inputs = tokenizer(text, return_tensors="pt").to(device)
 
+    # Generate waveform
     with torch.no_grad():
         output = model(**inputs)
-        waveform = output.waveform.squeeze().cpu().numpy()
 
-    # Mono conversion if multi-channel
-    if waveform.ndim > 1:
-        waveform = waveform.mean(axis=0)
+    print("Raw waveform shape:", output.waveform.shape)
+
+    waveform = output.waveform.cpu().numpy()
+
+    # Process waveform dimensions
+    if waveform.ndim == 3:
+        waveform = waveform[0]  # batch dimension
+    if waveform.ndim == 2:
+        waveform = waveform.mean(axis=0)  # average channels to mono
+
+    print("Processed waveform shape:", waveform.shape)
+    print("Waveform min/max before clip:", waveform.min(), waveform.max())
 
     waveform = waveform.astype(np.float32)
     waveform = np.clip(waveform, -1.0, 1.0)
 
     pcm_waveform = (waveform * 32767).astype(np.int16)
 
+    print("PCM waveform shape:", pcm_waveform.shape)
+    print("PCM waveform min/max:", pcm_waveform.min(), pcm_waveform.max())
+
     buf = io.BytesIO()
     sample_rate = getattr(model.config, "sampling_rate", 22050)
+    print("Sample rate:", sample_rate)
+
     scipy.io.wavfile.write(buf, rate=sample_rate, data=pcm_waveform)
     buf.seek(0)
 
-    print(f"Generated audio length: {pcm_waveform.shape[0]} samples, Sample rate: {sample_rate}")
-
     return StreamingResponse(buf, media_type="audio/wav")