243 lines
6.8 KiB
Python
243 lines
6.8 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):
|
|
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
|
|
|
|
log.write("\nSearching for workflow in order:\n")
|
|
|
|
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
|
|
|
|
#
|
|
#
|
|
#
|
|
|
|
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()
|
|
|
|
log.write("Merging settings in order:")
|
|
for path in setting_paths:
|
|
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}")
|
|
|
|
#
|
|
#
|
|
#
|
|
|
|
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:
|
|
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 = ""
|
|
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)
|
|
|
|
|
|
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()
|
|
|
|
# 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
|
|
|
|
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"]
|
|
# __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)
|
|
|
|
# Process 8 attachments at a time per one Discord message:
|
|
while True:
|
|
attachments_buffer = all_attachments[:8]
|
|
discord_files = []
|
|
|
|
for a in attachments_buffer:
|
|
#base64_str = a
|
|
#base64_bytes = base64_str.encode("ascii")
|
|
#attachment_bytes = base64.b64decode(base64_bytes)
|
|
|
|
f = discord.File(io.BytesIO(a), filename="file.png")
|
|
discord_files.append(f)
|
|
|
|
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
|