import gradio as gr
from transformers import ViTFeatureExtractor, ViTForImageClassification
import torch
import requests

# Load the pre-trained ViT model and feature extractor
model = ViTForImageClassification.from_pretrained("umitkantar/ViTBasePatch16Art")
feature_extractor = ViTFeatureExtractor.from_pretrained("umitkantar/ViTBasePatch16Art")

# Define the prediction function
def ai_image_detector(image):
    # Preprocess the input image
    inputs = feature_extractor(images=image, return_tensors="pt")
    
    # Make a prediction using the loaded model
    outputs = model(**inputs)
    predicted_class_idx = torch.argmax(outputs.logits).item()
    
    # Calculate the confidence score (percentage)
    confidence_score = torch.softmax(outputs.logits, dim=1)[0][predicted_class_idx].item() * 100
    
    return {f"AI-generated Probability (%): {confidence_score:.2f}": confidence_score}

# Set up the Gradio interface
image_input = gr.inputs.Image()
gr.Interface(ai_image_detector, image_input, "label", capture_session=True).launch()