diff --git a/.gitignore b/.gitignore index 3591c6c..6468feb 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,8 @@ __pycache__/ /venv/ /settings.json - /conf/* !/conf/example-workflow.json !/conf/example-workflow_settings.json +/log.txt +/log.txt.* diff --git a/bot.py b/bot.py index 0628170..d1411af 100644 --- a/bot.py +++ b/bot.py @@ -1,5 +1,6 @@ import discord +import lib.log as log from lib.settings import * from lib.events import * @@ -22,13 +23,14 @@ async def on_raw_reaction_add(rxn): @client.event async def on_ready(): - print("READY.") + log.write("READY.") def main(): global settings settings = get_settings(initialize=True) - client.run(settings["token"]) + client.run(settings["token"], log_handler=None) if __name__ == "__main__": + log.write("test") main() diff --git a/conf/example-workflow.json b/conf/example-workflow.json index 92442eb..000828c 100644 --- a/conf/example-workflow.json +++ b/conf/example-workflow.json @@ -31,7 +31,7 @@ }, "4": { "inputs": { - "ckpt_name": "anything-v3-fp16-pruned.safetensors" + "ckpt_name": "__CHECKPOINT_1__" }, "class_type": "CheckpointLoaderSimple", "_meta": { diff --git a/lib/events.py b/lib/events.py index f515a45..3ec8af9 100644 --- a/lib/events.py +++ b/lib/events.py @@ -4,7 +4,9 @@ import json import time import io import random +import logging +import lib.log as log from lib.helpers import * from lib.settings import * from lib.parser import * @@ -67,21 +69,21 @@ async def on_message_or_reaction(client, obj): 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 + log.write("\nSearching for workflow in order:\n") - try_path = f"{path}{extension}" - print(f"Looking for workflow {try_path}") - if os.path.isfile(try_path): - using_workflow_path = try_path + for path in workflow_paths: + check_text = " " + try_path = f"{path}.json" + + if os.path.isfile(try_path): + using_workflow_path = try_path + check_text = "x" + + log.write(f"[ {check_text} ] Selecting: {try_path}") if using_workflow_path is None: return - print(f"Using workflow {using_workflow_path}") - # # # @@ -101,12 +103,16 @@ async def on_message_or_reaction(client, obj): settings = get_settings() + log.write("Merging settings in order:") 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, {})) + check_text = " " + try_path = f"{path}_settings.json" + + if os.path.isfile(try_path): + settings = merge_dicts(settings, read_json(try_path, {})) + check_text = "x" + + log.write(f"[ {check_text} ] Merging: {try_path}") # # @@ -121,6 +127,20 @@ async def on_message_or_reaction(client, obj): if not ALLOW_REPEAT: return + # If allowed_users is specified, only listen to listed users: + if "allowed_users" in settings.keys(): + allowed_users = settings["allowed_users"] + if isinstance(allowed_users, list): + if str(user.id) not in allowed_users and str(user.name) not in allowed_users: + return + + # If ignored_users is specified, ignore listed users: + if "ignored_users" in settings.keys(): + ignored_users = settings["ignored_users"] + if isinstance(ignored_users, list): + if str(user.id) in ignored_users or str(user.name) in ignored_users: + return + # Read the found .json file: workflow_json = "" @@ -131,10 +151,25 @@ async def on_message_or_reaction(client, obj): params = get_prompt_parameters(msg.content, settings) + max_key_size = 0 + for k in params.keys(): + new_key_size = len(k) + 4 + if new_key_size > max_key_size: + max_key_size = new_key_size + + log.write("Replacing in workflow JSON body:") + for k in params.keys(): + v = params[k] + k = k.upper() + + label = f"__{k}__" + while len(label) < max_key_size: + label = label + " " + + log.write(f"[ - ] Replacing: {label} with {v}") - - time_start = time.perf_counter() + time_start = time.perf_counter() # Indicate to the user something is happening: await msg.add_reaction(WAITING_EMOJI) @@ -148,8 +183,15 @@ async def on_message_or_reaction(client, obj): 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) + if "seed" not in params.keys() or params["seed"] is None: + new_seed = random.randint(1, 999_999_999_999_999) + params_clone["seed"] = new_seed + + label = f"__SEED__" + while len(label) < max_key_size: + label = label + " " + + log.write(f"[ - ] Replacing: {label} with {new_seed}") # Make replacements, e.g. # __WIDTH__ to value of params["width"] @@ -157,6 +199,10 @@ async def on_message_or_reaction(client, obj): for k in params_clone.keys(): v = params_clone[k] k = k.upper() + + if v is None: + continue + workflow_json_clone = re.sub(rf"__{k}__", str(v), workflow_json_clone) # Must be valid JSON: diff --git a/lib/parser.py b/lib/parser.py index 734d714..e13c158 100644 --- a/lib/parser.py +++ b/lib/parser.py @@ -7,7 +7,7 @@ def get_prompt_parameters(user_message, settings): params = settings["defaults"].copy() mutable_string = user_message - mutable_string = re.sub(r"[^A-Za-z0-9 \-,\.>]", "", mutable_string) + 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: @@ -65,7 +65,4 @@ def get_prompt_parameters(user_message, settings): params_clone[target_key] = limit_value params = params_clone - - print(params) - return params diff --git a/lib/settings.py b/lib/settings.py index 4e9557b..844bfb4 100644 --- a/lib/settings.py +++ b/lib/settings.py @@ -12,13 +12,15 @@ def get_settings(initialize=False): "repeat_emoji": "♻️", "waiting_emoji": "🕒", "defaults": { + "checkpoint_1": "v1-5-pruned-emaonly.safetensors", "width": 512, "height": 512, "cfg": 7, "steps": 20, - "seed": 1, - "positive": "Test", - "negative": "Test", + "seed": None, + "positive": "", + "negative": "", + "negative_prefix": "((nsfw, nudity))" }, }