Update app.py

app.py

@@ -38,7 +38,15 @@ def _get_args():
     args = parser.parse_args()
     return args
 
-
+def handle_image_submission(_chatbot, task_history, file):
+    if file is None:
+        return _chatbot, task_history  # Return if no file is uploaded
+    file_path = save_image(file)
+    history_item = ((file_path,), None)
+    _chatbot.append(history_item)
+    task_history.append(history_item)
+    return predict(_chatbot, task_history)
+    
 def _load_model_tokenizer(args):
     model_id = args.checkpoint_path
     model_dir = snapshot_download(model_id, revision=args.revision)
@@ -230,13 +238,7 @@ Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder
         )
 
         submit_file_btn.click(
-            fn=predict,
-            inputs=[chatbot, task_history],
-            outputs=[chatbot]
-        )
-
-        submit_file_btn.click(
-            fn=add_file,
+            fn=handle_image_submission,
             inputs=[chatbot, task_history, file_upload],
             outputs=[chatbot, task_history]
         )