changed fastapi script to use onnx int8

src/api/app.py

@@ -5,39 +5,45 @@ import os
 import torch
 import librosa
 import numpy as np
-from transformers import AutoFeatureExtractor, Wav2Vec2ForSequenceClassification
+from optimum.onnxruntime import ORTModelForAudioClassification
+from transformers import AutoFeatureExtractor
 from typing import List, Dict
 import tempfile
 
-app = FastAPI(title="VigilAudio Emotion API")
+app = FastAPI(title="VigilAudio Optimized API")
 
-MODEL_PATH = "models/wav2vec2-finetuned"
-DEVICE = torch.device("cpu")
+# --- CONFIG ---
+# We use the INT8 model which proved to be the fastest in benchmarks
+MODEL_PATH = "models/onnx_quantized"
 
-print(f"Loading model into API memory...")
+# --- MODEL LOADING (Optimized with ONNX) ---
+print(f"Loading OPTIMIZED INT8 ONNX model into memory...")
 try:
     feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_PATH)
-    model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_PATH)
-    model.to(DEVICE)
-    model.eval()
+    # Note: we explicitly pass file_name since optimum expects model.onnx by default
+    model = ORTModelForAudioClassification.from_pretrained(MODEL_PATH, file_name="model_quantized.onnx")
+    
+    # Label mapping from config
     id2label = model.config.id2label
-    print(f"API Model Ready. Labels: {list(id2label.values())}")
+    print(f"Optimized API Ready. Speedup expected: ~1.8x")
 except Exception as e:
     print(f"API Failed to load model: {e}")
     model = None
 
+# --- UTILS ---
 def segment_audio(audio, sr, window_size=3.0):
-    """Splits audio into fixed-size windows."""
     chunk_len = int(window_size * sr)
     for i in range(0, len(audio), chunk_len):
         yield audio[i:i + chunk_len]
 
+# --- ENDPOINTS ---
 @app.get("/health")
 def health():
     return {
         "status": "online",
+        "engine": "ONNX Runtime (INT8)",
         "model_loaded": model is not None,
-        "device": str(DEVICE)
+        "labels": list(id2label.values()) if model else []
     }
 
 @app.post("/predict")
@@ -56,13 +62,14 @@ async def predict_emotion(file: UploadFile = File(...)):
         timeline = []
         
         for i, chunk in enumerate(segment_audio(speech, sr, window_size=3.0)):
-            if len(chunk) < 8000: 
-                continue
+            if len(chunk) < 8000: continue
                 
             inputs = feature_extractor(chunk, sampling_rate=16000, return_tensors="pt", padding=True)
             
+            # ONNX Inference
             with torch.no_grad():
-                logits = model(inputs.input_values.to(DEVICE)).logits
+                outputs = model(**inputs)
+                logits = outputs.logits
                 probs = torch.nn.functional.softmax(logits, dim=-1)
                 pred_id = torch.argmax(logits, dim=-1).item()
             
@@ -78,6 +85,7 @@ async def predict_emotion(file: UploadFile = File(...)):
 
         return {
             "filename": file.filename,
+            "engine": "ONNX_INT8",
             "duration_seconds": round(duration, 2),
             "dominant_emotion": dominant,
             "timeline": timeline
@@ -87,9 +95,8 @@ async def predict_emotion(file: UploadFile = File(...)):
         print(f"Prediction error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
     finally:
-        # Cleanup temp file
         if os.path.exists(tmp_path):
             os.remove(tmp_path)
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)