Update app.py

app.py

@@ -14,8 +14,16 @@ import gradio as gr
 from scipy.io import wavfile
 from scipy.io.wavfile import write
 
+from speechbrain.pretrained import SpectralMaskEnhancement
+
+enhance_model = SpectralMaskEnhancement.from_hparams(
+    source="speechbrain/metricgan-plus-voicebank",
+    savedir="pretrained_models/metricgan-plus-voicebank",
+)
+
 knn_vc = torch.hub.load('bshall/knn-vc', 'knn_vc', prematched=True, trust_repo=True, pretrained=True, device='cpu')
 
+
 def voice_change(audio_in, audio_ref):
     samplerate1, data1 = wavfile.read(audio_in)
     samplerate2, data2 = wavfile.read(audio_ref)
@@ -53,7 +61,14 @@ def tts(text, model, voice, api_key):
 
     # Get the file path of the temp file
     temp_file_path = temp_file.name
-
+    
+    noisy = enhance_model.load_audio(
+        temp_file_path
+    ).unsqueeze(0)
+    enhanced = enhance_model.enhance_batch(noisy, lengths=torch.tensor([1.]))
+    torchaudio.save('enhanced.wav', enhanced.cpu(), 16000)
+    return 'enhanced.wav'
+    
     return temp_file_path