Update app.py

app.py

@@ -90,8 +90,13 @@ def chat_handler(prompt, user_image, chat_history, output_filename="generated_im
     user_message_content = []
     if prompt:
         user_message_content.append(prompt)
-    if user_image:
-        user_message_content.append("Uploaded Image") # Placeholder for user image in chat history
+    if user_image is not None: # Handle case where no image is uploaded initially
+        # Convert user image to base64 for chatbot display
+        buffered = io.BytesIO()
+        user_image.save(buffered, format="PNG")
+        user_image_base64 = base64.b64encode(buffered.getvalue()).decode()
+        user_image_data_uri = f"data:image/png;base64,{user_image_base64}"
+        user_message_content.append(user_image_data_uri) # Use data URI for user image in chat history
     if user_message_content:
         chat_history.append({"role": "user", "content": user_message_content if len(user_message_content) > 1 else user_message_content[0]})
 
@@ -103,9 +108,7 @@ def chat_handler(prompt, user_image, chat_history, output_filename="generated_im
     # Generate image based on user input
     img, status = generate_image(prompt or "Generate an image", user_image, output_filename)
 
-    assistant_message_content = []
-    assistant_message_content.append(status) # Add text status
-
+    assistant_message_content = None # Initialize to None
     if img:
         # Create thumbnail for chatbot
         thumbnail_size = (100, 100) # Define thumbnail size
@@ -117,10 +120,12 @@ def chat_handler(prompt, user_image, chat_history, output_filename="generated_im
         thumbnail.save(buffered, format="PNG")
         thumbnail_base64 = base64.b64encode(buffered.getvalue()).decode()
         thumbnail_data_uri = f"data:image/png;base64,{thumbnail_base64}"
-        assistant_message_content.append(thumbnail_data_uri) # Add thumbnail to message
+        assistant_message_content = thumbnail_data_uri # ONLY data URI as assistant message
+    else:
+        assistant_message_content = status # If no image, send text status
 
     # Add assistant's response to chat history
-    chat_history.append({"role": "assistant", "content": assistant_message_content if len(assistant_message_content) > 1 else assistant_message_content[0]})
+    chat_history.append({"role": "assistant", "content": assistant_message_content})
 
     return chat_history, user_image, img, ""