Update app.py

app.py

@@ -5,8 +5,6 @@ import os
 import torch
 import librosa
 from transformers import AutoModel
-from huggingface_hub import hf_hub_download
-
 import net
 import utils
 
@@ -17,15 +15,8 @@ ser_model = None
 wav_mean = None
 wav_std = None
 
-# HF Model repo path (you uploaded the model files here)
-HF_MODEL_REPO = "Samara369/SER_1"
-MODEL_FILES = {
-    "ssl": "trained_models/final_ssl.pt",
-    "pool": "trained_models/final_pool.pt",
-    "ser": "trained_models/final_ser.pt",
-    "stat": "trained_models/train_norm_stat.pkl"
-}
-
+# Configuration
+MODEL_PATH = "trained_models"
 SSL_TYPE = utils.get_ssl_type("wavlm-large")
 POOLING_TYPE = "AttentiveStatisticsPooling"
 HEAD_DIM = 1024
@@ -35,36 +26,43 @@ EMOTION_NAMES = {
     'F': 'Fear', 'D': 'Disgust', 'C': 'Contempt', 'N': 'Neutral'
 }
 
+from huggingface_hub import hf_hub_download
+
 def load_models():
     global ssl_model, pool_model, ser_model, wav_mean, wav_std
 
     if ssl_model is None:
-        print("🔄 Downloading and loading models on CPU...")
-
-        # Download all model files using huggingface_hub
-        ssl_path = hf_hub_download(HF_MODEL_REPO, MODEL_FILES["ssl"])
-        pool_path = hf_hub_download(HF_MODEL_REPO, MODEL_FILES["pool"])
-        ser_path = hf_hub_download(HF_MODEL_REPO, MODEL_FILES["ser"])
-        stat_path = hf_hub_download(HF_MODEL_REPO, MODEL_FILES["stat"])
-
-        ssl_model = AutoModel.from_pretrained(SSL_TYPE)
-        ssl_model.freeze_feature_encoder()
-        ssl_model.load_state_dict(torch.load(ssl_path, map_location='cpu'))
+        print("Downloading and loading models from Hugging Face...")
+
+        # Paths to files in the repo
+        repo_id = "Samara369/SER_1"
+        ssl_path = hf_hub_download(repo_id=repo_id, filename="trained_models/final_ssl.pt")
+        pool_path = hf_hub_download(repo_id=repo_id, filename="trained_models/final_pool.pt")
+        ser_path = hf_hub_download(repo_id=repo_id, filename="trained_models/final_ser.pt")
+        norm_path = hf_hub_download(repo_id=repo_id, filename="trained_models/train_norm_stat.pkl")
+
+        # Load SSL model
+        ssl_model = AutoModel.from_pretrained("microsoft/wavlm-large")
+        ssl_model.load_state_dict(torch.load(ssl_path, map_location="cpu"))
         ssl_model.eval()
 
+        # Load pooling model
         feat_dim = ssl_model.config.hidden_size
         pool_net = getattr(net, POOLING_TYPE)
         pool_model = pool_net(feat_dim)
-        pool_model.load_state_dict(torch.load(pool_path, map_location='cpu'))
+        pool_model.load_state_dict(torch.load(pool_path, map_location="cpu"))
         pool_model.eval()
 
+        # Load regression head
         dh_input_dim = feat_dim * 2 if POOLING_TYPE == "AttentiveStatisticsPooling" else feat_dim
         ser_model = net.EmotionRegression(dh_input_dim, HEAD_DIM, 1, 8, dropout=0.5)
-        ser_model.load_state_dict(torch.load(ser_path, map_location='cpu'))
+        ser_model.load_state_dict(torch.load(ser_path, map_location="cpu"))
         ser_model.eval()
 
-        wav_mean, wav_std = utils.load_norm_stat(stat_path)
-        print("✅ Models loaded.")
+        # Load normalization stats
+        wav_mean, wav_std = utils.load_norm_stat(norm_path)
+
+        print("Models loaded from Hugging Face.")
 
 def process_single_audio(wav_path):
     wav, _ = librosa.load(wav_path, sr=16000)
@@ -89,7 +87,7 @@ def predict_emotion(audio_path):
         return predicted_emotion, confidence, logits
 
     except Exception as e:
-        print("❌ Error during inference:", e)
+        print("Error during inference:", e)
         return "Error", 0.0, None
 
 def process_audio_file(audio_file):