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

# -------------------------------
# Load model once (global)
# -------------------------------
model_id = "jacoballessio/ai-image-detect-distilled"

processor = ViTImageProcessor.from_pretrained(model_id)
model = AutoModelForImageClassification.from_pretrained(
    model_id,
    dtype=torch.float32,
    low_cpu_mem_usage=True
)

model.eval()
device = "cpu"
model.to(device)


# -------------------------------
# Prediction function
# -------------------------------
def predict(image: Image.Image):
    if image is None:
        return "Please upload an image", None

    # Preprocess
    inputs = processor(image, return_tensors="pt").to(device)

    # Inference
    with torch.no_grad():
        outputs = model(**inputs)

    # Probabilities
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    confidence = probs.max().item()
    predicted_label = model.config.id2label[probs.argmax().item()]

    # Convert to dict for Gradio Label
    labels = model.config.id2label
    scores = probs.squeeze().tolist()

    confidence_dict = {
        labels[i]: float(scores[i]) for i in range(len(scores))
    }

    # Result text
    if predicted_label.lower() == "fake":
        result = f"⚠️ AI-GENERATED\nConfidence: {confidence:.3f}"
    else:
        result = f"✅ REAL IMAGE\nConfidence: {confidence:.3f}"

    return result, confidence_dict


# -------------------------------
# UI
# -------------------------------
app = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Textbox(label="Prediction"),
        gr.Label(label="Confidence Scores")
    ],
    title="🖼️ AI vs Real Image Detector",
    description="Upload an image to check if it's AI-generated or real."
)


# -------------------------------
# Run app
# -------------------------------
if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)