From aa6d0bee6e528afa7cc7e684b9991bf41887fd27 Mon Sep 17 00:00:00 2001 From: Benjamyn Love Date: Thu, 15 Dec 2022 20:36:19 +1100 Subject: [PATCH] biiig changes --- commands/diffuse_cog.py | 153 ++++++++++++++++++++++++++++++++++++++++ modules/diffuseapi.py | 36 +++++----- slash_commands.py | 45 ++++++++++++ 3 files changed, 218 insertions(+), 16 deletions(-) create mode 100644 commands/diffuse_cog.py create mode 100644 slash_commands.py diff --git a/commands/diffuse_cog.py b/commands/diffuse_cog.py new file mode 100644 index 0000000..51608aa --- /dev/null +++ b/commands/diffuse_cog.py @@ -0,0 +1,153 @@ +import discord +import hashlib +from discord import app_commands +from discord.ext import commands +from modules.diffuseapi import DiffuseAPI +from base64 import b64decode, b64encode +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", ["Bot"], True, 28) + self.user_settings = {} + + def _get_user_settings(self, user_id: int): + user_data = self.user_settings.get(user_id) + if user_data == None: + self.user_settings[user_id] = { + "prompt": "", + "styles": ["Bot"], + "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="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) + + # 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() + + if digest == b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~': + 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 + ### /Bonk check + + # Decode the image to bytes and send it in the response as a discord.File object + data = io.BytesIO(decoded_image) + await interaction.followup.send(prompt, file=discord.File(data, "prompt.png"), view=view) + + @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") + async def change_orientation(self, interaction: discord.Interaction, orientation: str): + 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) + + +async def setup(bot: commands.Bot): + await bot.add_cog(DiffuseCog(bot), guild=discord.Object(id=775297479783874570)) + + diff --git a/modules/diffuseapi.py b/modules/diffuseapi.py index e65d812..9b1d081 100644 --- a/modules/diffuseapi.py +++ b/modules/diffuseapi.py @@ -18,7 +18,7 @@ class DiffuseAPI(): def set_steps(self, steps): try: new_steps = int(steps) - if 0 > new_steps < 50: + if new_steps > 0 and new_steps <= 50: self.num_steps = new_steps return True return False @@ -83,22 +83,26 @@ class DiffuseAPI(): def get_nsfw_filter(self): return self.nsfw_enabled - async def generate_image(self, prompt, neg_prompt=""): - payload = { - "prompt": prompt, - "styles": self.styles, - "steps": self.num_steps, - "seed": self.seed, - "n_iter": 1, - "height": self.height, - "width": self.width, - "negative_prompts": neg_prompt, - "cfg_scale": self.cfg_scale - } + async def generate_image(self, prompt, neg_prompt="", payload=None): + if payload is None: + payload = { + "prompt": prompt, + "styles": self.styles, + "steps": self.num_steps, + "seed": self.seed, + "n_iter": 1, + "height": self.height, + "width": self.width, + "negative_prompts": neg_prompt, + "cfg_scale": self.cfg_scale + } + + payload['prompt'] = prompt settings = { "filter_nsfw": not self.nsfw_enabled, - "enable_pnginfo": False + "enable_pnginfo": False, +# "sd_hypernetwork": "anime_3" } override_payload = { @@ -122,7 +126,7 @@ class DiffuseAPI(): "gfpgan_visibility": 0, "codeformer_visibility": 0, "codeformer_weight": 0, - "upscaling_resize": 4, + "upscaling_resize": 2, "upscaling_resize_w": 512, "upscaling_resize_h": 1024, "upscaling_crop": True, @@ -139,4 +143,4 @@ class DiffuseAPI(): return None async with session.post("/sdapi/v1/extra-single-image", json=payload) as image_json: image_data = await image_json.json() - return image_data["image"] \ No newline at end of file + return image_data["image"] diff --git a/slash_commands.py b/slash_commands.py new file mode 100644 index 0000000..e6a09ff --- /dev/null +++ b/slash_commands.py @@ -0,0 +1,45 @@ +import discord +from discord import app_commands +from discord.ext import commands +import os +import asyncio + + +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) + + +class aclient(commands.Bot): + def __init__(self): + activity = discord.Activity(type=discord.ActivityType.listening, name="Waiting for your prompt :3 (Main API)") + super().__init__(intents=discord.Intents.default(), command_prefix="!", activity=activity) + + self.synced = False + + async def load(self): + for filename in os.listdir("./commands"): + if filename.endswith(".py"): + print(f"Loading commands.{filename[:-3]}") + await self.load_extension(f"commands.{filename[:-3]}") + + async def on_ready(self): + await self.wait_until_ready() + if not self.synced: + await tree.sync(guild=discord.Object(id=775297479783874570)) + self.synced = True + print(f"We have logged in as {self.user}.") + pass + +client = aclient() +tree = client.tree + +async def main(): + await client.load() + await client.start(discord_client_token) + +# client.run(discord_client_token) +asyncio.run(main()) \ No newline at end of file