DiffuseAPI.py

@@ -0,0 +1,110 @@
+import aiohttp
+from sys import exit
+
+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=""):
+        settings = {
+        "filter_nsfw": not self.nsfw_enabled
+        }
+
+        override_payload = {
+            "override_settings": settings
+        }
+
+        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} | override_payload
+
+        sess = aiohttp.ClientSession(self.url)
+        alive = await sess.head('/')
+
+        if alive.status != 200:
+            alive.close()
+            await sess.close()
+            return None
+        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]

modules/diffuseapi.py

@@ -1,7 +1,5 @@
 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):
@@ -97,8 +95,7 @@ class DiffuseAPI():
         }
 
         settings = {
-        "filter_nsfw": not self.nsfw_enabled,
+        "filter_nsfw": not self.nsfw_enabled
-        "enable_pnginfo": False
         }
 
         override_payload = {
@@ -107,36 +104,18 @@ class DiffuseAPI():
 
         payload.update(override_payload)
 
-        async with aiohttp.ClientSession(self.url) as session:
+        sess = aiohttp.ClientSession(self.url)
-            async with session.head('/') as alive:
+        alive = await sess.head('/')
-                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):
+        if alive.status != 200:
-        payload = {
+            alive.close()
-            "resize_mode": 0,
+            await sess.close()
-            "show_extras_results": True,
+            return None
-            "gfpgan_visibility": 0,
+        request  = await sess.post("/sdapi/v1/txt2img", json=payload)
-            "codeformer_visibility": 0,
+        try:
-            "codeformer_weight": 0,
+            req_json = await request.json()
-            "upscaling_resize": 4,
+            request.close()
-            "upscaling_resize_w": 512,
+            await sess.close()
-            "upscaling_resize_h": 1024,
+            return req_json["images"][0]
-            "upscaling_crop": True,
+        except:
-            "upscaler_1": "R-ESRGAN 4x+ Anime6B",
+            return None
-            "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"]

weeeabot.py

@@ -1,11 +1,9 @@
 import discord
 from discord.ext import commands
 import aiohttp
-from base64 import b64decode, b64encode
+from base64 import b64decode
 from sys import exit
 from modules import diffuseapi
-import io
-import hashlib
 
 try:
     with open("token.secret", 'r') as f:
@@ -17,39 +15,10 @@ except FileNotFoundError:
 intents = discord.Intents.default()
 intents.message_content = True
 
-api = diffuseapi.DiffuseAPI("https://art.jurydoak.com", ["Bot"], False, 28)
+api = diffuseapi.DiffuseAPI("http://localhost:7860", ["default"], 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"]):
 
@@ -61,32 +30,79 @@ class Settings:
 
 
 main_settings = Settings()
-activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")
+activity = discord.Activity(type=discord.ActivityType.listening, name="!help")
 
 
 bot = commands.Bot(intents=intents, command_prefix="!", activity=activity, help_command=commands.DefaultHelpCommand())
 
+async def send_request(ctx, endpoint, payload):
+    global num_steps
+    global ai_seed
+    sess = aiohttp.ClientSession(main_settings.url)
+    alive = await sess.head('/')
+
+    if alive.status != 200:
+        await ctx.reply("I'm sorry but it appears my brain is offline")
+        alive.close()
+        await sess.close()
+        return {
+            'images': [None]
+        }
+
+    print(payload)
+    request  = await sess.post(endpoint, json=payload)
+    try:
+        req_json = await request.json()
+    except:
+        print(await request.text())
+        return {
+            'images': [None]
+        }
+    request.close()
+    await sess.close()
+    return req_json
+
 @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)
+    payload = {
-    view = Confirm()
+        "prompt": prompt,
-    if image_data is None:
+        "styles": main_settings.styles,
+        "steps": main_settings.num_steps,
+        "seed": main_settings.ai_seed,
+        "n_iter": 1,
+        "height": 1024,
+        "negative_prompts": "nsfw, not safe for work, nudity, multiple keyboards",
+        "cfg_scale": 12
+    }
+
+    settings = {
+        "filter_nsfw": not main_settings.nsfw_enabled,
+        "samples_save": True,
+    }
+
+    override_payload = {
+        "override_settings": settings
+    }
+
+    payload.update(override_payload)
+    image_data = await send_request(ctx, endpoint="/sdapi/v1/txt2img", payload=payload)
+
+    if image_data['images'][0] 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~':
+    with open('/tmp/image.png', 'wb') as f:
-        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"))
+        f.write(b64decode(image_data['images'][0]))
-        return
+
-    print(h.digest())
+    embed = discord.Embed()
-    data = io.BytesIO(decoded_image)
+    upload_file = discord.File("/tmp/image.png", filename="image.png")
-    await ctx.reply("", file=discord.File(data, "_".join(args) + ".png"), view=view)
+    embed.set_image(url="attachment://image.png")
+    # embed.title = prompt
+
+    await ctx.reply("", file=upload_file, embed=embed)
 
 
 @bot.command()
@@ -98,9 +114,9 @@ async def seed(ctx, arg):
         pass
 
     # global ai_seed
-    api.set_seed(arg)
+    main_settings.ai_seed = arg
 
-    await ctx.reply(f"I have updated the seed to {api.seed} for you my master.")
+    await ctx.reply(f"I have updated the seed to {main_settings.ai_seed} for you my master.")
 
 @bot.command()
 async def steps(ctx, arg):
@@ -115,9 +131,9 @@ async def steps(ctx, arg):
         return
 
     # global num_steps
-    api.set_steps(arg)
+    main_settings.num_steps = arg
 
-    await ctx.reply(f"I have updated the steps to {api.num_steps} for you my master.")
+    await ctx.reply(f"I have updated the steps to {main_settings.num_steps} for you my master.")
 
 @bot.command()
 async def settings(ctx):
@@ -125,37 +141,37 @@ async def settings(ctx):
     global ai_seed, steps
     settings = f"""
     ```
-    seed: {api.seed}
+    seed: {main_settings.ai_seed}
-    steps: {api.steps}
+    steps: {main_settings.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":
+    if main_settings.url != "http://localhost:7860":
-        api.url = "http://localhost:7860"
+        main_settings.url = "http://localhost:7860"
-        api.styles = ["default"]
+        main_settings.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"
+        main_settings.url = "https://art.jurydoak.com"
-        api.styles = ["Bot"]
+        main_settings.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.command()
+async def test2(ctx):
+    image_data = await api.generate_image(prompt="cute bot doing bot things")
+    if image_data is None:
+        return
+    with open('/tmp/image.png', 'wb') as f:
+        f.write(b64decode(image_data))
 
+    embed = discord.Embed()
+    upload_file = discord.File("/tmp/image.png", filename="image.png")
+    embed.set_image(url="attachment://image.png")
+    # embed.title = prompt
 
+    await ctx.reply("", file=upload_file, embed=embed)
 
 bot.run(discord_client_token)