comfyui-discord/lib/events.py

import discord
import os
import json
import time
import io
import random

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]

    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)

    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:
        for extension in ["", ".json"]:
            if using_workflow_path is not None:
                continue

            try_path = f"{path}{extension}"
            print(f"Looking for workflow {try_path}")
            if os.path.isfile(try_path):
                using_workflow_path = try_path

    if using_workflow_path is None:
        return

    print(f"Using workflow {using_workflow_path}")

    #
    #
    #

    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:
        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, {}))

    #
    #
    #

    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


    # 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)





    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():
            params_clone["seed"] = random.randint(1, 999_999_999_999_999)

        # 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()
            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