Push me

And then just touch me
So I can get my
Satisfaction

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)