Update app.py

app.py

@@ -1,14 +1,29 @@
 import gradio as gr
-from speechbrain.inference import SpectralMaskGAN
+import torch
+import torchaudio
+from speechbrain.inference.enhancement import SpectralMaskEnhancement
 
 # Load the MetricGAN model
-model = SpectralMaskGAN.from_hparams(source="speechbrain/metricgan-plus-voicebank", savedir="tmpdir_metricgan")
+enhance_model = SpectralMaskEnhancement.from_hparams(
+    source="speechbrain/metricgan-plus-voicebank",
+    savedir="tmpdir_metricgan",
+)
 
 # Define a function to enhance speech
 def enhance_speech(audio):
-    # Process the uploaded audio file through the model
-    enhanced_audio = model.enhance_file(audio)
-    return enhanced_audio
+    # Load the audio and add fake batch dimension
+    noisy = enhance_model.load_audio(audio).unsqueeze(0)
+    
+    # Add relative length tensor (assuming full length)
+    lengths = torch.tensor([1.])
+    
+    # Enhance the audio
+    enhanced = enhance_model.enhance_batch(noisy, lengths)
+    
+    # Save enhanced audio to a temporary file
+    output_path = "enhanced.wav"
+    torchaudio.save(output_path, enhanced.cpu(), 16000)
+    return output_path
 
 # Set up the Gradio interface
 iface = gr.Interface(