Update app.py

app.py

@@ -1,3 +1,27 @@
+from transformers import pipeline
+from custom_pipeline import CustomSpeechEnhancementPipeline
+from sgmse.model import ScoreModel
+from argparse import Namespace
 import gradio as gr
 
-gr.load("models/sp-uhh/speech-enhancement-sgmse").launch()
+# Define the arguments (as per your model configuration)
+args = Namespace(
+ device="cuda", # Use "cuda" if you have GPU support, otherwise "cpu"
+ corrector="ald", # Options: "ald", "langevin", "none"
+ corrector_steps=1, # Number of corrector steps
+ snr=0.5, # Signal-to-noise ratio for Langevin dynamics
+ N=30 # Number of reverse steps
+)
+
+# Load the speech enhancement model (provide the correct path to the model checkpoint)
+model = ScoreModel.load_from_checkpoint("path_to_your_model_checkpoint", map_location=args.device)
+
+# Create an instance of the custom pipeline
+enhancer = CustomSpeechEnhancementPipeline(model=model, target_sr=16000, pad_mode="zero_pad", args=args)
+
+# Define the Gradio interface using the custom pipeline
+def enhance_audio(audio):
+ return enhancer(audio)
+
+# Launch the Gradio interface
+gr.Interface(fn=enhance_audio, inputs="audio", outputs="audio").launch()