mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-11 06:40:48 +08:00
97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import TypedDict, Union, Optional, Sequence, Dict, ClassVar, Protocol, Tuple, TypeVar, Any, Literal, \
|
|
Callable
|
|
|
|
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) -> InputTypes: ...
|
|
|
|
# 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: Optional[ClassVar[bool]]
|
|
|
|
def __call__(self) -> T:
|
|
...
|
|
|
|
|
|
@dataclass
|
|
class ExportedNodes:
|
|
NODE_CLASS_MAPPINGS: Dict[str, CustomNode] = field(default_factory=dict)
|
|
NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = field(default_factory=dict)
|
|
EXTENSION_WEB_DIRS: Dict[str, str] = field(default_factory=dict)
|
|
|
|
def update(self, exported_nodes: ExportedNodes) -> ExportedNodes:
|
|
self.NODE_CLASS_MAPPINGS.update(exported_nodes.NODE_CLASS_MAPPINGS)
|
|
self.NODE_DISPLAY_NAME_MAPPINGS.update(exported_nodes.NODE_DISPLAY_NAME_MAPPINGS)
|
|
self.EXTENSION_WEB_DIRS.update(exported_nodes.EXTENSION_WEB_DIRS)
|
|
return self
|
|
|
|
def __len__(self):
|
|
return len(self.NODE_CLASS_MAPPINGS)
|
|
|
|
def __sub__(self, other: ExportedNodes):
|
|
exported_nodes = ExportedNodes().update(self)
|
|
for self_key in exported_nodes.NODE_CLASS_MAPPINGS:
|
|
if self_key in other.NODE_CLASS_MAPPINGS:
|
|
exported_nodes.NODE_CLASS_MAPPINGS.pop(self_key)
|
|
if self_key in other.NODE_DISPLAY_NAME_MAPPINGS:
|
|
exported_nodes.NODE_DISPLAY_NAME_MAPPINGS.pop(self_key)
|
|
for self_key in exported_nodes.EXTENSION_WEB_DIRS:
|
|
if self_key in other.EXTENSION_WEB_DIRS:
|
|
exported_nodes.EXTENSION_WEB_DIRS.pop(self_key)
|
|
return exported_nodes
|
|
|
|
def __add__(self, other):
|
|
exported_nodes = ExportedNodes().update(self)
|
|
return exported_nodes.update(other)
|