import os import gradio as gr import numpy as np import torch from lavis.models import load_model_and_preprocess from PIL import Image 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 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) if torch.cuda.is_available(): torch.cuda.empty_cache() return caption def chat(input_image, question, history): history = history or [] question = question.lower() image = vis_processors["eval"](input_image).unsqueeze(0).to(device) clean = lambda x: x.replace("
", "").replace("
", "").replace("\n", "") clean_h = lambda x: (clean(x[0]), clean(x[1])) context = list(map(clean_h, history)) template = "Question: {} Answer: {}." prompt = ( " ".join( [template.format(context[i][0], context[i][1]) for i in range(len(context))] ) + " Question: " + question + " Answer:" ) response = model.generate({"image": image, "prompt": prompt}) history.append((question, response[0])) return history, history def clear_chat(history): return [], [] with gr.Blocks() as demo: gr.Markdown( "### BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models" ) gr.Markdown( "This demo uses the `pretrain_opt2.7b` weights. For more information please visit [Github](https://github.com/salesforce/LAVIS/tree/main/projects/blip2) or [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 Decoding Strategy", value="Beam Search", ) btn_caption = gr.Button("Generate Caption") output_text = gr.Textbox(label="Answer", lines=5) with gr.Column(): chatbot = gr.Chatbot() chat_state = gr.State() question_txt = gr.Textbox(label="Question", lines=1) btn_answer = gr.Button("Generate Answer") btn_clear = gr.Button("Clear Chat") btn_caption.click( generate_caption, inputs=[input_image, caption_type], outputs=[output_text] ) btn_answer.click( chat, inputs=[input_image, question_txt, chat_state], outputs=[chatbot, chat_state], ) btn_clear.click(clear_chat, inputs=[chat_state], outputs=[chatbot, chat_state]) gr.Examples( [ ["./merlion.png", "Beam Search", "which city is this?"], [ "./Blue_Jay_0044_62759.jpg", "Beam Search", "what is the name of this bird?", ], ["./5kstbz-0001.png", "Beam Search", "where is the man standing?"], [ "ILSVRC2012_val_00000008.JPEG", "Beam Search", "Name the colors of macarons you see in the image.", ], ], inputs=[input_image, caption_type, question_txt], ) gr.Markdown( "Sample images are taken from [ImageNet](https://paperswithcode.com/sota/image-classification-on-imagenet), [CUB](https://paperswithcode.com/dataset/cub-200-2011) and [GamePhysics](https://asgaardlab.github.io/CLIPxGamePhysics/) datasets." ) demo.launch()