import gradio as gr
from transformers import pipeline

# Load the model using pipeline
pipe = pipeline("audio-classification", model="MelodyMachine/Deepfake-audio-detection-V2")

# Define the prediction function
def predict(audio):
    print("Audio file received:", audio)  # Debugging statement
    try:
        result = pipe(audio)
        print("Raw prediction result:", result)  # Debugging statement
        # Convert the result to the expected format
        output = {item['label']: item['score'] for item in result}
        print("Formatted prediction result:", output)  # Debugging statement
        return output
    except Exception as e:
        print("Error during prediction:", e)  # Debugging statement
        return {"error": str(e)}

# Create the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Audio(type="filepath"),
    outputs=gr.Label(),
    title="Testing Deepfake Audio Detection Simple Interface",
    description="Upload an audio file or record your voice to detect if the audio is a deepfake."
)

# Launch the interface
iface.launch()