#import gradio as gr
#gr.load("models/mistralai/Mistral-7B-Instruct-v0.3").launch()
import os
import requests
import discord
from discord.ext import commands
from dotenv import load_dotenv
import logging
import io
import time
from PIL import Image
import torch
from diffusers import (
    StableDiffusionXLPipeline, 
    KDPM2AncestralDiscreteScheduler,
    AutoencoderKL
)

# Load environment variables from the .env file
load_dotenv()

DISCORD_TOKEN = os.getenv('dsTOK')
HF_API_KEY = os.getenv('HFREAD')

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load VAE component and configure the pipeline
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipe = StableDiffusionXLPipeline.from_pretrained("Corcelio/mobius", vae=vae, torch_dtype=torch.float16)
pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.to('cuda')

# Initialize the Discord bot with a default prefix
intents = discord.Intents.default()
intents.messages = True
intents.message_content = True
default_prefix = "!"
bot = commands.Bot(command_prefix=commands.when_mentioned_or(default_prefix), intents=intents)

@bot.event
async def on_ready():
    logger.info(f'Bot is ready. Logged in as {bot.user}')

@bot.command(name='gen')
@commands.cooldown(rate=1, per=10, type=commands.BucketType.user)
async def generate(ctx, *description):
    """
    Command to generate an image based on a description using the Stable Diffusion XL model.
    """
    if not description:
        await ctx.send("You need to provide a description for the image.")
        return

    prompt = " ".join(description)
    negative_prompt = "bad quality, low resolution, ugly, worse quality, grayscale, blurry"

    try:
        # Generate the image
        image = pipe(
            prompt, 
            negative_prompt=negative_prompt, 
            width=512,
            height=512,
            guidance_scale=5,
            num_inference_steps=25,
            clip_skip=3
        ).images[0]

        with io.BytesIO() as image_binary:
            image.save(image_binary, 'PNG')
            image_binary.seek(0)
            await ctx.send(file=discord.File(fp=image_binary, filename='generated_image.png'))
    except Exception as e:
        logger.error(f"Error processing the generate command: {e}")
        await ctx.send("An error occurred while processing your request.")

@generate.error
async def generate_error(ctx, error):
    if isinstance(error, commands.CommandOnCooldown):
        await ctx.send(f"This command is on cooldown. Please try again after {int(error.retry_after)} seconds.")
    else:
        logger.error(f"Error in generate command: {error}")
        await ctx.send("An error occurred. Please try again later.")

@bot.command(name='setprefix')
@commands.has_permissions(administrator=True)
async def set_prefix(ctx, prefix: str):
    bot.command_prefix = commands.when_mentioned_or(prefix)
    await ctx.send(f"Command prefix changed to: {prefix}")

@bot.command(name='help_custom')
async def help_custom(ctx):
    help_text = (
        "Here are the commands you can use:\n"
        "!generate <description> - Generate an image based on the description.\n"
        "!setprefix <prefix> - Change the command prefix (admin only).\n"
        "!help_custom - Display this help message."
    )
    await ctx.send(help_text)

# Run the bot
bot.run(DISCORD_TOKEN)