Update app.py

app.py

@@ -4,38 +4,50 @@ from sgmse.model import ScoreModel
 import gradio as gr
 from sgmse.util.other import pad_spec
 
+# Define parameters based on the argparse configuration in enhancement.py
+args = {
+    "test_dir": "./test_data",  # example directory, adjust as needed
+    "enhanced_dir": "./enhanced_data",  # example directory, adjust as needed
+    "ckpt": "https://huggingface.co/sp-uhh/speech-enhancement-sgmse/resolve/main/train_vb_29nqe0uh_epoch%3D115.ckpt",
+    "corrector": "ald",
+    "corrector_steps": 1,
+    "snr": 0.5,
+    "N": 30,
+    "device": "cuda" if torch.cuda.is_available() else "cpu"
+}
+
 # Load the pre-trained model
-model = ScoreModel.load_from_checkpoint("https://huggingface.co/sp-uhh/speech-enhancement-sgmse/resolve/main/train_vb_29nqe0uh_epoch%3D115.ckpt")
+model = ScoreModel.load_from_checkpoint(args["ckpt"])
 
 def enhance_speech(audio_file):
     # Load and process the audio file
     y, sr = torchaudio.load(audio_file)
-
-    T_orig = y.size(1)   
+    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)
-    
+    Y = torch.unsqueeze(model._forward_transform(model._stft(y.to(args["device"]))), 0)
+    Y = pad_spec(Y, mode="constant")  # Ensure pad_mode is defined; replace with actual pad_mode if needed
+
     # 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)
+        '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, x_hat.cpu().numpy(), sr)
+    torchaudio.save(output_file, x_hat.cpu(), sr)
     
     return output_file