Spaces:
Sleeping
Sleeping
| 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("<p>", "").replace("</p>", "").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().style(color_map=("green", "pink")) | |
| 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() | |