.gitignore

@@ -1,5 +1,4 @@
 .vscode/
 venv/
 *.secret
-__pycache__/
-jmh/
+__pycache__/

DiffuseAPI.py

@@ -115,7 +115,7 @@ class DiffuseAPI():
         async with aiohttp.ClientSession(self.url) as session:
             async with session.head('/') as alive:
                 if alive.status != 200:
-                    return None
+                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]

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\": self.styles,\n",
+    "        \"steps\": self.num_steps,\n",
+    "        \"seed\": self.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(self.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
+}

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)