Spaces:
Runtime error
Runtime error
import os | |
import torch | |
from PIL import Image | |
import numpy as np | |
from PIL import Image | |
from lavis.models import load_model_and_preprocess | |
import gradio as gr | |
device = torch.device("cuda") if torch.cuda.is_available() else "cpu" | |
model, vis_processors, _ = load_model_and_preprocess( | |
name="blip2_opt", model_type="pretrain_opt2.7b", is_eval=True, device=device | |
) | |
def answer_question(image, prompt): | |
image = vis_processors["eval"](image).unsqueeze(0).to(device) | |
response = model.generate({"image": image, "prompt": f"Question: {prompt} Answer:"}) | |
response = '\n'.join(response) | |
return response | |
def generate_caption(image, caption_type): | |
image = vis_processors["eval"](image).unsqueeze(0).to(device) | |
if caption_type == "Beam Search": | |
caption = model.generate({"image": image}) | |
else: | |
caption = model.generate({"image": image}, use_nucleus_sampling=True, num_captions=3) | |
caption = '\n'.join(caption) | |
return caption | |
with gr.Blocks() as demo: | |
gr.Markdown("## BLIP-2 Demo") | |
gr.Markdown("Using `OPT2.7B` - [Github](https://github.com/salesforce/LAVIS/tree/main/projects/blip2) - [Paper](https://arxiv.org/abs/2301.12597)") | |
with gr.Row(): | |
with gr.Column(): | |
input_image = gr.Image(label="Image", type="pil") | |
caption_type = gr.Radio(["Beam Search", "Nucleus Sampling"], label="Caption Type", value="Beam Search") | |
btn_caption = gr.Button("Generate Caption") | |
question_txt = gr.Textbox(label="Question", lines=1) | |
btn_answer = gr.Button("Generate Answer") | |
with gr.Column(): | |
output_text = gr.Textbox(label="Answer", lines=5) | |
btn_caption.click(generate_caption, inputs=[input_image, caption_type], outputs=[output_text]) | |
btn_answer.click(answer_question, inputs=[input_image, question_txt], outputs=[output_text]) | |
gr.Examples([['./merlion.png', 'Beam Search', 'which city is this?']], inputs=[input_image, caption_type, question_txt]) | |
demo.launch() |