import gradio as gr
import torch
import torchaudio
from df import enhance, init_df

# Initialize DeepFilterNet model
model, df_state, _ = init_df()

def denoise_audio(audio):
    # Load the input audio file
    waveform, sample_rate = torchaudio.load(audio)
    
    # Denoise the audio
    enhanced_audio = enhance(model, df_state, waveform)
    
    # Save and return the enhanced audio file
    output_file = "enhanced_output.wav"
    torchaudio.save(output_file, enhanced_audio, sample_rate)
    
    return output_file

# Gradio interface
iface = gr.Interface(
    fn=denoise_audio, 
    inputs=gr.Audio(type="filepath"),  # Remove 'source' argument
    outputs="file",
    title="DeepFilterNet Audio Denoising",
    description="Upload an audio file to remove noise using DeepFilterNet."
)

iface.launch()