diff --git a/.gitignore b/.gitignore index 4ffff98..f7d7393 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ __pycache__/ /log.txt.* /*.log.txt /*.log.txt.* +/database.db diff --git a/bot.py b/bot.py index a5ede3c..5bf4b76 100644 --- a/bot.py +++ b/bot.py @@ -1,6 +1,7 @@ import discord import lib.log as log +import lib.db as db from lib.settings import * from lib.events import * diff --git a/lib/comfyui.py b/lib/comfyui.py index c9f0614..7db71c0 100644 --- a/lib/comfyui.py +++ b/lib/comfyui.py @@ -22,6 +22,8 @@ async def get_comfyui_generations(api_url, workflow): continue break + print(resp_json) + # Read the output history anmd fetch each image: history = resp_json[prompt_id] output_images = [] @@ -42,4 +44,20 @@ async def get_comfyui_generations(api_url, workflow): async with session.get(f"{api_url}/view?{url_params}") as resp: output_images.append(await resp.content.read()) - return output_images + + messages = resp_json[prompt_id]["status"]["messages"] + timestamp_start = 0 + timestamp_finish = 0 + + for m in messages: + label = m[0] + obj = m[1] + if label == "execution_start": timestamp_start = obj["timestamp"] + if label == "execution_success": timestamp_finish = obj["timestamp"] + + execution_time = (timestamp_finish - timestamp_start) / 1000.0 + + return { + "execution_time": execution_time, + "images": output_images + } diff --git a/lib/db.py b/lib/db.py new file mode 100644 index 0000000..d1d6863 --- /dev/null +++ b/lib/db.py @@ -0,0 +1,48 @@ +import sys +import logging +import logging.handlers +import re +import sqlite3 + +con = None + +def query(query_string, values=[]): + global con + + if con is None: + con = sqlite3.connect("database.db") + con.row_factory = sqlite3.Row + + cur = con.cursor() + cur.execute(query_string, values) + rows = cur.fetchall() + con.commit() + + return rows + +def insert(table_name, obj): + columns = "" + q_marks = "" + values = [] + typed_columns = "'id' integer primary key autoincrement, " + + for k in obj.keys(): + value = obj[k] + columns = columns + f"'{k}', " + q_marks = q_marks + "?, " + values.append(value) + + type = "text" + if isinstance(type, int): type = "integer" + if isinstance(type, float): type = "real" + + typed_columns = typed_columns + f"'{k}' {type} null, " + + typed_columns = typed_columns + "'inserted_on' timestamp default current_timestamp, " + columns = columns.rstrip(", ") + q_marks = q_marks.rstrip(", ") + typed_columns = typed_columns.rstrip(", ") + query_str = f"INSERT into '{table_name}' ({columns}) values ({q_marks})" + + query(f"CREATE table if not exists '{table_name}' ({typed_columns})") + query(query_str, values) diff --git a/lib/events.py b/lib/events.py index 8d01402..73aaa73 100644 --- a/lib/events.py +++ b/lib/events.py @@ -7,6 +7,8 @@ import random import logging import lib.log as log +import lib.db as db + from lib.helpers import * from lib.settings import * from lib.parser import * @@ -158,6 +160,7 @@ async def on_message_or_reaction(client, obj): repeat_n_times = 1 if "repeat_n_times" not in params.keys() else params["repeat_n_times"] all_attachments = [] + execution_time = 0 for i in range(0, repeat_n_times): workflow_json_clone = workflow_json @@ -182,8 +185,11 @@ async def on_message_or_reaction(client, obj): 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) + workflow = json.loads(workflow_json_clone) + response = await get_comfyui_generations(settings["api_url"], workflow) + + all_attachments = all_attachments + response["images"] + execution_time = execution_time + response["execution_time"] bytes_generated = 0 images_generated = 0 @@ -224,20 +230,26 @@ async def on_message_or_reaction(client, obj): # # - 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}") - - - 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")) - + db.insert("prompts", { + "handler_name": settings["handler_name"], + "api_url": settings["api_url"], + "bot_id": client.user.id, + "bot_name": get_legacy_username(client.user), + "guild_id": chl.guild.id, + "guild_name": chl.guild.name, + "channel_id": chl.id, + "channel_name": chl.name, + "author_id": author.id, + "author_name": get_legacy_username(author), + "user_id": user.id, + "user_name": get_legacy_username(user), + "prompt": msg.content, + "workflow_path": using_workflow_path, + "wait_time": time_taken, + "execution_time": execution_time, + "images_generated": images_generated, + "bytes_generated": bytes_generated + }) finally: lock = False diff --git a/lib/helpers.py b/lib/helpers.py index 30f3afb..6bc5991 100644 --- a/lib/helpers.py +++ b/lib/helpers.py @@ -46,3 +46,11 @@ def write_json(path, o): file.write(json.dumps(o, indent=4)) except: pass + +def get_legacy_username(obj): + username = obj.name + discriminator = 0 + if hasattr(obj, "discriminator"): + if obj.discriminator is not None: + discriminator = obj.discriminator + return f"{username}#{discriminator}" diff --git a/lib/log.py b/lib/log.py index 9c0fa42..a843a0b 100644 --- a/lib/log.py +++ b/lib/log.py @@ -60,31 +60,19 @@ def write(input_str, color=colors.fg.lightgrey, logger_name="discord.bot"): 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 = logging.handlers.RotatingFileHandler(filename="log.txt", encoding="utf-8", maxBytes=32 * 1024 * 1024, backupCount=10) + file_formatter = PlainTextFormatter(line_format, timestamp_format, style="{") file_handler.setFormatter(file_formatter) - console_handler = logging.StreamHandler() + console_handler = logging.StreamHandler() + console_formatter = logging.Formatter(colors.fg.darkgrey + line_format, timestamp_format, style="{") console_handler.setFormatter(console_formatter) root_logger = logging.getLogger() root_logger.addHandler(file_handler) root_logger.addHandler(console_handler) - 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) - - 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(logger_name).info(f"{color}{input_str}{colors.reset}") diff --git a/lib/settings.py b/lib/settings.py index 844bfb4..3518a3a 100644 --- a/lib/settings.py +++ b/lib/settings.py @@ -5,6 +5,7 @@ from lib.helpers import * def get_settings(initialize=False): settings_path = "settings.json" default_settings = { + "handler_name": "default", "token": "PASTE_TOKEN_HERE", "api_url": "http://127.0.0.1:8188", "allow_repeat": True,