Initial code commit, basic functionality working

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+/venv/
+/settings.json
+
+/conf/*
+!/conf/example-workflow.json
+!/conf/example-workflow_settings.json

bot.py

@@ -0,0 +1,34 @@
+import discord
+
+from lib.settings import *
+from lib.events import *
+
+intents     = discord.Intents(messages=True, guilds=True, message_content=True, reactions=True)
+client      = discord.Client(intents=intents)
+settings    = {}
+
+@client.event
+async def on_message(msg):
+    await on_message_or_reaction(client, msg)
+
+@client.event
+async def on_raw_reaction_add(rxn):
+    ALLOW_REPEAT    = settings["allow_repeat"]
+    REPEAT_EMOJI    = settings["repeat_emoji"]
+
+    if ALLOW_REPEAT:
+        if rxn.emoji.name == REPEAT_EMOJI:
+            await on_message_or_reaction(client, rxn)
+
+@client.event
+async def on_ready():
+    print("READY.")
+
+def main():
+    global settings
+    settings = get_settings(initialize=True)
+    client.run(settings["token"])
+
+if __name__ == "__main__":
+    main()
+

comfyui-discord.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+INSTALLATION=0
+
+if [[ ! -d venv ]]
+then
+    python3 -m venv venv || python -m venv venv || (
+        echo "Could not create a Python virtual environment."
+        exit 1
+    )
+
+    INSTALLATION=1
+fi
+
+if [[ -f ./venv/bin/activate ]]; then source ./venv/bin/activate; fi
+if [[ -f ./venv/Scripts/activate ]]; then source ./venv/Scripts/activate; fi
+
+
+if [[ $INSTALLATION -eq 1 ]]
+then
+    pip install discord.py
+fi
+
+python -u bot.py

conf/example-workflow.json

@@ -0,0 +1,107 @@
+{
+  "3": {
+    "inputs": {
+      "seed": __SEED__,
+      "steps": __STEPS__,
+      "cfg": __CFG__,
+      "sampler_name": "euler",
+      "scheduler": "normal",
+      "denoise": 1,
+      "model": [
+        "4",
+        0
+      ],
+      "positive": [
+        "6",
+        0
+      ],
+      "negative": [
+        "7",
+        0
+      ],
+      "latent_image": [
+        "5",
+        0
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "KSampler"
+    }
+  },
+  "4": {
+    "inputs": {
+      "ckpt_name": "anything-v3-fp16-pruned.safetensors"
+    },
+    "class_type": "CheckpointLoaderSimple",
+    "_meta": {
+      "title": "Load Checkpoint"
+    }
+  },
+  "5": {
+    "inputs": {
+      "width": __WIDTH__,
+      "height": __HEIGHT__,
+      "batch_size": 1
+    },
+    "class_type": "EmptyLatentImage",
+    "_meta": {
+      "title": "Empty Latent Image"
+    }
+  },
+  "6": {
+    "inputs": {
+      "text": "__POSITIVE__",
+      "clip": [
+        "4",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "7": {
+    "inputs": {
+      "text": "__NEGATIVE_PREFIX__ __NEGATIVE__",
+      "clip": [
+        "4",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP Text Encode (Prompt)"
+    }
+  },
+  "8": {
+    "inputs": {
+      "samples": [
+        "3",
+        0
+      ],
+      "vae": [
+        "4",
+        2
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE Decode"
+    }
+  },
+  "9": {
+    "inputs": {
+      "filename_prefix": "ComfyUI",
+      "images": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "SaveImage",
+    "_meta": {
+      "title": "Save Image"
+    }
+  }
+}

conf/example-workflow_settings.json

@@ -0,0 +1,9 @@
+{
+    "defaults": {
+        "widthy": 384,
+        "heighty": 512,
+        "max_width": 256,
+        "max_height": 256,
+        "seed": null
+    }
+}

lib/comfyui.py

@@ -0,0 +1,45 @@
+import time
+
+import asyncio
+import aiohttp
+
+from urllib.request import urlopen as urlopen
+from urllib.parse import urlencode as urlencode
+
+async def get_comfyui_generations(api_url, workflow):
+    async with aiohttp.ClientSession() as session:
+        async with session.post(f"{api_url}/prompt", json={"prompt": workflow}) as resp:
+            resp_json = await resp.json()
+
+        prompt_id = resp_json["prompt_id"]
+
+        # Loop endlessly until ComfyUI confirms our prompt's completion:
+        while True:
+            async with session.get(f"{api_url}/history/{prompt_id}") as resp:
+                resp_json = await resp.json()
+            if not resp_json:
+                time.sleep(1)
+                continue
+            break
+
+        # Read the output history anmd fetch each image:
+        history         = resp_json[prompt_id]
+        output_images   = []
+
+        for o in history["outputs"]:
+            for node_id in history["outputs"]:
+                node_output     = history["outputs"][node_id]
+                files_output    = []
+
+                if "images" in node_output:
+                    for image in node_output["images"]:
+                        url_params = urlencode({
+                            "filename":     image["filename"],
+                            "subfolder":    image["subfolder"],
+                            "type":         image["type"]
+                        })
+
+                        async with session.get(f"{api_url}/view?{url_params}") as resp:
+                            output_images.append(await resp.content.read())
+
+        return output_images

lib/events.py

@@ -0,0 +1,194 @@
+import discord
+import os
+import json
+import time
+import io
+import random
+
+from lib.helpers import *
+from lib.settings import *
+from lib.parser import *
+from lib.comfyui import *
+
+async def on_message_or_reaction(client, obj):
+    msg     = None
+    chl     = None
+    user    = None
+    author  = None
+    roles   = None
+    rxn     = None
+
+    msg_types = [discord.MessageType.default]
+
+    if isinstance(obj, discord.RawReactionActionEvent):
+        chl     = await client.fetch_channel(obj.channel_id)
+        msg     = await chl.fetch_message(obj.message_id)
+        user    = await client.fetch_user(obj.user_id)
+        author  = await client.fetch_user(obj.message_author_id)
+        roles   = obj.member.roles
+        rxn     = obj
+
+        msg_types.append(discord.MessageType.reply)
+
+    if isinstance(obj, discord.Message):
+        msg     = obj
+        chl     = obj.channel
+        user    = msg.author
+        author  = msg.author
+        roles   = msg.author.roles
+        rxn     = None
+
+    if user == client.user:         return
+    if msg.type not in msg_types:   return
+    if user.bot:                    return
+
+    #
+    #
+    #
+
+    chl_topic_parts     = []
+    chl_topic_part_1    = ""
+    if chl.topic is not None:
+        chl_topic_parts     = chl.topic.split(",")
+        chl_topic_part_1    = chl_topic_parts[0]
+
+
+    # Try different paths to find workflow .json file:
+    workflow_paths = [
+        f"conf/{chl.category.name}/{chl_topic_part_1}",
+        f"conf/{chl.category.name}/{chl.name}",
+        f"conf/{chl_topic_part_1}",
+        f"conf/{chl.name}",
+        f"conf/{chl.category.name}",
+    ]
+
+    using_workflow_path = None
+    using_settings_path = None
+
+    for path in workflow_paths:
+        for extension in ["", ".json"]:
+            if using_workflow_path is not None:
+                continue
+
+            try_path = f"{path}{extension}"
+            print(f"Looking for workflow {try_path}")
+            if os.path.isfile(try_path):
+                using_workflow_path = try_path
+
+    if using_workflow_path is None:
+        return
+
+    print(f"Using workflow {using_workflow_path}")
+
+    #
+    #
+    #
+
+    setting_paths = [
+        f"conf/{chl.category.name}",
+        f"conf/{chl.name}",
+    ]
+
+    for i in chl_topic_parts:
+        setting_paths = setting_paths + ["conf/" + i.strip()]
+
+    setting_paths = setting_paths + [f"conf/{chl.category.name}/{chl.name}"]
+
+    for i in chl_topic_parts:
+        setting_paths = setting_paths + [f"conf/{chl.category.name}/" + i.strip()]
+
+    settings = get_settings()
+
+    for path in setting_paths:
+        for extension in ["", ".json"]:
+            try_path = f"{path}{extension}_settings.json"
+            print(f"Looking for setting {try_path}")
+            if os.path.isfile(try_path):
+                settings = merge_dicts(settings, read_json(try_path, {}))
+
+    #
+    #
+    #
+
+    ALLOW_REPEAT    = settings["allow_repeat"]
+    SHOW_REPEAT     = settings["show_repeat"]
+    REPEAT_EMOJI    = settings["repeat_emoji"]
+    WAITING_EMOJI   = settings["waiting_emoji"]
+
+    if rxn is not None:
+        if not ALLOW_REPEAT:
+            return
+
+
+    # Read the found .json file:
+    workflow_json = ""
+    with open(using_workflow_path, "r") as file:
+        workflow_json = file.read()
+
+    # Break the user message into prompt parameters:
+    params = get_prompt_parameters(msg.content, settings)
+
+
+
+
+
+    time_start  = time.perf_counter()
+
+    # Indicate to the user something is happening:
+    await msg.add_reaction(WAITING_EMOJI)
+    await chl.typing()
+
+    repeat_n_times  = 1 if "repeat_n_times" not in params.keys() else params["repeat_n_times"]
+    all_attachments = []
+
+    for i in range(0, repeat_n_times):
+        workflow_json_clone = workflow_json
+        params_clone        = params.copy()
+
+        # If the seed is not specified, generate one:
+        if "seed" not in params.keys():
+            params_clone["seed"] = random.randint(1, 999_999_999_999_999)
+
+        # Make replacements, e.g.
+        # __WIDTH__ to value of params["width"]
+        # __POSITIVE__ to value of params["positive"]
+        for k in params_clone.keys():
+            v = params_clone[k]
+            k = k.upper()
+            workflow_json_clone = re.sub(rf"__{k}__", str(v), workflow_json_clone)
+
+        # Must be valid JSON:
+        workflow        = json.loads(workflow_json_clone)
+        all_attachments = all_attachments + await get_comfyui_generations(settings["api_url"], workflow)
+
+    # Process 8 attachments at a time per one Discord message:
+    while True:
+        attachments_buffer  = all_attachments[:8]
+        discord_files       = []
+        
+        for a in attachments_buffer:
+            #base64_str          = a
+            #base64_bytes        = base64_str.encode("ascii")
+            #attachment_bytes    = base64.b64decode(base64_bytes)
+
+            f = discord.File(io.BytesIO(a), filename="file.png")
+            discord_files.append(f)
+
+        if len(discord_files) < 1: 
+            discord_files = None
+
+        post = await chl.send(files=discord_files, content=msg.content, reference=msg)
+        del all_attachments[:8]
+
+        if len(all_attachments) < 1:
+            break
+
+    if ALLOW_REPEAT:
+        if SHOW_REPEAT:
+            await msg.add_reaction(REPEAT_EMOJI)
+            await post.add_reaction(REPEAT_EMOJI)
+    
+    await msg.remove_reaction(WAITING_EMOJI, client.user)
+
+    time_end    = time.perf_counter()
+    time_taken  = time_end - time_start

lib/helpers.py

@@ -0,0 +1,51 @@
+import json
+import os
+
+
+# Make it easy to read a json file
+def read_json(path, default_value={}):
+    if os.path.isfile(path):
+        try:
+            with open(path, "r") as file:
+                json_string = file.read()
+                json_dict   = json.loads(json_string)
+                return json_dict
+        except:
+            pass
+    return default_value
+
+
+# Merge dictionary b into a
+def merge_dicts(a, b):
+    if b is None:
+        return a
+
+    output = a
+
+    for k in b.keys():
+        if k not in output.keys():
+            if b[k] is not None:
+                output[k] = b[k]
+
+        if b[k] is None:
+            if k in output.keys():
+                output.pop(k, None)
+
+        if isinstance(b[k], dict):
+            if isinstance(output[k], dict):
+                output[k] = merge_dicts(output[k], b[k])
+                continue
+
+        output[k] = b[k]
+        if output[k] is None:
+            del output[k]
+            
+    return output
+
+
+def write_json(path, o):
+    try:
+        with open(path, "w") as file:
+            file.write(json.dumps(o, indent=4))
+    except:
+        pass

lib/parser.py

@@ -0,0 +1,71 @@
+import random
+import re
+import copy
+
+def get_prompt_parameters(user_message, settings):
+    params          = {}
+    params          = settings["defaults"].copy()
+
+    mutable_string = user_message
+    mutable_string = re.sub(r"[^A-Za-z0-9 \-,\.>]", "", mutable_string)
+    mutable_string = re.sub(r"\s{1,}", " ", mutable_string)
+
+    # Get the options from the mutable string:
+    options     = ""
+    options_re  = r"^((.+)>)"
+    match       = re.search(options_re, mutable_string)
+    if match is not None:
+        options         = match.group(2).strip()
+        mutable_string  = re.sub(options_re, "", mutable_string)
+
+
+    # Get the repeat_n_times from the mutable string:
+    repeat_re = r"( x([0-9]{1,2}))$"
+    match       = re.search(repeat_re, mutable_string)
+    if match is not None:
+        params["repeat_n_times"]    = int(match.group(2).strip())
+        mutable_string              = re.sub(repeat_re, "", mutable_string)
+
+
+    # Get the negative prompt from the mutable string:
+    negative_re = r"( -(.+))$"
+    match       = re.search(negative_re, mutable_string)
+    if match is not None:
+        params["negative"]  = match.group(2).strip()
+        mutable_string      = re.sub(negative_re, "", mutable_string)
+
+    params["positive"] = mutable_string.strip()
+
+    overwrite_re = r"([A-Za-z_]{1,})=?([0-9\.]{1,})"
+    for (k, v) in re.findall(overwrite_re, options):
+        params[k] = float(v)
+
+    # Process limitations, e.g. max_width and max_height:
+    params_clone = copy.deepcopy(params)
+
+    for key in params.keys():
+        limit_value = params[key]
+
+        term = "min_"
+        if re.search(rf"^{term}", key):
+            target_key = re.sub(rf"^{term}", "", key)
+            if target_key not in params.keys():
+                params_clone[target_key] = limit_value
+            params_clone[target_key] = limit_value if params[target_key] < limit_value else params[target_key]
+
+        term = "max_"
+        if re.search(rf"^{term}", key):
+            target_key = re.sub(rf"^{term}", "", key)
+            if target_key in params.keys():
+                params_clone[target_key] = limit_value if params[target_key] > limit_value else params[target_key]
+
+        term = "force_"
+        if re.search(rf"^{term}", key):
+            target_key = re.sub(rf"^{term}", "", key)
+            params_clone[target_key] = limit_value
+
+    params = params_clone
+
+    print(params)
+
+    return params

lib/settings.py

@@ -0,0 +1,37 @@
+import os
+
+from lib.helpers import *
+
+def get_settings(initialize=False):
+    settings_path       = "settings.json"
+    default_settings    = {
+        "token":            "PASTE_TOKEN_HERE",
+        "api_url":          "http://127.0.0.1:8188",
+        "allow_repeat":     True,
+        "show_repeat":      True,
+        "repeat_emoji":     "♻️",
+        "waiting_emoji":    "🕒",
+        "defaults": {
+            "width":    512,
+            "height":   512,
+            "cfg":      7,
+            "steps":    20,
+            "seed":     1,
+            "positive": "Test",
+            "negative": "Test",
+        },
+    }
+
+    if initialize:
+        if not os.path.isfile(settings_path):
+            write_json(settings_path, default_settings)
+            print(f"A bot token is required, please modify {settings_path} to include the token and rerun this Python script.")
+            exit()
+
+    settings = default_settings.copy()
+    settings = merge_dicts(settings, read_json(settings_path))
+
+    if initialize:
+        write_json(settings_path, settings)
+
+    return settings