Update app.py

app.py

@@ -1,7 +1,72 @@
+import torch
+import torchaudio
 import gradio as gr
+from sgmse.model import ScoreModel
+from sgmse.util.other import pad_spec
+import time  # Import the time module
+import os
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Define parameters based on the 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"
+}
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# Ensure the model is loaded to the correct device
+model = ScoreModel.load_from_checkpoint(args["ckpt"]).to(args["device"])
+
+def enhance_speech(audio_file):
+    start_time = time.time()  # Start the timer
+
+    # Load and process the audio file
+    y, sr = torchaudio.load(audio_file)  # Gradio passes the file path
+    print(f"Loaded audio in {time.time() - start_time:.2f}s")
+    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)
+    print(f"Transformed input in {time.time() - start_time:.2f}s")
+
+    Y = pad_spec(Y, mode="zero_pad")  # Use "zero_pad" mode for padding
+
+    # 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
+
+    # Create a temporary path for saving the enhanced audio in Hugging Face Space
+    output_file = "/tmp/enhanced_output.wav"  # Use a temporary directory
+    torchaudio.save(output_file, x_hat.cpu(), sr)
+
+    print(f"Processed audio in {time.time() - start_time:.2f}s")
+    
+    # Return the path to the enhanced file for Gradio to handle
+    return output_file
+
+# Gradio interface setup
+inputs = gr.Audio(label="Input Audio", type="filepath")  # Adjusted to 'filepath'
+outputs = gr.Audio(label="Enhanced Audio", type="filepath")  # Output as filepath
+title = "Speech Enhancement using SGMSE"
+description = "This Gradio demo uses the SGMSE model for speech enhancement. Upload your audio file to enhance it."
+article = "<p style='text-align: center'><a href='https://huggingface.co/SP-UHH/speech-enhancement-sgmse' target='_blank'>Model Card</a></p>"
+
+# Launch the Gradio interface
+gr.Interface(fn=enhance_speech, inputs=inputs, outputs=outputs, title=title, description=description, article=article).launch()