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


device = 0 if torch.cuda.is_available() else -1

# Loading in Model
model_name = "dima806/ai_vs_real_image_detection"
pipe = pipeline("image-classification", model=model_name, device = device)
    


# Classification function
def classify_image(img: Image.Image):
    results = pipe(img)
    top = results[0]
    label = top["label"]
    score = top["score"]
    return f"Prediction: {label} (Confidence: {score:.2f})"

# Gradio interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Real vs AI Image Detection",
    description="Upload an image to see if it's REAL or AI-generated."
)

interface.launch()