Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import base64
 import io
 import logging
@@ -5,15 +6,22 @@ import numpy as np
 import torch
 import librosa
 import uvicorn
+
 from fastapi import FastAPI, HTTPException, Security, Depends
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.security.api_key import APIKeyHeader
 from pydantic import BaseModel
 from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
 
-# Config
+# ======================================================
+# CONFIG & SECRETS
+# ======================================================
 API_KEY_NAME = "access_token"
 API_KEY_VALUE = "HCL_SECURE_KEY_2026"
+
+# Get your Hugging Face token from the Space's Secret settings
+HF_TOKEN = os.getenv("HF_Token") 
+
 MODEL_ID = "melba-t/wav2vec2-fake-speech-detection" 
 TARGET_SR = 16000
 LABEL_MAP = {0: "HUMAN", 1: "AI_GENERATED"}
@@ -21,13 +29,34 @@ LABEL_MAP = {0: "HUMAN", 1: "AI_GENERATED"}
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("hcl-api")
 
-# Initialize Model
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_ID)
-model = AutoModelForAudioClassification.from_pretrained(MODEL_ID).to(DEVICE)
-model.eval()
 
+# ======================================================
+# MODEL INITIALIZATION (WITH AUTH)
+# ======================================================
+try:
+    logger.info(f"Loading private model {MODEL_ID}...")
+    # Passing the token allows access to the private/restricted repo
+    feature_extractor = AutoFeatureExtractor.from_pretrained(
+        MODEL_ID, 
+        token=HF_TOKEN
+    )
+    model = AutoModelForAudioClassification.from_pretrained(
+        MODEL_ID, 
+        token=HF_TOKEN
+    ).to(DEVICE)
+    model.eval()
+    logger.info("Model loaded successfully.")
+except Exception as e:
+    logger.error(f"Error loading model: {e}")
+    # Fallback to prevent app crash if token is missing
+    model = None 
+
+# ======================================================
+# FASTAPI APP
+# ======================================================
 app = FastAPI(title="HCL AI Voice Detection API")
+api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
 
 app.add_middleware(
     CORSMiddleware,
@@ -39,8 +68,6 @@ app.add_middleware(
 class AudioRequest(BaseModel):
     audio_base64: str
 
-api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
-
 async def verify_api_key(api_key: str = Security(api_key_header)):
     if api_key != API_KEY_VALUE:
         raise HTTPException(status_code=403, detail="Invalid API Key")
@@ -51,14 +78,14 @@ def preprocess_audio(b64_string: str):
         if "," in b64_string:
             b64_string = b64_string.split(",")[1]
         
-        # Correct padding
+        # Standardize padding
         missing_padding = len(b64_string) % 4
         if missing_padding:
             b64_string += "=" * (4 - missing_padding)
 
         audio_bytes = base64.b64decode(b64_string)
         
-        # Load via librosa for better MP3 compatibility
+        # Load audio using librosa (requires ffmpeg in packages.txt)
         with io.BytesIO(audio_bytes) as bio:
             audio, sr = librosa.load(bio, sr=TARGET_SR)
 
@@ -72,13 +99,21 @@ def preprocess_audio(b64_string: str):
 
 @app.get("/")
 def home():
-    return {"message": "API is running. Visit /docs for Swagger UI"}
+    return {"message": "HCL Voice Detection API Active. Visit /docs"}
 
 @app.post("/predict")
 async def predict(request: AudioRequest, _: str = Depends(verify_api_key)):
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded. Check HF_Token.")
+        
     try:
-        waveform = preprocess_audio(request.audio_base_64)
-        inputs = feature_extractor(waveform, sampling_rate=TARGET_SR, return_tensors="pt").to(DEVICE)
+        waveform = preprocess_audio(request.audio_base64)
+        
+        inputs = feature_extractor(
+            waveform, 
+            sampling_rate=TARGET_SR, 
+            return_tensors="pt"
+        ).to(DEVICE)
 
         with torch.inference_mode():
             logits = model(**inputs).logits
@@ -86,15 +121,19 @@ async def predict(request: AudioRequest, _: str = Depends(verify_api_key)):
 
         confidence, pred_idx = torch.max(probs, dim=-1)
         
+        # Map prediction to required hackathon labels
+        label = LABEL_MAP.get(int(pred_idx.item()), "UNKNOWN")
+
         return {
-            "classification": LABEL_MAP.get(int(pred_idx.item()), "UNKNOWN"),
+            "classification": label,
             "confidence_score": round(float(confidence.item()), 4)
         }
     except ValueError as ve:
         raise HTTPException(status_code=400, detail=str(ve))
     except Exception as e:
         logger.error(f"Prediction error: {e}")
-        raise HTTPException(status_code=500, detail="Internal Server Error")
+        raise HTTPException(status_code=500, detail="Inference Error")
 
 if __name__ == "__main__":
-    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)
+    # Port 7860 is required for Hugging Face Spaces
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)