Imagine/imagine-v5/vyro_workflows/__init__.py

import yaml
import shutil
import json
import os
import sys

import server
from server import web
    
from .nodes.pipe_input import NODE_CLASS_MAPPINGS as PipeInput_MAPPING
from .nodes.perlin import NODE_CLASS_MAPPINGS as Perlin_MAPPING
from .nodes.loaders import NODE_CLASS_MAPPINGS as Loaders_MAPPING
from .nodes.prompt_analyzer import NODE_CLASS_MAPPINGS as PromptAnalyzer_MAPPING
from .nodes.routing import NODE_CLASS_MAPPINGS as Routing_MAPPING
from .nodes.samplers import NODE_CLASS_MAPPINGS as Samplers_MAPPING
from .nodes.interposer import NODE_CLASS_MAPPINGS as Interposer_MAPPING
from .nodes.IPAdapterPlus import NODE_CLASS_MAPPINGS as IPAdapterPlus_MAPPING
from .nodes.util import NODE_CLASS_MAPPINGS as Utils_MAPPING



NODE_CLASS_MAPPINGS = {
    **PipeInput_MAPPING,
    **Perlin_MAPPING,
    **Loaders_MAPPING,
    **PromptAnalyzer_MAPPING,
    **Routing_MAPPING,
    **Samplers_MAPPING,
    **Interposer_MAPPING,
    **IPAdapterPlus_MAPPING,

    **Utils_MAPPING,
   
}


try:
    dir_path = os.path.dirname(os.path.abspath(__file__))
    sys.path.append(dir_path)
    from directory_paths import ext_path, default_nested_nodes_path, config_path, js_extensions_repo_path
except ImportError as e:
    print(f"[NestedNodeBuilder] Error importing directory_paths.py: {e}")


# Config keys
config_nested_nodes_path = "nested_nodes_path"

# Load config
config = {}
with open(config_path, 'r') as config_file:
    try:
        config = yaml.safe_load(config_file)
    except yaml.YAMLError as yaml_e:
        print("[NestedNodeBuilder] Error loading NestedNodeBuilder config:", yaml_e)


def load_nested_node_defs():
    defs_list = []
    if config_nested_nodes_path not in config:
        print(
            f'[NestedNodeBuilder] missing entry "{config_nested_nodes_path}" in config.yaml, \
            using default path "{default_nested_nodes_path}".'
        )
    nested_nodes_path = config.get(config_nested_nodes_path, default_nested_nodes_path)
    if not os.path.isabs(nested_nodes_path):
        nested_nodes_path = os.path.join(ext_path, nested_nodes_path)
    if not os.path.exists(nested_nodes_path):
        os.makedirs(nested_nodes_path)

    # Load each json file in the nested_nodes_path and add it to the defs list
    for file_name in os.listdir(nested_nodes_path):
        if file_name.endswith(".json"):
            file_path = os.path.join(nested_nodes_path, file_name)
            with open(file_path, 'r') as json_file:
                try:
                    node_def = json.load(json_file)
                except json.JSONDecodeError as json_e:
                    print(f"[NestedNodeBuilder] Error loading {file_name}:", json_e)
                    continue
                if "name" not in node_def:
                    print("[NestedNodeBuilder] missing property \"name\" in node definition:", file_name)
                    continue
                defs_list.append(node_def)
    defs = {}
    for node_def in defs_list:
        key = node_def["name"]
        defs[key] = node_def

    # Add the nested node defs
    # import nodes
    # from .nested_nodes import createNestedNodeClass
    # for node_def in defs_list:
    #     nodes.NODE_CLASS_MAPPINGS[node_def["name"]] = createNestedNodeClass(node_def)
    #     nodes.NODE_DISPLAY_NAME_MAPPINGS[node_def["name"]] = node_def["display_name"]

    return defs


def save_nested_def(node_def):
    if config_nested_nodes_path not in config:
        print(
            f'[NestedNodeBuilder] missing entry "{config_nested_nodes_path}" in config.yaml, \
            using default path "{default_nested_nodes_path}".'
        )
    nested_nodes_path = config.get(config_nested_nodes_path, default_nested_nodes_path)
    if not os.path.isabs(nested_nodes_path):
        nested_nodes_path = os.path.join(ext_path, nested_nodes_path)
    if not os.path.exists(nested_nodes_path):
        os.makedirs(nested_nodes_path)
    file_name = node_def["name"] + ".json"
    file_path = os.path.join(nested_nodes_path, file_name)
    # Raise error if file already exists
    # if os.path.exists(file_path):
    #     raise FileExistsError(f"[NestedNodeBuilder] {file_name} already exists.")
    with open(file_path, 'w') as json_file:
        json.dump(node_def, json_file, indent=4)


def place_js():
    src = os.path.join(ext_path, "js")
    dst = js_extensions_repo_path
    shutil.copytree(src, dst, dirs_exist_ok=True)

try: 
    @server.PromptServer.instance.routes.get('/nested_node_builder/nested_defs')
    async def server_add_def_route(request):
        nested_node_defs = load_nested_node_defs()
        return web.json_response(nested_node_defs)


    @server.PromptServer.instance.routes.post('/nested_node_builder/nested_defs')
    async def server_add_def_route(request):
        nested_def = await request.json()
        save_nested_def(nested_def)
        return web.Response(text="ok")

    #
    # Setup
    #
    WEB_DIRECTORY = "./js"
    skip_js_copy = hasattr(server.PromptServer.instance, "supports") and "custom_nodes_from_web" in server.PromptServer.instance.supports
    if not skip_js_copy:
        print("[NestedNodeBuilder] Installed ComfyUI version doesn't support in-place web extension loading. Copying files to web directory...")
        place_js()
except:
    print("Prompt Server not initliazed yet, Skipping server routes.")