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()