from transformers import ViTForImageClassification
from PIL import Image
import torch
import gradio as gr
from transformers import pipeline


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Loading in Model
model_name = "dima806/ai_vs_real_image_detection"
model = ViTForImageClassification.from_pretrained(model_name).to(device)
model.to(device)


#Classification function
def classify_image(img: Image.Image):
    inputs = model(images=img, return_tensors="pt").to(device)
    results = model(inputs)
    top = results[0]
    label = top["label"]
    score = top["score"]
    return f"Prediction: {label} (Confidence: {score:.2f})"



# Interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Real vs AI Image detection",
    description="Check if your image is Real or AI"
)