Update app.py

app.py

@@ -5,25 +5,55 @@ import requests
 import asyncio
 import aiohttp  # Use aiohttp for asynchronous HTTP requests
 import gradio as gr  # Import Gradio
-
+import google.generativeai as genai
+from io import BytesIO
 # --- Environment Variables & Setup ---
 DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
-
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 # --- Discord Bot Setup ---
 intents = discord.Intents.default()
 client = discord.Client(intents=intents)
 tree = app_commands.CommandTree(client)
 
-
-
-
-
-
-
-
-
-
-
+if not DISCORD_BOT_TOKEN or not GEMINI_API_KEY:
+    raise ValueError("Both DISCORD_BOT_TOKEN and GEMINI_API_KEY must be set.")
+
+
+@tree.command(name="imagen3", description="Generates iamge(s) using imagen3 model")
+@app_commands.choices(
+    aspect_ratio=[
+        app_commands.Choice(name="1:1 (Square)", value="1:1"),
+        app_commands.Choice(name="9:16 (Vertical)", value="9:16"),
+        app_commands.Choice(name="16:9 (Horizontal)", value="16:9"),
+        app_commands.Choice(name="3:4", value="3:4"),
+        app_commands.Choice(name="4:3", value="4:3"),
+    ]
+)
+async def generate_command(
+    interaction: discord.Interaction,
+    prompt: str,
+    img_count:int,
+    aspect_ratio: app_commands.Choice[str],
+):
+    genai.configure(api_key=GEMINI_API_KEY)
+    imagen = genai.ImageGenerationModel("imagen-3.0-generate-001")
+
+    result = imagen.generate_images(
+        prompt=prompt,
+        number_of_images=img_count,
+        safety_filter_level="block_only_high",
+        person_generation="allow_adult",
+        aspect_ratio=aspect_ratio,
+    )
+    files = []
+    for i, image in enumerate(result.generated_images):
+        # Assuming `image` is a byte array or binary stream of the image
+        image_file = BytesIO(image)  # Convert to BytesIO
+        files.append(discord.File(fp=image_file, filename=f"image_{i+1}.png"))
+
+    # Send images to Discord
+    await interaction.followup.send(content=f"Here are your {img_count} images for prompt: `{prompt}`", files=files)
+    
 
 # --- Gradio Interface ---
 def echo_text(text):