imageai / app.py
Groundspout's picture
Update app.py
9d3ded5
import discord
from discord.ext import commands
import os
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch
import asyncio
import random
import threading
from concurrent.futures import ThreadPoolExecutor
from discord import app_commands
executor = ThreadPoolExecutor()
bot = commands.Bot(command_prefix='!', intents=discord.Intents.all())
model_id = "stabilityai/stable-diffusion-2-base"
queue = []
# Use the Euler scheduler here instead
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler)
pipe = pipe.to("cpu")
semaphore = threading.Semaphore(value=1)
def generate_image(prompt, id):
semaphore.acquire()
try:
image = pipe(prompt).images[0]
path = f"{prompt}{random.randint(1,10000000)}.png"
image.save(path)
file = discord.File(path)
finally:
semaphore.release()
queue.popleft()
asyncio.run_coroutine_threadsafe(update_queue_messages(), bot.loop)
return (file,path)
@bot.event
async def on_ready():
await bot.tree.sync()
@bot.tree.command(name="ai", description="Generate an image using AI, describe the image with up to 100 characters.")
async def generate(interaction:discord.Interaction, prompt: str=None):
if not prompt or len(prompt)>100:
await interaction.response.send_message("Please write a description of the image up to 100 characters.")
position = len(queue) + 1
await interaction.response.send_message(f"Generating image... may take some time")
response = await interaction.original_response()
value = await bot.loop.run_in_executor(executor, generate_image, prompt, interaction.user.id)
image = value[0]
if image:
await interaction.followup.send(content=prompt, file=image)
os.remove(value[1])
else:
await interaction.followup.send("there was an error, try again later...")
await response.edit(content="Finished generating")
bot.run(
'MTExMzEwNDIyNzIyMTM4MTE1MA.GoPxZb.d0w449b8n3DMxasP1-HHiJo7nRCT1YnHQ20hYY')