Update app.py

app.py

@@ -382,11 +382,14 @@ def visualize_all_bbox_together(image, generation):
     return pil_image, generation_colored
 
 
-def gradio_reset(chat_state, img_list):
+def gradio_reset(chat_state, img_list, path_list):
     if chat_state is not None:
         chat_state.messages = []
     if img_list is not None:
         img_list = []
+    for path in path_list:
+        os.remove(path)
+    path_list.clear()
     return None, gr.update(value=None, interactive=True), gr.update(placeholder='Upload your image and chat',
                                                                     interactive=True), chat_state, img_list
 
@@ -410,7 +413,7 @@ def example_trigger(text_input, image, upload_flag, replace_flag, img_list):
     return upload_flag, replace_flag
 
 
-def gradio_ask(user_message, chatbot, chat_state, gr_img, img_list, upload_flag, replace_flag):
+def gradio_ask(user_message, chatbot, chat_state, gr_img, img_list, upload_flag, replace_flag, path_list):
     if len(user_message) == 0:
         text_box_show = 'Input should not be empty!'
     else:
@@ -448,6 +451,7 @@ def gradio_ask(user_message, chatbot, chat_state, gr_img, img_list, upload_flag,
         visual_img, _ = visualize_all_bbox_together(gr_img, user_message)
         if visual_img is not None:
             file_path = save_tmp_img(visual_img)
+            path_list.append(file_path)
             chatbot = chatbot + [[(file_path,), None]]
 
     return text_box_show, chatbot, chat_state, img_list, upload_flag, replace_flag
@@ -483,7 +487,7 @@ def gradio_stream_answer(chatbot, chat_state, img_list, temperature):
     return chatbot, chat_state
 
 
-def gradio_visualize(chatbot, gr_img):
+def gradio_visualize(chatbot, gr_img, path_list):
     if isinstance(gr_img, dict):
         gr_img, mask = gr_img['image'], gr_img['mask']
 
@@ -493,6 +497,7 @@ def gradio_visualize(chatbot, gr_img):
         if len(generation_color):
             chatbot[-1][1] = generation_color
         file_path = save_tmp_img(visual_img)
+        path_list.append(file_path)
         chatbot = chatbot + [[None, (file_path,)]]
 
     return chatbot
@@ -581,6 +586,7 @@ with gr.Blocks() as demo:
 
     upload_flag = gr.State(value=0)
     replace_flag = gr.State(value=0)
+    path_list = gr.State(value=[])
     image.upload(image_upload_trigger, [upload_flag, replace_flag, img_list], [upload_flag, replace_flag])
 
     with gr.Row():
@@ -617,7 +623,7 @@ with gr.Blocks() as demo:
 
     text_input.submit(
         gradio_ask,
-        [text_input, chatbot, chat_state, image, img_list, upload_flag, replace_flag],
+        [text_input, chatbot, chat_state, image, img_list, upload_flag, replace_flag, path_list],
         [text_input, chatbot, chat_state, img_list, upload_flag, replace_flag], queue=False
     ).success(
         gradio_stream_answer,
@@ -625,14 +631,14 @@ with gr.Blocks() as demo:
         [chatbot, chat_state]
     ).success(
         gradio_visualize,
-        [chatbot, image],
+        [chatbot, image, path_list],
         [chatbot],
         queue=False,
     )
 
     send.click(
         gradio_ask,
-        [text_input, chatbot, chat_state, image, img_list, upload_flag, replace_flag],
+        [text_input, chatbot, chat_state, image, img_list, upload_flag, replace_flag, path_list],
         [text_input, chatbot, chat_state, img_list, upload_flag, replace_flag]
     ).success(
         gradio_stream_answer,
@@ -640,10 +646,10 @@ with gr.Blocks() as demo:
         [chatbot, chat_state]
     ).success(
         gradio_visualize,
-        [chatbot, image],
+        [chatbot, image, path_list],
         [chatbot],
     )
 
-    clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, chat_state, img_list], queue=False)
+    clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, chat_state, img_list, path_list], queue=False)
 
 demo.launch(enable_queue=True)