Added better logging, added allowed and ignored users options

This commit is contained in:
Conner Harkness 2025-03-04 14:57:40 -07:00
parent e7915ad2b6
commit 6196f9cc88
6 changed files with 78 additions and 30 deletions

3
.gitignore vendored
View File

@ -1,7 +1,8 @@
__pycache__/
/venv/
/settings.json
/conf/*
!/conf/example-workflow.json
!/conf/example-workflow_settings.json
/log.txt
/log.txt.*

6
bot.py
View File

@ -1,5 +1,6 @@
import discord
import lib.log as log
from lib.settings import *
from lib.events import *
@ -22,13 +23,14 @@ async def on_raw_reaction_add(rxn):
@client.event
async def on_ready():
print("READY.")
log.write("READY.")
def main():
global settings
settings = get_settings(initialize=True)
client.run(settings["token"])
client.run(settings["token"], log_handler=None)
if __name__ == "__main__":
log.write("test")
main()

View File

@ -31,7 +31,7 @@
},
"4": {
"inputs": {
"ckpt_name": "anything-v3-fp16-pruned.safetensors"
"ckpt_name": "__CHECKPOINT_1__"
},
"class_type": "CheckpointLoaderSimple",
"_meta": {

View File

@ -4,7 +4,9 @@ import json
import time
import io
import random
import logging
import lib.log as log
from lib.helpers import *
from lib.settings import *
from lib.parser import *
@ -67,21 +69,21 @@ async def on_message_or_reaction(client, obj):
using_workflow_path = None
using_settings_path = None
for path in workflow_paths:
for extension in ["", ".json"]:
if using_workflow_path is not None:
continue
log.write("\nSearching for workflow in order:\n")
try_path = f"{path}{extension}"
print(f"Looking for workflow {try_path}")
if os.path.isfile(try_path):
using_workflow_path = try_path
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}")
if using_workflow_path is None:
return
print(f"Using workflow {using_workflow_path}")
#
#
#
@ -101,12 +103,16 @@ async def on_message_or_reaction(client, obj):
settings = get_settings()
log.write("Merging settings in order:")
for path in setting_paths:
for extension in ["", ".json"]:
try_path = f"{path}{extension}_settings.json"
print(f"Looking for setting {try_path}")
if os.path.isfile(try_path):
settings = merge_dicts(settings, read_json(try_path, {}))
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}")
#
#
@ -121,6 +127,20 @@ async def on_message_or_reaction(client, obj):
if not ALLOW_REPEAT:
return
# If allowed_users is specified, only listen to listed users:
if "allowed_users" in settings.keys():
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:
return
# If ignored_users is specified, ignore listed users:
if "ignored_users" in settings.keys():
ignored_users = settings["ignored_users"]
if isinstance(ignored_users, list):
if str(user.id) in ignored_users or str(user.name) in ignored_users:
return
# Read the found .json file:
workflow_json = ""
@ -131,10 +151,25 @@ async def on_message_or_reaction(client, obj):
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}")
time_start = time.perf_counter()
time_start = time.perf_counter()
# Indicate to the user something is happening:
await msg.add_reaction(WAITING_EMOJI)
@ -148,8 +183,15 @@ async def on_message_or_reaction(client, obj):
params_clone = params.copy()
# If the seed is not specified, generate one:
if "seed" not in params.keys():
params_clone["seed"] = random.randint(1, 999_999_999_999_999)
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}")
# Make replacements, e.g.
# __WIDTH__ to value of params["width"]
@ -157,6 +199,10 @@ async def on_message_or_reaction(client, obj):
for k in params_clone.keys():
v = params_clone[k]
k = k.upper()
if v is None:
continue
workflow_json_clone = re.sub(rf"__{k}__", str(v), workflow_json_clone)
# Must be valid JSON:

View File

@ -7,7 +7,7 @@ def get_prompt_parameters(user_message, settings):
params = settings["defaults"].copy()
mutable_string = user_message
mutable_string = re.sub(r"[^A-Za-z0-9 \-,\.>]", "", mutable_string)
mutable_string = re.sub(r"[^A-Za-z0-9 \-,\.>\(\)]", "", mutable_string)
mutable_string = re.sub(r"\s{1,}", " ", mutable_string)
# Get the options from the mutable string:
@ -65,7 +65,4 @@ def get_prompt_parameters(user_message, settings):
params_clone[target_key] = limit_value
params = params_clone
print(params)
return params

View File

@ -12,13 +12,15 @@ def get_settings(initialize=False):
"repeat_emoji": "♻️",
"waiting_emoji": "🕒",
"defaults": {
"checkpoint_1": "v1-5-pruned-emaonly.safetensors",
"width": 512,
"height": 512,
"cfg": 7,
"steps": 20,
"seed": 1,
"positive": "Test",
"negative": "Test",
"seed": None,
"positive": "",
"negative": "",
"negative_prefix": "((nsfw, nudity))"
},
}