Merge pull request 'Merge bot code into module branch' (#1) from discord-bot-addition into stable-diffuse-api

Reviewed-on: #1

.gitignore

@@ -0,0 +1,3 @@
+.vscode/
+venv/
+*.secret

requirements.txt

@@ -0,0 +1,2 @@
+discord.py
+aiohttp

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)