Update app.py

app.py

@@ -1,168 +1,29 @@
-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
-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()
-
-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 get_file_url(filepath):
-    # Convert the relative file path to a URL that Gradio can serve
-    # Gradio serves files using the /file/ prefix
-    return f"/file/{filepath}"
-
-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
-            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(user_input, user_image, chat_history):
-    # Add user message to chat history
-    if user_image:
-        # Save the uploaded image to the static directory
-        img_path = save_image_to_static(user_image, prefix="uploaded")
-        # Convert the file path to a URL that Gradio can serve
-        img_url = get_file_url(img_path)
-        # Add the image URL to the chat history
-        chat_history.append({"role": "user", "content": img_url})
-    
-    # Add the text prompt to the chat history
-    if user_input:
-        chat_history.append({"role": "user", "content": user_input})
-    
-    # If no input (neither text nor image), return early
-    if not user_input and not user_image:
-        chat_history.append({"role": "assistant", "content": "Please provide a prompt or an image."})
-        return chat_history, None, ""
-
-    # Generate image based on user input
-    img, status = generate_image(user_input or "Generate an image", user_image)
-    
-    # Add AI response to chat history
-    if img:
-        # Save the generated image to the static directory
-        img_path = save_image_to_static(img, prefix="generated")
-        # Convert the file path to a URL that Gradio can serve
-        img_url = get_file_url(img_path)
-        # Add the image URL to the chat history
-        chat_history.append({"role": "assistant", "content": img_url})
-    
-    # Add the status message
-    chat_history.append({"role": "assistant", "content": status})
-    
-    return chat_history, None, ""
-
-# Create Gradio interface with chatbot layout
+# Alternative layout using a separate gr.Image component
 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 the conversation thread
+    # Chatbot display area for text messages
     chatbot = gr.Chatbot(
         label="Chat",
         height=300,
-        type="messages",  # Explicitly set to 'messages' format
-        avatar_images=(None, None)  # No avatars for simplicity
+        type="messages",
+        avatar_images=(None, None)
     )
     
+    # Separate image outputs
+    with gr.Row():
+        uploaded_image_output = gr.Image(label="Uploaded Image")
+        generated_image_output = gr.Image(label="Generated Image")
+    
     # Input area
     with gr.Row():
-        # Image upload button
         image_input = gr.Image(
             label="Upload Image",
             type="pil",
             scale=1,
             height=100
         )
-        # Text input
         prompt_input = gr.Textbox(
             label="",
             placeholder="Type something",
@@ -170,24 +31,41 @@ with gr.Blocks(title="Image Editing Chatbot") as demo:
             container=False,
             scale=3
         )
-        # Run button
         run_btn = gr.Button("Run", scale=1)
     
     # State to maintain chat history
     chat_state = gr.State([])
 
+    def chat_handler(user_input, user_image, chat_history):
+        # Add user message to chat history
+        if user_input:
+            chat_history.append({"role": "user", "content": user_input})
+        
+        # If no input, return early
+        if not user_input and not user_image:
+            chat_history.append({"role": "assistant", "content": "Please provide a prompt or an image."})
+            return chat_history, None, user_image, None, ""
+
+        # Generate image based on user input
+        img, status = generate_image(user_input or "Generate an image", user_image)
+        
+        # Add AI response to chat history
+        chat_history.append({"role": "assistant", "content": status})
+        
+        return chat_history, None, user_image, img, ""
+
     # Connect the button to the chat handler
     run_btn.click(
         fn=chat_handler,
         inputs=[prompt_input, image_input, chat_state],
-        outputs=[chatbot, image_input, prompt_input]
+        outputs=[chatbot, image_input, uploaded_image_output, generated_image_output, prompt_input]
     )
 
     # Also allow Enter key to submit
     prompt_input.submit(
         fn=chat_handler,
         inputs=[prompt_input, image_input, chat_state],
-        outputs=[chatbot, image_input, prompt_input]
+        outputs=[chatbot, image_input, uploaded_image_output, generated_image_output, prompt_input]
     )
 
 if __name__ == "__main__":