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