Update app.py

app.py

@@ -8,6 +8,7 @@ import gradio as gr  # Import Gradio
 import google.generativeai as genai
 from io import BytesIO
 import PIL.Image
+import base64
 # --- Environment Variables & Setup ---
 DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
@@ -36,15 +37,30 @@ async def hello_command(interaction):
 async def generate_command(
     interaction: discord.Interaction,
     prompt: str,
-    # img_count: int,
-    # aspect_ratio: app_commands.Choice[str],
+    images: list[discord.Attachment] = None,
 ):
     try:
         await interaction.response.defer()  # Defer the interaction
         genai.configure(api_key=GEMINI_API_KEY)
         model = genai.GenerativeModel("gemini-2.0-flash-exp")
-        response = model.generate_content(prompt, stream=True)
-        current_message = None  # Initialize outside the loop to keep track of current message
+        content = []  # List to store the content parts
+        if images:
+            for image in images:
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(image.url) as resp:
+                         if resp.status == 200:
+                            image_bytes = await resp.read()
+                            base64_image = base64.b64encode(image_bytes).decode('utf-8')
+                            content.append({'mime_type':'image/jpeg', 'data': base64_image})
+
+                         else:
+                            await interaction.followup.send(f"Failed to download image: {image.filename}, status code: {resp.status}")
+                            return
+        content.append(prompt)
+
+        response = model.generate_content(content, stream=True)
+
+        current_message = None
         current_message_content = ""
 
         for part in response:
@@ -57,23 +73,21 @@ async def generate_command(
                 else:
                     current_message = await interaction.followup.send(content=current_message_content)
             else:
-                # Send the current message
                 if current_message:
                   await current_message.edit(content=current_message_content)
                 else:
                   await interaction.followup.send(content=current_message_content)
-                current_message_content = text_chunk  # Start new message with the current chunk
+                current_message_content = text_chunk
                 current_message = await interaction.followup.send(content=current_message_content)
 
-
-        # Send any remaining content
         if current_message_content:
            if current_message:
              await current_message.edit(content=current_message_content)
            else:
-             await interaction.followup.send(content=current_message_content)
+            await interaction.followup.send(content=current_message_content)
     except Exception as e:
         print(e)
+        await interaction.followup.send(f"An error occurred: {e}")
 
         
 async def on_ready():