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, image_3):
+def chat(message, history, image_1, image_2, image_3, image_4):
     print(history)
     if history:
         if image_1 is not None and image_2 is None:
@@ -28,11 +28,18 @@ def chat(message, history, image_1, image_2, image_3):
                 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]
+                if image_4 is None:
+                    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:
+                    past_message = "USER: The first image: <|image|>\nThe second image: <|image|>\nThe third image:<|image|>\nThe fourth 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, image_4]
     else:  
         if image_1 is not None and image_2 is None:
             message = "USER: The image: <|image|> " + message + " ASSISTANT:"
@@ -42,17 +49,21 @@ def chat(message, history, image_1, image_2, image_3):
                 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]
+                if image_4 is None:
+                    message = "USER: The first image: <|image|>\nThe second image: <|image|>\nThe third image:<|image|>" + message + " ASSISTANT:"
+                    images = [image_1, image_2, image_3]
+                else:
+                    message = "USER: The first image: <|image|>\nThe second image: <|image|>\nThe third image:<|image|>\nThe fourth image:<|image|>" + message + " ASSISTANT:"
+                    images = [image_1, image_2, image_3, image_4]
     
     print(message)
     
-    return model.tokenizer.batch_decode(model.chat(message, images, max_new_tokens=150).clamp(0, 100000))[0].split("ASSISTANT:")[-1]
+    return model.tokenizer.batch_decode(model.chat(message, images, max_new_tokens=300).clamp(0, 100000))[0].split("ASSISTANT:")[-1]
 
 
 with gr.Blocks(title="img") as demo:
     title_markdown = ("""
-    <div align="center">*Preview Version (v1)! Now we support two images as inputs! Try it now!*</div>
+    <h3 align="center">*Super Version of Q-Instruct with Multi-image (up to 4, same as GPT-4V) Support!*</h3>
 <h1 align="center"><a href="https://github.com/Q-Future/Q-Instruct"><img src="https://github.com/Q-Future/Q-Instruct/blob/main/q_instruct_logo.png?raw=true", alt="Q-Instruct (mPLUG-Owl-2)" border="0" style="margin: 0 auto; height: 85px;" /></a> </h1>
 <h2 align="center">Q-Instruct: Improving Low-level Visual Abilities for Multi-modality Foundation Models</h2>
 <h5 align="center"> Please find our more accurate visual scoring demo on <a href='https://huggingface.co/spaces/teowu/OneScorer'>[OneScorer]</a>!</h2>
@@ -69,5 +80,6 @@ with gr.Blocks(title="img") as demo:
             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_3 = gr.Image(type='pil', label="Image 3 (Third image)")
-    gr.ChatInterface(fn = chat, additional_inputs=[input_img_1, input_img_2, input_img_3])
+            input_img_4 = gr.Image(type='pil', label="Image 4 (Third image)")
+    gr.ChatInterface(fn = chat, additional_inputs=[input_img_1, input_img_2, input_img_3, input_img_4])
     demo.launch(share=True)