244 lines
7.7 KiB
Python
244 lines
7.7 KiB
Python
import discord
|
|
import os
|
|
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 *
|
|
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
|
|
|
|
#
|
|
#
|
|
#
|
|
|
|
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]
|
|
|
|
# Try different paths to find workflow .json file:
|
|
workflow_paths = [
|
|
f"conf/{chl.category.name}/{chl_topic_part_1}",
|
|
f"conf/{chl.category.name}/{chl.name}",
|
|
f"conf/{chl_topic_part_1}",
|
|
f"conf/{chl.name}",
|
|
f"conf/{chl.category.name}",
|
|
]
|
|
|
|
using_workflow_path = None
|
|
using_settings_path = None
|
|
|
|
for path in workflow_paths:
|
|
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"conf/{chl.category.name}",
|
|
f"conf/{chl.name}",
|
|
]
|
|
|
|
for i in chl_topic_parts:
|
|
setting_paths = setting_paths + ["conf/" + i.strip()]
|
|
|
|
setting_paths = setting_paths + [f"conf/{chl.category.name}/{chl.name}"]
|
|
|
|
for i in chl_topic_parts:
|
|
setting_paths = setting_paths + [f"conf/{chl.category.name}/" + i.strip()]
|
|
|
|
settings = get_settings()
|
|
|
|
for path in setting_paths:
|
|
try_path = f"{path}_settings.json"
|
|
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 = []
|
|
|
|
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)
|
|
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:
|
|
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)
|
|
|
|
#
|
|
#
|
|
#
|
|
|
|
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
|