From 9d3eb747963fe927773902344959f0f7a1d3e545 Mon Sep 17 00:00:00 2001 From: doctorpangloss <@hiddenswitch.com> Date: Wed, 28 Feb 2024 14:11:34 -0800 Subject: [PATCH] Fix importing vanilla custom nodes --- README.md | 15 +++- comfy/cmd/execution.py | 4 + comfy/nodes/package.py | 32 ++++---- comfy/nodes/vanilla_node_importing.py | 106 ++++++++++++++++++++++++++ 4 files changed, 140 insertions(+), 17 deletions(-) create mode 100644 comfy/nodes/vanilla_node_importing.py diff --git a/README.md b/README.md index 462ea71b1..b49b09f3c 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,8 @@ A vanilla, up-to-date fork of [ComfyUI](https://github.com/comfyanonymous/comfyu - [Workflows](https://comfyanonymous.github.io/ComfyUI_examples/) - [Installing](#installing) - [Configuration](#command-line-arguments) -- [Custom Nodes Authoring](#custom-nodes) +- [Installing Custom Nodes](#installing-custom-nodes) +- [Authoring Custom Nodes](#custom-nodes) - [API](#using-comfyui-as-an-api--programmatically) - [Distributed](#distributed-multi-process-and-multi-gpu-comfy) @@ -196,6 +197,18 @@ For AMD 7600 and maybe other RDNA3 cards: ```HSA_OVERRIDE_GFX_VERSION=11.0.0 com Custom Nodes can be added to ComfyUI by copying and pasting Python files into your `./custom_nodes` directory. +## Installing Custom Nodes + +There are two kinds of custom nodes: vanilla custom nodes, which generally expect to be dropped into the `custom_nodes` directory and managed by a tool called the ComfyUI Extension manager ("vanilla" custom nodes) and this repository's opinionated, installable custom nodes ("installable"). + +### Vanilla Custom Nodes + +Clone the repository containing the custom nodes into `custom_nodes/` in your working directory. + +### Installable Custom Nodes + +Run `pip install git+https://github.com/owner/repository`, replacing the `git` repository with the installable custom nodes URL. This is just the GitHub URL. + ## Authoring Custom Nodes Create a `requirements.txt`: diff --git a/comfy/cmd/execution.py b/comfy/cmd/execution.py index 930248cb8..f9a0fa2d6 100644 --- a/comfy/cmd/execution.py +++ b/comfy/cmd/execution.py @@ -11,6 +11,10 @@ import typing from typing import List, Optional, Tuple, Union from typing_extensions import TypedDict +# Suppress warnings during import +import warnings +warnings.filterwarnings("ignore", message="torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.") + import torch from ..component_model.abstract_prompt_queue import AbstractPromptQueue diff --git a/comfy/nodes/package.py b/comfy/nodes/package.py index ed971d765..a7e6e5fb2 100644 --- a/comfy/nodes/package.py +++ b/comfy/nodes/package.py @@ -6,20 +6,16 @@ import os import pkgutil import time import types -import typing - -from . import base_nodes -from comfy_extras import nodes as comfy_extras_nodes - -try: - import custom_nodes -except: - custom_nodes: typing.Optional[types.ModuleType] = None -from .package_typing import ExportedNodes from functools import reduce -from pkg_resources import resource_filename from importlib.metadata import entry_points +from pkg_resources import resource_filename + +from comfy_extras import nodes as comfy_extras_nodes +from . import base_nodes +from .package_typing import ExportedNodes +from .vanilla_node_importing import mitigated_import_of_vanilla_custom_nodes + _comfy_nodes = ExportedNodes() @@ -42,7 +38,8 @@ def _import_nodes_in_module(exported_nodes: ExportedNodes, module: types.ModuleT return node_class_mappings and len(node_class_mappings) > 0 or web_directory -def _import_and_enumerate_nodes_in_module(module: types.ModuleType, print_import_times=False) -> ExportedNodes: +def _import_and_enumerate_nodes_in_module(module: types.ModuleType, print_import_times=False, + depth=100) -> ExportedNodes: exported_nodes = ExportedNodes() timings = [] if _import_nodes_in_module(exported_nodes, module): @@ -60,7 +57,8 @@ def _import_and_enumerate_nodes_in_module(module: types.ModuleType, print_import submodule = importlib.import_module(full_name) # Recursively call the function if it's a package exported_nodes.update( - _import_and_enumerate_nodes_in_module(submodule, print_import_times=print_import_times)) + _import_and_enumerate_nodes_in_module(submodule, print_import_times=print_import_times, + depth=depth - 1)) except KeyboardInterrupt as interrupted: raise interrupted except Exception as x: @@ -78,7 +76,8 @@ def _import_and_enumerate_nodes_in_module(module: types.ModuleType, print_import return exported_nodes -def import_all_nodes_in_workspace() -> ExportedNodes: +def import_all_nodes_in_workspace(vanilla_custom_nodes=True) -> ExportedNodes: + global _comfy_nodes if len(_comfy_nodes) == 0: base_and_extra = reduce(lambda x, y: x.update(y), map(_import_and_enumerate_nodes_in_module, [ @@ -88,8 +87,9 @@ def import_all_nodes_in_workspace() -> ExportedNodes: ]), ExportedNodes()) custom_nodes_mappings = ExportedNodes() - if custom_nodes is not None: - custom_nodes_mappings.update(_import_and_enumerate_nodes_in_module(custom_nodes, print_import_times=True)) + + if vanilla_custom_nodes: + custom_nodes_mappings += mitigated_import_of_vanilla_custom_nodes() # load from entrypoints for entry_point in entry_points().select(group='comfyui.custom_nodes'): diff --git a/comfy/nodes/vanilla_node_importing.py b/comfy/nodes/vanilla_node_importing.py new file mode 100644 index 000000000..b50f977ca --- /dev/null +++ b/comfy/nodes/vanilla_node_importing.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import importlib +import os +import sys +import time +import types +from typing import Dict + +from . import base_nodes +from .package_typing import ExportedNodes + + +def _vanilla_load_custom_nodes_1(module_path, ignore=set()) -> ExportedNodes: + exported_nodes = ExportedNodes() + module_name = os.path.basename(module_path) + if os.path.isfile(module_path): + sp = os.path.splitext(module_path) + module_name = sp[0] + try: + if os.path.isfile(module_path): + module_spec = importlib.util.spec_from_file_location(module_name, module_path) + module_dir = os.path.split(module_path)[0] + else: + module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py")) + module_dir = module_path + + module = importlib.util.module_from_spec(module_spec) + sys.modules[module_name] = module + module_spec.loader.exec_module(module) + + if hasattr(module, "WEB_DIRECTORY") and getattr(module, "WEB_DIRECTORY") is not None: + web_dir = os.path.abspath(os.path.join(module_dir, getattr(module, "WEB_DIRECTORY"))) + if os.path.isdir(web_dir): + exported_nodes.EXTENSION_WEB_DIRS[module_name] = web_dir + + if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None: + for name in module.NODE_CLASS_MAPPINGS: + if name not in ignore: + exported_nodes.NODE_CLASS_MAPPINGS[name] = module.NODE_CLASS_MAPPINGS[name] + if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, + "NODE_DISPLAY_NAME_MAPPINGS") is not None: + exported_nodes.NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS) + return exported_nodes + else: + print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.") + return exported_nodes + except Exception as e: + import traceback + print(traceback.format_exc()) + print(f"Cannot import {module_path} module for custom nodes:", e) + return exported_nodes + + +def _vanilla_load_custom_nodes_2() -> ExportedNodes: + from ..cmd import folder_paths + base_node_names = set(base_nodes.NODE_CLASS_MAPPINGS.keys()) + node_paths = folder_paths.get_folder_paths("custom_nodes") + node_import_times = [] + exported_nodes = ExportedNodes() + for custom_node_path in node_paths: + if not os.path.exists(custom_node_path) or not os.path.isdir(custom_node_path): + pass + possible_modules = os.listdir(os.path.realpath(custom_node_path)) + if "__pycache__" in possible_modules: + possible_modules.remove("__pycache__") + + for possible_module in possible_modules: + module_path = os.path.join(custom_node_path, possible_module) + if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue + if module_path.endswith(".disabled"): continue + time_before = time.perf_counter() + possible_exported_nodes = _vanilla_load_custom_nodes_1(module_path, base_node_names) + node_import_times.append( + (time.perf_counter() - time_before, module_path, len(possible_exported_nodes.NODE_CLASS_MAPPINGS) > 0)) + exported_nodes.update(possible_exported_nodes) + + if len(node_import_times) > 0: + print("\nImport times for custom nodes:") + for n in sorted(node_import_times): + if n[2]: + import_message = "" + else: + import_message = " (IMPORT FAILED)" + print("{:6.1f} seconds{}:".format(n[0], import_message), n[1]) + print() + return exported_nodes + + +def mitigated_import_of_vanilla_custom_nodes() -> ExportedNodes: + # only vanilla custom nodes will ever go into the custom_nodes directory + # this mitigation puts files that custom nodes expects are at the root of the repository back where they should be found + from ..cmd import cuda_malloc, folder_paths, execution, server, latent_preview + for module in (cuda_malloc, folder_paths, execution, server, latent_preview): + module_short_name = module.__name__.split(".")[-1] + sys.modules[module_short_name] = module + sys.modules['nodes'] = base_nodes + comfy_extras_mitigation: Dict[str, types.ModuleType] = {} + for module_name, module in sys.modules.items(): + if not module_name.startswith("comfy_extras.nodes"): + continue + module_short_name = module_name.split(".")[-1] + comfy_extras_mitigation[f'comfy_extras.{module_short_name}'] = module + sys.modules.update(comfy_extras_mitigation) + vanilla_custom_nodes = _vanilla_load_custom_nodes_2() + return vanilla_custom_nodes