257 lines
8.2 KiB
Python
257 lines
8.2 KiB
Python
import discord
|
|
import os
|
|
import json
|
|
import time
|
|
import io
|
|
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 *
|
|
from lib.comfyui import *
|
|
|
|
async def on_message_or_reaction(client, obj):
|
|
global RUN_TIME
|
|
|
|
msg = None
|
|
chl = None
|
|
user = None
|
|
author = None
|
|
roles = None
|
|
rxn = None
|
|
|
|
msg_types = [discord.MessageType.default]
|
|
|
|
# Handle reactions:
|
|
if isinstance(obj, discord.RawReactionActionEvent):
|
|
chl = await client.fetch_channel(obj.channel_id)
|
|
msg = await chl.fetch_message(obj.message_id)
|
|
user = await client.fetch_user(obj.user_id)
|
|
author = await client.fetch_user(obj.message_author_id)
|
|
roles = obj.member.roles
|
|
rxn = obj
|
|
|
|
msg_types.append(discord.MessageType.reply)
|
|
|
|
# Handle direct messages from users:
|
|
if isinstance(obj, discord.Message):
|
|
msg = obj
|
|
chl = obj.channel
|
|
user = msg.author
|
|
author = msg.author
|
|
roles = msg.author.roles
|
|
rxn = None
|
|
|
|
if user == client.user: return
|
|
if msg.type not in msg_types: return
|
|
if user.bot: return
|
|
|
|
#
|
|
#
|
|
#
|
|
|
|
settings = get_settings()
|
|
|
|
CONFIG_DIR = settings["config_dir"]
|
|
|
|
chl_topic_parts = []
|
|
chl_topic_part_1 = ""
|
|
if chl.topic is not None:
|
|
chl_topic_parts = chl.topic.split(",")
|
|
chl_topic_part_1 = chl_topic_parts[0]
|
|
|
|
using_workflow_path = None
|
|
|
|
for path in [f"{CONFIG_DIR}workflows/{chl_topic_part_1}"]:
|
|
try_path = f"{path}.json"
|
|
if os.path.isfile(try_path):
|
|
using_workflow_path = try_path
|
|
log.write(f"Workflow {log.colors.fg.lightcyan}{try_path}")
|
|
|
|
if using_workflow_path is None:
|
|
return
|
|
|
|
#
|
|
#
|
|
#
|
|
|
|
setting_paths = [
|
|
f"{CONFIG_DIR}guild_defaults/{chl.guild.id}",
|
|
f"{CONFIG_DIR}guild_defaults/{chl.guild.name}",
|
|
f"{CONFIG_DIR}category_defaults/{chl.category.id}",
|
|
f"{CONFIG_DIR}category_defaults/{chl.category.name}",
|
|
f"{CONFIG_DIR}channel_defaults/{chl.id}",
|
|
f"{CONFIG_DIR}channel_defaults/{chl.name}",
|
|
]
|
|
|
|
for i in chl_topic_parts:
|
|
chl_topic_part = i.strip()
|
|
setting_paths = setting_paths + [f"{CONFIG_DIR}defaults/{chl_topic_part}"]
|
|
|
|
for r in roles:
|
|
setting_paths = setting_paths + [f"{CONFIG_DIR}role_defaults/{r.name}"]
|
|
|
|
setting_paths = setting_paths + [
|
|
f"{CONFIG_DIR}user_defaults/{user.id}",
|
|
f"{CONFIG_DIR}user_defaults/{user.name}",
|
|
]
|
|
|
|
for path in setting_paths:
|
|
try_path = f"{path}.json"
|
|
log.write(f"Seeking {try_path}")
|
|
if os.path.isfile(try_path):
|
|
settings = merge_dicts(settings, read_json(try_path, {}))
|
|
log.write(f"Merging {log.colors.fg.lightcyan}{try_path}")
|
|
|
|
#
|
|
#
|
|
#
|
|
|
|
ALLOW_REPEAT = settings["allow_repeat"]
|
|
SHOW_REPEAT = settings["show_repeat"]
|
|
REPEAT_EMOJI = settings["repeat_emoji"]
|
|
WAITING_EMOJI = settings["waiting_emoji"]
|
|
|
|
if rxn is not None:
|
|
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:
|
|
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:
|
|
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:
|
|
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:
|
|
workflow_json = file.read()
|
|
|
|
# Break the user message into prompt parameters:
|
|
params = get_prompt_parameters(msg.content, settings)
|
|
|
|
for k in params.keys():
|
|
v = params[k]
|
|
k = k.upper()
|
|
label = f"__{k}__"
|
|
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:
|
|
await msg.add_reaction(WAITING_EMOJI)
|
|
await chl.typing()
|
|
|
|
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
|
|
params_clone = params.copy()
|
|
|
|
# If the seed is not specified, generate one:
|
|
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
|
|
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"]
|
|
# __POSITIVE__ to value of params["positive"]
|
|
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:
|
|
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
|
|
|
|
# Process 8 attachments at a time per one Discord message:
|
|
while True:
|
|
attachments_buffer = all_attachments[:8]
|
|
discord_files = []
|
|
|
|
for a in attachments_buffer:
|
|
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
|
|
|
|
post = await chl.send(files=discord_files, content=msg.content, reference=msg)
|
|
del all_attachments[:8]
|
|
|
|
if len(all_attachments) < 1:
|
|
break
|
|
|
|
if ALLOW_REPEAT:
|
|
if SHOW_REPEAT:
|
|
await msg.add_reaction(REPEAT_EMOJI)
|
|
await post.add_reaction(REPEAT_EMOJI)
|
|
|
|
await msg.remove_reaction(WAITING_EMOJI, client.user)
|
|
|
|
time_end = time.perf_counter()
|
|
time_taken = time_end - time_start
|
|
|
|
log.write(f"{time_taken:0.2f}s taken.", log.colors.fg.lightgreen)
|
|
|
|
#
|
|
#
|
|
#
|
|
|
|
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
|