Update app.py

app.py

@@ -40,13 +40,9 @@ def _get_args() -> ArgumentParser:
     return args
 
 def handle_image_submission(_chatbot, task_history, file) -> tuple:
-    print("handle_image_submission called")
     if file is None:
-        print("No file uploaded")
         return _chatbot, task_history
-    print("File received:", file)
     file_path = save_image(file, uploaded_file_dir)
-    print("File saved at:", file_path)
     history_item = ((file_path,), None)
     _chatbot.append(history_item)
     task_history.append(history_item)
@@ -112,12 +108,10 @@ def _parse_text(text: str) -> str:
     return text
 
 def save_image(image_file, upload_dir: str) -> str:
-    print("save_image called with:", image_file)
     Path(upload_dir).mkdir(parents=True, exist_ok=True)
     filename = secrets.token_hex(10) + Path(image_file.name).suffix
     file_path = Path(upload_dir) / filename
-    print("Saving to:", file_path)
-    with open(image_file.name, "rb") as f_input, open(file_path, "wb") as f_output:
+    with open(image_file, "rb") as f_input, open(file_path, "wb") as f_output:
         f_output.write(f_input.read())
     return str(file_path)
 
@@ -134,16 +128,10 @@ def predict(_chatbot, task_history) -> tuple:
     print("predict called")
     chat_query, chat_response = _chatbot[-1]
 
-    # Check if the chat query is a tuple (indicating an image), and if so, extract the image path
-    if isinstance(chat_query, tuple):
-        chat_query = chat_query[0]
-
-    # Main logic for processing the query
     if isinstance(chat_query, tuple):
         query = [{'image': chat_query[0]}]
     else:
         query = [{'text': _parse_text(chat_query)}]
-    print("Query for model:", query)
 
     inputs = tokenizer.from_list_format(query)
     tokenized_inputs = tokenizer(inputs, return_tensors='pt')
@@ -151,9 +139,10 @@ def predict(_chatbot, task_history) -> tuple:
 
     pred = model.generate(**tokenized_inputs)
     response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)
+
+    if 'image' in query[0]:
     print("Model response:", response)
 
-    # Handling the response based on the type of query (image or text)
     if 'image' in query[0]:
         image = tokenizer.draw_bbox_on_latest_picture(response)
         if image is not None:
@@ -164,7 +153,6 @@ def predict(_chatbot, task_history) -> tuple:
     else:
         _chatbot[-1] = (chat_query, response)
 
-    # Return the updated chatbot and task history
     return _chatbot, task_history
 
 def save_uploaded_image(image_file, upload_dir):