added chat

app.py

@@ -1,59 +1,105 @@
 import os
-import torch
-from PIL import Image
+
+import gradio as gr
 import numpy as np
-from PIL import Image
+import torch
 from lavis.models import load_model_and_preprocess
-import gradio as gr
+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 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)
-    
+        caption = model.generate(
+            {"image": image}, use_nucleus_sampling=True, num_captions=3
+        )
+
+    caption = "\n".join(caption)
+
     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 Demo")
-    gr.Markdown("Using `OPT2.7B` - [Github](https://github.com/salesforce/LAVIS/tree/main/projects/blip2) - [Paper](https://arxiv.org/abs/2301.12597)")
-    
-    
+    gr.Markdown("# BLIP-2")
+    gr.Markdown(
+        "## Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models"
+    )
+    gr.Markdown(
+        "This demo uses `OPT2.7B` weights. For more information please see [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 Type", value="Beam Search")
+            caption_type = gr.Radio(
+                ["Beam Search", "Nucleus Sampling"],
+                label="Caption Type",
+                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")
-        
-        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()
+            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?", None, None]],
+        inputs=[input_image, caption_type, question_txt, chat_state, chatbot],
+    )
+
+demo.launch()