From c941ee09fcaab5223ec39bf7e38dd406908b67a5 Mon Sep 17 00:00:00 2001
From: doctorpangloss <@hiddenswitch.com>
Date: Wed, 21 Feb 2024 23:44:16 -0800
Subject: [PATCH] Improve nodes handling

---
 comfy/cli_args_types.py       |  1 +
 comfy/nodes/package.py        |  2 ++
 comfy/nodes/package_typing.py | 59 +++++++++++++++++++++++++++++------
 setup.py                      |  5 ++-
 4 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/comfy/cli_args_types.py b/comfy/cli_args_types.py
index 9e10261f1..58a0789bd 100644
--- a/comfy/cli_args_types.py
+++ b/comfy/cli_args_types.py
@@ -75,6 +75,7 @@ class Configuration(dict):
         distributed_queue_worker (bool): Workers will pull requests off the AMQP URL.
         distributed_queue_name (str): This name will be used by the frontends and workers to exchange prompt requests and replies. Progress updates will be prefixed by the queue name, followed by a '.', then the user ID.
         external_address (str): Specifies a base URL for external addresses reported by the API, such as for image paths.
+        verbose (bool): Shows extra output for debugging purposes such as import errors of custom nodes.
     """
     def __init__(self, **kwargs):
         super().__init__()
diff --git a/comfy/nodes/package.py b/comfy/nodes/package.py
index 3dd4d370b..adabc8392 100644
--- a/comfy/nodes/package.py
+++ b/comfy/nodes/package.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import importlib
+import logging
 import os
 import pkgutil
 import time
@@ -63,6 +64,7 @@ def _import_and_enumerate_nodes_in_module(module: types.ModuleType, print_import
             except KeyboardInterrupt as interrupted:
                 raise interrupted
             except Exception as x:
+                logging.error(f"{full_name} import failed", exc_info=x)
                 success = False
             timings.append((time.perf_counter() - time_before, full_name, success))
 
diff --git a/comfy/nodes/package_typing.py b/comfy/nodes/package_typing.py
index 5f3be9f48..0d91b27db 100644
--- a/comfy/nodes/package_typing.py
+++ b/comfy/nodes/package_typing.py
@@ -1,23 +1,64 @@
 from __future__ import annotations
 
-import typing
-from typing import Protocol, ClassVar, Tuple, Dict
 from dataclasses import dataclass, field
+from typing import TypedDict, Union, Optional, Sequence, Dict, ClassVar, Protocol, Tuple, TypeVar, Any, Literal, \
+    Callable
 
-T = typing.TypeVar('T', bound='CustomNode')
+T = TypeVar('T')
+
+
+class NumberSpecOptions(TypedDict, total=False):
+    default: Union[int, float]
+    min: Union[int, float]
+    max: Union[int, float]
+    step: Union[int, float]
+    round: int
+
+
+IntSpec = Dict[str, Union[
+    Literal["INT"],
+    Tuple[Literal["INT"], Dict[str, Union[int, float, str]]]
+]]
+FloatSpec = Dict[str, Union[
+    Literal["FLOAT"],
+    Tuple[Literal["FLOAT"], Dict[str, Union[int, float, str]]]
+]]
+StringSpec = Dict[str, Union[
+    Literal["STRING"],
+    Tuple[Literal["STRING"], Dict[str, str]]
+]]
+ChoiceSpec = Dict[str, Union[
+    Sequence[str],  # Directly a list of choices
+    Tuple[Sequence[str], Dict[str, Any]]  # Choices with additional specifications
+]]
+
+ComplexInputSpec = Dict[str, Any]
+InputTypeSpec = Union[IntSpec, FloatSpec, StringSpec, ChoiceSpec, ComplexInputSpec]
+
+
+class InputTypes(Protocol):
+    required: Dict[str, InputTypeSpec]
+    optional: Optional[Dict[str, InputTypeSpec]]
+    hidden: Optional[Dict[str, InputTypeSpec]]
+
+
+ValidateInputsMethod = Optional[Callable[..., Union[bool, str]]]
 
 
 class CustomNode(Protocol):
     @classmethod
-    def INPUT_TYPES(cls) -> dict: ...
+    def INPUT_TYPES(cls) -> InputTypes: ...
 
-    RETURN_TYPES: ClassVar[typing.Sequence[str]]
-    RETURN_NAMES: typing.Optional[ClassVar[Tuple[str]]]
-    OUTPUT_IS_LIST: typing.Optional[ClassVar[typing.Sequence[bool]]]
-    INPUT_IS_LIST: typing.Optional[ClassVar[bool]]
+    # Optional method signature for VALIDATE_INPUTS
+    VALIDATE_INPUTS: ClassVar[ValidateInputsMethod] = None
+
+    RETURN_TYPES: ClassVar[Sequence[str]]
+    RETURN_NAMES: Optional[ClassVar[Tuple[str]]]
+    OUTPUT_IS_LIST: Optional[ClassVar[Sequence[bool]]]
+    INPUT_IS_LIST: Optional[ClassVar[bool]]
     FUNCTION: ClassVar[str]
     CATEGORY: ClassVar[str]
-    OUTPUT_NODE: typing.Optional[ClassVar[bool]]
+    OUTPUT_NODE: Optional[ClassVar[bool]]
 
     def __call__(self) -> T:
         ...
diff --git a/setup.py b/setup.py
index c37ee6cae..fd1785bfb 100644
--- a/setup.py
+++ b/setup.py
@@ -176,9 +176,8 @@ setup(
     author="",
     version=version,
     python_requires=">=3.9,<3.13",
-    # todo: figure out how to include the web directory to eventually let main live inside the package
-    # todo: see https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/ for more about adding plugins
-    packages=find_packages(exclude=[] if is_editable else ['custom_nodes']),
+    packages=find_packages(exclude=["tests"] + [] if is_editable else ['custom_nodes']),
+    package_dir={'': ''},
     include_package_data=True,
     install_requires=dependencies(),
     setup_requires=["pip", "wheel"],