Update app.py

app.py

@@ -103,12 +103,29 @@ def chat_handler(prompt, user_image, chat_history, output_filename="generated_im
     
     return chat_history, user_image, img, ""
 
+# Function to update chat history with thumbnails
+def update_chat_with_thumbnails(chat_history, uploaded_image_url, generated_image_url):
+    # Create a copy of the chat history to avoid modifying the input directly
+    updated_history = chat_history.copy()
+    
+    # If there's an uploaded image, add its thumbnail to the chat history
+    if uploaded_image_url:
+        thumbnail_html = f'<img src="{uploaded_image_url}" width="100px" style="margin: 5px;" />'
+        updated_history.append({"role": "user", "content": thumbnail_html})
+    
+    # If there's a generated image, add its thumbnail to the chat history
+    if generated_image_url:
+        thumbnail_html = f'<img src="{generated_image_url}" width="100px" style="margin: 5px;" />'
+        updated_history.append({"role": "assistant", "content": thumbnail_html})
+    
+    return updated_history
+
 # Create Gradio interface
 with gr.Blocks(title="Image Editing Chatbot") as demo:
     gr.Markdown("# Image Editing Chatbot")
     gr.Markdown("Upload an image and/or type a prompt to generate or edit an image using Google's Gemini model")
     
-    # Chatbot display area for text messages
+    # Chatbot display area for text messages and thumbnails
     chatbot = gr.Chatbot(
         label="Chat",
         height=300,
@@ -142,9 +159,43 @@ with gr.Blocks(title="Image Editing Chatbot") as demo:
             )
             generate_btn = gr.Button("Generate Image")
     
+    # Hidden components to store image URLs
+    uploaded_image_url = gr.State("")
+    generated_image_url = gr.State("")
+    
     # State to maintain chat history
     chat_state = gr.State([])
 
+    # JavaScript to extract image URLs and update the chat history
+    gr.HTML("""
+    <script>
+    async function updateImageURLs() {
+        // Get the image elements from the gr.Image components
+        const uploadedImage = document.querySelector('#component-3 img'); // Adjust component ID based on Gradio's generated IDs
+        const generatedImage = document.querySelector('#component-4 img'); // Adjust component ID based on Gradio's generated IDs
+        
+        // Extract the src attributes (URLs)
+        const uploadedImageURL = uploadedImage && uploadedImage.src ? uploadedImage.src : "";
+        const generatedImageURL = generatedImage && generatedImage.src ? generatedImage.src : "";
+        
+        // Call the Python function to update the chat history with the URLs
+        const chatHistory = await gradioApp().querySelector('#component-2').value; // Adjust component ID for chat_state
+        const result = await gradioApp().querySelector('#component-0').callFunction('update_chat_with_thumbnails', [chatHistory, uploadedImageURL, generatedImageURL]);
+        return result;
+    }
+
+    // Run the function when the Generate Image button is clicked
+    document.querySelector('#component-8').addEventListener('click', async () => { // Adjust component ID for generate_btn
+        // Wait for the images to update
+        setTimeout(async () => {
+            const updatedChat = await updateImageURLs();
+            // Update the chatbot component with the new history
+            document.querySelector('#component-1').value = updatedChat; // Adjust component ID for chatbot
+        }, 1000); // Delay to ensure images are loaded
+    });
+    </script>
+    """)
+
     # Connect the button to the chat handler
     generate_btn.click(
         fn=chat_handler,