Update main.py

main.py

@@ -1,9 +1,8 @@
 import os
 import json
 import torch
-import librosa
+import torchaudio
 import requests
-import soundfile as sf
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import FileResponse
 from transformers import (
@@ -24,82 +23,88 @@ VOICE_ID = config["eleven_voice_id"]
 LLM_URL = config["llm_url"]
 
 
-# STT Model
- 
+
+def load_audio(audio_path, target_sr=16000):
+    wav, sr = torchaudio.load(audio_path)
+
+    if wav.shape[0] > 1:
+        wav = wav.mean(dim=0, keepdim=True)
+
+
+    if sr != target_sr:
+        wav = torchaudio.functional.resample(wav, sr, target_sr)
+
+    return wav.squeeze().numpy(), target_sr
+
+
+#  STT MODEL 
 print("Loading STT model...")
 stt_processor = Wav2Vec2Processor.from_pretrained("facebook/mms-1b-all")
 stt_model = Wav2Vec2ForCTC.from_pretrained("facebook/mms-1b-all").to(DEVICE)
 stt_model.eval()
-print("STT loaded ")
-
+print("STT loaded")
 
 def transcribe(audio_path):
-    wav, sr = librosa.load(audio_path, sr=16000)
-    inputs = stt_processor(wav, sampling_rate=16000, return_tensors="pt", padding=True)
+    wav, sr = load_audio(audio_path)
+    inputs = stt_processor(wav, sampling_rate=sr, return_tensors="pt", padding=True)
     with torch.no_grad():
         logits = stt_model(inputs.input_values.to(DEVICE)).logits
     ids = torch.argmax(logits, dim=-1)
     return stt_processor.batch_decode(ids)[0].strip()
 
 
-#  Emotion Model
-
+# EMOTION MODEL  #
 print("Loading Emotion model...")
 emotion_extractor = AutoFeatureExtractor.from_pretrained("superb/hubert-base-superb-er")
 emotion_model = AutoModelForAudioClassification.from_pretrained(
     "superb/hubert-base-superb-er"
 ).to(DEVICE)
 emotion_model.eval()
-print("Emotion model loaded ")
-
+print("Emotion model loaded")
 
 def get_emotion(audio_path):
-    wav, sr = librosa.load(audio_path, sr=16000)
-    feats = emotion_extractor(wav, sampling_rate=16000, return_tensors="pt")
+    wav, sr = load_audio(audio_path)
+    feats = emotion_extractor(wav, sampling_rate=sr, return_tensors="pt")
     with torch.no_grad():
         out = emotion_model(feats["input_values"].to(DEVICE))
     pred = torch.argmax(out.logits, dim=-1).item()
     return emotion_model.config.id2label[pred]
 
 
-
-#  LLM Call
-
+#  LLM CALL  
 def ask_llm(text):
     payload = {"query": text}
     r = requests.post(LLM_URL, json=payload, timeout=200)
+
     try:
         return r.json()["answer"]
     except:
         return str(r.json())
 
 
-
-#  TTS
-
+# TTS 
 def tts_eleven(text, out_file="response.mp3"):
     url = f"https://api.elevenlabs.io/v1/text-to-speech/{VOICE_ID}"
     headers = {
         "xi-api-key": ELEVEN_API_KEY,
-        "Content-Type": "application/json"
+        "Content-Type": "application/json",
     }
     payload = {"text": text, "model_id": "eleven_multilingual_v2"}
 
     resp = requests.post(url, json=payload, headers=headers)
     if resp.status_code != 200:
-        raise Exception(f"ElevenLabs TTS Error: {resp.text}")
+        raise Exception(f"ElevenLabs API Error: {resp.text}")
 
     with open(out_file, "wb") as f:
         f.write(resp.content)
-    return out_file
-
 
+    return out_file
 
-# FastAPI App
 
+# FASTAPI APP 
 app = FastAPI(title="Voice AI API")
 
-# Enable CORS for Hugging Face Spaces frontend
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -108,23 +113,22 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-
 @app.post("/process-audio/")
 async def process_audio(file: UploadFile = File(...)):
     audio_path = f"temp_{file.filename}"
     with open(audio_path, "wb") as f:
         f.write(await file.read())
 
-   
     transcript = transcribe(audio_path)
     emotion = get_emotion(audio_path)
-    llm_out = ask_llm(transcript)
-    tts_file = tts_eleven(llm_out)
+    llm_response = ask_llm(transcript)
+    tts_file = tts_eleven(llm_response)
 
-    # Return TTS file as downloadable mp3
     return FileResponse(tts_file, media_type="audio/mpeg", filename="response.mp3")
 
 
 @app.get("/")
 async def root():
-    return {"message": "Voice AI API is running. Use /process-audio/ endpoint to upload audio."}
+    return {
+        "message": "Voice AI API is running. Use /process-audio/ to upload audio."
+    }