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

# Load the Hugging Face model
MODEL_ID = "jacoballessio/ai-image-detect-distilled"

pipe = pipeline("image-classification", model=MODEL_ID)

# Prediction function
def predict_image(image):
    try:
        results = pipe(image)
        
        # Extract top two predictions
        top_results = sorted(results, key=lambda x: x['score'], reverse=True)[:2]
        labels = [r['label'] for r in top_results]
        scores = [round(r['score'] * 100, 2) for r in top_results]

        # Determine final result
        final_label = labels[0]
        final_score = scores[0]

        # Display result
        return {
            "Model Results": top_results,
            "Final AI Probability (%)": final_score,
            "Overall Decision": "AI-generated" if "ai" in final_label.lower() else "Human"
        }
    except Exception as e:
        return {"error": str(e)}

# Gradio Interface
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs="json",
    title="AI Image Detector",
    description="Detect whether an image is AI-generated or real using Jacob Allessio's distilled model."
)

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