diff --git a/easy_nodes/__init__.py b/easy_nodes/__init__.py index 3320dcc..2283a71 100644 --- a/easy_nodes/__init__.py +++ b/easy_nodes/__init__.py @@ -20,11 +20,11 @@ TensorVerifier, TypeVerifier, create_field_setter_node, + get_node_mappings, initialize_easy_nodes, register_type, show_image, show_text, - get_node_mappings ) # For backwards compatibility with the original comfy_annotations module. diff --git a/easy_nodes/comfy_types.py b/easy_nodes/comfy_types.py index ba20348..cf76b8c 100644 --- a/easy_nodes/comfy_types.py +++ b/easy_nodes/comfy_types.py @@ -5,31 +5,19 @@ from easy_nodes.easy_nodes import AnythingVerifier, TensorVerifier, TypeVerifier, register_type, AnyType, any_type -class ImageTensor(torch.Tensor): - pass +class ImageTensor(torch.Tensor): pass +class MaskTensor(torch.Tensor): pass -class MaskTensor(torch.Tensor): - pass - -class LatentTensor(torch.Tensor): - pass - -class ConditioningTensor(torch.Tensor): - pass - -class ModelTensor(torch.Tensor): - pass - -class SigmasTensor(torch.Tensor): - pass +class LatentTensor(torch.Tensor): pass +class ConditioningTensor(torch.Tensor): pass +class ModelTensor(torch.Tensor): pass +class SigmasTensor(torch.Tensor): pass # Maybe there's an actual class for this? -class PhotoMaker: - pass +class PhotoMaker: pass # Abstract type, not for instantiating. -class NumberType: - pass +class NumberType: pass # ComfyUI will get the special string that anytype is registered with, which is hardcoded to match anything. diff --git a/easy_nodes/easy_nodes.py b/easy_nodes/easy_nodes.py index 389dd1e..34f9f6a 100644 --- a/easy_nodes/easy_nodes.py +++ b/easy_nodes/easy_nodes.py @@ -1,4 +1,5 @@ +import enum import functools import hashlib import importlib @@ -59,6 +60,7 @@ class EasyNodesConfig: NODE_CLASS_MAPPINGS: dict NODE_DISPLAY_NAME_MAPPINGS: dict num_registered: int = 0 + get_node_mappings_called: bool = False # Keep track of the config from the last init, because different custom_nodes modules @@ -66,8 +68,13 @@ class EasyNodesConfig: _current_config: EasyNodesConfig = None +# Changing the default of auto-register to false in 1.1, so catch if the user hasn't set it explicitly so we can give them a warning. +class AutoRegisterSentinel(enum.Enum): + DEFAULT = enum.auto() + + def initialize_easy_nodes(default_category: str = "EasyNodes", - auto_register: bool = True, + auto_register: bool = AutoRegisterSentinel.DEFAULT, docstring_mode: AutoDescriptionMode = AutoDescriptionMode.FULL, verify_level: CheckSeverityMode = CheckSeverityMode.WARN, auto_move_tensors: bool = False): @@ -78,7 +85,7 @@ def initialize_easy_nodes(default_category: str = "EasyNodes", Args: default_category (str, optional): The default category for nodes. Defaults to "EasyNodes". - auto_register (bool, optional): Whether to automatically register nodes with ComfyUI (so you don't have to export). Defaults to True. Experimental. + auto_register (bool, optional): Whether to automatically register nodes with ComfyUI (so you don't have to export). Defaults to False. Experimental. docstring_mode (AutoDescriptionMode, optional): The mode for generating node docstrings. Defaults to AutoDescriptionMode.FULL. verify_level (bool, optional): Whether to verify tensors for shape and data type according to ComfyUI type (MASK, IMAGE, etc). Runs on inputs and outputs. Defaults to False. auto_move_tensors (bool, optional): Whether to automatically move torch Tensors to the GPU before your function gets called, and then to the CPU on output. Defaults to False. @@ -94,28 +101,37 @@ def initialize_easy_nodes(default_category: str = "EasyNodes", assert _current_config.auto_register or not _current_config.NODE_CLASS_MAPPINGS, ( f"Auto-registration was turned off by previous initializer, but {len(_current_config.NODE_CLASS_MAPPINGS)} nodes were not picked up.") + NODE_CLASS_MAPPINGS = {} + NODE_DISPLAY_NAME_MAPPINGS = {} + logging.info(f"Initializing EasyNodes. Auto-registration: {auto_register}") - if auto_register: + if auto_register is AutoRegisterSentinel.DEFAULT: + logging.warning("Auto-registration not set explicitly. Will default to False in a future version. Call easy_nodes.initialize_easy_nodes(auto_register=True|False) to suppress this warning.") + + if auto_register is True: NODE_CLASS_MAPPINGS = comfyui_nodes.NODE_CLASS_MAPPINGS NODE_DISPLAY_NAME_MAPPINGS = comfyui_nodes.NODE_DISPLAY_NAME_MAPPINGS + + if auto_register is True or auto_register is AutoRegisterSentinel.DEFAULT: frame = sys._getframe(1).f_globals['__name__'] _ensure_package_dicts_exist(frame) - else: - NODE_CLASS_MAPPINGS = {} - NODE_DISPLAY_NAME_MAPPINGS = {} + _current_config = EasyNodesConfig(default_category, auto_register, docstring_mode, verify_level, auto_move_tensors, NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS) def get_node_mappings(): - assert _current_config is not None, "EasyNodes not initialized. Call easy_nodes.init() before using ComfyFunc." - assert not _current_config.auto_register, "Auto-node registration is on. Call easy_nodes.init(auto_register=False) if you want to export manually." - _current_config.initialized = False + assert _current_config is not None, "EasyNodes not initialized. Call easy_nodes.initialize_easy_nodes() before using ComfyNode." + assert _current_config.num_registered > 0, "No nodes registered. Use the @ComfyNode() decorator to register nodes after calling easy_nodes.initialize_easy_nodes()." + assert _current_config.auto_register is not True, "Auto-node registration is on. Call easy_nodes.initialize_easy_nodes(auto_register=False) if you want to export manually." + assert not _current_config.get_node_mappings_called, "get_node_mappings() already called. This function should only be called once." + _current_config.get_node_mappings_called = True return _current_config.NODE_CLASS_MAPPINGS, _current_config.NODE_DISPLAY_NAME_MAPPINGS -def _get_curr_config(): +def _get_curr_config() -> EasyNodesConfig: if _current_config is None: + logging.warning("easy_nodes.initialize_easy_nodes() should be called prior to any other EasyNodes activity. Initializing now with easy_nodes.initialize_easy_nodes() for backwards compatibility.") easy_nodes.initialize_easy_nodes() return _current_config @@ -737,7 +753,7 @@ def decorator(func: callable): filename = func.__code__.co_filename - wrapped_name = func.__qualname__ + "_comfyfunc_wrapper" + wrapped_name = func.__qualname__ + "_comfynode_wrapper" source_location = f"{filename}:{func.__code__.co_firstlineno}" code_origin_loc = f"\n Source: {func.__qualname__} {source_location}" original_is_changed = is_changed @@ -823,6 +839,11 @@ def wrapped_is_changed(*args, **kwargs): @functools.wraps(func) def wrapper(*args, **kwargs): + if curr_config.auto_register == AutoRegisterSentinel.DEFAULT and curr_config.get_node_mappings_called is False: + logging.warning("EasyNodes auto-registration not explicitly enabled, and easy_nodes.get_node_mappings() has not been called. " + + "In the future auto_register will default to False, so please set explicitly via easy_nodes.initialize_easy_nodes(auto_register=True), " + + "or use easy_nodes.get_node_mappings() to export to ComfyUI after all ComfyNodes have been created.") + if debug: logger.info( f"Calling {func.__name__} with {len(args)} args and {len(kwargs)} kwargs. Is class method: {is_cls_mth}" @@ -1181,6 +1202,12 @@ def _create_comfy_node( class_map[workflow_name] = node_class display_map[workflow_name] = display_name + + # Temporary for backwards compatibility. + if easy_nodes_config.auto_register is AutoRegisterSentinel.DEFAULT: + comfyui_nodes.NODE_CLASS_MAPPINGS[workflow_name] = node_class + comfyui_nodes.NODE_DISPLAY_NAME_MAPPINGS[workflow_name] = display_name + easy_nodes_config.num_registered += 1 diff --git a/example/__init__.py b/example/__init__.py index 553012e..0cdd011 100644 --- a/example/__init__.py +++ b/example/__init__.py @@ -1,12 +1,8 @@ -# Simply importing your module gives it a chance to add the @ComfyFunc nodes since -# EasyNodes will automatically export the NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS -# for you. -from .example_nodes import * +import easy_nodes +easy_nodes.initialize_easy_nodes(default_category="EasyNodes Examples", auto_register=False) +# Simply importing your module gives the ComfyNode decorator a chance to register your nodes. +from .example_nodes import * # noqa: F403, E402 -# Alternatively, to export yourself, you can do the following: -# import easy_nodes -# easy_nodes.init(auto_register=False) -# import example.example_nodes -# NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = easy_nodes.get_node_mappings() -# __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file +NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = easy_nodes.get_node_mappings() +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/example/example_nodes.py b/example/example_nodes.py index 49b3d6c..59ea4df 100644 --- a/example/example_nodes.py +++ b/example/example_nodes.py @@ -11,16 +11,11 @@ import easy_nodes import torch - -my_category = "EasyNodes Examples" - -# By default EasyNodes will auto-register any decorated function (automatically insert it into ComfyUI's node registry). -# If you want to manually register your nodes the regular way, turn off -# auto_register and call easy_nodes.get_node_mappings() -easy_nodes.initialize_easy_nodes(default_category=my_category, auto_register=True) +# Important! Make sure easy_nodes.initialize_easy_nodes is called before any nodes are defined. +# See __init__.py for an example of how to do this. # This is the converted example node from ComfyUI's example_node.py.example file. -@ComfyNode(my_category) +@ComfyNode() def annotated_example( image: ImageTensor, string_field: str = StringInput("Hello World!", multiline=False), @@ -47,7 +42,7 @@ def another_function(foo: float = 1.0): print("Hello World!", foo) -ComfyNode(my_category, is_changed=lambda: random.random())(another_function) +ComfyNode(is_changed=lambda: random.random())(another_function) # You can register arbitrary classes to be used as inputs or outputs. @@ -64,14 +59,14 @@ def __init__(self): easy_nodes.create_field_setter_node(MyFunClass) -@ComfyNode(my_category, is_output_node=True, color="#4F006F") +@ComfyNode(is_output_node=True, color="#4F006F") def my_fun_class_node_processor(fun_class: MyFunClass) -> ImageTensor: print(f"Processing MyFunClass: {fun_class.width} {fun_class.height} {fun_class.color}") my_image = torch.rand((1, fun_class.height, fun_class.width, 3)) * fun_class.color return my_image -@ComfyNode(my_category) +@ComfyNode() def create_random_image(width: int=NumberInput(128, 128, 1024), height: int=NumberInput(128, 128, 1024)) -> ImageTensor: return torch.rand((1, height, width, 3)) @@ -95,14 +90,13 @@ def my_is_changed_func(): return random() ComfyNode( - my_category, is_changed=my_is_changed_func, description="Descriptions can also be passed in manually. This operation increments a counter", )(ExampleClass.my_method) # Preview text and images right in the nodes. -@ComfyNode(my_category, is_output_node=True) +@ComfyNode(is_output_node=True) def preview_example(str2: str = StringInput("")) -> str: easy_nodes.show_text(f"hello: {str2}") return str2 @@ -118,7 +112,7 @@ def my_class_method(cls, foo: float): cls.class_counter += 1 -ComfyNode(my_category, is_changed=lambda: random.random())( +ComfyNode(is_changed=lambda: random.random())( AnotherExampleClass.my_class_method ) @@ -127,12 +121,12 @@ def my_class_method(cls, foo: float): # differentiate between images and masks in ComfyUI. This is purely cosmetic, and they # are interchangeable in Python. If you annotate the type of a parameter as torch.Tensor # it will be treated as an ImageTensor. -@ComfyNode(my_category, color="#00FF00") +@ComfyNode(color="#00FF00") def convert_to_image(mask: MaskTensor) -> ImageTensor: return mask -@ComfyNode(my_category) +@ComfyNode() def text_repeater(text: str=StringInput("Sample text"), times: int=NumberInput(10, 1, 100)) -> list[str]: return [text] * times @@ -141,7 +135,7 @@ def text_repeater(text: str=StringInput("Sample text"), # If you wrap your input types in list[], under the hood the decorator will make sure you get # everything in a single call with the list inputs passed to you as lists automatically. # If you don't, then you'll get multiple calls with a single item on each call. -@ComfyNode(my_category) +@ComfyNode() def combine_lists( image1: list[ImageTensor], image2: list[ImageTensor] ) -> list[ImageTensor]: @@ -150,7 +144,7 @@ def combine_lists( # Adding a default for a param makes it optional, so ComfyUI won't require it to run your node. -@ComfyNode(my_category) +@ComfyNode() def add_images( image1: ImageTensor, image2: ImageTensor, image3: ImageTensor = None ) -> ImageTensor: @@ -160,7 +154,7 @@ def add_images( return combined_tensors -@ComfyNode(my_category, is_output_node=True, color="#006600") +@ComfyNode(is_output_node=True, color="#006600") def example_show_mask(mask: MaskTensor) -> MaskTensor: easy_nodes.show_image(mask) logging.info("Showing mask") @@ -177,7 +171,7 @@ def threshold_image(image: ImageTensor, threshold_value: float = NumberInput(0.5 # ImageTensor and MaskTensor are just torch.Tensors, so you can treat them as such. -@ComfyNode(my_category, color="#0000FF") +@ComfyNode(color="#0000FF") def example_mask_image(image: ImageTensor, mask: MaskTensor, value: float=NumberInput(0, 0, 1, 0.0001, display="slider")) -> ImageTensor: @@ -189,13 +183,13 @@ def example_mask_image(image: ImageTensor, # As long as Python is happy, ComfyUI will be happy with whatever you tell it the return type is. # You can set the node color by passing in a color argument to the decorator. -@ComfyNode(my_category, color="#FF0000") +@ComfyNode(color="#FF0000") def convert_to_mask(image: ImageTensor, threshold: float = NumberInput(0.5, 0, 1, 0.0001, display="slider")) -> MaskTensor: return (image > threshold).float() # The decorated functions remain normal Python functions, so we can nest them inside each other too. -@ComfyNode(my_category) +@ComfyNode() def mask_image_with_image( image: ImageTensor, image_to_use_as_mask: ImageTensor ) -> ImageTensor: diff --git a/readme.md b/readme.md index 3642041..8207e25 100644 --- a/readme.md +++ b/readme.md @@ -31,6 +31,12 @@ Note that ImageTensor/MaskTensor are just syntactic sugar for semantically diffe For more control, you can call [easy_nodes.initialize_easy_nodes(...)](https://github.com/andrewharp/ComfyUI-EasyNodes?tab=readme-ov-file#initialization-options) before creating nodes and and turn on some advanced settings that will apply to all nodes you create. +## New in 1.1: + +- Custom verifiers for types on input and output for your nodes. For example, it will automatically verify that images always have 1, 3 or 4 channels (B&W, RGB and RGBA). Set `verify_level` when calling initialize_easy_nodes to either CheckSeverityMode OFF, WARN, or FATAL (default is WARN). You can write your own verifiers. See [comfy_types.py](comfy_types.py) for examples of types with verifiers. +- Expanded ComfyUI type support. See [comfy_types.py](comfy_types.py) for the full list of registered types. +- Added warnings if relying on node auto-registration without explicitly asking for it (while also supporting get_node_mappings() at the same time). This is because the default for auto_register will change to False in a future release, in order to make ComfyUI-EasyNodes more easily findable by indexers like ComfyUI-Manager, which expect your nodes to be found in your `__init__.py`. You can enable auto-registration explicitly with `easy_nodes.initialize_easy_nodes(auto_register=True)`. + ## New in 1.0: - Renamed to ComfyUI-EasyNodes from ComfyUI-Annotations to better reflect the package's goal (rather than the means)