Update app.py

app.py

@@ -1,117 +1,4 @@
-import base64
-import os
-import mimetypes
-from google import genai
-from google.genai import types
-import gradio as gr
-import io
-from PIL import Image
-
-def save_binary_file(file_name, data):
-    f = open(file_name, "wb")
-    f.write(data)
-    f.close()
-
-def generate_image(prompt, image=None, output_filename="generated_image"):
-    # Initialize client with the API key
-    client = genai.Client(
-        api_key="AIzaSyAQcy3LfrkMy6DqS_8MqftAXu1Bx_ov_E8",
-    )
-
-    model = "gemini-2.0-flash-exp-image-generation"
-    parts = [types.Part.from_text(text=prompt)]
-    
-    # If an image is provided, add it to the content
-    if image:
-        # Convert PIL Image to bytes
-        img_byte_arr = io.BytesIO()
-        image.save(img_byte_arr, format="PNG")
-        img_bytes = img_byte_arr.getvalue()
-        # Add the image as a Part with inline_data
-        parts.append({
-            "inline_data": {
-                "mime_type": "image/png",
-                "data": img_bytes
-            }
-        })
-
-    contents = [
-        types.Content(
-            role="user",
-            parts=parts,
-        ),
-    ]
-    generate_content_config = types.GenerateContentConfig(
-        temperature=1,
-        top_p=0.95,
-        top_k=40,
-        max_output_tokens=8192,
-        response_modalities=[
-            "image",
-            "text",
-        ],
-        safety_settings=[
-            types.SafetySetting(
-                category="HARM_CATEGORY_CIVIC_INTEGRITY",
-                threshold="OFF",
-            ),
-        ],
-        response_mime_type="text/plain",
-    )
-
-    # Generate the content
-    response = client.models.generate_content_stream(
-        model=model,
-        contents=contents,
-        config=generate_content_config,
-    )
-
-    # Process the response
-    for chunk in response:
-        if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
-            continue
-        if chunk.candidates[0].content.parts[0].inline_data:
-            inline_data = chunk.candidates[0].content.parts[0].inline_data
-            file_extension = mimetypes.guess_extension(inline_data.mime_type)
-            filename = f"{output_filename}{file_extension}"
-            save_binary_file(filename, inline_data.data)
-            
-            # Convert binary data to PIL Image for Gradio display
-            img = Image.open(io.BytesIO(inline_data.data))
-            return img, f"Image saved as {filename}"
-        else:
-            return None, chunk.text
-
-    return None, "No image generated"
-
-# Function to handle chat interaction
-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, image_history, ""
-
-    # Generate image based on user input
-    img, status = generate_image(prompt or "Generate an image", user_image, output_filename)
-    
-    # Add the status message to the chat history
-    chat_history.append({"role": "assistant", "content": status})
-    
-    # Add the generated image to the image history
-    if img:
-        image_history.append(img)
-    
-    return chat_history, user_image, img, image_history, ""
-
-# Create Gradio interface
+# Fallback: Use a gr.Gallery component for thumbnails
 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")
@@ -124,16 +11,10 @@ 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
-    )
+    # Gallery for thumbnails
+    thumbnail_gallery = gr.Gallery(label="Conversation Thumbnails", columns=4, height=150)
     
-    # Separate image outputs
+    # Separate image outputs (gallery)
     with gr.Row():
         uploaded_image_output = gr.Image(label="Uploaded Image")
         generated_image_output = gr.Image(label="Generated Image")
@@ -159,22 +40,48 @@ with gr.Blocks(title="Image Editing Chatbot") as demo:
             )
             generate_btn = gr.Button("Generate Image")
     
-    # State to maintain chat history and image history
+    # State to maintain chat history and thumbnail list
     chat_state = gr.State([])
-    image_history_state = gr.State([])
+    thumbnail_list = gr.State([])
+
+    def chat_handler(prompt, user_image, chat_history, thumbnails, 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 thumbnail list
+        if user_image:
+            thumbnails.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, thumbnails, ""
+
+        # Generate image based on user input
+        img, status = generate_image(prompt or "Generate an image", user_image, output_filename)
+        
+        # Add the status message to the chat history
+        chat_history.append({"role": "assistant", "content": status})
+        
+        # Add the generated image to the thumbnail list
+        if img:
+            thumbnails.append(img)
+        
+        return chat_history, user_image, img, thumbnails, ""
 
     # Connect the button to the chat handler
     generate_btn.click(
         fn=chat_handler,
-        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]
+        inputs=[prompt_input, image_input, chat_state, thumbnail_list, filename_input],
+        outputs=[chatbot, uploaded_image_output, generated_image_output, thumbnail_gallery, prompt_input]
     )
 
     # Also allow Enter key to submit
     prompt_input.submit(
         fn=chat_handler,
-        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]
+        inputs=[prompt_input, image_input, chat_state, thumbnail_list, filename_input],
+        outputs=[chatbot, uploaded_image_output, generated_image_output, thumbnail_gallery, prompt_input]
     )
 
 if __name__ == "__main__":