import gradio as gr
import torch
from model import pred_audio

# Define the audio input component
input_audio = gr.Audio(sources=["upload"], waveform_options=dict(waveform_color="#01C6FF"), type="filepath")
output1 = gr.Number(label="ShallowCNN", show_label=True)
output2 = gr.Number(label="SimpleLSTM", show_label=True)
output3 = gr.Number(label="DTDNN", show_label=True)
output4 = gr.Number(label="Ensemble", show_label=True)

# Create the Gradio interface
interface = gr.Interface(
    fn=pred_audio,
    inputs=input_audio,
    outputs=[output1, output2, output3, output4],
    title="Audio DeepFake Detection"
)

interface.launch()