import gradio as gr
from transformers import pipeline

clf = pipeline("image-classification",
               model="google/vit-base-patch16-224",
               device_map="auto")

def predict(img):
    out = clf(img)[:5]
    return {o['label']: float(o['score']) for o in out}

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=5),
    title="Image Classification (ViT)"
)

if __name__ == "__main__":
    demo.launch(share=True)  # share=True gives you a temporary Colab link