import gradio as gr
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import torch

MODEL_NAME = "prithivMLmods/Deep-Fake-Detector-v2-Model"

processor = ViTImageProcessor.from_pretrained(MODEL_NAME)
model = ViTForImageClassification.from_pretrained(MODEL_NAME)

def classify_image(image):
    image = Image.fromarray(image).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
    labels = model.config.id2label
    return { labels[i]: float(probs[i]) for i in range(len(probs)) }

demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Label(num_top_classes=2, label="Prediction (Real vs Deepfake)"),
    title="Deepfake Detector (ViT)",
    description="Upload an image — model classifies it as Real or Deepfake."
)

if __name__ == "__main__":
    demo.launch()