Upload gradio_web_server.py

gradio_web_server.py

@@ -1,3 +1,4 @@
+
 import argparse
 import datetime
 import json
@@ -19,9 +20,9 @@ logger = build_logger("gradio_web_server", "gradio_web_server.log")
 
 headers = {"User-Agent": "LLaVA Client"}
 
-no_change_btn = gr.Button.update()
-enable_btn = gr.Button.update(interactive=True)
-disable_btn = gr.Button.update(interactive=False)
+no_change_btn = gr.Button()
+enable_btn = gr.Button(interactive=True)
+disable_btn = gr.Button(interactive=False)
 
 priority = {
     "vicuna-13b": "aaaaaaa",
@@ -58,12 +59,11 @@ function() {
 def load_demo(url_params, request: gr.Request):
     logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
 
-    dropdown_update = gr.Dropdown.update(visible=True)
+    dropdown_update = gr.Dropdown(visible=True)
     if "model" in url_params:
         model = url_params["model"]
         if model in models:
-            dropdown_update = gr.Dropdown.update(
-                value=model, visible=True)
+            dropdown_update = gr.Dropdown(value=model, visible=True)
 
     state = default_conversation.copy()
     return state, dropdown_update
@@ -73,7 +73,7 @@ def load_demo_refresh_model_list(request: gr.Request):
     logger.info(f"load_demo. ip: {request.client.host}")
     models = get_model_list()
     state = default_conversation.copy()
-    dropdown_update = gr.Dropdown.update(
+    dropdown_update = gr.Dropdown(
         choices=models,
         value=models[0] if len(models) > 0 else ""
     )
@@ -124,8 +124,7 @@ 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) + (disable_btn,) * 5
-	
-	
+
 
 def add_text(state, text, image, image_process_mode, request: gr.Request):
     logger.info(f"add_text. ip: {request.client.host}. len: {len(text)}")
@@ -153,65 +152,14 @@ def add_text(state, text, image, image_process_mode, request: gr.Request):
     state.skip_next = False
     return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
 
-	
-def batch_process_images(folder_path, textbox, model_selector, temperature, top_p, max_output_tokens, request: gr.Request):
-    print("Starting batch processing of images")
-
-    # Initialize counters and timer
-    image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
-    total_images = len(image_files)
-    processed_images = 0
-    total_processing_time = 0
-
-    # Process each image file
-    for filename in image_files:
-        image_path = os.path.join(folder_path, filename)
-        start_time = time.time()
-        
-        with Image.open(image_path) as image:
-            state = default_conversation.copy()
-            state, _, _, _, _, _, _, _, _ = add_text(state, textbox, image, "Default", request)
-
-            # Call http_bot and iterate over the generator
-            response_text = ""
-            for state_update in http_bot(state, model_selector, temperature, top_p, max_output_tokens, request):
-                # Update state and extract response text
-                state, chatbot_output, *_ = state_update
-                response_text = chatbot_output
-
-            # Save the final response to a file
-            try:
-                with open(os.path.splitext(image_path)[0] + '.txt', 'w') as f:
-                    f.write(response_text[0][1])
-            except Exception as e:
-                print(f"An error occurred: {e}")
-
-        # Update processing information
-        processed_images += 1
-        processing_time = time.time() - start_time
-        total_processing_time += processing_time
-        average_processing_time = total_processing_time / processed_images
-        images_left = total_images - processed_images
-        eta_seconds = average_processing_time * images_left
-        eta = datetime.timedelta(seconds=int(eta_seconds))
-
-        # Display progress information
-        print(f"{processed_images}/{total_images} images processed, {images_left} left, average process time {average_processing_time :.2f} seconds, ETA: {str(eta)}")
-
-    return "Batch processing completed."
-
-	
-
 
 def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request: gr.Request):
     logger.info(f"http_bot. ip: {request.client.host}")
-    print(f"model_selector {model_selector}")
     start_tstamp = time.time()
     model_name = model_selector
-    
+
     if state.skip_next:
         # This generate call is skipped due to invalid inputs
-        print("invalid input state.skip_next")
         yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
         return
 
@@ -220,6 +168,15 @@ def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request:
         if "llava" in model_name.lower():
             if 'llama-2' in model_name.lower():
                 template_name = "llava_llama_2"
+            elif "mistral" in model_name.lower() or "mixtral" in model_name.lower():
+                if 'orca' in model_name.lower():
+                    template_name = "mistral_orca"
+                elif 'hermes' in model_name.lower():
+                    template_name = "chatml_direct"
+                else:
+                    template_name = "mistral_instruct"
+            elif 'llava-v1.6-34b' in model_name.lower():
+                template_name = "chatml_direct"
             elif "v1" in model_name.lower():
                 if 'mmtag' in model_name.lower():
                     template_name = "v1_mmtag"
@@ -242,7 +199,6 @@ def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request:
             template_name = "llama_2"
         else:
             template_name = "vicuna_v1"
-        print(f"template_name {template_name}")
         new_state = conv_templates[template_name].copy()
         new_state.append_message(new_state.roles[0], state.messages[-2][1])
         new_state.append_message(new_state.roles[1], None)
@@ -258,7 +214,6 @@ def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request:
     # No available worker
     if worker_addr == "":
         state.messages[-1][-1] = server_error_msg
-        print(f"error No available worker")
         yield (state, state.to_gradio_chatbot(), disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
         return
 
@@ -285,12 +240,12 @@ def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request:
         "images": f'List of {len(state.get_images())} images: {all_image_hash}',
     }
     logger.info(f"==== request ====\n{pload}")
-    
+
     pload['images'] = state.get_images()
 
     state.messages[-1][-1] = "▌"
     yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
-    print(f"entering Stream output")
+
     try:
         # Stream output
         response = requests.post(worker_addr + "/worker_generate_stream",
@@ -334,7 +289,8 @@ def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request:
 
 title_markdown = ("""
 Most Up To Date Scripts On : https://www.patreon.com/posts/sota-very-best-90744385 \n
-Original Project :  https://llava-vl.github.io
+Original Project :  https://llava-vl.github.io\n
+REFRESH PAGE AFTER PART 3 TO SEE LOADED MODEL
 """)
 
 tos_markdown = ("""
@@ -359,17 +315,56 @@ block_css = """
 
 """
 
+def batch_process_images(folder_path, textbox, model_selector, temperature, top_p, max_output_tokens, request: gr.Request):
+    print("Starting batch processing of images")
+
+    # Initialize counters and timer
+    image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
+    total_images = len(image_files)
+    processed_images = 0
+    total_processing_time = 0
+
+    # Process each image file
+    for filename in image_files:
+        image_path = os.path.join(folder_path, filename)
+        start_time = time.time()
+        
+        with Image.open(image_path) as image:
+            state = default_conversation.copy()
+            state, _, _, _, _, _, _, _, _ = add_text(state, textbox, image, "Default", request)
+
+            # Call http_bot and iterate over the generator
+            response_text = ""
+            for state_update in http_bot(state, model_selector, temperature, top_p, max_output_tokens, request):
+                # Update state and extract response text
+                state, chatbot_output, *_ = state_update
+                response_text = chatbot_output
+
+            # Save the final response to a file
+            try:
+                with open(os.path.splitext(image_path)[0] + '.txt', 'w') as f:
+                    f.write(response_text[0][1])
+            except Exception as e:
+                print(f"An error occurred: {e}")
 
+        # Update processing information
+        processed_images += 1
+        processing_time = time.time() - start_time
+        total_processing_time += processing_time
+        average_processing_time = total_processing_time / processed_images
+        images_left = total_images - processed_images
+        eta_seconds = average_processing_time * images_left
+        eta = datetime.timedelta(seconds=int(eta_seconds))
 
+        # Display progress information
+        print(f"{processed_images}/{total_images} images processed, {images_left} left, average process time {average_processing_time :.2f} seconds, ETA: {str(eta)}")
 
+    return "Batch processing completed."
 
-def build_demo(embed_mode):
+def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
     textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False)
-    
-    # New components for batch processing
     folder_input = gr.Textbox(label="Enter Folder Path for Batch Processing")
     batch_btn = gr.Button("Batch Process")
-
     with gr.Blocks(title="LLaVA", theme=gr.themes.Default(), css=block_css) as demo:
         state = gr.State()
 
@@ -392,7 +387,8 @@ def build_demo(embed_mode):
                     value="Default",
                     label="Preprocess for non-square image", visible=False)
 
-                cur_dir = os.path.dirname(os.path.abspath(__file__))
+                if cur_dir is None:
+                    cur_dir = os.path.dirname(os.path.abspath(__file__))
                 gr.Examples(examples=[
                     [f"{cur_dir}/examples/extreme_ironing.jpg", "just caption the image with details, colors, items, objects, emotions, art style, drawing style and objects but do not add any description or comment. do not miss any item in the given image"],
                 ], inputs=[imagebox, textbox])
@@ -403,7 +399,12 @@ def build_demo(embed_mode):
                     max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",)
 
             with gr.Column(scale=8):
-                chatbot = gr.Chatbot(elem_id="chatbot", label="LLaVA Chatbot", height=550)
+                chatbot = gr.Chatbot(
+                    elem_id="chatbot",
+                    label="LLaVA Chatbot",
+                    height=650,
+                    layout="panel",
+                )
                 with gr.Row():
                     with gr.Column(scale=8):
                         textbox.render()
@@ -413,9 +414,13 @@ def build_demo(embed_mode):
                     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", interactive=False)
 
+        if not embed_mode:
+            gr.Markdown(tos_markdown)
+            gr.Markdown(learn_more_markdown)
         url_params = gr.JSON(visible=False)
 
         # Add new components for batch processing
@@ -435,31 +440,28 @@ def build_demo(embed_mode):
         upvote_btn.click(
             upvote_last_response,
             [state, model_selector],
-            [textbox, upvote_btn, downvote_btn, flag_btn],
-            queue=False
+            [textbox, upvote_btn, downvote_btn, flag_btn]
         )
         downvote_btn.click(
             downvote_last_response,
             [state, model_selector],
-            [textbox, upvote_btn, downvote_btn, flag_btn],
-            queue=False
+            [textbox, upvote_btn, downvote_btn, flag_btn]
         )
         flag_btn.click(
             flag_last_response,
             [state, model_selector],
-            [textbox, upvote_btn, downvote_btn, flag_btn],
-            queue=False
+            [textbox, upvote_btn, downvote_btn, flag_btn]
         )
 
         regenerate_btn.click(
             regenerate,
             [state, image_process_mode],
-            [state, chatbot, textbox, imagebox] + btn_list,
-            queue=False
+            [state, chatbot, textbox, imagebox] + btn_list
         ).then(
             http_bot,
             [state, model_selector, temperature, top_p, max_output_tokens],
-            [state, chatbot] + btn_list
+            [state, chatbot] + btn_list,
+            concurrency_limit=concurrency_count
         )
 
         clear_btn.click(
@@ -477,18 +479,19 @@ def build_demo(embed_mode):
         ).then(
             http_bot,
             [state, model_selector, temperature, top_p, max_output_tokens],
-            [state, chatbot] + btn_list
+            [state, chatbot] + btn_list,
+            concurrency_limit=concurrency_count
         )
 
         submit_btn.click(
             add_text,
             [state, textbox, imagebox, image_process_mode],
-            [state, chatbot, textbox, imagebox] + btn_list,
-            queue=False
+            [state, chatbot, textbox, imagebox] + btn_list
         ).then(
             http_bot,
             [state, model_selector, temperature, top_p, max_output_tokens],
-            [state, chatbot] + btn_list
+            [state, chatbot] + btn_list,
+            concurrency_limit=concurrency_count
         )
 
         if args.model_list_mode == "once":
@@ -496,8 +499,7 @@ def build_demo(embed_mode):
                 load_demo,
                 [url_params],
                 [state, model_selector],
-                _js=get_window_url_params,
-                queue=False
+                _js=get_window_url_params
             )
         elif args.model_list_mode == "reload":
             demo.load(
@@ -517,8 +519,8 @@ if __name__ == "__main__":
     parser.add_argument("--host", type=str, default="0.0.0.0")
     parser.add_argument("--port", type=int)
     parser.add_argument("--controller-url", type=str, default="http://localhost:10000")
-    parser.add_argument("--concurrency-count", type=int, default=10)
-    parser.add_argument("--model-list-mode", type=str, default="reload",
+    parser.add_argument("--concurrency-count", type=int, default=16)
+    parser.add_argument("--model-list-mode", type=str, default="once",
         choices=["once", "reload"])
     parser.add_argument("--share", action="store_true")
     parser.add_argument("--moderate", action="store_true")
@@ -529,12 +531,12 @@ if __name__ == "__main__":
     models = get_model_list()
 
     logger.info(args)
-    demo = build_demo(args.embed)
+    demo = build_demo(args.embed, concurrency_count=args.concurrency_count)
     demo.queue(
-        concurrency_count=args.concurrency_count,
         api_open=False
     ).launch(
         server_name=args.host,
         server_port=args.port,
-        share=args.share
+        share=args.share,
+        inbrowser=True
     )