Update app.py

app.py

@@ -11,7 +11,7 @@ model = AutoModelForCausalLM.from_pretrained("q-future/co-instruct-preview",
                                              attn_implementation="eager",
                                              device_map={"":"cuda:0"})
 
-def chat(message, history, image_1, image_2):
+def chat(message, history, image_1, image_2, image_3):
     print(history)
     if history:
         if image_1 is not None and image_2 is None:
@@ -21,18 +21,29 @@ def chat(message, history, image_1, image_2):
             message = past_message + "USER:" + message + " ASSISTANT:"
             images = [image_1]
         if image_1 is not None and image_2 is not None:
-            past_message = "USER: The first image: <|image|>\nThe second image: <|image|>" + history[0][0] + " ASSISTANT:" + history[0][1] + "</s>"
-            for i in range((len(history) - 1)):
-                past_message += "USER:" + history[i][0] + " ASSISTANT:" + history[i][1] + "</s>"
-            message = past_message + "USER:" + message + " ASSISTANT:"
-            images = [image_1, image_2]
+            if image_3 is None:
+                past_message = "USER: The first image: <|image|>\nThe second image: <|image|>" + history[0][0] + " ASSISTANT:" + history[0][1] + "</s>"
+                for i in range((len(history) - 1)):
+                    past_message += "USER:" + history[i][0] + " ASSISTANT:" + history[i][1] + "</s>"
+                message = past_message + "USER:" + message + " ASSISTANT:"
+                images = [image_1, image_2]
+            else:
+                past_message = "USER: The first image: <|image|>\nThe second image: <|image|>\nThe third image:<|image|>" + history[0][0] + " ASSISTANT:" + history[0][1] + "</s>"
+                for i in range((len(history) - 1)):
+                    past_message += "USER:" + history[i][0] + " ASSISTANT:" + history[i][1] + "</s>"
+                message = past_message + "USER:" + message + " ASSISTANT:"
+                images = [image_1, image_2, image_3]
     else:  
         if image_1 is not None and image_2 is None:
             message = "USER: The image: <|image|> " + message + " ASSISTANT:"
             images = [image_1]
         if image_1 is not None and image_2 is not None:
-            message = "USER: The first image: <|image|>\nThe second image: <|image|>" + message + " ASSISTANT:"
-            images = [image_1, image_2]
+            if image_3 is None:
+                message = "USER: The first image: <|image|>\nThe second image: <|image|>" + message + " ASSISTANT:"
+                images = [image_1, image_2]
+            else:
+                message = "USER: The first image: <|image|>\nThe second image: <|image|>\nThe third image:<|image|>" + message + " ASSISTANT:"
+                images = [image_1, image_2, image_3]
     
     print(message)
     
@@ -55,7 +66,8 @@ with gr.Blocks(title="img") as demo:
 """)
     gr.Markdown(title_markdown)
     with gr.Row():
-            input_img_1 = gr.Image(type='pil', label="Image 1 (The first image)")
-            input_img_2 = gr.Image(type='pil', label="Image 2 (The second image)")
-    gr.ChatInterface(fn = chat, additional_inputs=[input_img_1, input_img_2])
+            input_img_1 = gr.Image(type='pil', label="Image 1 (First image)")
+            input_img_2 = gr.Image(type='pil', label="Image 2 (Second image)")
+            input_img_2 = gr.Image(type='pil', label="Image 3 (Third image)")
+    gr.ChatInterface(fn = chat, additional_inputs=[input_img_1, input_img_2, input_img_3])
     demo.launch(share=True)