caharkness/comfyui-discord
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added better logging

Browse Source
This commit is contained in:
Conner Harkness 2025-03-05 14:59:24 -07:00
parent 25f05b5c4f
commit 02da214790
5 changed files with 200 additions and 124 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

@ -6,3 +6,5 @@ __pycache__/
!/conf/example-workflow_settings.json
/log.txt
/log.txt.*
/*.log.txt
/*.log.txt.*

12
bot.py
View File

@ -6,7 +6,6 @@ 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):
@ -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)
@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()

81
lib/events.py
View File

@ -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"
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"
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}")
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,25 +141,15 @@ 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()
label = f"__{k}__"
while len(label) < max_key_size:
label = label + " "
log.write(f"[ - ] Replacing: {label} with {v}")
log.write(f"Replacing {log.colors.fg.yellow}{label}{log.colors.reset} with {log.colors.fg.white}{v}")
log.write_prompt(client, msg, chl, user, author, roles, rxn)
try:
time_start = time.perf_counter()
# Indicate to the user something is happening:
@ -186,12 +167,7 @@ async def on_message_or_reaction(client, obj):
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}")
log.write(f"Replacing {log.colors.fg.yellow}__SEED__{log.colors.reset} with {log.colors.fg.white}{new_seed}")
# Make replacements, e.g.
# __WIDTH__ to value of params["width"]
@ -209,18 +185,19 @@ async def on_message_or_reaction(client, obj):
workflow = json.loads(workflow_json_clone)
all_attachments = all_attachments + await get_comfyui_generations(settings["api_url"], workflow)
bytes_generated = 0
images_generated = 0
# 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)
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)
if len(discord_files) < 1:
discord_files = None
@ -240,3 +217,27 @@ async def on_message_or_reaction(client, obj):
time_end = time.perf_counter()
time_taken = time_end - time_start
log.write(f"{time_taken:0.2f}s taken.", log.colors.fg.lightgreen)
#
#
#
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"))
finally:
lock = False

126
lib/log.py
View File

@ -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)
logFormatter = logging.Formatter(
"[{asctime}] [{threadName:<16}] [{levelname:<8}] {name}: {message}",
"%Y-%m-%d %H:%M:%S",
style="{")
console_handler = logging.StreamHandler()
console_handler.setFormatter(console_formatter)
rootLogger = logging.getLogger()
root_logger = logging.getLogger()
root_logger.addHandler(file_handler)
root_logger.addHandler(console_handler)
#fileHandler = logging.FileHandler("log.txt")
fileHandler = logging.handlers.RotatingFileHandler(
filename="log.txt",
encoding="utf-8",
maxBytes=32 * 1024 * 1024,
backupCount=10,
)
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.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

1
lib/parser.py
View File

@ -4,6 +4,7 @@ import copy
def get_prompt_parameters(user_message, settings):
params = {}
if "defaults" in settings.keys():
params = settings["defaults"].copy()
mutable_string = user_message