Update app.py

app.py

@@ -1,71 +1,71 @@
-import gradio as gr
-import torch
-from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor
-import numpy as np
-import torchaudio
-
-# =========================
-# CONFIG
-# =========================
-MODEL_NAME = "your-username/Audio-Emotion-Detection"  # <- replace with your repo name
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# =========================
-# LOAD MODEL & PROCESSOR
-# =========================
-processor = Wav2Vec2Processor.from_pretrained(MODEL_NAME)
-model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_NAME).to(DEVICE)
-
-# Emotion labels in same order used during training
-LABELS = ["Angry", "Disgusted", "Fearful", "Happy", "Neutral", "Sad", "Surprised"]
-
-# =========================
-# PREDICTION PIPELINE
-# =========================
-def predict(audio):
-    # audio: tuple (sample_rate, numpy array)
-    sr, data = audio
-
-    # Resample to 16k if necessary
-    if sr != 16000:
-        data = torchaudio.functional.resample(torch.tensor(data), sr, 16000).numpy()
-        sr = 16000
-
-    # Process input
-    inputs = processor(
-        data,
-        sampling_rate=sr,
-        return_tensors="pt",
-        padding=True,
-        truncation=True
-    ).to(DEVICE)
-
-    # Forward pass
-    with torch.no_grad():
-        logits = model(**inputs).logits
-        probs = torch.nn.functional.softmax(logits, dim=-1)[0]
-        pred_idx = torch.argmax(probs).item()
-
-    return {LABELS[i]: float(probs[i]) for i in range(len(LABELS))}
-
-# =========================
-# GRADIO INTERFACE
-# =========================
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Audio(sources=["upload", "microphone"], type="numpy", label="Upload or Record Audio"),
-    outputs=gr.Label(num_top_classes=3),
-    title="Audio Emotion Detection 🎧",
-    description=(
-        "Fine-tuned Wav2Vec2 model for detecting emotions from voice. "
-        "Supports 7 emotions: Angry, Disgusted, Fearful, Happy, Neutral, Sad, and Surprised. "
-        "All audio should be 16kHz."
-    ),
-    allow_flagging="never",
-)
-
-# =========================
-# LAUNCH APP
-# =========================
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+import torch
+from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor
+import numpy as np
+import torchaudio
+
+# =========================
+# CONFIG
+# =========================
+MODEL_NAME = "Hatman/audio-emotion-detection"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# =========================
+# LOAD MODEL & PROCESSOR
+# =========================
+processor = Wav2Vec2Processor.from_pretrained(MODEL_NAME)
+model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_NAME).to(DEVICE)
+
+# Emotion labels in same order used during training (matches the model card)
+LABELS = ["angry", "disgust", "fear", "happy", "neutral", "sad", "surprised"]
+
+# =========================
+# PREDICTION PIPELINE
+# =========================
+def predict(audio):
+    # audio: tuple (sample_rate, numpy array)
+    sr, data = audio
+
+    # Resample to 16k if needed
+    if sr != 16000:
+        data = torchaudio.functional.resample(torch.tensor(data), sr, 16000).numpy()
+        sr = 16000
+
+    # Process input
+    inputs = processor(
+        data,
+        sampling_rate=sr,
+        return_tensors="pt",
+        padding=True,
+        truncation=True
+    ).to(DEVICE)
+
+    # Forward pass
+    with torch.no_grad():
+        logits = model(**inputs).logits
+        probs = torch.nn.functional.softmax(logits, dim=-1)[0]
+        pred_idx = torch.argmax(probs).item()
+
+    return {LABELS[i]: float(probs[i]) for i in range(len(LABELS))}
+
+# =========================
+# GRADIO INTERFACE
+# =========================
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Audio(sources=["upload", "microphone"], type="numpy", label="Upload or Record Audio"),
+    outputs=gr.Label(num_top_classes=3),
+    title="Audio Emotion Detection 🎧",
+    description=(
+        "Wav2Vec2 emotion classification model. "
+        "Supports 7 emotions: Angry, Disgust, Fear, Happy, Neutral, Sad, and Surprised. "
+        "Upload audio or use your microphone."
+    ),
+    allow_flagging="never",
+)
+
+# =========================
+# LAUNCH APP
+# =========================
+if __name__ == "__main__":
+    demo.launch()