biiig changes

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

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=""):
+    async def generate_image(self, prompt, neg_prompt="", payload=None):
-        payload = {
+        if payload is None:
-        "prompt": prompt,
+            payload = {
-        "styles": self.styles,
+            "prompt": prompt,
-        "steps": self.num_steps,
+            "styles": self.styles,
-        "seed": self.seed,
+            "steps": self.num_steps,
-        "n_iter": 1,
+            "seed": self.seed,
-        "height": self.height,
+            "n_iter": 1,
-        "width": self.width,
+            "height": self.height,
-        "negative_prompts": neg_prompt,
+            "width": self.width,
-        "cfg_scale": self.cfg_scale
+            "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,

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