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

.gitignore

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

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 *
 

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

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)

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)
+            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}")
+        db.insert("prompts", {
-        log.write_stat(f"user.{user.name}.bytes_generated", f"{bytes_generated}")
+            "handler_name":     settings["handler_name"],
-        log.write_stat(f"user.{user.name}.images_generated", f"{images_generated}")
+            "api_url":          settings["api_url"],
-
+            "bot_id":           client.user.id,
-
+            "bot_name":         get_legacy_username(client.user),
-        t_time_taken = log.get_stat(r".*\.time_taken\: (.+)$")
+            "guild_id":         chl.guild.id,
-        t_bytes_generated = log.get_stat(r".*\.bytes_generated\: (.+)$")
+            "guild_name":       chl.guild.name,
-        t_images_generated = log.get_stat(r".*\.images_generated\: (.+)$")
+            "channel_id":       chl.id,
-
+            "channel_name":     chl.name,
-        kb_gen = t_bytes_generated / 1024
+            "author_id":        author.id,
-        mb_gen = kb_gen / 1024
+            "author_name":      get_legacy_username(author),
-
+            "user_id":          user.id,
-        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"))
+            "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

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

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="{")
+        file_handler    = logging.handlers.RotatingFileHandler(filename="log.txt", encoding="utf-8", maxBytes=32 * 1024 * 1024, backupCount=10)
-        prompt_formatter    = PlainTextFormatter("[{asctime}] {message}", timestamp_format, style="{")
+        file_formatter  = PlainTextFormatter(line_format, 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     = 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}")

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,