Misc changes for vision arena: change log, fix image moderation (lm-sys#3341)

Author: BabyChouSr

fastchat/conversation.py

@@ -69,6 +69,8 @@ class Conversation:
     stop_str: Union[str, List[str]] = None
     # Stops generation if meeting any token in this list
     stop_token_ids: List[int] = None
+    # The maximum image size in megabytes that this model takes in. None means we do not resize the image.
+    max_image_size_mb: int = None
 
     def get_prompt(self) -> str:
         """Get the prompt for generation."""
@@ -351,10 +353,11 @@ def update_last_message(self, message: str):
         """
         self.messages[-1][1] = message
 
-    def convert_image_to_base64(self, image, resize_image=False):
+    def convert_image_to_base64(self, image):
         """Given an image, return the base64 encoded image string."""
         from PIL import Image
         import requests
+        from fastchat.utils import resize_image_and_return_image_in_bytes
 
         # Load image if it has not been loaded in yet
         if type(image) == str:
@@ -367,23 +370,10 @@ def convert_image_to_base64(self, image, resize_image=False):
             else:
                 image = Image.open(image).convert("RGB")
 
-        if resize_image:
-            max_hw, min_hw = max(image.size), min(image.size)
-            aspect_ratio = max_hw / min_hw
-            max_len, min_len = 2048, 2048
-            shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
-            longest_edge = int(shortest_edge * aspect_ratio)
-            W, H = image.size
-            if longest_edge != max(image.size):
-                if H > W:
-                    H, W = longest_edge, shortest_edge
-                else:
-                    H, W = shortest_edge, longest_edge
-                image = image.resize((W, H))
-
-        buffered = BytesIO()
-        image.save(buffered, format="PNG")
-        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+        image_bytes = resize_image_and_return_image_in_bytes(
+            image, self.max_image_size_mb
+        )
+        img_b64_str = base64.b64encode(image_bytes.getvalue()).decode()
 
         return img_b64_str
 
@@ -400,7 +390,7 @@ def to_gradio_chatbot(self):
                     ):
                         img_str = f'<img src="{img_b64_str}" alt="user upload image" />'
                     else:
-                        img_str = f'<img src="data:image/jpeg;base64,{img_b64_str}" alt="user upload image" />'
+                        img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
                     msg = img_str + msg.replace("<image>\n", "").strip()
 
                 ret.append([msg, None])
@@ -429,7 +419,7 @@ def to_openai_image_format(self, image_urls):
                         base64.b64encode(base64.b64decode(image_url))
                         == image_url.encode()
                     ), "The image data is not a valid base64 encoded string"
-                    openai_images.append(f"data:image/jpeg;base64,{image_url}")
+                    openai_images.append(f"data:image/png;base64,{image_url}")
                 except:
                     raise ValueError(
                         f"This file is not valid or not currently supported by the OpenAI API: {image_url}"
@@ -438,12 +428,16 @@ def to_openai_image_format(self, image_urls):
 
     def to_openai_vision_api_messages(self):
         """Convert the conversation to OpenAI vision api completion format"""
-        ret = [
-            {
-                "role": "system",
-                "content": [{"type": "text", "text": self.system_message}],
-            }
-        ]
+        if self.system_message == "":
+            ret = []
+        else:
+            ret = [
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": self.system_message}],
+                }
+            ]
+
         for i, (_, msg) in enumerate(self.messages[self.offset :]):
             if i % 2 == 0:
                 if type(msg) is tuple:
@@ -598,7 +592,7 @@ def to_reka_api_messages(self):
                                 {
                                     "type": "human",
                                     "text": text,
-                                    "media_url": f"data:image/jpeg;base64,{image}",
+                                    "media_url": f"data:image/png;base64,{image}",
                                 }
                             )
                 else:
@@ -680,6 +674,7 @@ def copy(self):
             sep2=self.sep2,
             stop_str=self.stop_str,
             stop_token_ids=self.stop_token_ids,
+            max_image_size_mb=self.max_image_size_mb,
         )
 
     def dict(self):
@@ -1078,6 +1073,7 @@ def get_conv_template(name: str) -> Conversation:
         roles=("user", "assistant"),
         sep_style=SeparatorStyle.DEFAULT,
         sep=None,
+        max_image_size_mb=None,  # OpenAI does auto-resizing
     )
 )
 
@@ -1115,6 +1111,7 @@ def get_conv_template(name: str) -> Conversation:
         roles=("Human", "Assistant"),
         sep_style=SeparatorStyle.ADD_COLON_SINGLE,
         sep="\n\n",
+        max_image_size_mb=5 / 1.35,
     )
 )
 
@@ -1138,6 +1135,7 @@ def get_conv_template(name: str) -> Conversation:
         roles=("user", "assistant"),
         sep_style=SeparatorStyle.DEFAULT,
         sep=None,
+        max_image_size_mb=5 / 1.35,
     )
 )
 
@@ -1161,6 +1159,7 @@ def get_conv_template(name: str) -> Conversation:
         roles=("user", "assistant"),
         sep_style=SeparatorStyle.DEFAULT,
         sep=None,
+        max_image_size_mb=5 / 1.35,
     )
 )
 
@@ -1193,6 +1192,7 @@ def get_conv_template(name: str) -> Conversation:
         roles=("user", "assistant"),
         sep_style=SeparatorStyle.DEFAULT,
         sep=None,
+        max_image_size_mb=5 / 1.35,
     )
 )
 
@@ -1288,6 +1288,7 @@ def get_conv_template(name: str) -> Conversation:
         roles=("user", "model"),
         sep_style=SeparatorStyle.DEFAULT,
         sep=None,
+        max_image_size_mb=20,
     )
 )
 

fastchat/serve/gradio_block_arena_vision.py

@@ -40,7 +40,7 @@
     image_moderation_filter,
 )
 
-logger = build_logger("gradio_web_server_multi", "gradio_web_server_multi.log")
+logger = build_logger("gradio_web_server", "gradio_web_server.log")
 
 no_change_btn = gr.Button()
 enable_btn = gr.Button(interactive=True, visible=True)
@@ -82,7 +82,7 @@ def add_image(textbox):
 
 
 def vote_last_response(state, vote_type, model_selector, request: gr.Request):
-    filename = get_conv_log_filename(state.is_vision)
+    filename = get_conv_log_filename(state.is_vision, state.has_csam_image)
     with open(filename, "a") as fout:
         data = {
             "tstamp": round(time.time(), 4),
@@ -158,7 +158,7 @@ def moderate_input(text, all_conv_text, model_list, images, ip):
         elif text_flagged and image_flagged:
             text = MODERATION_MSG
 
-    return text, csam_flagged
+    return text, image_flagged, csam_flagged
 
 
 def add_text(state, model_selector, chat_input, request: gr.Request):
@@ -176,10 +176,17 @@ def add_text(state, model_selector, chat_input, request: gr.Request):
     all_conv_text = state.conv.get_prompt()
     all_conv_text = all_conv_text[-2000:] + "\nuser: " + text
 
-    text, csam_flag = moderate_input(
+    text, image_flagged, csam_flag = moderate_input(
         text, all_conv_text, [state.model_name], images, ip
     )
 
+    if image_flagged:
+        logger.info(f"image flagged. ip: {ip}. text: {text}")
+        state.skip_next = True
+        return (state, state.to_gradio_chatbot(), {"text": IMAGE_MODERATION_MSG}) + (
+            no_change_btn,
+        ) * 5
+
     if (len(state.conv.messages) - state.conv.offset) // 2 >= CONVERSATION_TURN_LIMIT:
         logger.info(f"conversation turn limit. ip: {ip}. text: {text}")
         state.skip_next = True

fastchat/serve/gradio_block_arena_vision_anony.py

@@ -168,7 +168,7 @@ def clear_history_example(request: gr.Request):
 
 
 def vote_last_response(states, vote_type, model_selectors, request: gr.Request):
-    filename = get_conv_log_filename(states[0].is_vision)
+    filename = get_conv_log_filename(states[0].is_vision, states[0].has_csam_image)
 
     with open(filename, "a") as fout:
         data = {
@@ -309,7 +309,7 @@ def add_text(
         )
 
     model_list = [states[i].model_name for i in range(num_sides)]
-    text, csam_flag = moderate_input(text, text, model_list, images, ip)
+    text, image_flagged, csam_flag = moderate_input(text, text, model_list, images, ip)
 
     conv = states[0].conv
     if (len(conv.messages) - conv.offset) // 2 >= CONVERSATION_TURN_LIMIT:
@@ -327,6 +327,21 @@ def add_text(
             + [""]
         )
 
+    if image_flagged:
+        logger.info(f"image flagged. ip: {ip}. text: {text}")
+        for i in range(num_sides):
+            states[i].skip_next = True
+        return (
+            states
+            + [x.to_gradio_chatbot() for x in states]
+            + [{"text": IMAGE_MODERATION_MSG}]
+            + [
+                no_change_btn,
+            ]
+            * 6
+            + [""]
+        )
+
     text = text[:INPUT_CHAR_LEN_LIMIT]  # Hard cut-off
     for i in range(num_sides):
         post_processed_text = _prepare_text_with_image(

fastchat/serve/gradio_block_arena_vision_named.py

@@ -53,9 +53,7 @@
 )
 
 
-logger = build_logger(
-    "gradio_web_server_vision_multi", "gradio_web_server_vision_multi.log"
-)
+logger = build_logger("gradio_web_server_multi", "gradio_web_server_multi.log")
 
 num_sides = 2
 enable_moderation = False
@@ -72,7 +70,7 @@ def clear_history_example(request: gr.Request):
 
 
 def vote_last_response(states, vote_type, model_selectors, request: gr.Request):
-    filename = get_conv_log_filename(states[0].is_vision)
+    filename = get_conv_log_filename(states[0].is_vision, states[0].has_csam_image)
     with open(filename, "a") as fout:
         data = {
             "tstamp": round(time.time(), 4),
@@ -189,7 +187,9 @@ def add_text(
         all_conv_text_left[-1000:] + all_conv_text_right[-1000:] + "\nuser: " + text
     )
 
-    text, csam_flag = moderate_input(text, all_conv_text, model_list, images, ip)
+    text, image_flagged, csam_flag = moderate_input(
+        text, all_conv_text, model_list, images, ip
+    )
 
     conv = states[0].conv
     if (len(conv.messages) - conv.offset) // 2 >= CONVERSATION_TURN_LIMIT:
@@ -206,6 +206,20 @@ def add_text(
             * 6
         )
 
+    if image_flagged:
+        logger.info(f"image flagged. ip: {ip}. text: {text}")
+        for i in range(num_sides):
+            states[i].skip_next = True
+        return (
+            states
+            + [x.to_gradio_chatbot() for x in states]
+            + [{"text": IMAGE_MODERATION_MSG}]
+            + [
+                no_change_btn,
+            ]
+            * 6
+        )
+
     text = text[:INPUT_CHAR_LEN_LIMIT]  # Hard cut-off
     for i in range(num_sides):
         post_processed_text = _prepare_text_with_image(

fastchat/serve/gradio_web_server.py

@@ -131,9 +131,11 @@ def dict(self):
             {
                 "conv_id": self.conv_id,
                 "model_name": self.model_name,
-                "has_csam_image": self.has_csam_image,
             }
         )
+
+        if self.is_vision:
+            base.update({"has_csam_image": self.has_csam_image})
         return base
 
 
@@ -144,11 +146,13 @@ def set_global_vars(controller_url_, enable_moderation_, use_remote_storage_):
     use_remote_storage = use_remote_storage_
 
 
-def get_conv_log_filename(is_vision=False):
+def get_conv_log_filename(is_vision=False, has_csam_image=False):
     t = datetime.datetime.now()
     conv_log_filename = f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json"
-    if is_vision:
+    if is_vision and not has_csam_image:
         name = os.path.join(LOGDIR, f"vision-tmp-{conv_log_filename}")
+    elif is_vision and has_csam_image:
+        name = os.path.join(LOGDIR, f"vision-csam-{conv_log_filename}")
     else:
         name = os.path.join(LOGDIR, conv_log_filename)
 
@@ -306,10 +310,8 @@ def _prepare_text_with_image(state, text, images, csam_flag):
             # reset convo with new image
             state.conv = get_conversation_template(state.model_name)
 
-        resize_image = "llava" in state.model_name
         image = state.conv.convert_image_to_base64(
-            image,
-            resize_image=resize_image,
+            image
         )  # PIL type is not JSON serializable
 
         if csam_flag:
@@ -572,7 +574,9 @@ def bot_response(
         has_csam_images=state.has_csam_image, use_remote_storage=use_remote_storage
     )
 
-    filename = get_conv_log_filename(is_vision=state.is_vision)
+    filename = get_conv_log_filename(
+        is_vision=state.is_vision, has_csam_image=state.has_csam_image
+    )
 
     with open(filename, "a") as fout:
         data = {