Update app.py

app.py

@@ -86,14 +86,20 @@ def generate_image(prompt, image=None, output_filename="generated_image"):
 
 # Function to handle chat interaction
 def chat_handler(prompt, user_image, chat_history, output_filename="generated_image"):
-    # Add the user prompt and image to the chat history
-    user_message_content = []
+    # Add the user prompt to the chat history - ONLY TEXT PROMPT for user message
     if prompt:
-        user_message_content.append(prompt)
+        chat_history.append({"role": "user", "content": prompt})
     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]})
+        # If there's a user image, add a separate message for the *thumbnail* to chat history
+        user_thumbnail_size = (100, 100)
+        user_thumbnail = user_image.copy()
+        user_thumbnail.thumbnail(user_thumbnail_size)
+        buffered = io.BytesIO()
+        user_thumbnail.save(buffered, format="PNG")
+        user_thumbnail_base64 = base64.b64encode(buffered.getvalue()).decode()
+        user_thumbnail_data_uri = f"data:image/png;base64,{user_thumbnail_base64}"
+        chat_history.append({"role": "user", "content": gr.HTML(f'<img src="{user_thumbnail_data_uri}" alt="Uploaded Image Thumbnail">')})
+
 
     # If no input, return early
     if not prompt and not user_image:
@@ -115,13 +121,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}"
-        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
 
 
-    # Update chat history
+    # Update chat history - Assistant message is now EITHER gr.HTML or text
     chat_history.append({"role": "assistant", "content": assistant_message_content})
 
     return chat_history, user_image, img, ""