initial commit

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+merlion.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,59 @@
+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()

requirements.txt

@@ -0,0 +1,3 @@
+torch
+torchvision
+salesforce-lavis