import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification

# Load the model and processor
print("Loading model...")
processor = AutoImageProcessor.from_pretrained("Organika/sdxl-detector")
model = AutoModelForImageClassification.from_pretrained("Organika/sdxl-detector")
print("Model loaded successfully!")

def detect_ai(image):
    """
    Detect if an image is AI-generated or real.
    
    Args:
        image: PIL Image object
        
    Returns:
        dict: Probabilities for each class (AI-generated vs Real)
    """
    if image is None:
        return {}
    
    try:
        # Direct model inference
        inputs = processor(images=image, return_tensors="pt")
        outputs = model(**inputs)
        logits = outputs.logits
        probs = logits.softmax(dim=-1)[0].tolist()
        id2label = model.config.id2label
        
        # Create result dictionary
        result = {
            id2label[0]: probs[0],
            id2label[1]: probs[1],
        }
        
        return result
    
    except Exception as e:
        print(f"Error processing image: {e}")
        return {"Error": "Failed to process image"}

# Create the Gradio interface
demo = gr.Interface(
    fn=detect_ai,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Label(num_top_classes=2, label="AI vs Real Probability"),
    title="🤖 AI‑Generated Image Detector",
    description="""
    Upload an image to detect whether it's AI-generated or real.
    
    This model can help identify images generated by AI systems like DALL-E, Midjourney, Stable Diffusion, and others.
    
    **How to use:**
    1. Upload an image (JPG, PNG, etc.)
    2. The model will analyze it and return probabilities
    3. Higher probability for "AI-generated" suggests the image was created by AI
    """,
    article="""
    ### About the Model
    The model has been trained to detect various AI-generated images 
    with a focus on SDXL and similar diffusion models.
    
    ### Limitations
    - The model may not be 100% accurate on all images
    - Performance may vary depending on the AI model used to generate the image
    - Very high-quality AI images or heavily post-processed real images might be misclassified
    """,
    examples=[
        # You can add example images here if you have them
    ],
    cache_examples=False,
)

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