Update app.py

app.py

@@ -6,19 +6,25 @@ from google.genai import types
 import gradio as gr
 import io
 from PIL import Image
+import uuid
+
+# Create a static directory to store images
+STATIC_DIR = "static"
+if not os.path.exists(STATIC_DIR):
+    os.makedirs(STATIC_DIR)
 
 def save_binary_file(file_name, data):
     f = open(file_name, "wb")
     f.write(data)
     f.close()
 
-# Function to convert a PIL Image to a base64 string
-def image_to_base64(img):
-    img_byte_arr = io.BytesIO()
-    img.save(img_byte_arr, format="PNG")
-    img_bytes = img_byte_arr.getvalue()
-    base64_string = base64.b64encode(img_bytes).decode("utf-8")
-    return f"data:image/png;base64,{base64_string}"
+def save_image_to_static(img, prefix="image"):
+    # Generate a unique filename
+    filename = f"{prefix}_{uuid.uuid4().hex}.png"
+    filepath = os.path.join(STATIC_DIR, filename)
+    img.save(filepath, format="PNG")
+    # Return the relative path
+    return filepath
 
 def generate_image(prompt, image=None, output_filename="generated_image"):
     # Initialize client with the API key
@@ -96,10 +102,10 @@ def generate_image(prompt, image=None, output_filename="generated_image"):
 def chat_handler(user_input, user_image, chat_history):
     # Add user message to chat history
     if user_image:
-        # Convert the uploaded image to base64
-        img_base64 = image_to_base64(user_image)
+        # Save the uploaded image to the static directory
+        img_path = save_image_to_static(user_image, prefix="uploaded")
         # Add the image to the chat history
-        chat_history.append({"role": "user", "content": img_base64})
+        chat_history.append({"role": "user", "content": img_path})
     
     # Add the text prompt to the chat history
     if user_input:
@@ -115,10 +121,10 @@ def chat_handler(user_input, user_image, chat_history):
     
     # Add AI response to chat history
     if img:
-        # Convert the generated image to base64
-        img_base64 = image_to_base64(img)
+        # Save the generated image to the static directory
+        img_path = save_image_to_static(img, prefix="generated")
         # Add the image to the chat history
-        chat_history.append({"role": "assistant", "content": img_base64})
+        chat_history.append({"role": "assistant", "content": img_path})
     
     # Add the status message
     chat_history.append({"role": "assistant", "content": status})
@@ -133,7 +139,7 @@ with gr.Blocks(title="Image Editing Chatbot") as demo:
     # Chatbot display area for the conversation thread
     chatbot = gr.Chatbot(
         label="Chat",
-        height=300,  # Already reduced to 300
+        height=300,
         type="messages",  # Explicitly set to 'messages' format
         avatar_images=(None, None)  # No avatars for simplicity
     )