custom_nodes_checker.py

import json
import os
import re
import sys

IGNORED_NODES = {"Note", "PrimitiveNode", "Reroute"}


def get_nodes_from_python_file(filepath):
    # Extract node names from a `nodes*.py` file based on its NODE_CLASS_MAPPINGS dictionary.
    if not os.path.exists(filepath):
        print(f"Error: '{filepath}' not found.")
        sys.exit(1)
    if os.path.isdir(filepath):
        print(f"Error: '{filepath}' is a directory, expected a Python file.")
        sys.exit(1)

    nodes = set()
    node_class_mappings_found = False
    node_pattern = re.compile(r'"([^"]+)"')

    with open(filepath, "r") as file:
        for line in file:
            line = line.strip()
            if line.startswith("NODE_CLASS_MAPPINGS"):
                node_class_mappings_found = True
                continue
            if node_class_mappings_found:
                if line.startswith("}"):
                    break
                # Find all matches between quotes and add to the set
                matches = node_pattern.findall(line)
                nodes.update(matches)

    return nodes


def get_built_in_nodes(comfyui_path):
    # Collect all built-in nodes by parsing `nodes.py` file and any` nodes_*.py` files
    # in `comfy_extras`.

    main_nodes_py_path = os.path.join(comfyui_path, "nodes.py")
    built_in_nodes = get_nodes_from_python_file(main_nodes_py_path)

    comfy_extras_dir = os.path.join(comfyui_path, "comfy_extras")
    if os.path.exists(comfy_extras_dir) and os.path.isdir(comfy_extras_dir):
        for filename in os.listdir(comfy_extras_dir):
            if filename.startswith("nodes_") and filename.endswith(".py"):
                extras_file_path = os.path.join(comfy_extras_dir, filename)
                extras_nodes = get_nodes_from_python_file(extras_file_path)
                built_in_nodes.update(extras_nodes)

    return built_in_nodes


def get_nodes_from_workflow_json_file(filepath):
    # Extract node types from a workflow JSON file.

    types = set()

    with open(filepath, "r") as file:
        data = json.load(file)
        nodes = data.get("nodes", [])
        for node in nodes:
            node_type = node.get("type")
            if node_type:
                types.add(node_type.strip())

    return list(types)


def compare_nodes(built_in_nodes, workflow_nodes):
    # Compare the built-in nodes with the nodes used in the workflow.

    workflow_nodes = [node for node in workflow_nodes if node not in IGNORED_NODES]
    missing_nodes = sorted(
        [
            workflow_node
            for workflow_node in workflow_nodes
            if workflow_node not in built_in_nodes
        ]
    )
    return missing_nodes


def load_extension_node_map(extension_node_map_path):
    # Load ComfyUI Manager `extension-node-map.json` file if it exists.

    if not os.path.exists(extension_node_map_path):
        print(
            "ComfyUI Manager is not installed. Install it to get additional custom node information."
        )
        return None

    with open(extension_node_map_path, "r") as file:
        extension_node_map = json.load(file)

    return extension_node_map


def find_node_in_extension_map(node_name, extension_node_map):
    # Find the URL corresponding to a node name in the `extension-node-map.json`.

    for url, (nodes, _) in extension_node_map.items():
        if node_name in nodes:
            return url
    return None


def process_workflow_json(file_path, built_in_nodes, extension_node_map):
    # Process a workflow JSON file, comparing its nodes with the built-in nodes and
    # optionally checking the `extension-node-map.json`.

    workflow_nodes = get_nodes_from_workflow_json_file(file_path)
    missing_nodes = compare_nodes(built_in_nodes, workflow_nodes)
    json_filename = os.path.basename(file_path)  # Get the filename from the path
    print(f"\n{json_filename}")
    if missing_nodes:
        for missing_node in missing_nodes:
            if extension_node_map:
                extension_url = find_node_in_extension_map(
                    missing_node, extension_node_map
                )
                if extension_url:
                    print(f"  {missing_node} / {extension_url}")
                else:
                    print(f"  {missing_node} / NOT FOUND")
            else:
                print(f"  {missing_node} / NOT FOUND")
    else:
        print("  Workflow does not use custom nodes.")


def process_workflow_dir(
    dir_path, built_in_nodes, extension_node_map, current_level=0, max_levels=2
):
    # Recursively process a directory of workflow JSON files.

    if current_level > max_levels:
        return

    if not os.path.exists(dir_path):
        print(f"Error: Directory '{dir_path}' not found.")
        sys.exit(1)

    for entry in os.listdir(dir_path):
        entry_path = os.path.join(dir_path, entry)
        if os.path.isfile(entry_path) and entry_path.endswith(".json"):
            process_workflow_json(entry_path, built_in_nodes, extension_node_map)
        elif os.path.isdir(entry_path):
            process_workflow_dir(
                entry_path, built_in_nodes, extension_node_map, current_level + 1
            )


def print_help():

    help_message = """
    [ComfyUI Workflow Custom Nodes Checker]

    This script reads ComfyUI workflows (either a single workflow or workflows in a
    directory) and displays the custom nodes the workflows require.

    Usage:
    python custom_nodes_checker.py [<comfyui_path>] [<workflow_path>]

    Arguments:
    <comfyui_path>   - Optional. Path to ComfyUI root directory or directly to `nodes.py` file.
                       If not provided, assumes `nodes.py` is in the same directory
                       as the script.
    <workflow_path>  - Optional. Path to a JSON workflow file or workflows directory.
                       If not provided, assumes a `workflows` directory in the same
                       directory as the script.

    Options:
    -h, --help       - Display this help message.

    Example:
    python custom_nodes_checker.py /path/to/comfyui /path/to/workflow.json
    python custom_nodes_checker.py /path/to/comfyui /path/to/directory/
    python custom_nodes_checker.py
    """
    print(help_message)


def main(comfyui_path=None, workflow_path=None):
    """
    Main function to orchestrate the processing of workflows and comparison of nodes.

    Args:
    - comfyui_path (str or None): The path to the ComfyUI root directory.
    - workflow_path (str or None): The path to the workflow JSON file or directory.
    """
    script_dir = os.path.dirname(os.path.abspath(__file__))
    comfyui_path = comfyui_path if comfyui_path else script_dir
    workflow_dir_path = (
        workflow_path if workflow_path else os.path.join(script_dir, "workflows")
    )

    built_in_nodes = get_built_in_nodes(comfyui_path)

    extension_node_map_path = os.path.join(
        comfyui_path,
        "custom_nodes",
        "ComfyUI-Manager",
        "node_db",
        "new",
        "extension-node-map.json",
    )
    extension_node_map = load_extension_node_map(extension_node_map_path)

    if os.path.isfile(workflow_dir_path) and workflow_dir_path.endswith(".json"):
        process_workflow_json(workflow_dir_path, built_in_nodes, extension_node_map)
    elif os.path.isdir(workflow_dir_path):
        process_workflow_dir(workflow_dir_path, built_in_nodes, extension_node_map)
    else:
        print(f"Error: Provided workflow path '{workflow_dir_path}' is invalid.")
        sys.exit(1)


if __name__ == "__main__":

    if len(sys.argv) == 2 and sys.argv[1] in ("--help", "-h"):
        print_help()
        sys.exit(0)

    comfyui_path = sys.argv[1] if len(sys.argv) > 1 else None
    workflow_path = sys.argv[2] if len(sys.argv) > 2 else None

    main(comfyui_path, workflow_path)