diff --git a/images/404.jpg b/images/404.jpg new file mode 100644 index 0000000..716b2d7 Binary files /dev/null and b/images/404.jpg differ diff --git a/modules/diffuseapi.py b/modules/diffuseapi.py index e34b8da..e65d812 100644 --- a/modules/diffuseapi.py +++ b/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): @@ -40,7 +42,7 @@ class DiffuseAPI(): return False except: return False - + def set_styles(self, styles): if type(styles) == list: self.styles = styles @@ -64,20 +66,20 @@ class DiffuseAPI(): 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 @@ -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) - alive = await sess.head('/') + 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] - if alive.status != 200: - alive.close() - await sess.close() - return None - request = await sess.post("/sdapi/v1/txt2img", json=payload) - try: - req_json = await request.json() - request.close() - await sess.close() - return req_json["images"][0] - except: - return None + 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"] \ No newline at end of file diff --git a/weeeabot.py b/weeeabot.py index b339e73..6dd2165 100644 --- a/weeeabot.py +++ b/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 = { - "prompt": prompt, - "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: + 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() - with open('/tmp/image.png', 'wb') as f: - f.write(b64decode(image_data['images'][0])) - - 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) + 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() @@ -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} - steps: {main_settings.num_steps} + seed: {api.seed} + steps: {api.steps} ``` """ await ctx.message.channel.send(settings) @bot.command() -async def test(ctx): - '''Test function, currently changes the URL and bot settings''' - if main_settings.url != "http://localhost:7860": - main_settings.url = "http://localhost:7860" - main_settings.styles = ["default"] - await ctx.reply("Set to fastboi") - else: - main_settings.url = "https://art.jurydoak.com" - main_settings.styles = ["Bot"] - await ctx.reply("Set to main api") +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 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)) +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") + - 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) \ No newline at end of file