import gradio as gr
from df.enhance import enhance, init_df, load_audio, save_audio
import torch
import tempfile

# Model ko initialize karein
model, df_state, _ = init_df()

def clean_audio(audio_path):
    if audio_path is None:
        return None
    try:
        # Audio load karein
        audio, _ = load_audio(audio_path, sr=df_state.sr())
        # Enhance (Clean) karein
        enhanced_audio = enhance(model, df_state, audio)
        # Save karein
        output_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
        save_audio(output_file, enhanced_audio, df_state.sr())
        return output_file
    except Exception as e:
        return None

# Interface
iface = gr.Interface(
    fn=clean_audio,
    inputs=gr.Audio(type="filepath", label="Upload Audio"),
    outputs=gr.Audio(type="filepath", label="Clean Audio"),
    title="Voice Cleaner"
)

iface.launch()