import os
import io
import IPython.display
from PIL import Image
import base64
from transformers import pipeline
import gradio as gr

import warnings
warnings.filterwarnings("ignore")

task = "image-to-text"
model = "nlpconnect/vit-gpt2-image-captioning"
image_captioner = pipeline("image-to-text", model = model)

def captioner(image):
    result = image_captioner(image)
    return result[0]['generated_text']

gr.close_all()
demo = gr.Interface(fn=captioner,
                    inputs=[gr.Image(label="Upload image", type="pil")],
                    outputs=[gr.Textbox(label="Caption")],
                    title="Image Captioning with {model}",
                    description=f"Caption any image using the {model} model",
                    allow_flagging="never",
                   )

demo.launch()