Merge pull request 'stable-diffuse-api' (#2) from stable-diffuse-api into master

Reviewed-on: #2

.gitignore

@@ -1,3 +1,5 @@
 .vscode/
 venv/
 *.secret
+__pycache__/
+jmh/

modules/diffuseapi.py

@@ -0,0 +1,142 @@
+import aiohttp
+from base64 import b64decode
+from sys import exit
+from pprint import pprint
+
+class DiffuseAPI():
+    def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):
+
+        self.nsfw_enabled = nsfw_enabled
+        self.num_steps = num_steps
+        self.url = url
+        self.styles = styles
+        self.seed = -1
+        self.width = 512
+        self.height = 1024
+        self.cfg_scale = 12
+
+    def set_steps(self, steps):
+        try:
+            new_steps = int(steps)
+            if 0 > new_steps < 50:
+                self.num_steps = new_steps
+                return True
+            return False
+        except:
+            return False
+
+    def set_seed(self, seed):
+        try:
+            new_seed = int(seed)
+            self.seed = new_seed
+            return True
+        except:
+            return False
+    
+    def set_cfg_scale(self, scale):
+        try:
+            new_scale = int(scale)
+            if 0 > new_scale < 30:
+                self.cfg_scale = new_scale
+                return True
+            return False
+        except:
+            return False
+
+    def set_styles(self, styles):
+        if type(styles) == list:
+            self.styles = styles
+            return True
+        return False
+
+    def set_orientation(self, orientation):
+        new_orientation = str(orientation)
+        match new_orientation:
+            case "portrait":
+                self.width = 512
+                self.height = 1024
+                return True
+            case "landscape":
+                self.width = 1024
+                self.height = 512
+                return True
+            case "square":
+                self.width = 512
+                self.height = 512
+                return True
+            case _:
+                return False
+
+    def get_orientation(self):
+        if self.width == 512 and self.height == 512:
+            return ("square", self.width, self.height)
+        elif self.width == 1024:
+            return ("landscape", self.width, self.height)
+        return ("portrait", self.width, self.height)
+
+    def set_nsfw_filter(self, filter_state):
+        if type(filter_state) == bool:
+            self.nsfw_enabled = filter_state
+            return True
+        return False
+
+    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
+        }
+
+        settings = {
+        "filter_nsfw": not self.nsfw_enabled,
+        "enable_pnginfo": False
+        }
+
+        override_payload = {
+            "override_settings": settings
+        }
+
+        payload.update(override_payload)
+
+        async with aiohttp.ClientSession(self.url) as session:
+            async with session.head('/') as alive:
+                if alive.status != 200:
+                    return None
+            async with session.post("/sdapi/v1/txt2img", json=payload) as image_json:
+                image_data = await image_json.json()
+                return image_data["images"][0]
+
+    async def generate_upscale(self, image):
+        payload = {
+            "resize_mode": 0,
+            "show_extras_results": True,
+            "gfpgan_visibility": 0,
+            "codeformer_visibility": 0,
+            "codeformer_weight": 0,
+            "upscaling_resize": 4,
+            "upscaling_resize_w": 512,
+            "upscaling_resize_h": 1024,
+            "upscaling_crop": True,
+            "upscaler_1": "R-ESRGAN 4x+ Anime6B",
+            "upscaler_2": "None",
+            "extras_upscaler_2_visibility": 0,
+            "upscale_first": False,
+            "image": image
+        }
+
+        async with aiohttp.ClientSession(self.url) as session:
+            async with session.head('/') as alive:
+                if alive.status != 200:
+                    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"]

modules/diffuseapi.py_ref

@@ -0,0 +1,41 @@
+import aiohttp
+import asyncio
+
+class DiffuseAPI:
+    def __init__(self, url, nsfw_enabled, styles, steps):
+        self.url = url
+        self.nsfw_enabled = nsfw_enabled
+        self.styles = styles
+        self.steps = steps
+        self.seed = -1
+
+    def _generate_payload(self):
+        payload = {
+            "prompt": "",
+            "negative_prompt": "",
+            "steps": self.steps,
+            "seed": self.seed,
+            "styles": self.styles,
+            "height": 1024,
+            "width": 512
+        }
+        settings = {
+            "filter_nsfw": not self.nsfw_enabled,
+            "samples_save": True,
+        }
+
+        override_payload = {
+            "override_settings": settings
+        }
+        payload.update(override_payload)
+        return payload
+
+    async def generate_image(self, prompt, negative_prompt=""):
+        async with aiohttp.ClientSession(self.url) as session:
+            payload = self._generate_payload()
+            payload.update({"prompt": prompt, "negative_prompt": negative_prompt})
+            print(payload)
+            async with session.post('/sdapi/v1/txt2img', json=payload) as image_handler:
+                image_data = await image_handler.json()
+
+        print(image_data)

requirements.txt

@@ -0,0 +1,2 @@
+discord.py
+aiohttp

weeeabot.py

@@ -0,0 +1,161 @@
+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"], False, 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)
+    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)
+
+
+@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.set_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.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)