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

# Load the AI image detection model pipeline
detector = pipeline("image-classification", model="capcheck/ai-image-detection")

@spaces.GPU
def analyze_image(image: Image.Image):
    # Run prediction on the uploaded image
    results = detector(image)
    
    # Format the outputs into a dictionary format for Gradio
    return {res["label"]: float(res["score"]) for res in results}

demo = gr.Interface(
    fn=analyze_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2),
    title="IsRealOrNot Forensic Engine"
)

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