Update app.py

app.py

@@ -138,16 +138,24 @@ def predict(_chatbot, task_history) -> tuple:
     tokenized_inputs = tokenizer(inputs, return_tensors='pt')
     tokenized_inputs = tokenized_inputs.to(model.device)
     pred = model.generate(**tokenized_inputs)
-    response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)    
+    response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)
     if 'image' in query[0]:
         print("Model response:", response)
-        image_path = response.strip()  # Ensure it's just a string path or text
-        _chatbot[-1] = (chat_query, image_path)
-        task_history[-1] = (chat_query, image_path)
+        image = tokenizer.draw_bbox_on_latest_picture(response)
+        if image is not None:
+            image_path = save_image(image, uploaded_file_dir)
+            formatted_response = (chat_query, image_path)
+        else:
+            formatted_response = (chat_query, response)
+    else:
+        text_response = response.strip()
+        formatted_response = (chat_query, text_response)
+
+    _chatbot[-1] = formatted_response
+    if task_history:
+        task_history[-1] = formatted_response
     else:
-        text_response = response.strip() 
-        _chatbot[-1] = (chat_query, text_response)
-        task_history[-1] = (chat_query, text_response)
+        task_history.append(formatted_response)
 
     return _chatbot, task_history