Update app.py

app.py

@@ -20,6 +20,12 @@ from minigpt4.runners import *
 from minigpt4.tasks import *
 
 def parse_args():
+    """
+    Parse command line arguments.
+
+    Returns:
+    argparse.Namespace: Parsed command line arguments.
+    """
     parser = argparse.ArgumentParser(description="Demo")
     parser.add_argument("--cfg-path", type=str, default='eval_configs/minigpt4.yaml', help="path to configuration file.")
     parser.add_argument(
@@ -32,8 +38,13 @@ def parse_args():
     args = parser.parse_args()
     return args
 
-
 def setup_seeds(config):
+    """
+    Set up random seeds for reproducibility.
+
+    Parameters:
+    config (Config): Configuration object.
+    """
     seed = config.run_cfg.seed + get_rank()
 
     random.seed(seed)
@@ -42,37 +53,39 @@ def setup_seeds(config):
 
     cudnn.benchmark = False
     cudnn.deterministic = True
-    
-# ========================================
-#             Model Initialization
-# ========================================
 
-SHARED_UI_WARNING = f'''### [NOTE] It is possible that you are waiting in a lengthy queue.
+def initialize_chat():
+    """
+    Initialize the chat model.
 
-You can duplicate and use it with a paid private GPU.
+    Returns:
+    Chat: Initialized chat model.
+    """
+    print('Initializing Chat')
+    config = Config(parse_args())
 
-<a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/Vision-CAIR/minigpt4?duplicate=true"><img style="margin-top:0;margin-bottom:0" src="https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-xl-dark.svg" alt="Duplicate Space"></a>
+    model_config = config.model_cfg
+    model_cls = registry.get_model_class(model_config.arch)
+    model = model_cls.from_config(model_config).to('cuda:0')
 
-Alternatively, you can also use the demo on our [project page](https://minigpt-4.github.io).
-'''
+    vis_processor_cfg = config.datasets_cfg.cc_align.vis_processor.train
+    vis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg)
+    chat = Chat(model, vis_processor)
+    print('Initialization Finished')
 
-print('Initializing Chat')
-cfg = Config(parse_args())
+    return chat
 
-model_config = cfg.model_cfg
-model_cls = registry.get_model_class(model_config.arch)
-model = model_cls.from_config(model_config).to('cuda:0')
-
-vis_processor_cfg = cfg.datasets_cfg.cc_align.vis_processor.train
-vis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg)
-chat = Chat(model, vis_processor)
-print('Initialization Finished')
+def gradio_reset(chat_state, img_list):
+    """
+    Reset the Gradio interface.
 
-# ========================================
-#             Gradio Setting
-# ========================================
+    Parameters:
+    chat_state (gr.State): The current state of the chat.
+    img_list (gr.State): The current list of images.
 
-def gradio_reset(chat_state, img_list):
+    Returns:
+    tuple: Updated Gradio interface elements.
+    """
     if chat_state is not None:
         chat_state.messages = []
     if img_list is not None:
@@ -80,6 +93,17 @@ def gradio_reset(chat_state, img_list):
     return None, gr.update(value=None, interactive=True), gr.update(placeholder='Please upload your image first', interactive=False), gr.update(value="Upload & Start Chat", interactive=True), chat_state, img_list
 
 def upload_img(gr_img, text_input, chat_state):
+    """
+    Upload an image and update the Gradio interface.
+
+    Parameters:
+    gr_img (gr.Image): The uploaded image.
+    text_input (gr.Textbox): The text input box.
+    chat_state (gr.State): The current state of the chat.
+
+    Returns:
+    tuple: Updated Gradio interface elements.
+    """
     if gr_img is None:
         return None, None, gr.update(interactive=True), chat_state, None
     chat_state = CONV_VISION.copy()
@@ -88,67 +112,99 @@ def upload_img(gr_img, text_input, chat_state):
     return gr.update(interactive=False), gr.update(interactive=True, placeholder='Type and press Enter'), gr.update(value="Start Chatting", interactive=False), chat_state, img_list
 
 def gradio_ask(user_message, chatbot, chat_state):
+    """
+    Process user message and update the Gradio interface.
+
+    Parameters:
+    user_message (str): The message input by the user.
+    chatbot (list): The current state of the chatbot.
+    chat_state (gr.State): The current state of the chat.
+
+    Returns:
+    tuple: Updated Gradio interface elements.
+    """
     if len(user_message) == 0:
         return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
     chat.ask(user_message, chat_state)
     chatbot = chatbot + [[user_message, None]]
     return '', chatbot, chat_state
 
-
 def gradio_answer(chatbot, chat_state, img_list, num_beams, temperature):
+    """
+    Generate a chatbot answer and update the Gradio interface.
+
+    Parameters:
+    chatbot (list): The current state of the chatbot.
+    chat_state (gr.State): The current state of the chat.
+    img_list (gr.State): The current list of images.
+    num_beams (int): The number of beams for the beam search.
+    temperature (float): The temperature for the generation.
+
+    Returns:
+    tuple: Updated Gradio interface elements.
+    """
     llm_message = chat.answer(conv=chat_state, img_list=img_list, max_new_tokens=300, num_beams=1, temperature=temperature, max_length=2000)[0]
     chatbot[-1][1] = llm_message
     return chatbot, chat_state, img_list
 
-title = """<h1 align="center">Demo of MiniGPT-4</h1>"""
-description = """<h3>This is the demo of MiniGPT-4. Upload your images and start chatting!</h3>"""
-article = """<div style='display:flex; gap: 0.25rem; '><a href='https://minigpt-4.github.io'><img src='https://img.shields.io/badge/Project-Page-Green'></a><a href='https://github.com/Vision-CAIR/MiniGPT-4'><img src='https://img.shields.io/badge/Github-Code-blue'></a><a href='https://github.com/TsuTikgiau/blip2-llm/blob/release_prepare/MiniGPT_4.pdf'><img src='https://img.shields.io/badge/Paper-PDF-red'></a></div>
-"""
-
-#TODO show examples below
-
-with gr.Blocks() as demo:
-    gr.Markdown(title)
-    gr.Markdown(SHARED_UI_WARNING)
-    gr.Markdown(description)
-    gr.Markdown(article)
-
-    with gr.Row():
-        with gr.Column(scale=0.5):
-            image = gr.Image(type="pil")
-            upload_button = gr.Button(value="Upload & Start Chat", interactive=True, variant="primary")
-            clear = gr.Button("Restart")
-            
-            num_beams = gr.Slider(
-                minimum=1,
-                maximum=5,
-                value=1,
-                step=1,
-                interactive=True,
-                label="beam search numbers)",
-            )
-            
-            temperature = gr.Slider(
-                minimum=0.1,
-                maximum=2.0,
-                value=1.0,
-                step=0.1,
-                interactive=True,
-                label="Temperature",
-            )
-            
-
-        with gr.Column():
-            chat_state = gr.State()
-            img_list = gr.State()
-            chatbot = gr.Chatbot(label='MiniGPT-4')
-            text_input = gr.Textbox(label='User', placeholder='Please upload your image first', interactive=False)
-    
-    upload_button.click(upload_img, [image, text_input, chat_state], [image, text_input, upload_button, chat_state, img_list])
-    
-    text_input.submit(gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state]).then(
-        gradio_answer, [chatbot, chat_state, img_list, num_beams, temperature], [chatbot, chat_state, img_list]
-    )
-    clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, upload_button, chat_state, img_list], queue=False)
+def main():
+    """
+    Main function to run the Gradio interface.
+    """
+    # Initialize the chat model
+    chat = initialize_chat()
+
+    # Set up the Gradio interface
+    title = """<h1 align="center">Demo of MiniGPT-4</h1>"""
+    description = """<h3>This is the demo of MiniGPT-4. Upload your images and start chatting!</h3>"""
+    article = """<div style='display:flex; gap: 0.25rem; '><a href='https://minigpt-4.github.io'><img src='https://img.shields.io/badge/Project-Page-Green'></a><a href='https://github.com/Vision-CAIR/MiniGPT-4'><img src='https://img.shields.io/badge/Github-Code-blue'></a><a href='https://github.com/TsuTikgiau/blip2-llm/blob/release_prepare/MiniGPT_4.pdf'><img src='https://img.shields.io/badge/Paper-PDF-red'></a></div>
+    """
+
+    with gr.Blocks() as demo:
+        gr.Markdown(title)
+        gr.Markdown(SHARED_UI_WARNING)
+        gr.Markdown(description)
+        gr.Markdown(article)
+
+        with gr.Row():
+            with gr.Column(scale=0.5):
+                image = gr.Image(type="pil")
+                upload_button = gr.Button(value="Upload & Start Chat", interactive=True, variant="primary")
+                clear = gr.Button("Restart")
+                
+                num_beams = gr.Slider(
+                    minimum=1,
+                    maximum=5,
+                    value=1,
+                    step=1,
+                    interactive=True,
+                    label="beam search numbers)",
+                )
+                
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=2.0,
+                    value=1.0,
+                                        step=0.1,
+                    interactive=True,
+                    label="Temperature",
+                )
+
+            with gr.Column():
+                chat_state = gr.State()
+                img_list = gr.State()
+                chatbot = gr.Chatbot(label='MiniGPT-4')
+                text_input = gr.Textbox(label='User', placeholder='Please upload your image first', interactive=False)
+        
+        upload_button.click(upload_img, [image, text_input, chat_state], [image, text_input, upload_button, chat_state, img_list])
+        
+        text_input.submit(gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state]).then(
+            gradio_answer, [chatbot, chat_state, img_list, num_beams, temperature], [chatbot, chat_state, img_list]
+        )
+        clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, upload_button, chat_state, img_list], queue=False)
+
+    demo.launch(enable_queue=True)
+
+if __name__ == "__main__":
+    main()
 
-demo.launch(enable_queue=True)