Update app.py

app.py

@@ -85,15 +85,19 @@ def generate_image(prompt, image=None, output_filename="generated_image"):
     return None, "No image generated"
 
 # Function to handle chat interaction
-def chat_handler(prompt, user_image, chat_history, output_filename="generated_image"):
+def chat_handler(prompt, user_image, chat_history, image_history, output_filename="generated_image"):
     # Add the prompt to the chat history
     if prompt:
         chat_history.append({"role": "user", "content": prompt})
     
+    # Add the uploaded image to the image history
+    if user_image:
+        image_history.append(user_image)
+    
     # If no input, return early
     if not prompt and not user_image:
         chat_history.append({"role": "assistant", "content": "Please provide a prompt or an image."})
-        return chat_history, user_image, None, ""
+        return chat_history, user_image, None, image_history, ""
 
     # Generate image based on user input
     img, status = generate_image(prompt or "Generate an image", user_image, output_filename)
@@ -101,31 +105,18 @@ def chat_handler(prompt, user_image, chat_history, output_filename="generated_im
     # Add the status message to the chat history
     chat_history.append({"role": "assistant", "content": status})
     
-    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})
+    # Add the generated image to the image history
+    if img:
+        image_history.append(img)
     
-    # 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
+    return chat_history, user_image, img, image_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 and thumbnails
+    # Chatbot display area for text messages
     chatbot = gr.Chatbot(
         label="Chat",
         height=300,
@@ -133,6 +124,15 @@ with gr.Blocks(title="Image Editing Chatbot") as demo:
         avatar_images=(None, None)
     )
     
+    # Gallery to display image history as thumbnails
+    image_history_display = gr.Gallery(
+        label="Image History",
+        height=200,
+        columns=4,
+        object_fit="contain",
+        preview=True
+    )
+    
     # Separate image outputs
     with gr.Row():
         uploaded_image_output = gr.Image(label="Uploaded Image")
@@ -159,55 +159,22 @@ 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
+    # State to maintain chat history and image 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>
-    """)
+    image_history_state = gr.State([])
 
     # Connect the button to the chat handler
     generate_btn.click(
         fn=chat_handler,
-        inputs=[prompt_input, image_input, chat_state, filename_input],
-        outputs=[chatbot, uploaded_image_output, generated_image_output, prompt_input]
+        inputs=[prompt_input, image_input, chat_state, image_history_state, filename_input],
+        outputs=[chatbot, uploaded_image_output, generated_image_output, image_history_display, prompt_input]
     )
 
     # Also allow Enter key to submit
     prompt_input.submit(
         fn=chat_handler,
-        inputs=[prompt_input, image_input, chat_state, filename_input],
-        outputs=[chatbot, uploaded_image_output, generated_image_output, prompt_input]
+        inputs=[prompt_input, image_input, chat_state, image_history_state, filename_input],
+        outputs=[chatbot, uploaded_image_output, generated_image_output, image_history_display, prompt_input]
     )
 
 if __name__ == "__main__":