import gradio as gr
from transformers import pipeline

# Load the models using pipeline
audio_model = pipeline("audio-classification", model="MelodyMachine/Deepfake-audio-detection-V2")
image_model = pipeline("image-classification", model="dima806/deepfake_vs_real_image_detection")

# Define the prediction function
def predict(data, model_choice):
    print("Data received:", data)  # Debugging statement
    try:
        if model_choice == "Audio Deepfake Detection":
            result = audio_model(data)
        elif model_choice == "Image Deepfake Detection":
            result = image_model(data)
        else:
            return {"error": "Invalid model choice"}
        
        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)}

# Define the interface based on the selected model
def update_interface(model_choice):
    if model_choice == "Audio Deepfake Detection":
        return gr.Audio(type="filepath"), gr.Label()
    elif model_choice == "Image Deepfake Detection":
        return gr.Image(type="filepath"), gr.Label()
    else:
        return None, None

# Create the Gradio interface
with gr.Blocks() as iface:
    model_choice = gr.Radio(choices=["Audio Deepfake Detection", "Image Deepfake Detection"], label="Select Model", value="Audio Deepfake Detection")
    input_component, output_component = update_interface(model_choice.value)
    
    def update_inputs(model_choice):
        input_component, output_component = update_interface(model_choice)
        input_placeholder.update(visible=False)
        output_placeholder.update(visible=False)
        input_placeholder.update(visible=True, component=input_component)
        output_placeholder.update(visible=True, component=output_component)
        
    input_placeholder = gr.Placeholder(gr.Component, visible=True)
    output_placeholder = gr.Placeholder(gr.Component, visible=True)
    
    model_choice.change(fn=update_inputs, inputs=model_choice, outputs=[input_placeholder, output_placeholder])
    
    submit_button = gr.Button("Submit")
    submit_button.click(fn=predict, inputs=[input_placeholder, model_choice], outputs=output_placeholder)
    
    iface.launch()