Update app.py

app.py

@@ -156,6 +156,22 @@ class InferenceDemo(object):
         self.conversation = conv_templates[args.conv_mode].copy()
         self.num_frames = args.num_frames
 
+class ChatSessionManager:
+    def __init__(self):
+        self.chatbot_instance = None
+
+    def initialize_chatbot(self, args, model_path, tokenizer, model, image_processor, context_len):
+        self.chatbot_instance = InferenceDemo(args, model_path, tokenizer, model, image_processor, context_len)
+        print(f"Initialized Chatbot instance with ID: {id(self.chatbot_instance)}")
+
+    def reset_chatbot(self):
+        self.chatbot_instance = None
+
+    def get_chatbot(self, args, model_path, tokenizer, model, image_processor, context_len):
+        if self.chatbot_instance is None:
+            self.initialize_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+        return self.chatbot_instance
+        
 
 def is_valid_video_filename(name):
     video_extensions = ["avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg"]
@@ -209,13 +225,6 @@ def load_image(image_file):
     return image
 
 
-def clear_history(history):
-
-    our_chatbot.conversation = conv_templates[our_chatbot.conv_mode].copy()
-
-    return None
-
-
 def clear_response(history):
     for index_conv in range(1, len(history)):
         # loop until get a text response from our model.
@@ -226,60 +235,80 @@ def clear_response(history):
     history = history[:-index_conv]
     return history, question
 
+chat_manager = ChatSessionManager()
+
+
+def clear_history(history):
+    chatbot_instance = chat_manager.get_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+    chatbot_instance.conversation = conv_templates[chatbot_instance.conv_mode].copy()
+    return None
 
-# def print_like_dislike(x: gr.LikeData):
-#     print(x.index, x.value, x.liked)
 
 
 def add_message(history, message):
-    # history=[]
-    global our_chatbot
-    if len(history) == 0:
-        our_chatbot = InferenceDemo(
-            args, model_path, tokenizer, model, image_processor, context_len
-        )
+    global chat_image_num
+    if not history:
+        history = []
+        our_chatbot = chat_manager.get_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+        chat_image_num = 0
+        
+    if len(message["files"]) <= 1:
+        for x in message["files"]:
+            history.append(((x,), None))
+            chat_image_num += 1
+        if chat_image_num > 1:
+            history = []
+            chat_manager.reset_chatbot()
+            our_chatbot = chat_manager.get_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+            chat_image_num = 0
+            for x in message["files"]:
+                history.append(((x,), None))
+                chat_image_num += 1
+                
+        if message["text"] is not None:
+            history.append((message["text"], None))
+        
+        print(f"### Chatbot instance ID: {id(our_chatbot)}")
+        return history, gr.MultimodalTextbox(value=None, interactive=False)
+    else:
+        for x in message["files"]:
+            history.append(((x,), None))
+        if message["text"] is not None:
+            history.append((message["text"], None))
 
-    for x in message["files"]:
-        history.append(((x,), None))
-    if message["text"] is not None:
-        history.append((message["text"], None))
-    return history, gr.MultimodalTextbox(value=None, interactive=False)
+        return history, gr.MultimodalTextbox(value=None, interactive=False)
 
 
 @spaces.GPU
 def bot(history, temperature, top_p, max_output_tokens):
-    print("### turn start history",history)
-    print("### turn start conv",our_chatbot.conversation)
+    our_chatbot = chat_manager.get_chatbot(args, model_path, tokenizer, model, image_processor, context_len)
+    print(f"### Chatbot instance ID: {id(our_chatbot)}")
     text = history[-1][0]
     images_this_term = []
     text_this_term = ""
-    # import pdb;pdb.set_trace()
+    
     num_new_images = 0
+    previous_image = False
     for i, message in enumerate(history[:-1]):
         if type(message[0]) is tuple:
-            if len(message[0])>1:
+            if previous_image:
                 gr.Warning("Only one image can be uploaded in a conversation. Please reduce the number of images and start a new conversation.")
-                return history 
+                our_chatbot.conversation = conv_templates[our_chatbot.conv_mode].copy()
+                return None
+           
+            images_this_term.append(message[0][0])
+            if is_valid_video_filename(message[0][0]):
+                raise ValueError("Video is not supported")
+                num_new_images += our_chatbot.num_frames
+            elif is_valid_image_filename(message[0][0]):
+                print("#### Load image from local file",message[0][0])
+                num_new_images += 1
             else:
-                images_this_term.append(message[0][0])
-                if is_valid_video_filename(message[0][0]):
-                    # 不接受视频
-                    raise ValueError("Video is not supported")
-                    num_new_images += our_chatbot.num_frames
-                elif is_valid_image_filename(message[0][0]):
-                    print("#### Load image from local file",message[0][0])
-                    num_new_images += 1
-                else:
-                    raise ValueError("Invalid image file")
+                raise ValueError("Invalid image file")
+            previous_image = True
         else:
             num_new_images = 0
-
-    # for message in history[-i-1:]:
-    #     images_this_term.append(message[0][0])
-
-    # assert len(images_this_term) > 0, "must have an image"
-    # image_files = (args.image_file).split(',')
-    # image = [load_image(f) for f in images_this_term if f]
+            previous_image = False
 
     all_image_hash = []
     all_image_path = []
@@ -323,9 +352,7 @@ def bot(history, temperature, top_p, max_output_tokens):
 
     image_tensor = torch.stack(image_tensor)
     image_token = DEFAULT_IMAGE_TOKEN * num_new_images
-    # if our_chatbot.model.config.mm_use_im_start_end:
-    #     inp = DEFAULT_IM_START_TOKEN + image_token + DEFAULT_IM_END_TOKEN + "\n" + inp
-    # else:
+    
     inp = text
     inp = image_token + "\n" + inp
     our_chatbot.conversation.append_message(our_chatbot.conversation.roles[0], inp)
@@ -333,18 +360,6 @@ def bot(history, temperature, top_p, max_output_tokens):
     our_chatbot.conversation.append_message(our_chatbot.conversation.roles[1], None)
     prompt = our_chatbot.conversation.get_prompt()
 
-    
-    if len(images_this_term) > 1:
-        gr.Warning("Only one image can be uploaded in a conversation. Please reduce the number of images and start a new conversation.")
-        return history
-
-    # input_ids = (
-    #     tokenizer_image_token(
-    #         prompt, our_chatbot.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
-    #     )
-    #     .unsqueeze(0)
-    #     .to(our_chatbot.model.device)
-    # )
     input_ids = tokenizer_image_token(
             prompt, our_chatbot.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
         ).unsqueeze(0).to(our_chatbot.model.device)
@@ -358,9 +373,7 @@ def bot(history, temperature, top_p, max_output_tokens):
     stopping_criteria = KeywordsStoppingCriteria(
         keywords, our_chatbot.tokenizer, input_ids
     )
-    # streamer = TextStreamer(
-    #     our_chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True
-    # )
+    
     streamer = TextIteratorStreamer(
         our_chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True
     )
@@ -368,27 +381,6 @@ def bot(history, temperature, top_p, max_output_tokens):
     print(input_ids.device)
     print(image_tensor.device)
     
-    # with torch.inference_mode():
-    #     output_ids = our_chatbot.model.generate(
-    #         input_ids,
-    #         images=image_tensor,
-    #         do_sample=True,
-    #         temperature=0.7,
-    #         top_p=1.0,
-    #         max_new_tokens=4096,
-    #         streamer=streamer,
-    #         use_cache=False,
-    #         stopping_criteria=[stopping_criteria],
-    #     )
-
-    # outputs = our_chatbot.tokenizer.decode(output_ids[0]).strip()
-    # if outputs.endswith(stop_str):
-    #     outputs = outputs[: -len(stop_str)]
-    # our_chatbot.conversation.messages[-1][-1] = outputs
-
-    # history[-1] = [text, outputs]
-    
-    # return history
     generate_kwargs = dict(
         inputs=input_ids,
         streamer=streamer,
@@ -407,13 +399,12 @@ def bot(history, temperature, top_p, max_output_tokens):
     outputs = []
     for stream_token in streamer:
         outputs.append(stream_token)
-        # print("### stream_token",stream_token)
-        # our_chatbot.conversation.messages[-1][-1] = "".join(outputs)
+        
         history[-1] = [text, "".join(outputs)]
         yield history
     our_chatbot.conversation.messages[-1][-1] = "".join(outputs)
-    print("### turn end history", history)
-    print("### turn end conv",our_chatbot.conversation)
+    # print("### turn end history", history)
+    # print("### turn end conv",our_chatbot.conversation)
         
     with open(get_conv_log_filename(), "a") as fout:
         data = {
@@ -677,11 +668,10 @@ with gr.Blocks(
         gr.Markdown(learn_more_markdown)
         gr.Markdown(bibtext)
 
-    chat_msg = chat_input.submit(
-        add_message, [chatbot, chat_input], [chatbot, chat_input]
-    )
-    bot_msg = chat_msg.then(bot, [chatbot,temperature, top_p, max_output_tokens], chatbot, api_name="bot_response")
-    bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
+    chat_input.submit(
+    add_message, [chatbot, chat_input], [chatbot, chat_input]
+    ).then(bot, [chatbot, temperature, top_p, max_output_tokens], chatbot, api_name="bot_response").then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
+  
 
     # chatbot.like(print_like_dislike, None, None)
     clear_btn.click(
@@ -727,5 +717,5 @@ if __name__ == "__main__":
     model_name = get_model_name_from_path(args.model_path)
     tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit)
     model=model.to(torch.device('cuda'))
-    our_chatbot = None
+    chat_image_num = 0
     demo.launch()