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

# Load model
captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")

# Inference function
def generate_caption(image):
    result = captioner(image)
    return result[0]["generated_text"]

# Gradio UI
iface = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="🖼️ Image Caption Generator",
    description="Upload an image and the model will describe it in a sentence.",
)

iface.launch()