Update app.py

app.py

@@ -98,6 +98,13 @@ def _parse_text(text):
     text = "".join(lines)
     return text
 
+def add_file(history, task_history, file):
+    if file is None:
+        return history, task_history  # Return if no file is uploaded
+    file_path = save_image(file)
+    history = history + [((file_path,), None)]
+    task_history = task_history + [((file_path,), None)]
+    return history, task_history
 
 def _launch_demo(args, model, tokenizer):
     uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(
@@ -122,13 +129,27 @@ def _launch_demo(args, model, tokenizer):
 
         if 'image' in query[0]:
             image = tokenizer.draw_bbox_on_latest_picture(response)
-            image_path = save_image(image)  # Ensure you have a function to save the image
-            _chatbot[-1] = (chat_query, (image_path,))
+               if image is not None:
+                image_path = save_image(image)
+                _chatbot[-1] = (chat_query, (image_path,))
+            else:
+                _chatbot[-1] = (chat_query, "No image to display.")
         else:
             _chatbot[-1] = (chat_query, response)
-
+    
         return _chatbot
 
+    def save_image(image):
+        if image is None:
+            return None
+        temp_dir = secrets.token_hex(20)
+        temp_dir = Path(uploaded_file_dir) / temp_dir
+        temp_dir.mkdir(exist_ok=True, parents=True)
+        name = f"tmp{secrets.token_hex(5)}.jpg"
+        filename = temp_dir / name
+        image.save(str(filename))
+        return str(filename)
+
     def regenerate(_chatbot, task_history):
         if not task_history:
             return _chatbot
@@ -180,15 +201,6 @@ def _launch_demo(args, model, tokenizer):
     def reset_state(task_history):
         task_history.clear()
         return []
-        
-    def save_image(image):
-        temp_dir = secrets.token_hex(20)
-        temp_dir = Path(uploaded_file_dir) / temp_dir
-        temp_dir.mkdir(exist_ok=True, parents=True)
-        name = f"tmp{secrets.token_hex(5)}.jpg"
-        filename = temp_dir / name
-        image.save(str(filename))
-        return str(filename)
 
 
     with gr.Blocks() as demo:
@@ -219,6 +231,12 @@ Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder
             outputs=[chatbot]
         )
 
+        submit_file_btn.click(
+            fn=predict,
+            inputs=[chatbot, task_history],
+            outputs=[chatbot]
+        )
+
         submit_file_btn.click(
             fn=add_file,
             inputs=[chatbot, task_history, file_upload],