weeabot/commands/diffuse_cog.py

import discord
import hashlib
from discord import app_commands
from discord.ext import commands
from modules.diffuseapi import DiffuseAPI
from base64 import b64decode, b64encode
from typing import Literal
import io





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

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

    # When the confirm function is run we will send the image to the upscale endpoint
    @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 self.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()

    # The cancel function will delete the message
    @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 DiffuseCog(commands.Cog):
    def __init__(self, bot: commands.Bot):
        self.bot = bot
        self.api = DiffuseAPI("https://art.jurydoak.com", ["default"], True, 28)
        self.user_settings = {}
        self.prev_prompt = None


    def _get_user_settings(self, user_id: int):
        user_data = self.user_settings.get(user_id)
        if user_data is None:
            self.user_settings[user_id] = {
                "prompt": "",
                "styles": ["default"],
                "steps": 28,
                "seed": -1,
                "n_iter": 1,
                "height": 768,
                "width": 512,
                "negative_prompts": "",
                "cfg_scale": 12
            }

        user_data = self.user_settings.get(user_id)
        return user_data


    



    @app_commands.command(name="pregen", description="Generate an image appending your prompt to the end")
    async def pregenerate_image(self, interaction: discord.Interaction, prompt: str):
        """/pregen prompt (Generates an image using the prompt)"""
        # This can take a while so we defer the response, this also tells the client the bot is thinking
        await interaction.response.defer()

        # Create an instance of the Confirm view so we can use it in our response
        view = Confirm(self.api)
        new_prompt = f"{self.prev_prompt}, {prompt}"
        # Get any custom settings the user has entered
        user_data = self._get_user_settings(interaction.user.id)

        print(f"SENDING PROMPT: {new_prompt} from User: {interaction.user.display_name}")

        # Send the prompt + the users custom payload to the API
        image_data = await self.api.generate_image(prompt=new_prompt, payload=user_data)

        # print(image_data)
        # Verify that an image was returned
        if image_data is None:
            await interaction.followup.send("Failed to generate image, please bitch at admin :3")


        # ### Bonk check
        decoded_image = b64decode(image_data)
        # h = hashlib.md5()
        # h.update(decoded_image)
        # digest = h.digest()

        # return await self.bonk_check(interaction, digest)
        ### /Bonk check

        # Decode the image to bytes and send it in the response as a discord.File object
        data = io.BytesIO(decoded_image)
        self.prev_prompt = new_prompt
        await interaction.followup.send(new_prompt, file=discord.File(data, "prompt.png"), view=view)

    @app_commands.command(name="gen", description="Generate an image with the following prompt")
    async def generate_image(self, interaction: discord.Interaction, prompt: str):
        """/gen prompt (Generates an image using the prompt)"""
        # This can take a while so we defer the response, this also tells the client the bot is thinking
        await interaction.response.defer()

        # Create an instance of the Confirm view so we can use it in our response
        view = Confirm(self.api)

        # Get any custom settings the user has entered
        user_data = self._get_user_settings(interaction.user.id)

        print(f"SENDING PROMPT: {prompt} from User: {interaction.user.display_name}")

        # Send the prompt + the users custom payload to the API
        image_data = await self.api.generate_image(prompt=prompt, payload=user_data)

        # print(image_data)
        # Verify that an image was returned
        if image_data is None:
            await interaction.followup.send("Failed to generate image, please bitch at admin :3")


        # ### Bonk check
        decoded_image = b64decode(image_data)
        # h = hashlib.md5()
        # h.update(decoded_image)
        # digest = h.digest()

        # return await self.bonk_check(interaction, digest)
        ### /Bonk check

        # Decode the image to bytes and send it in the response as a discord.File object
        data = io.BytesIO(decoded_image)
        self.prev_prompt = prompt
        await interaction.followup.send(prompt, file=discord.File(data, "prompt.png"), view=view)

    # async def bonk_check(self, interaction, digest):
    #     digests = [b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~']
    #     if digest in digests :
    #         await interaction.followup.send("Were no stranger to lewds, but you know the rules, and so do I", file=discord.File("images/404.jpg", "404.jpg"))
    #         return

    @app_commands.command(name="steps", description="Set the amount of steps the bot will run")
    async def set_steps(self, interaction: discord.Interaction, steps: int):
        user_settings = self._get_user_settings(interaction.user.id)
        
        try:
            user_settings['steps'] = int(steps)
            await interaction.response.send_message("Set steps uWu", ephemeral=True)
        except Exception:
            await interaction.response.send_message("Failed to set steps, baka", ephemeral=True)


    @app_commands.command(name="api", description="Swaps between the main and test API", )
    async def swap_api(self, interaction: discord.Interaction):
        user_settings = self._get_user_settings(interaction.user.id)

        '''Test function, currently changes the URL and bot settings'''
        if self.api.url != "http://localhost:7860":
            self.api.url = "http://localhost:7860"
            user_settings['style'] = ["default"]
            self.api.styles = ["default"]
            activity = discord.Activity(type=discord.ActivityType.listening, name="Waiting for your prompt :3 (Test API)")
            await self.bot.change_presence(activity=activity)
            await interaction.response.send_message("Set to fastboi")
        else:
            self.api.url = "https://art.jurydoak.com"
            self.api.styles = ["Bot"]
            user_settings['style'] = ["Bot"]
            activity = discord.Activity(type=discord.ActivityType.listening, name="Waiting for your prompt :3 (Main API)")
            await self.bot.change_presence(activity=activity)
            await interaction.response.send_message("Set to main api")

    @app_commands.command(name="orientation", description="Set the orientation of the generated image")
    # @app_commands.choices([app_commands.Choice("portrait"), app_commands.Choice("landscape"), app_commands.Choice("square")])
    async def change_orientation(self, interaction: discord.Interaction, orientation: Literal["portrait", "landscape", "square"]):
        if orientation.lower() not in ["portrait", "landscape", "square"]:
            await interaction.response.send_message("You can only set the orientation to portrait, landscape or square")
            return
        
        user_data = self._get_user_settings(interaction.user.id)
        match orientation.lower():
            case "portrait":
                user_data['height'] = 768
                user_data['width'] = 512
            case "landscape":
                user_data['height'] = 512
                user_data['width'] = 768
            case "square":
                user_data['height'] = 512
                user_data['width'] = 512
        
        await interaction.response.send_message(f"Set the orientation to {orientation}", ephemeral=True)
    
    @app_commands.command(name="seed", description="Set a static seed (-1 is random)")
    async def set_seed(self, interaction: discord.Interaction, seed: int):
        try:
            seed = int(seed)
            user_settings = self._get_user_settings(interaction.user.id)
            user_settings["seed"] = seed
            await interaction.response.send_message(f"Set the seed to {seed}", ephemeral=True)

        except:
            await interaction.response.send_message("Failed to set the seed")

async def setup(bot: commands.Bot):
    await bot.add_cog(DiffuseCog(bot), guild=discord.Object(id=775297479783874570))