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

# Load pre-trained model and processor
model_name = "facebook/deit-base-distilled-patch16-224"
processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

def detect_deepfake(image):
    # Preprocess the image
    inputs = processor(images=image, return_tensors="pt")
    
    # Make prediction
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    
    # For demonstration, we'll assume class 0 is real and 1 is fake
    # (In a real project, you'd need to verify this with your model)
    return "Real" if predicted_class_idx == 0 else "Fake (Possible Deepfake)"

# Create a simple interface
iface = gr.Interface(
    fn=detect_deepfake,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Deepfake Detection",
    description="Upload an image to check if it might be a deepfake."
)

iface.launch()