import gradio as gr
from transformers import pipeline
from optimum.pipelines import pipeline as onnx_pipeline

# Load Models
clean_pipe = pipeline("image-classification", model="WinKawaks/vit-small-patch16-224")
mal_pipe = onnx_pipeline("image-classification", model="onnx/model.onnx", accelerator="ort")

# Interface Functions

def classify_image(model_type, image):
    if model_type == "Clean Model":
        return clean_pipe(image)
    elif model_type == "Malicious Model":
        return mal_pipe(image)
    else:
        return "Invalid model type"

# Gradio Interface
inputs = [
    gr.inputs.Radio(choices=["Clean Model", "Malicious Model"], label="Select Model"),
    gr.inputs.Image(type="filepath", label="Upload Image")
]

outputs = gr.outputs.Label(num_top_classes=1, label="Classification Result")

app = gr.Interface(fn=classify_image, inputs=inputs, outputs=outputs, title="Model Comparison: Clean vs Malicious", description="Compare the behavior of a clean model and a potentially malicious model using the same image input.")

app.launch()