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

def wav_to_mp3(wav_path, mp3_path):
    # Load the WAV file
    audio = AudioSegment.from_wav(wav_path)
    # Export as MP3
    audio.export(mp3_path, format="mp3")
    
# 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)
    wav_to_mp3(output_file,"enhanced.mp3")
    output_file="enhanced.mp3"
    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()