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

MODEL_ID = "enoch-alterego/anime-character-classifier"

# Load model once when the app starts
print("Loading model...")
proc  = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
model.eval()
print("Model loaded!")

def predict(image):
    if image is None:
        return {}

    # Convert and run prediction
    inputs = proc(images=image.convert("RGB"), return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits

    probs = torch.softmax(logits, dim=-1)[0]
    topk  = torch.topk(probs, k=5)

    # Return top 5 results
    return {
        model.config.id2label[idx.item()]: float(score.item())
        for score, idx in zip(topk.values, topk.indices)
    }

# Build the interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an anime character image"),
    outputs=gr.Label(num_top_classes=5, label="Which anime is this from?"),
    title="🎌 Guess the Anime Character",
    description="Upload any anime character image and the AI will guess which series they are from.",
)

demo.launch()