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

Reduced text logging, added SQLite3 prompt and statistics logging, reading and measuring execution time from ComfyUI

Browse Source
This commit is contained in:
Conner Harkness 2025-03-10 12:58:11 -06:00
parent 02da214790
commit 3e2f8e61b2
8 changed files with 111 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -8,3 +8,4 @@ __pycache__/
/log.txt.*
/*.log.txt
/*.log.txt.*
/database.db

1
bot.py
View File

@ -1,6 +1,7 @@
import discord
import lib.log as log
import lib.db as db
from lib.settings import *
from lib.events import *

20
lib/comfyui.py
View File

@ -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
}

48
lib/db.py Normal file
View File

@ -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)

42
lib/events.py
View File

@ -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
@ -183,7 +186,10 @@ async def on_message_or_reaction(client, obj):
# Must be valid JSON:
workflow = json.loads(workflow_json_clone)
all_attachments = all_attachments + await get_comfyui_generations(settings["api_url"], workflow)
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

8
lib/helpers.py
View File

@ -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}"

16
lib/log.py
View File

@ -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_formatter = PlainTextFormatter(line_format, timestamp_format, style="{")
file_handler.setFormatter(file_formatter)
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}")

1
lib/settings.py
View File

@ -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,