diff --git a/DiffuseAPI.py b/DiffuseAPI.py new file mode 100644 index 0000000..5164f75 --- /dev/null +++ b/DiffuseAPI.py @@ -0,0 +1,128 @@ +import aiohttp +from base64 import b64decode +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=""): + payload = { + "prompt": prompt, + "styles": main_settings.styles, + "steps": main_settings.num_steps, + "seed": main_settings.ai_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 + } + + override_payload = { + "override_settings": settings + } + + payload.update(override_payload) + + sess = aiohttp.ClientSession(main_settings.url) + alive = await sess.head('/') + + 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 + + + + + + diff --git a/classweeb.ipynb b/classweeb.ipynb new file mode 100644 index 0000000..87826af --- /dev/null +++ b/classweeb.ipynb @@ -0,0 +1,183 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import aiohttp\n", + "from base64 import b64decode\n", + "from sys import exit" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class DiffuseAPI():\n", + " def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):\n", + "\n", + " self.nsfw_enabled = nsfw_enabled\n", + " self.num_steps = num_steps\n", + " self.url = url\n", + " self.styles = styles\n", + " self.seed = -1\n", + " self.width = 512\n", + " self.height = 1024\n", + " self.cfg_scale = 12\n", + "\n", + " def set_steps(self, steps):\n", + " try:\n", + " new_steps = int(steps)\n", + " if 0 > new_steps < 50:\n", + " self.num_steps = new_steps\n", + " return True\n", + " return False\n", + " except:\n", + " return False\n", + "\n", + " def set_seed(self, seed):\n", + " try:\n", + " new_seed = int(seed)\n", + " self.seed = new_seed\n", + " return True\n", + " except:\n", + " return False\n", + " \n", + " def set_cfg_scale(self, scale):\n", + " try:\n", + " new_scale = int(scale)\n", + " if 0 > new_scale < 30:\n", + " self.cfg_scale = new_scale\n", + " return True\n", + " return False\n", + " except:\n", + " return False\n", + " \n", + " def set_styles(self, styles):\n", + " if type styles = list:\n", + " self.styles = styles\n", + " return True\n", + " return False\n", + "\n", + " def set_orientation(self, orientation):\n", + " new_orientation = str(orientation)\n", + " match new_orientation:\n", + " case \"portrait\":\n", + " self.width = 512\n", + " self.height = 1024\n", + " return True\n", + " case \"landscape\":\n", + " self.width = 1024\n", + " self.height = 512\n", + " return True\n", + " case \"square\":\n", + " self.width = 512\n", + " self.height = 512\n", + " return True\n", + " case _:\n", + " return False\n", + " \n", + " def get_orientation(self):\n", + " if self.width == 512 and self.height == 512:\n", + " return (\"square\", self.width, self.height)\n", + " elif self.width == 1024:\n", + " return (\"landscape\", self.width, self.height)\n", + " return (\"portrait\", self.width, self.height)\n", + " \n", + " def set_nsfw_filter(self, filter_state):\n", + " if type filter_state = bool:\n", + " self.nsfw_enabled = filter_state\n", + " return True\n", + " return False\n", + " \n", + " def get_nsfw_filter(self):\n", + " return self.nsfw_enabled\n", + "\n", + " async def generate_image(self, prompt, neg_prompt=\"\"):\n", + " payload = {\n", + " \"prompt\": prompt,\n", + " \"styles\": main_settings.styles,\n", + " \"steps\": main_settings.num_steps,\n", + " \"seed\": main_settings.ai_seed,\n", + " \"n_iter\": 1,\n", + " \"height\": self.height,\n", + " \"width\": self.width,\n", + " \"negative_prompts\": neg_prompt,\n", + " \"cfg_scale\": self.cfg_scale\n", + " }\n", + "\n", + " settings = {\n", + " \"filter_nsfw\": not self.nsfw_enabled\n", + " }\n", + "\n", + " override_payload = {\n", + " \"override_settings\": settings\n", + " }\n", + "\n", + " payload.update(override_payload)\n", + "\n", + " sess = aiohttp.ClientSession(main_settings.url)\n", + " alive = await sess.head('/')\n", + "\n", + " if alive.status != 200:\n", + " alive.close()\n", + " await sess.close()\n", + " return None\n", + " request = await sess.post(\"/sdapi/v1/txt2img\", json=payload)\n", + " try:\n", + " req_json = await request.json()\n", + " request.close()\n", + " await sess.close()\n", + " return req_json[\"images\", [0]]\n", + " except:\n", + " return None\n", + "\n", + "\n", + " \n", + "\n", + " \n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "return image[\"images\", [0]]" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3.11.0 ('venv': venv)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.0" + }, + "orig_nbformat": 4, + "vscode": { + "interpreter": { + "hash": "aa19d1669203bff990bc2820f3654fa686baacd13bba0cbe075930f54e624f06" + } + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/test.py b/test.py new file mode 100644 index 0000000..7b04ed9 --- /dev/null +++ b/test.py @@ -0,0 +1,159 @@ +import discord +from discord.ext import commands +import aiohttp +from base64 import b64decode +from sys import exit + +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 + +# client = discord.Client(intents=intents) + +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") + + +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: + await ctx.reply("Something went wrong, please report this to the admin so it can be ignored") + return + + 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) + + +@bot.command() +async def seed(ctx, arg): + '''Set the seed for the image generation''' + try: + arg = int(arg) + except: + pass + + # global ai_seed + main_settings.ai_seed = arg + + 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): + '''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 + main_settings.num_steps = arg + + await ctx.reply(f"I have updated the steps to {main_settings.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: {main_settings.ai_seed} + steps: {main_settings.num_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") + +bot.run(discord_client_token) \ No newline at end of file