Removed old python file

And then just touch me
So I can get my
Satisfaction

DiffuseAPI.py

@@ -1,110 +0,0 @@
-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,5 +1,7 @@
 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):
@@ -95,7 +97,8 @@ class DiffuseAPI():
         }
 
         settings = {
-        "filter_nsfw": not self.nsfw_enabled
+        "filter_nsfw": not self.nsfw_enabled,
+        "enable_pnginfo": False
         }
 
         override_payload = {
@@ -104,18 +107,36 @@ class DiffuseAPI():
 
         payload.update(override_payload)
 
-        sess = aiohttp.ClientSession(self.url)
+        async with aiohttp.ClientSession(self.url) as session:
-        alive = await sess.head('/')
+            async with session.head('/') as alive:
-
                 if alive.status != 200:
-            alive.close()
-            await sess.close()
                     return None
-        request  = await sess.post("/sdapi/v1/txt2img", json=payload)
+            async with session.post("/sdapi/v1/txt2img", json=payload) as image_json:
-        try:
+                image_data = await image_json.json()
-            req_json = await request.json()
+                return image_data["images"][0]
-            request.close()
+
-            await sess.close()
+    async def generate_upscale(self, image):
-            return req_json["images"][0]
+        payload = {
-        except:
+            "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"]

weeeabot.py

@@ -1,9 +1,11 @@
 import discord
 from discord.ext import commands
 import aiohttp
-from base64 import b64decode
+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:
@@ -15,10 +17,39 @@ except FileNotFoundError:
 intents = discord.Intents.default()
 intents.message_content = True
 
-api = diffuseapi.DiffuseAPI("http://localhost:7860", ["default"], True, 28)
+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"]):
 
@@ -30,79 +61,32 @@ class Settings:
 
 
 main_settings = Settings()
-activity = discord.Activity(type=discord.ActivityType.listening, name="!help")
+activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")
 
 
 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")
-    payload = {
+    image_data = await api.generate_image(prompt)
-        "prompt": prompt,
+    view = Confirm()
-        "styles": main_settings.styles,
+    if image_data is None:
-        "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()
 
-    with open('/tmp/image.png', 'wb') as f:
+    if digest == b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~':
-        f.write(b64decode(image_data['images'][0]))
+        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
-    embed = discord.Embed()
+    print(h.digest())
-    upload_file = discord.File("/tmp/image.png", filename="image.png")
+    data = io.BytesIO(decoded_image)
-    embed.set_image(url="attachment://image.png")
+    await ctx.reply("", file=discord.File(data, "_".join(args) + ".png"), view=view)
-    # embed.title = prompt
-
-    await ctx.reply("", file=upload_file, embed=embed)
 
 
 @bot.command()
@@ -114,9 +98,9 @@ async def seed(ctx, arg):
         pass
 
     # global ai_seed
-    main_settings.ai_seed = arg
+    api.set_seed(arg)
 
-    await ctx.reply(f"I have updated the seed to {main_settings.ai_seed} for you my master.")
+    await ctx.reply(f"I have updated the seed to {api.seed} for you my master.")
 
 @bot.command()
 async def steps(ctx, arg):
@@ -131,9 +115,9 @@ async def steps(ctx, arg):
         return
 
     # global num_steps
-    main_settings.num_steps = arg
+    api.set_steps(arg)
 
-    await ctx.reply(f"I have updated the steps to {main_settings.num_steps} for you my master.")
+    await ctx.reply(f"I have updated the steps to {api.num_steps} for you my master.")
 
 @bot.command()
 async def settings(ctx):
@@ -141,37 +125,37 @@ async def settings(ctx):
     global ai_seed, steps
     settings = f"""
     ```
-    seed: {main_settings.ai_seed}
+    seed: {api.seed}
-    steps: {main_settings.num_steps}
+    steps: {api.steps}
     ```
     """
     await ctx.message.channel.send(settings)
 
 @bot.command()
-async def test(ctx):
+async def upscale(ctx):
-    '''Test function, currently changes the URL and bot settings'''
+    """Upscale the attached image"""
-    if main_settings.url != "http://localhost:7860":
+    orig_image_data = await ctx.message.attachments[0].read()
-        main_settings.url = "http://localhost:7860"
+    new_data = await api.generate_upscale(b64encode(orig_image_data).decode('utf-8'))
-        main_settings.styles = ["default"]
+    data = io.BytesIO(b64decode(new_data))
-        await ctx.reply("Set to fastboi")
+    await ctx.reply("", file=discord.File(data, "upscaled_img.png"))
-    else:
-        main_settings.url = "https://art.jurydoak.com"
-        main_settings.styles = ["Bot"]
-        await ctx.reply("Set to main api")
 
 @bot.command()
-async def test2(ctx):
+async def test(ctx):
-    image_data = await api.generate_image(prompt="cute bot doing bot things")
+    '''Test function, currently changes the URL and bot settings'''
-    if image_data is None:
+    if api.url != "http://localhost:7860":
-        return
+        api.url = "http://localhost:7860"
-    with open('/tmp/image.png', 'wb') as f:
+        api.styles = ["default"]
-        f.write(b64decode(image_data))
+        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")
+
 
-    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)