Update app.py

app.py

@@ -90,13 +90,8 @@ 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 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_image is not None:
+        user_message_content.append("Uploaded Image") # Simple text 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]})
 
@@ -108,7 +103,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 = None # Initialize to None
+    thumbnail_data_uri = None # Initialize to None
     if img:
         # Create thumbnail for chatbot
         thumbnail_size = (100, 100) # Define thumbnail size
@@ -120,11 +115,13 @@ 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 = thumbnail_data_uri # ONLY data URI as assistant message
+        print("Thumbnail Data URI:", thumbnail_data_uri) # Print to console
+        assistant_message_content = gr.HTML(f'<img src="{thumbnail_data_uri}" alt="Generated Image Thumbnail">') # Use gr.HTML
     else:
         assistant_message_content = status # If no image, send text status
 
-    # Add assistant's response to chat history
+
+    # Update chat history
     chat_history.append({"role": "assistant", "content": assistant_message_content})
 
     return chat_history, user_image, img, ""