Add v2 pipeline model (prithivMLmods), route by model param

app.py

@@ -13,7 +13,7 @@ import torch.nn.functional as F
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from safetensors.torch import load_file
-from transformers import Wav2Vec2Model, Wav2Vec2Config
+from transformers import Wav2Vec2Model, Wav2Vec2Config, pipeline
 
 app = FastAPI(title="Speech Emotion Recognition")
 
@@ -25,10 +25,14 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-MODEL_CLASSES = ["angry", "disgust", "fear", "happy", "neutral", "pleasant_surprise", "sad"]
-MODEL_CLASS_INDEX = {c: i for i, c in enumerate(MODEL_CLASSES)}
+# ── v1: Custom Wav2Vec2 + SERHead (7 classes) ──────────────────────────
+MODEL_CLASSES_V1 = ["angry", "disgust", "fear", "happy", "neutral", "pleasant_surprise", "sad"]
 
-_model = None
+# ── v2: HuggingFace pipeline (8 classes) ────────────────────────────────
+V2_LABELS = ["ANG", "CAL", "DIS", "FEA", "HAP", "NEU", "SAD", "SUR"]
+
+_model_v1 = None
+_model_v2 = None
 
 
 class SERHead(nn.Module):
@@ -46,72 +50,67 @@ class SERHead(nn.Module):
         return self.classifier(F.relu(self.projector(pooled)))
 
 
-def load_model():
-    global _model
+def load_model_v1():
+    global _model_v1
     device = "cuda" if torch.cuda.is_available() else "cpu"
-
     config = Wav2Vec2Config.from_pretrained("facebook/wav2vec2-base")
     backbone = Wav2Vec2Model(config).to(device).eval()
     head = SERHead().to(device).eval()
 
     model_path = "model.safetensors"
     if not os.path.exists(model_path):
-        print(f"[WARN] {model_path} not found — using dummy predictor")
-        _model = "dummy"
+        print("[WARN] model.safetensors not found — v1 unavailable")
+        _model_v1 = "unavailable"
         return
 
     state = load_file(model_path)
-
-    # Strip "wav2vec2." prefix from backbone keys
     backbone_prefix = "wav2vec2."
-    backbone_state = {
-        k[len(backbone_prefix):]: v
-        for k, v in state.items()
-        if k.startswith(backbone_prefix)
-    }
-    missing, unexpected = backbone.load_state_dict(backbone_state, strict=False)
-    if missing:
-        print(f"[WARN] Backbone missing keys: {missing[:5]}...")
-    if unexpected:
-        print(f"[WARN] Backbone unexpected keys: {unexpected[:5]}...")
-
+    backbone_state = {k[len(backbone_prefix):]: v for k, v in state.items() if k.startswith(backbone_prefix)}
+    backbone.load_state_dict(backbone_state, strict=False)
     head.load_state_dict(state, strict=False)
+    _model_v1 = {"backbone": backbone, "head": head, "device": device}
+    print("[INFO] v1 (Wav2Vec2) loaded")
 
-    # Sanity check: run a tiny forward pass
-    dummy_input = torch.randn(1, 16000)
-    with torch.no_grad():
-        out = backbone(dummy_input, output_hidden_states=True)
-        logits = head(out.hidden_states)
-        probs = F.softmax(logits, dim=-1).squeeze(0)
-    print(f"[INFO] Sanity check — probs: {probs[:3].tolist()} ... {probs[-3:].tolist()}")
-    print(f"[INFO] Sanity check — argmax: {MODEL_CLASSES[int(probs.argmax())]}")
 
-    _model = {"backbone": backbone, "head": head, "device": device}
-    print(f"[INFO] Model loaded on {device}")
+def load_model_v2():
+    global _model_v2
+    try:
+        _model_v2 = pipeline(
+            "audio-classification",
+            model="prithivMLmods/Speech-Emotion-Classification",
+        )
+        print("[INFO] v2 (prithivMLmods) loaded")
+    except Exception as e:
+        print(f"[WARN] v2 failed to load: {e}")
+        _model_v2 = "unavailable"
 
 
 @app.on_event("startup")
 async def startup():
-    load_model()
+    load_model_v1()
+    load_model_v2()
 
 
-def _extract_features(waveform: np.ndarray, sr: int) -> np.ndarray:
+def _decode_audio(wav_bytes: bytes):
+    import soundfile as sf
     import librosa
-    mfcc = librosa.feature.mfcc(y=waveform, sr=sr, n_mfcc=40, n_fft=512, hop_length=256)
-    mfcc_delta = librosa.feature.delta(mfcc)
-    mfcc_delta2 = librosa.feature.delta(mfcc, order=2)
-    features = np.concatenate([mfcc, mfcc_delta, mfcc_delta2], axis=0)
-    return features.mean(axis=1)
-
-
-def _predict(waveform: np.ndarray, sr: int) -> dict:
-    global _model
-    if _model is None or _model == "dummy":
-        return _dummy_predict(waveform, sr)
+    buf = io.BytesIO(wav_bytes)
+    waveform, sr = sf.read(buf)
+    if waveform.ndim > 1:
+        waveform = waveform.mean(axis=1)
+    if sr != 16000:
+        waveform = librosa.resample(y=waveform, orig_sr=sr, target_sr=16000)
+        sr = 16000
+    return waveform, sr
+
+
+def _predict_v1(waveform: np.ndarray) -> dict:
+    if _model_v1 is None or _model_v1 == "unavailable":
+        return {"emotion": "neutral", "confidence": 0.0, "probabilities": {}}
 
-    backbone = _model["backbone"]
-    head = _model["head"]
-    device = _model["device"]
+    backbone = _model_v1["backbone"]
+    head = _model_v1["head"]
+    device = _model_v1["device"]
 
     wav_t = torch.from_numpy(waveform).float().unsqueeze(0).to(device)
     with torch.no_grad():
@@ -121,57 +120,29 @@ def _predict(waveform: np.ndarray, sr: int) -> dict:
 
     probs_np = probs.cpu().numpy()
     pred_idx = int(probs_np.argmax())
-    emotion = MODEL_CLASSES[pred_idx]
-    prob_map = {c: round(float(probs_np[i]), 4) for i, c in enumerate(MODEL_CLASSES)}
+    emotion = MODEL_CLASSES_V1[pred_idx]
+    prob_map = {c: round(float(probs_np[i]), 4) for i, c in enumerate(MODEL_CLASSES_V1)}
+    return {"emotion": emotion, "confidence": round(float(probs_np[pred_idx]), 4), "probabilities": prob_map}
 
-    return {
-        "emotion": emotion,
-        "confidence": round(float(probs_np[pred_idx]), 4),
-        "probabilities": prob_map,
-    }
 
+def _predict_v2(waveform: np.ndarray, sr: int) -> dict:
+    if _model_v2 is None or _model_v2 == "unavailable":
+        return {"emotion": "neutral", "confidence": 0.0, "probabilities": {}}
 
-def _dummy_predict(waveform: np.ndarray, sr: int) -> dict:
-    rng = np.random.default_rng(42)
-    probs = rng.dirichlet(np.ones(len(MODEL_CLASSES)) * 0.5)
-    probs = probs / probs.sum()
-    prob_map = {c: round(float(p), 4) for c, p in zip(MODEL_CLASSES, probs)}
+    result = _model_v2(waveform, top_k=8)
+    probs = {r["label"]: r["score"] for r in result}
+    top = result[0]
     return {
-        "emotion": "neutral",
-        "confidence": round(float(probs[MODEL_CLASS_INDEX["neutral"]]), 4),
-        "probabilities": prob_map,
+        "emotion": top["label"],
+        "confidence": round(float(top["score"]), 4),
+        "probabilities": probs,
     }
 
 
-def _predict_from_bytes(wav_bytes: bytes) -> dict:
-    try:
-        import soundfile as sf
-        import librosa
-
-        buf = io.BytesIO(wav_bytes)
-        waveform, sr = sf.read(buf)
-        if waveform.ndim > 1:
-            waveform = waveform.mean(axis=1)
-        if sr != 16000:
-            waveform = librosa.resample(y=waveform, orig_sr=sr, target_sr=16000)
-            sr = 16000
-    except Exception as e:
-        print(f"[WARN] Failed to decode audio: {e}")
-        return {"emotion": "neutral", "confidence": 0.0, "probabilities": {}}
-
-    return _predict(waveform, sr)
-
-
 @app.get("/")
 @app.get("/health")
 async def health():
-    return {
-        "status": "ok",
-        "model_loaded": _model is not None and _model != "dummy",
-    }
-
-
-
+    return {"status": "ok", "v1_loaded": _model_v1 is not None and _model_v1 != "unavailable", "v2_loaded": _model_v2 is not None and _model_v2 != "unavailable"}
 
 
 @app.post("/predict_b64")
@@ -183,8 +154,8 @@ async def predict_b64(request: Request):
         if "application/json" in content_type or body.startswith(b"{"):
             payload = json.loads(body)
             b64_str = payload.get("audio") or payload.get("image", "")
+            model_ver = payload.get("model", "v1")
         else:
-            # form-encoded with key "data"
             import urllib.parse
             parsed = urllib.parse.parse_qs(body.decode())
             raw = parsed.get("data", [None])[0]
@@ -192,12 +163,20 @@ async def predict_b64(request: Request):
                 raise HTTPException(status_code=400, detail="Missing 'data' field")
             payload = json.loads(raw)
             b64_str = payload.get("audio") or payload.get("image", "") or raw
+            model_ver = payload.get("model", "v1")
 
         if not b64_str:
             raise HTTPException(status_code=400, detail="No audio data found")
 
         wav_bytes = base64.b64decode(b64_str)
-        result = _predict_from_bytes(wav_bytes)
+        waveform, sr = _decode_audio(wav_bytes)
+
+        if model_ver == "v2":
+            result = _predict_v2(waveform, sr)
+        else:
+            result = _predict_v1(waveform)
+
+        result["model"] = model_ver
         return result
     except HTTPException:
         raise