app.py

#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
from PIL import Image

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

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

API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
headers = {"Authorization": f"Bearer {HF_API_KEY}"}

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

# Function to query the Hugging Face model with a structured prompt
def query(payload):
    try:
        response = requests.post(API_URL, headers=headers, json=payload)
        response.raise_for_status()
        return response.content
    except requests.exceptions.RequestException as e:
        logger.error(f"Error querying the API: {e}")
        return None

# 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='generate')
@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 Hugging Face model.
    """
    if not description:
        await ctx.send("You need to provide a description for the image.")
        return

    description = " ".join(description)
    try:
        payload = {"inputs": description}
        image_bytes = query(payload)

        if image_bytes:
            image = Image.open(io.BytesIO(image_bytes))
            image.save("generated_image.png")
            await ctx.send(file=discord.File("generated_image.png"))
        else:
            await ctx.send("An error occurred while generating the image.")
    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)