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')