Update app.py
Browse files
app.py
CHANGED
@@ -1,3 +1,27 @@
|
|
|
|
|
|
|
|
|
|
1 |
import gradio as gr
|
2 |
|
3 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from transformers import pipeline
|
2 |
+
from custom_pipeline import CustomSpeechEnhancementPipeline
|
3 |
+
from sgmse.model import ScoreModel
|
4 |
+
from argparse import Namespace
|
5 |
import gradio as gr
|
6 |
|
7 |
+
# Define the arguments (as per your model configuration)
|
8 |
+
args = Namespace(
|
9 |
+
device="cuda", # Use "cuda" if you have GPU support, otherwise "cpu"
|
10 |
+
corrector="ald", # Options: "ald", "langevin", "none"
|
11 |
+
corrector_steps=1, # Number of corrector steps
|
12 |
+
snr=0.5, # Signal-to-noise ratio for Langevin dynamics
|
13 |
+
N=30 # Number of reverse steps
|
14 |
+
)
|
15 |
+
|
16 |
+
# Load the speech enhancement model (provide the correct path to the model checkpoint)
|
17 |
+
model = ScoreModel.load_from_checkpoint("path_to_your_model_checkpoint", map_location=args.device)
|
18 |
+
|
19 |
+
# Create an instance of the custom pipeline
|
20 |
+
enhancer = CustomSpeechEnhancementPipeline(model=model, target_sr=16000, pad_mode="zero_pad", args=args)
|
21 |
+
|
22 |
+
# Define the Gradio interface using the custom pipeline
|
23 |
+
def enhance_audio(audio):
|
24 |
+
return enhancer(audio)
|
25 |
+
|
26 |
+
# Launch the Gradio interface
|
27 |
+
gr.Interface(fn=enhance_audio, inputs="audio", outputs="audio").launch()
|