import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration

processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

def get_completion(raw_image):

  inputs = processor(raw_image, return_tensors="pt")
  out = model.generate(**inputs)

  return processor.decode(out[0], skip_special_tokens=True)

demo = gr.Interface(fn=get_completion,
                    inputs=[gr.Image(label="Upload image", type="pil")],
                    outputs=[gr.Textbox(label="Caption")],
                    title="Image Captioning with BLIP",
                    description="Caption any image using the BLIP model",
                    allow_flagging="never",
                    examples=["https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg", "https://free-images.com/sm/9596/dog_animal_greyhound_983023.jpg"])

demo.launch()