Spaces:
Runtime error
Runtime error
File size: 2,142 Bytes
0daa710 55e02db 10d711e 0daa710 9d3ded5 0daa710 9d3ded5 0daa710 9d3ded5 0daa710 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
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')
|