import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

model_name = "prithivMLmods/open-deepfake-detection"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
model.eval()

def predict(img):
    inputs = processor(images=img, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=1).squeeze()
        is_ai = bool(torch.argmax(probs).item())
        confidence = float(probs[1].item()) if is_ai else float(probs[0].item())
    message = "AI-generated image detected." if is_ai else "Image appears original/authentic."
    return {"is_ai_generated": is_ai, "confidence": confidence, "message": message}

iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs="json", allow_flagging="never")
iface.launch()