diff --git a/.gitignore b/.gitignore index b34f717..e8c9d7d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ .vscode/ venv/ *.secret +__pycache__/ +jmh/ 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/__init__.py b/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/modules/diffuseapi.py b/modules/diffuseapi.py new file mode 100644 index 0000000..e65d812 --- /dev/null +++ b/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"] \ No newline at end of file diff --git a/modules/diffuseapi.py_ref b/modules/diffuseapi.py_ref new file mode 100644 index 0000000..0fa1c52 --- /dev/null +++ b/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) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..3aed5b2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +discord.py +aiohttp \ No newline at end of file diff --git a/testing.ipynb b/testing.ipynb index f488cc2..d91af8c 100644 --- a/testing.ipynb +++ b/testing.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 32, + "execution_count": 5, "metadata": {}, "outputs": [], "source": [ @@ -13,7 +13,7 @@ }, { "cell_type": "code", - "execution_count": 57, + "execution_count": 6, "metadata": {}, "outputs": [], "source": [ @@ -40,24 +40,24 @@ "payload.update(override_payload)\n", "\n", "# url = \"https://art.jurydoak.com/sdapi/v1\"\n", - "url = \"http://localhost:7860/sdapi/v1\"" + "url = \"http://10.6.9.69:7860/sdapi/v1\"" ] }, { "cell_type": "code", - "execution_count": 58, + "execution_count": 7, "metadata": {}, "outputs": [ { "data": { "text/html": [ - "" + "" ], "text/plain": [ "" ] }, - "execution_count": 58, + "execution_count": 7, "metadata": {}, "output_type": "execute_result" } @@ -81,7 +81,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3.10.8 64-bit", + "display_name": "Python 3.11.0 ('venv': venv)", "language": "python", "name": "python3" }, @@ -95,12 +95,12 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.8" + "version": "3.11.0" }, "orig_nbformat": 4, "vscode": { "interpreter": { - "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a" + "hash": "aa19d1669203bff990bc2820f3654fa686baacd13bba0cbe075930f54e624f06" } } }, diff --git a/weeeabot.py b/weeeabot.py new file mode 100644 index 0000000..6dd2165 --- /dev/null +++ b/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) \ No newline at end of file