File size: 3,681 Bytes
d48ca0f
 
 
 
211da2c
 
d48ca0f
df9d99f
ae1e871
11be39b
ae1e871
d48ca0f
211da2c
d48ca0f
 
c0fe403
211da2c
d48ca0f
36b3803
211da2c
 
df9d99f
 
 
 
b05ac74
11be39b
 
 
 
 
 
 
 
 
 
 
 
 
211da2c
df9d99f
211da2c
 
 
df9d99f
 
211da2c
 
 
df9d99f
211da2c
ae1e871
df9d99f
36b3803
211da2c
ae1e871
211da2c
ae1e871
 
03f3768
 
36b3803
e4934d1
ae1e871
 
 
 
 
36b3803
 
 
 
ae1e871
06b452a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11be39b
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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
#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

# 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/Corcelio/mobius"
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, max_retries=5, delay=5):
    retries = 0
    while retries < max_retries:
        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}. Retrying in {delay} seconds...")
            time.sleep(delay)
            retries += 1
    logger.error(f"Failed to query the API after {max_retries} attempts.")
    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))
            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'))
        else:
            await ctx.send("An error occurred while generating the image. Please try again later.")
    except Exception as e:
        logger.error(f"Error processing the generate command: {e}")
        await ctx.send("An error occurred while processing your request. Please try again later.")

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