Update app.py to load model checkpoint

app.py

@@ -1,7 +1,34 @@
+import torch
 import gradio as gr
+from model import SGMSE  # Adjust according to your model's definition
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Load your model
+model_path = "https://huggingface.co/sp-uhh/speech-enhancement-sgmse/resolve/main/pretrained_checkpoints/speech_enhancement/train_vb_29nqe0uh_epoch%3D115.ckpt"
+model = SGMSE()  # Initialize your model class
+model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+model.eval()  # Set the model to evaluation mode
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+def enhance_audio(input_audio):
+    import torchaudio
+
+    # Load the input audio file
+    waveform, sample_rate = torchaudio.load(input_audio)
+
+    with torch.no_grad():
+        enhanced_waveform = model(waveform)
+
+    output_path = "enhanced_audio.wav"
+    torchaudio.save(output_path, enhanced_waveform.cpu(), sample_rate)
+    return output_path
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=enhance_audio,
+    inputs=gr.Audio(source="upload", type="filepath"),
+    outputs=gr.Audio(type="file"),
+    title="Speech Enhancement Model",
+    description="Upload a noisy audio file to enhance it using the SGMSE model."
+)
+
+if __name__ == "__main__":
+    iface.launch()