Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import torch
 from transformers import VitsModel, AutoTokenizer
 from fastapi import FastAPI, HTTPException, Header
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 import scipy.io.wavfile
 import uuid
@@ -10,7 +10,7 @@ import numpy as np
 
 app = FastAPI()
 
-# --- CORS Permissions ---
+# --- CORS Permissions (Zaroori) ---
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -23,61 +23,79 @@ OUTPUT_DIR = "/tmp"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 API_KEY = os.getenv("API_KEY", "MySecretPassword123")
 
-print("⏳ Loading AI Models... (Thora time lagega)")
+print("⏳ Loading AI Models... (1-2 minute lagenge)")
 
-# --- LOADING REAL AI MODELS (Facebook MMS) ---
-# Ye models pehli baar run honay par download honge (10-20 seconds)
+# --- GLOBAL VARIABLES FOR MODELS ---
+model_ur = None
+tokenizer_ur = None
+model_hi = None
+tokenizer_hi = None
+
+# --- LOADING REAL AI MODELS (Fixed Names) ---
 try:
-    # Urdu Model
-    model_ur = VitsModel.from_pretrained("facebook/mms-tts-urd-script-arabic")
-    tokenizer_ur = AutoTokenizer.from_pretrained("facebook/mms-tts-urd-script-arabic")
+    # 1. URDU MODEL (Correct Name: facebook/mms-tts-urd)
+    print("Downloading Urdu Model...")
+    model_ur = VitsModel.from_pretrained("facebook/mms-tts-urd")
+    tokenizer_ur = AutoTokenizer.from_pretrained("facebook/mms-tts-urd")
     
-    # Hindi Model
+    # 2. HINDI MODEL (Correct Name: facebook/mms-tts-hin)
+    print("Downloading Hindi Model...")
     model_hi = VitsModel.from_pretrained("facebook/mms-tts-hin")
     tokenizer_hi = AutoTokenizer.from_pretrained("facebook/mms-tts-hin")
     
-    print("✅ AI Models Loaded Successfully!")
+    print("✅ All AI Models Loaded Successfully!")
+
 except Exception as e:
-    print(f"❌ Model Loading Error: {e}")
+    print(f"❌ CRITICAL ERROR LOADING MODELS: {e}")
 
 @app.get("/")
 def home():
-    return {"status": "Online", "message": "Real AI VITS Model Running"}
+    if model_ur is None:
+        return {"status": "Error", "message": "Models failed to load. Check Logs."}
+    return {"status": "Online", "message": "Real AI VITS Model Ready (Fixed)"}
 
 @app.post("/generate")
 async def generate_tts(
     text: str,
-    voice_id: str = "urdu", # urdu or hindi
+    voice_id: str = "urdu", 
     x_api_key: str = Header(None)
 ):
-    # 1. Security Check
+    # 1. Security
     if x_api_key != API_KEY:
         raise HTTPException(status_code=401, detail="Invalid API Key")
 
+    # 2. Check if models exist
+    if model_ur is None or model_hi is None:
+         return JSONResponse(status_code=500, content={"error": "Models not loaded yet. Check Server Logs."})
+
     filename = f"{uuid.uuid4()}.wav"
     filepath = os.path.join(OUTPUT_DIR, filename)
 
     try:
-        # 2. Select Language Model
+        # 3. Generate Logic
+        inputs = None
+        model = None
+
         if "hindi" in voice_id.lower():
             inputs = tokenizer_hi(text, return_tensors="pt")
-            with torch.no_grad():
-                output = model_hi(**inputs).waveform
+            model = model_hi
         else:
-            # Default to Urdu
+            # Urdu Default
             inputs = tokenizer_ur(text, return_tensors="pt")
-            with torch.no_grad():
-                output = model_ur(**inputs).waveform
+            model = model_ur
+
+        # 4. Create Waveform (No Internet Needed)
+        with torch.no_grad():
+            output = model(**inputs).waveform
 
-        # 3. Save Audio File (WAV format)
-        # Convert PyTorch tensor to audio file
+        # 5. Save File
         audio_data = output.cpu().numpy().squeeze()
-        scipy.io.wavfile.write(filepath, rate=model_ur.config.sampling_rate, data=audio_data)
+        scipy.io.wavfile.write(filepath, rate=model.config.sampling_rate, data=audio_data)
 
         return FileResponse(filepath, media_type="audio/wav", filename="ai_audio.wav")
 
     except Exception as e:
-        return {"error": str(e)}
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
 if __name__ == "__main__":
     import uvicorn