File size: 1,090 Bytes
573fd5b dcea6f9 573fd5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
from transformers import pipeline
from custom_pipeline import CustomSpeechEnhancementPipeline
from sgmse.model import ScoreModel
from argparse import Namespace
import gradio as gr
# 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()
|