Exception handling

app.py

@@ -39,39 +39,52 @@ def load_model():
             logger.info("Model loaded successfully...")
         except Exception as e:
             logger.error(f"Model loading failed: {e}")
-
-            raise gr.Error(
-                "Failed to load model.........."
-            )
+            raise gr.Error( "Failed to load model...")
     return feature_extractor, model
 
 
-
 def predict_audio(audio_path):
     logger.info(f"inside predict_audio : {audio_path}")
     
     feature_extractor, model = load_model()
     id2label = model.config.id2label
 
-    waveform, sr = librosa.load(audio_path, sr=AST_SR, mono=True)
-
-    waveform = torch.tensor(waveform)
-    max_val = waveform.abs().max()
-    if max_val > 0:
-        waveform = waveform / max_val
-
-    inputs = feature_extractor(
-        waveform.numpy(),
-        sampling_rate=sr,
-        return_tensors="pt"
-    )
-
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    with torch.no_grad():
-        logits = model(**inputs).logits.squeeze(0)
-
-    probs = torch.softmax(logits, dim=0).cpu().numpy()
-
+    try:
+        waveform, sr = librosa.load(audio_path, sr=AST_SR, mono=True)
+    except Exception as e:
+        logger.error(f"Audio loading failed : {e}")
+        raise.gr.Error("Failed to load the audio ...")
+
+    try:
+        waveform = torch.tensor(waveform)
+        max_val = waveform.abs().max()
+        if max_val > 0:
+            waveform = waveform / max_val
+    except Exception as e:
+        logger.error(f"Preprocessing failed: {e}")
+        raise.gr.Error("Preprocessing failed ...")
+
+    try:
+        inputs = feature_extractor(
+            waveform.numpy(),
+            sampling_rate=sr,
+            return_tensors="pt"
+        )
+    
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+    except Exception as e:
+        logger.error(f"Feature extraction failed : {e}")
+        raise.gr.Error("Feature extraction failed ...")
+
+    try:
+        with torch.no_grad():
+            logits = model(**inputs).logits.squeeze(0)
+    
+        probs = torch.softmax(logits, dim=0).cpu().numpy()
+    except Exception as e:
+        logger.error(f"Inference failed : {e}")
+        raise.gr.Error("Inference failed ...")
+        
     return waveform.numpy(), probs, id2label
 
 
@@ -99,7 +112,7 @@ with gr.Blocks(title="AST Model") as demo:
         img = librosa.display.specshow(
             mel_db, sr = AST_SR, 
             x_axis = 'time', y_axis = 'mel',
-            ax = ax[1]
+            cmap = 'viridis', ax = ax[1]
         )
 
         ax[1].set_title("Mel Spectrogram")