weeabot/weeeabot.py

import discord
from discord.ext import commands
import aiohttp
from base64 import b64decode, b64encode
from sys import exit
from modules import diffuseapi
import io
import hashlib

try:
    with open("token.secret", 'r') as f:
        discord_client_token = f.read()
except FileNotFoundError:
    print("Cannot locate token.secret please generate a token")
    exit(69)

intents = discord.Intents.default()
intents.message_content = True

api = diffuseapi.DiffuseAPI("https://art.jurydoak.com", ["Bot"], True, 28)

# client = discord.Client(intents=intents)

class Confirm(discord.ui.View):
    def __init__(self):
        super().__init__()
        self.value = None

    async def callback(self, interaction: discord.Interaction):
        await interaction.response.send_message("You wish to upscale this image!")

    # When the confirm button is pressed, set the inner value to `True` and
    # stop the View from listening to more input.
    # We also send the user an ephemeral message that we're confirming their choice.
    @discord.ui.button(label='Upscale', style=discord.ButtonStyle.green)
    async def confirm(self, interaction: discord.Interaction, button: discord.ui.Button):
        image_data = await interaction.message.attachments[0].read()
        await interaction.response.defer()
        upscaled_image = await api.generate_upscale(b64encode(image_data).decode('utf-8'))
        data = io.BytesIO(b64decode(upscaled_image))

        await interaction.followup.send("", file=discord.File(data, "upscaled.png"))
        self.stop()


    # This one is similar to the confirmation button except sets the inner value to `False`
    @discord.ui.button(label='Delete', style=discord.ButtonStyle.grey)
    async def cancel(self, interaction: discord.Interaction, button: discord.ui.Button):
        await interaction.message.delete()

        self.stop()

class Settings:
    def __init__(self, nsfw_enabled=True, num_steps=28, ai_seed=-1, url="https://art.jurydoak.com", styles=["Bot"]):

        self.nsfw_enabled = nsfw_enabled
        self.num_steps = num_steps
        self.ai_seed = ai_seed
        self.url = url
        self.styles = styles


main_settings = Settings()
activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")


bot = commands.Bot(intents=intents, command_prefix="!", activity=activity, help_command=commands.DefaultHelpCommand())

@bot.command()
async def prompt(ctx, *args):
    '''Generate an image with the provided prompt'''
    prompt = " ".join(args)
    user_roles = [role.name for role in ctx.message.author.roles]

    async with ctx.typing():
        if "Bot Ban" not in user_roles:
            # await ctx.reply("Generating your image boo")
            image_data = await api.generate_image(prompt)
            view = Confirm()
            if image_data is None:
                await ctx.reply("Something went wrong, please report this to the admin so it can be ignored")
                return
            decoded_image = b64decode(image_data)
            h = hashlib.md5()
            h.update(decoded_image)
            digest = h.digest()

            if digest == b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~':
                await ctx.reply("Were no stranger to lewds, but you know the rules, and so do I", file=discord.File("images/404.jpg", "404.jpg"))
                return
            print(h.digest())
            data = io.BytesIO(decoded_image)
            await ctx.reply("", file=discord.File(data, "_".join(args) + ".png"), view=view)
        else:
            await ctx.reply("NO!")


@bot.command()
async def seed(ctx, arg):
    '''Set the seed for the image generation'''
    try:
        arg = int(arg)
    except:
        pass

    # global ai_seed
    api.set_seed(arg)

    await ctx.reply(f"I have updated the seed to {api.seed} for you my master.")

@bot.command()
async def steps(ctx, arg):
    '''Set how many steps the AI will run (max 50)'''
    try:
        arg = int(arg)
    except:
        pass

    if arg > 50:
        await ctx.reply("I'm sorry Dave, I can't do that")
        return

    # global num_steps
    api.num_steps = arg

    await ctx.reply(f"I have updated the steps to {api.num_steps} for you my master.")

@bot.command()
async def settings(ctx):
    '''See the currently configured settings (BROKEN)'''
    global ai_seed, steps
    settings = f"""
    ```
    seed: {api.seed}
    steps: {api.num_steps}
    ```
    """
    await ctx.message.channel.send(settings)

@bot.command()
async def upscale(ctx):
    """Upscale the attached image"""
    orig_image_data = await ctx.message.attachments[0].read()
    new_data = await api.generate_upscale(b64encode(orig_image_data).decode('utf-8'))
    data = io.BytesIO(b64decode(new_data))
    await ctx.reply("", file=discord.File(data, "upscaled_img.png"))

@bot.command()
async def test(ctx):
    '''Test function, currently changes the URL and bot settings'''
    if api.url != "http://localhost:7860":
        api.url = "http://localhost:7860"
        api.styles = ["default"]
        activity = discord.Activity(type=discord.ActivityType.listening, name="!help fast")
        await bot.change_presence(activity=activity)
        await ctx.reply("Set to fastboi")
    else:
        api.url = "https://art.jurydoak.com"
        api.styles = ["Bot"]
        activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")
        await bot.change_presence(activity=activity)
        await ctx.reply("Set to main api")




bot.run(discord_client_token)