diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b34f717 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +.vscode/ +venv/ +*.secret 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..84e488f 100644 --- a/testing.ipynb +++ b/testing.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 32, + "execution_count": 1, "metadata": {}, "outputs": [], "source": [ @@ -13,14 +13,14 @@ }, { "cell_type": "code", - "execution_count": 57, + "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "payload = {\n", - " \"prompt\": \"view from behind web developer sitting in a chair, sfw\",\n", - " \"styles\": [\"default\"],\n", - " \"steps\": 28,\n", + " \"prompt\": \"cute girlhand on cheek eyes closed thinking hard, sfw\",\n", + " \"styles\": [\"Bot\"],\n", + " \"steps\": 70,\n", " \"seed\": -1,\n", " \"n_iter\": 1,\n", " \"height\": 1024,\n", @@ -39,25 +39,25 @@ "\n", "payload.update(override_payload)\n", "\n", - "# url = \"https://art.jurydoak.com/sdapi/v1\"\n", - "url = \"http://localhost:7860/sdapi/v1\"" + "url = \"https://art.jurydoak.com/sdapi/v1\"\n", + "# url = \"http://localhost:7860/sdapi/v1\"" ] }, { "cell_type": "code", - "execution_count": 58, + "execution_count": 3, "metadata": {}, "outputs": [ { "data": { "text/html": [ - "" + "" ], "text/plain": [ "" ] }, - "execution_count": 58, + "execution_count": 3, "metadata": {}, "output_type": "execute_result" } @@ -95,7 +95,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.8" + "version": "3.10.8 (main, Nov 1 2022, 14:18:21) [GCC 12.2.0]" }, "orig_nbformat": 4, "vscode": { diff --git a/weeeabot.py b/weeeabot.py new file mode 100644 index 0000000..7b04ed9 --- /dev/null +++ b/weeeabot.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