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

# Initialize Model
get_completion = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base", device=0)

@spaces.GPU(duration=120)
def captioner(input: Image.Image) -> str:
    """
    Generate a caption for the given image using the BLIP-IMAGE-CAPTIONING-BASE model.

    Args:
        input (Image.Image): The input image for which to generate a caption.

    Returns:
        str: The generated caption text.
    """
    output = get_completion(input)
    return output[0]['generated_text']

####### GRADIO APP #######
title = """<h1 id="title"> Image Captioning with BLIP model </h1>"""

description = """
- The model used for Generating Captions [BLIP-IMAGE-CAPTIONING-BASE](https://huggingface.co./Salesforce/blip-image-captioning-base).
"""

css = '''
h1#title {
  text-align: center;
}
'''

theme = gr.themes.Soft()
demo = gr.Blocks(css=css, theme=theme)

with demo:
  gr.Markdown(title)
  gr.Markdown(description)
  interface = gr.Interface(fn=captioner,
                    inputs=[gr.Image(label="Upload image", type="pil")],
                    outputs=[gr.Textbox(label="Caption")],
                    allow_flagging="never")

demo.launch()