Update app.py

app.py

@@ -8,21 +8,14 @@ import onnxruntime as ort
 import soundfile as sf
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from huggingface_hub import hf_hub_download
 
-# ==========================================
-# 1. Setup Logging
-# ==========================================
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(message)s",
-    handlers=[logging.StreamHandler()]
-)
+# Setup Logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 logger = logging.getLogger("LID_Engine")
 
 app = FastAPI(title="Pakistani LID AI Engine (SOTA V3)")
 
-# ✅ CORS Fix for your HTML frontend
+# CORS Fix
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -31,54 +24,41 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ==========================================
-# 2. Model Initialization
-# ==========================================
-logger.info("Initializing SOTA Engine...")
+# Load Model (Baked into the Docker image)
+MODEL_DIR = "local_model"
+MODEL_PATH = os.path.join(MODEL_DIR, "pakistani_lid_v3.onnx")
+
+logger.info("🚀 Loading pre-baked ONNX model...")
 try:
-    os.makedirs("local_model", exist_ok=True)
+    # Check if files exist just in case
+    if not os.path.exists(MODEL_PATH):
+        raise FileNotFoundError(f"Model not found at {MODEL_PATH}")
     
-    logger.info("Downloading ONNX files...")
-    hf_hub_download(repo_id="Hammad712/pakistani-lid-v3-sota", filename="pakistani_lid_v3.onnx.data", local_dir="local_model")
-    model_path = hf_hub_download(repo_id="Hammad712/pakistani-lid-v3-sota", filename="pakistani_lid_v3.onnx", local_dir="local_model")
-    
-    session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
-    logger.info("✅ Model loaded successfully!")
+    session = ort.InferenceSession(MODEL_PATH, providers=['CPUExecutionProvider'])
+    logger.info("✅ Engine is LIVE and Ready!")
 except Exception as e:
-    logger.error(f"❌ Initialization failed: {e}")
+    logger.error(f"❌ Failed to load model: {e}")
     raise e
 
 labels = ("balochi", "english", "pashto", "sindhi", "urdu")
 id2label = {i: label for i, label in enumerate(labels)}
 
-# ==========================================
-# 3. Inference Logic
-# ==========================================
 def predict_audio(audio_path):
-    # 🚨 Using Soundfile to avoid Torchaudio backend errors
     data, sr = sf.read(audio_path)
-    
-    # Convert to torch tensor [channels, samples]
     waveform = torch.from_numpy(data).float()
     if waveform.ndim == 2:
-        waveform = waveform.T  # soundfile uses [samples, channels]
-        waveform = waveform.mean(dim=0, keepdim=True)
+        waveform = waveform.T.mean(dim=0, keepdim=True)
     else:
         waveform = waveform.unsqueeze(0)
     
-    # Resample to 16kHz
     if sr != 16000:
         waveform = torchaudio.functional.resample(waveform, sr, 16000)
     
-    # Normalize & Clip to 15s
     target_frames = 16000 * 15
-    if waveform.shape[1] > target_frames:
-        waveform = waveform[:, :target_frames]
-    
+    waveform = waveform[:, :target_frames]
     waveform = (waveform / waveform.abs().max().clamp(min=1e-6)) - waveform.mean()
     waveform = waveform / waveform.std().clamp(min=1e-6)
     
-    # Create Mask
     length = waveform.shape[1]
     mask = torch.zeros(target_frames, dtype=torch.long)
     if length < target_frames:
@@ -87,7 +67,6 @@ def predict_audio(audio_path):
     else:
         mask[:] = 1
 
-    # ONNX Inference
     ort_inputs = {
         "input_values": waveform.numpy(),
         "attention_mask": mask.unsqueeze(0).numpy()
@@ -99,29 +78,19 @@ def predict_audio(audio_path):
     
     return id2label[pred_id], float(probs[0][pred_id])
 
-# ==========================================
-# 4. API Endpoint
-# ==========================================
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    logger.info(f"Inference request: {file.filename}")
     temp_path = f"temp_{file.filename}"
-    
     try:
         with open(temp_path, "wb") as f:
             f.write(await file.read())
-
         lang, conf = predict_audio(temp_path)
         os.remove(temp_path)
-        
-        logger.info(f"Result: {lang} ({conf:.2%})")
         return {"success": True, "language": lang.upper(), "confidence": round(conf * 100, 2)}
-    
     except Exception as e:
-        logger.error(f"Prediction error: {e}")
         if os.path.exists(temp_path): os.remove(temp_path)
         return {"success": False, "error": str(e)}
 
 @app.get("/")
-def health_check():
-    return {"status": "online", "model": "Pakistani LID V3 SOTA"}
+def health():
+    return {"status": "online"}