# First, make sure you've installed required packages
# !pip install -U gradio transformers torch torchvision

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

# Load the pipeline (auto-detects CUDA if available)
device = 0 if torch.cuda.is_available() else -1
pipe = pipeline("image-classification", model="prithivMLmods/Deep-Fake-Detector-v2-Model", device=device)

def classify_image(image=None, url=None):
    if image is None and not url:
        return "Skill issue: You gave me nothing to work with."

    try:
        if url:
            image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
        elif image:
            image = Image.fromarray(image).convert("RGB")
    except Exception as e:
        return f"Bro... that ain't an image: {str(e)}"

    result = pipe(image)
    return {entry["label"]: round(entry["score"], 3) for entry in result}

# Set up the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# 🔍 DeepFake Detector\nUpload an image or paste a URL. Let's see if you're being catfished.")

    with gr.Row():
        image_input = gr.Image(type="numpy", label="Upload Image")
        url_input = gr.Textbox(label="Or Enter Image URL")

    submit_btn = gr.Button("🚨 Detect")

    output = gr.Label(num_top_classes=2)

    submit_btn.click(fn=classify_image, inputs=[image_input, url_input], outputs=output)

# Launch the app
demo.launch()