Update app.py

app.py

@@ -128,47 +128,48 @@ def flag_last_response(state, model_selector, request: gr.Request):
     return ("",) + (disable_btn,) * 3
 
 
-def regenerate(state, image_process_mode1, image_process_mode2, request: gr.Request):
+def regenerate(state, image_process_mode, request: gr.Request):
     logger.info(f"regenerate. ip: {request.client.host}")
     state.messages[-1][-1] = None
     prev_human_msg = state.messages[-2]
     if type(prev_human_msg[1]) in (tuple, list):
-        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode1, image_process_mode2)
+        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
     state.skip_next = False
-    return (state, state.to_gradio_chatbot(), "", None, None) + (disable_btn,) * 5
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
 
 
 def clear_history(request: gr.Request):
     logger.info(f"clear_history. ip: {request.client.host}")
     state = default_conversation.copy()
-    return (state, state.to_gradio_chatbot(), "", None, None) + (disable_btn,) * 5
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
 
 
-def add_text(state, text, image1, image2, image_process_mode1, image_process_mode2, request: gr.Request):
+def add_text(state, text, image, image_process_mode, request: gr.Request):
     logger.info(f"add_text. ip: {request.client.host}. len: {len(text)}")
-    if len(text) <= 0 and image1 is None and image2 is None:
+    if len(text) <= 0 and image is None:
         state.skip_next = True
-        return (state, state.to_gradio_chatbot(), "", None, None) + (no_change_btn,) * 5
+        return (state, state.to_gradio_chatbot(), "", None) + (no_change_btn,) * 5
     if args.moderate:
         flagged = violates_moderation(text)
         if flagged:
             state.skip_next = True
-            return (state, state.to_gradio_chatbot(), moderation_msg, None, None) + (
+            return (state, state.to_gradio_chatbot(), moderation_msg, None) + (
                 no_change_btn,
             ) * 5
 
     text = text[:1536]  # Hard cut-off
-    if image1 is not None or image2 is not None:
+    if image is not None:
         text = text[:1200]  # Hard cut-off for images
         if "<image>" not in text:
+            # text = '<Image><image></Image>' + text
             text = text + "\n<image>"
-        text = (text, image1, image2, image_process_mode1, image_process_mode2)
+        text = (text, image, image_process_mode)
         if len(state.get_images(return_pil=True)) > 0:
             state = default_conversation.copy()
     state.append_message(state.roles[0], text)
     state.append_message(state.roles[1], None)
     state.skip_next = False
-    return (state, state.to_gradio_chatbot(), "", None, None) + (disable_btn,) * 5
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
 
 
 def http_bot(
@@ -179,10 +180,12 @@ def http_bot(
     model_name = model_selector
 
     if state.skip_next:
+        # This generate call is skipped due to invalid inputs
         yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
         return
 
     if len(state.messages) == state.offset + 2:
+        # First round of conversation
         if "llava" in model_name.lower():
             if "llama-2" in model_name.lower():
                 template_name = "llava_llama_2"
@@ -219,6 +222,7 @@ def http_bot(
         new_state.append_message(new_state.roles[1], None)
         state = new_state
 
+    # Query worker address
     controller_url = args.controller_url
     ret = requests.post(
         controller_url + "/get_worker_address", json={"model": model_name}
@@ -226,6 +230,7 @@ def http_bot(
     worker_addr = ret.json()["address"]
     logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")
 
+    # No available worker
     if worker_addr == "":
         state.messages[-1][-1] = server_error_msg
         yield (
@@ -239,6 +244,7 @@ def http_bot(
         )
         return
 
+    # Construct prompt
     prompt = state.get_prompt()
 
     all_images = state.get_images(return_pil=True)
@@ -252,6 +258,7 @@ def http_bot(
             os.makedirs(os.path.dirname(filename), exist_ok=True)
             image.save(filename)
 
+    # Make requests
     pload = {
         "model": model_name,
         "prompt": prompt,
@@ -271,6 +278,7 @@ def http_bot(
     yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
 
     try:
+        # Stream output
         response = requests.post(
             worker_addr + "/worker_generate_stream",
             headers=headers,
@@ -331,13 +339,17 @@ def http_bot(
 title_markdown = """
 # 🌋 AI Tutor Vision: Large Language and Vision Assistant
 [[website]](https://myapps.ai) [[Paper]](https://arxiv.org/abs/2304.08485) [[Model]](https://github.com/haotian-liu/LLaVA/blob/main/docs/MODEL_ZOO.md)
+
 ONLY WORKS WITH GPU!
+
 You can load the model with 4-bit or 8-bit quantization to make it fit in smaller hardwares. Setting the environment variable `bits` to control the quantization.
 *Note: 8-bit seems to be slower than both 4-bit/16-bit. Although it has enough VRAM to support 8-bit, until we figure out the inference speed issue, we recommend 4-bit for A10G for the best efficiency.*
+
 Recommended configurations:
 | Hardware          | T4-Small (16G)  | A10G-Small (24G) | A100-Large (40G) |
 |-------------------|-----------------|------------------|------------------|
 | **Bits**          | 4 (default)     | 4                | 16               |
+
 """
 
 tos_markdown = """
@@ -355,9 +367,11 @@ The service is a research preview intended for non-commercial use only, subject
 """
 
 block_css = """
+
 #buttons button {
     min-width: min(120px,100%);
 }
+
 """
 
 
@@ -384,15 +398,8 @@ def build_demo(embed_mode):
                         container=False,
                     )
 
-                imagebox1 = gr.Image(type="pil")
-                imagebox2 = gr.Image(type="pil")
-                image_process_mode1 = gr.Radio(
-                    ["Crop", "Resize", "Pad", "Default"],
-                    value="Default",
-                    label="Preprocess for non-square image",
-                    visible=False,
-                )
-                image_process_mode2 = gr.Radio(
+                imagebox = gr.Image(type="pil")
+                image_process_mode = gr.Radio(
                     ["Crop", "Resize", "Pad", "Default"],
                     value="Default",
                     label="Preprocess for non-square image",
@@ -411,7 +418,7 @@ def build_demo(embed_mode):
                             "What are the things I should be cautious about when I visit here?",
                         ],
                     ],
-                    inputs=[imagebox1, textbox, imagebox2],
+                    inputs=[imagebox, textbox],
                 )
 
                 with gr.Accordion("Parameters", open=False) as parameter_row:
@@ -442,18 +449,20 @@ def build_demo(embed_mode):
 
             with gr.Column(scale=8):
                 chatbot = gr.Chatbot(
-                    elem_id="chatbot", label="AI Tutor Vision Chatbot", height=550
+                    elem_id="chatbot", label="AI Tutor Vision", height=550
                 )
                 with gr.Row():
                     with gr.Column(scale=8):
                         textbox.render()
                     with gr.Column(scale=1, min_width=50):
                         submit_btn = gr.Button(
-                            value="Send", variant="primary", interactive=False)
+                            value="Send", variant="primary", interactive=False
+                        )
                 with gr.Row(elem_id="buttons") as button_row:
                     upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
                     downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
                     flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
+                    # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
                     regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
                     clear_btn = gr.Button(value="🗑️  Clear history", interactive=False)
 
@@ -462,6 +471,7 @@ def build_demo(embed_mode):
             gr.Markdown(learn_more_markdown)
         url_params = gr.JSON(visible=False)
 
+        # Register listeners
         btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
         upvote_btn.click(
             upvote_last_response,
@@ -480,21 +490,21 @@ def build_demo(embed_mode):
         )
         regenerate_btn.click(
             regenerate,
-            [state, image_process_mode1, image_process_mode2],
-            [state, chatbot, textbox, imagebox1, imagebox2] + btn_list,
+            [state, image_process_mode],
+            [state, chatbot, textbox, imagebox] + btn_list,
         ).then(
             http_bot,
             [state, model_selector, temperature, top_p, max_output_tokens],
             [state, chatbot] + btn_list,
         )
         clear_btn.click(
-            clear_history, None, [state, chatbot, textbox, imagebox1, imagebox2] + btn_list
+            clear_history, None, [state, chatbot, textbox, imagebox] + btn_list
         )
 
         textbox.submit(
             add_text,
-            [state, textbox, imagebox1, imagebox2, image_process_mode1, image_process_mode2],
-            [state, chatbot, textbox, imagebox1, imagebox2] + btn_list,
+            [state, textbox, imagebox, image_process_mode],
+            [state, chatbot, textbox, imagebox] + btn_list,
         ).then(
             http_bot,
             [state, model_selector, temperature, top_p, max_output_tokens],
@@ -502,8 +512,8 @@ def build_demo(embed_mode):
         )
         submit_btn.click(
             add_text,
-            [state, textbox, imagebox1, imagebox2, image_process_mode1, image_process_mode2],
-            [state, chatbot, textbox, imagebox1, imagebox2] + btn_list,
+            [state, textbox, imagebox, image_process_mode],
+            [state, chatbot, textbox, imagebox] + btn_list,
         ).then(
             http_bot,
             [state, model_selector, temperature, top_p, max_output_tokens],
@@ -600,6 +610,7 @@ if __name__ == "__main__":
     controller_proc = start_controller()
     worker_proc = start_worker(model_path, bits=bits)
 
+    # Wait for worker and controller to start
     time.sleep(10)
 
     exit_status = 0
@@ -612,4 +623,4 @@ if __name__ == "__main__":
         worker_proc.kill()
         controller_proc.kill()
 
-        sys.exit(exit_status)
+        sys.exit(exit_status)