Update app.py

app.py

@@ -2,6 +2,7 @@ import torch
 import torchaudio
 from sgmse.model import ScoreModel
 import gradio as gr
+from sgmse.util.other import pad_spec
 
 # Load the pre-trained model
 model = ScoreModel.load_from_checkpoint("pretrained_checkpoints/speech_enhancement/train_vb_29nqe0uh_epoch=115.ckpt")
@@ -11,12 +12,34 @@ def enhance_speech(audio_file):
     noisy, sr = torchaudio.load(audio_file)
     noisy = noisy.unsqueeze(0)  # Add fake batch dimension if needed
     
-    # Run the speech enhancement model
-    enhanced = model.predict(noisy)
+    if sr != target_sr:
+        y = torch.tensor(resample(y.numpy(), orig_sr=sr, target_sr=target_sr))
+
+    T_orig = y.size(1)   
+
+    # Normalize
+    norm_factor = y.abs().max()
+    y = y / norm_factor
+    
+    # Prepare DNN input
+    Y = torch.unsqueeze(model._forward_transform(model._stft(y.to(args.device))), 0)
+    Y = pad_spec(Y, mode=pad_mode)
+    
+    # Reverse sampling
+    sampler = model.get_pc_sampler(
+        'reverse_diffusion', args.corrector, Y.to(args.device), N=args.N, 
+        corrector_steps=args.corrector_steps, snr=args.snr)
+    sample, _ = sampler()
+    
+    # Backward transform in time domain
+    x_hat = model.to_audio(sample.squeeze(), T_orig)
+
+    # Renormalize
+    x_hat = x_hat * norm_factor
     
     # Save the enhanced audio
     output_file = 'enhanced_output.wav'
-    torchaudio.save(output_file, enhanced.cpu().squeeze(0), sr)
+    torchaudio.save(output_file, x_hat.cpu().numpy(), sr)
     
     return output_file