diff --git a/.gitignore b/.gitignore index 6468feb..4ffff98 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,5 @@ __pycache__/ !/conf/example-workflow_settings.json /log.txt /log.txt.* +/*.log.txt +/*.log.txt.* diff --git a/bot.py b/bot.py index d1411af..a5ede3c 100644 --- a/bot.py +++ b/bot.py @@ -4,9 +4,8 @@ import lib.log as log 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 = {} +intents = discord.Intents(messages=True, guilds=True, message_content=True, reactions=True) +client = discord.Client(intents=intents) @client.event async def on_message(msg): @@ -14,23 +13,16 @@ async def on_message(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) + await on_message_or_reaction(client, rxn) @client.event async def on_ready(): - log.write("READY.") + log.write("READY.", log.colors.fg.lightgreen) def main(): - global settings + log.write("Starting up...", log.colors.fg.yellow) settings = get_settings(initialize=True) client.run(settings["token"], log_handler=None) if __name__ == "__main__": - log.write("test") main() - diff --git a/lib/events.py b/lib/events.py index 3ec8af9..8d01402 100644 --- a/lib/events.py +++ b/lib/events.py @@ -13,6 +13,8 @@ from lib.parser import * from lib.comfyui import * async def on_message_or_reaction(client, obj): + global RUN_TIME + msg = None chl = None user = None @@ -56,7 +58,6 @@ async def on_message_or_reaction(client, obj): 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}", @@ -69,17 +70,11 @@ async def on_message_or_reaction(client, obj): using_workflow_path = None using_settings_path = None - log.write("\nSearching for workflow in order:\n") - for path in workflow_paths: - check_text = " " - try_path = f"{path}.json" - + 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}") + log.write(f"Workflow {log.colors.fg.lightcyan}{try_path}") if using_workflow_path is None: return @@ -103,16 +98,11 @@ async def on_message_or_reaction(client, obj): settings = get_settings() - log.write("Merging settings in order:") for path in setting_paths: - check_text = " " - try_path = f"{path}_settings.json" - + 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}") + settings = merge_dicts(settings, read_json(try_path, {})) + log.write(f"Merging {log.colors.fg.lightcyan}{try_path}") # # @@ -132,6 +122,7 @@ async def on_message_or_reaction(client, obj): 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: + log.write(f"User {user.name} or their ID not in list of allowed users", log.colors.fg.lightred) return # If ignored_users is specified, ignore listed users: @@ -139,9 +130,9 @@ async def on_message_or_reaction(client, obj): ignored_users = settings["ignored_users"] if isinstance(ignored_users, list): if str(user.id) in ignored_users or str(user.name) in ignored_users: + log.write(f"User {user.name} or their ID is explicitly being ignored", log.colors.fg.lightred) return - # Read the found .json file: workflow_json = "" with open(using_workflow_path, "r") as file: @@ -150,93 +141,103 @@ async def on_message_or_reaction(client, obj): # Break the user message into prompt parameters: 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() - + v = params[k] + k = k.upper() label = f"__{k}__" - while len(label) < max_key_size: - label = label + " " + log.write(f"Replacing {log.colors.fg.yellow}{label}{log.colors.reset} with {log.colors.fg.white}{v}") - log.write(f"[ - ] Replacing: {label} with {v}") + log.write_prompt(client, msg, chl, user, author, roles, rxn) + try: + time_start = time.perf_counter() - time_start = time.perf_counter() + # Indicate to the user something is happening: + await msg.add_reaction(WAITING_EMOJI) + await chl.typing() - # 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 = [] - 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() - 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() or params["seed"] is None: + new_seed = random.randint(1, 999_999_999_999_999) + params_clone["seed"] = new_seed + log.write(f"Replacing {log.colors.fg.yellow}__SEED__{log.colors.reset} with {log.colors.fg.white}{new_seed}") - # If the seed is not specified, generate one: - 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 + # 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() - label = f"__SEED__" - while len(label) < max_key_size: - label = label + " " + if v is None: + continue - log.write(f"[ - ] Replacing: {label} with {new_seed}") + workflow_json_clone = re.sub(rf"__{k}__", str(v), workflow_json_clone) - # 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() + # Must be valid JSON: + workflow = json.loads(workflow_json_clone) + all_attachments = all_attachments + await get_comfyui_generations(settings["api_url"], workflow) - if v is None: - continue + bytes_generated = 0 + images_generated = 0 - workflow_json_clone = re.sub(rf"__{k}__", str(v), workflow_json_clone) + # Process 8 attachments at a time per one Discord message: + while True: + attachments_buffer = all_attachments[:8] + discord_files = [] + + for a in attachments_buffer: + bytes_generated = bytes_generated + len(a) + images_generated = images_generated + 1 + discord_file = discord.File(io.BytesIO(a), filename="file.png") + discord_files.append(discord_file) - # Must be valid JSON: - workflow = json.loads(workflow_json_clone) - all_attachments = all_attachments + await get_comfyui_generations(settings["api_url"], workflow) + if len(discord_files) < 1: + discord_files = None - # Process 8 attachments at a time per one Discord message: - while True: - attachments_buffer = all_attachments[:8] - discord_files = [] + 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) - for a in attachments_buffer: - #base64_str = a - #base64_bytes = base64_str.encode("ascii") - #attachment_bytes = base64.b64decode(base64_bytes) + await msg.remove_reaction(WAITING_EMOJI, client.user) - f = discord.File(io.BytesIO(a), filename="file.png") - discord_files.append(f) + time_end = time.perf_counter() + time_taken = time_end - time_start - if len(discord_files) < 1: - discord_files = None + log.write(f"{time_taken:0.2f}s taken.", log.colors.fg.lightgreen) - post = await chl.send(files=discord_files, content=msg.content, reference=msg) - del all_attachments[:8] + # + # + # - if len(all_attachments) < 1: - break + log.write_stat(f"user.{user.name}.time_taken", f"{time_taken:0.2f}") + log.write_stat(f"user.{user.name}.bytes_generated", f"{bytes_generated}") + log.write_stat(f"user.{user.name}.images_generated", f"{images_generated}") - 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 + t_time_taken = log.get_stat(r".*\.time_taken\: (.+)$") + t_bytes_generated = log.get_stat(r".*\.bytes_generated\: (.+)$") + t_images_generated = log.get_stat(r".*\.images_generated\: (.+)$") + + kb_gen = t_bytes_generated / 1024 + mb_gen = kb_gen / 1024 + + await client.change_presence(activity=discord.Game(name=f"{t_time_taken:0.1f}s, {mb_gen:0.1f} MiB, {t_images_generated:0.0f}x")) + + + finally: + lock = False diff --git a/lib/log.py b/lib/log.py index 7d418ae..9c0fa42 100644 --- a/lib/log.py +++ b/lib/log.py @@ -1,41 +1,121 @@ import sys import logging import logging.handlers +import re LOG_SETUP = False +# The colors class below is from: +# https://www.geeksforgeeks.org/print-colors-python-terminal +class colors: + reset = "\033[0m" + bold = "\033[01m" + disable = "\033[02m" + underline = "\033[04m" + reverse = "\033[07m" + strikethrough = "\033[09m" + invisible = "\033[08m" + + class fg: + white = "\033[97m" + black = "\033[30m" + red = "\033[31m" + green = "\033[32m" + orange = "\033[33m" + blue = "\033[34m" + purple = "\033[35m" + cyan = "\033[36m" + lightgrey = "\033[37m" + darkgrey = "\033[90m" + lightred = "\033[91m" + lightgreen = "\033[92m" + yellow = "\033[93m" + lightblue = "\033[94m" + pink = "\033[95m" + lightcyan = "\033[96m" + + class bg: + black = "\033[40m" + red = "\033[41m" + green = "\033[42m" + orange = "\033[43m" + blue = "\033[44m" + purple = "\033[45m" + cyan = "\033[46m" + lightgrey = "\033[47m" + +class PlainTextFormatter(logging.Formatter): + def format(self, record): + log_message = super().format(record) + log_message = re.sub(r"\033\[[0-9]{1,2}m", "", log_message) + return log_message + # Make it easy to log text: -def write(input_str): +def write(input_str, color=colors.fg.lightgrey, logger_name="discord.bot"): global LOG_SETUP + if not LOG_SETUP: + logging.getLogger("discord").setLevel(logging.DEBUG) + logging.getLogger("discord.http").setLevel(logging.INFO) - logger = logging.getLogger('discord') - logger.setLevel(logging.DEBUG) - logging.getLogger('discord.http').setLevel(logging.INFO) + line_format = "[{asctime}] [{threadName}] [{levelname}] {name}: {message}" + timestamp_format = "%Y-%m-%d %H:%M:%S" + console_formatter = logging.Formatter(colors.fg.darkgrey + line_format, timestamp_format, style="{") + file_formatter = PlainTextFormatter(line_format, timestamp_format, style="{") + prompt_formatter = PlainTextFormatter("[{asctime}] {message}", timestamp_format, style="{") + stats_formatter = PlainTextFormatter("{message}", timestamp_format, style="{") + + file_handler = logging.handlers.RotatingFileHandler(filename="log.txt", encoding="utf-8", maxBytes=32 * 1024 * 1024, backupCount=10) + file_handler.setFormatter(file_formatter) + console_handler = logging.StreamHandler() + console_handler.setFormatter(console_formatter) - logFormatter = logging.Formatter( - "[{asctime}] [{threadName:<16}] [{levelname:<8}] {name}: {message}", - "%Y-%m-%d %H:%M:%S", - style="{") + root_logger = logging.getLogger() + root_logger.addHandler(file_handler) + root_logger.addHandler(console_handler) - rootLogger = logging.getLogger() + prompt_logger = logging.getLogger("discord.bot.prompts") + prompt_handler = logging.handlers.RotatingFileHandler(filename="prompts.log.txt", encoding="utf-8", maxBytes=32 * 1024 * 1024, backupCount=10) + prompt_handler.setFormatter(prompt_formatter) + prompt_logger.addHandler(prompt_handler) - #fileHandler = logging.FileHandler("log.txt") - fileHandler = logging.handlers.RotatingFileHandler( - filename="log.txt", - encoding="utf-8", - maxBytes=32 * 1024 * 1024, - backupCount=10, - ) - - fileHandler.setFormatter(logFormatter) - rootLogger.addHandler(fileHandler) - - consoleHandler = logging.StreamHandler() - consoleHandler.setFormatter(logFormatter) - rootLogger.addHandler(consoleHandler) + stats_logger = logging.getLogger("discord.bot.stats") + stats_handler = logging.handlers.RotatingFileHandler(filename="stats.log.txt", encoding="utf-8", maxBytes=32 * 1024 * 1024, backupCount=10) + stats_handler.setFormatter(stats_formatter) + stats_logger.addHandler(stats_handler) LOG_SETUP = True - logging.getLogger("discord").info(input_str) + logging.getLogger(logger_name).info(f"{color}{input_str}{colors.reset}") + +def write_prompt(client=None, msg=None, chl=None, user=None, author=None, roles=None, rxn=None): + bot_name = client.user.name + "#" + client.user.discriminator + server_name = msg.channel.guild.name + category_name = msg.channel.category.name + channel_name = msg.channel.name + user_name = user.name + + discrim = 0 + if user.discriminator is not None: + discrim = user.discriminator + + user_name = f"{user_name}#{discrim}" + message_content = msg.content + message_text = f"[{bot_name}/{server_name}/{category_name}/{channel_name}/{user_name}] {colors.fg.white}{message_content}" + + write(message_text, logger_name="discord.bot.prompts") + +def write_stat(key, value): + write(f"{key}: {value}", logger_name="discord.bot.stats") + +def get_stat(key, sum=True): + value = 0 + with open("stats.log.txt") as file: + lines = [line.strip() for line in file] + for line in lines: + match = re.search(key, line) + if match is not None: + value = value + float(match.group(1).strip()) + + return value diff --git a/lib/parser.py b/lib/parser.py index e13c158..eefde7c 100644 --- a/lib/parser.py +++ b/lib/parser.py @@ -3,8 +3,9 @@ import re import copy def get_prompt_parameters(user_message, settings): - params = {} - params = settings["defaults"].copy() + params = {} + if "defaults" in settings.keys(): + params = settings["defaults"].copy() mutable_string = user_message mutable_string = re.sub(r"[^A-Za-z0-9 \-,\.>\(\)]", "", mutable_string)