imagen3 / app.py
artintel235's picture
Update app.py
1ce3dab verified
raw
history blame
6.92 kB
import discord
from discord import app_commands
import os
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
import PIL.Image
import base64
from typing import List
# --- 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)
genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel("gemini-2.0-flash-exp")
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="hello", description="Says hello!")
async def hello_command(interaction):
await interaction.response.send_message("Hello there!")
@tree.command(name="gemini", description="Chat with Gemini")
async def generate_command(
interaction: discord.Interaction,
prompt: str,
):
try:
# Defer the interaction to allow time for processing
await interaction.response.defer()
genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel("gemini-2.0-flash-exp")
content = [] # List to store the content parts
# Handle attachments if the command is used in reply to a message
if interaction.message.reference:
ref_message = await interaction.channel.fetch_message(interaction.message.reference.message_id)
if ref_message.attachments: # Check if the replied-to message has attachments
for attachment in ref_message.attachments:
if not attachment.content_type or not attachment.content_type.startswith("image/"):
await interaction.followup.send(f"The attachment {attachment.filename} is not an image.", ephemeral=True)
continue
# Download each attachment
async with aiohttp.ClientSession() as session:
async with session.get(attachment.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: {attachment.filename}, status code: {resp.status}"
)
return
# Add the text prompt to the content
content.append(prompt)
# Call the Generative AI model
response = model.generate_content(content, stream=True)
current_message = None
current_message_content = ""
# Process streaming responses
for part in response:
text_chunk = part.text
if len(current_message_content) + len(text_chunk) <= 2000:
current_message_content += text_chunk
if current_message:
await current_message.edit(content=current_message_content)
else:
current_message = await interaction.followup.send(content=current_message_content)
else:
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
current_message = await interaction.followup.send(content=current_message_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)
except Exception as e:
print(e)
await interaction.followup.send(f"An error occurred: {e}")
@tree.command(name="discribe", description="Process an image from a message")
async def process_image(interaction: discord.Interaction):
try:
# Ensure the command is replying to a message
if not interaction.message.reference:
await interaction.response.send_message("Please reply to a message with an image to use this command.", ephemeral=True)
return
# Fetch the original message being replied to
ref_message = await interaction.channel.fetch_message(interaction.message.reference.message_id)
# Check if the original message has attachments
if not ref_message.attachments:
await interaction.response.send_message("The replied-to message does not contain any attachments.", ephemeral=True)
return
# Process the first attachment (assuming it's an image)
attachment = ref_message.attachments[0]
if not attachment.content_type.startswith("image/"):
await interaction.response.send_message("The attachment is not an image.", ephemeral=True)
return
# Defer the response to allow time for processing
await interaction.response.defer()
# Download and process the image as needed
async with aiohttp.ClientSession() as session:
async with session.get(attachment.url) as resp:
if resp.status == 200:
image_bytes = await resp.read()
# Process the image bytes as needed
# For example, pass them to your AI model
else:
await interaction.followup.send(f"Failed to download the image: {attachment.filename}")
return
# After processing, send a response
await interaction.followup.send("Image processed successfully.")
except Exception as e:
print(f"An error occurred: {e}")
await interaction.followup.send(f"An error occurred: {e}")
async def on_ready():
await tree.sync()
print("Bot is ready!")
client.event(on_ready)
# --- Gradio Interface ---
def echo_text(text):
return text
def run_gradio():
gr.Interface(
fn=echo_text,
inputs="text",
outputs="text",
live=False,
title="Minimal Gradio Interface",
).launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=True)
async def main():
bot_task = asyncio.create_task(client.start(DISCORD_BOT_TOKEN))
gradio_task = asyncio.to_thread(run_gradio)
await asyncio.gather(bot_task, gradio_task)
if __name__ == "__main__":
asyncio.run(main())