diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3591c6c --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +__pycache__/ +/venv/ +/settings.json + +/conf/* +!/conf/example-workflow.json +!/conf/example-workflow_settings.json diff --git a/bot.py b/bot.py new file mode 100644 index 0000000..0628170 --- /dev/null +++ b/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() + diff --git a/comfyui-discord.sh b/comfyui-discord.sh new file mode 100644 index 0000000..f6e11b4 --- /dev/null +++ b/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 diff --git a/conf/example-workflow.json b/conf/example-workflow.json new file mode 100644 index 0000000..92442eb --- /dev/null +++ b/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" + } + } +} diff --git a/conf/example-workflow_settings.json b/conf/example-workflow_settings.json new file mode 100644 index 0000000..09486af --- /dev/null +++ b/conf/example-workflow_settings.json @@ -0,0 +1,9 @@ +{ + "defaults": { + "widthy": 384, + "heighty": 512, + "max_width": 256, + "max_height": 256, + "seed": null + } +} diff --git a/lib/comfyui.py b/lib/comfyui.py new file mode 100644 index 0000000..c9f0614 --- /dev/null +++ b/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 diff --git a/lib/events.py b/lib/events.py new file mode 100644 index 0000000..555f2f8 --- /dev/null +++ b/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 diff --git a/lib/helpers.py b/lib/helpers.py new file mode 100644 index 0000000..8e91513 --- /dev/null +++ b/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 diff --git a/lib/parser.py b/lib/parser.py new file mode 100644 index 0000000..734d714 --- /dev/null +++ b/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 diff --git a/lib/settings.py b/lib/settings.py new file mode 100644 index 0000000..4e9557b --- /dev/null +++ b/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