From e6623a13597fb3622ef2a109b7b8783729e85cd9 Mon Sep 17 00:00:00 2001 From: Benjamin Berman Date: Thu, 29 Feb 2024 15:54:42 -0800 Subject: [PATCH] Fix CLIPLoader node, fix CustomNode typing, improve digest --- comfy/digest.py | 15 ++++- comfy/nodes/base_nodes.py | 4 +- comfy/nodes/package_typing.py | 112 ++++++++++++++++++++++++---------- 3 files changed, 93 insertions(+), 38 deletions(-) diff --git a/comfy/digest.py b/comfy/digest.py index 298247209..478be4f9f 100644 --- a/comfy/digest.py +++ b/comfy/digest.py @@ -2,10 +2,19 @@ from __future__ import annotations import hashlib import json +from typing_extensions import Buffer -def digest(data: dict | str) -> str: - json_str = data if isinstance(data, str) else json.dumps(data, separators=(',', ':')) +def digest(data: dict | str | Buffer) -> str: hash_object = hashlib.sha256() - hash_object.update(json_str.encode()) + if isinstance(data, Buffer): + hash_object.update(data) + else: + if isinstance(data, str): + json_str = data + elif isinstance(data, dict): + json_str = json.dumps(data, separators=(',', ':')) + else: + raise RuntimeError("invalid data type") + hash_object.update(json_str.encode()) return hash_object.hexdigest() diff --git a/comfy/nodes/base_nodes.py b/comfy/nodes/base_nodes.py index 37f653e61..3d8aad92a 100644 --- a/comfy/nodes/base_nodes.py +++ b/comfy/nodes/base_nodes.py @@ -833,9 +833,9 @@ class CLIPLoader: CATEGORY = "advanced/loaders" def load_clip(self, clip_name, type="stable_diffusion"): - clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION + clip_type = sd.CLIPType.STABLE_DIFFUSION if type == "stable_cascade": - clip_type = comfy.sd.CLIPType.STABLE_CASCADE + clip_type = sd.CLIPType.STABLE_CASCADE clip_path = folder_paths.get_full_path("clip", clip_name) clip = sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type) diff --git a/comfy/nodes/package_typing.py b/comfy/nodes/package_typing.py index 0d91b27db..5d9ced150 100644 --- a/comfy/nodes/package_typing.py +++ b/comfy/nodes/package_typing.py @@ -1,50 +1,96 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import TypedDict, Union, Optional, Sequence, Dict, ClassVar, Protocol, Tuple, TypeVar, Any, Literal, \ - Callable +from typing_extensions import TypedDict, NotRequired, Generic +from typing import Union, Optional, Sequence, Dict, ClassVar, Protocol, Tuple, TypeVar, Any, Literal, \ + Callable, List 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 +class IntSpecOptions(TypedDict, total=True): + default: int + min: int + max: int + step: NotRequired[int] + display: NotRequired[Literal["number", "slider"]] -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 FloatSpecOptions(TypedDict, total=True): + default: float + min: float + max: float + step: NotRequired[float] + round: NotRequired[float] + display: NotRequired[Literal["number", "slider"]] -class InputTypes(Protocol): +class StringSpecOptions(TypedDict, total=True): + multiline: NotRequired[bool] + default: NotRequired[str] + + +# todo: analyze the base_nodes for these types +CommonReturnTypes = Union[ + Literal["IMAGE", "STRING", "INT", "FLOAT", "CONDITIONING", "LATENT", "MASK", "MODEL", "VAE", "CLIP"], str] + +IntSpec = Tuple[Literal["INT"], IntSpecOptions] + +FloatSpec = Tuple[Literal["FLOAT"], FloatSpecOptions] + +StringSpec = Tuple[Literal["STRING"], StringSpecOptions] + +ChoiceSpec = Tuple[Union[Sequence[str], Sequence[float], Sequence[int]]] + +NonPrimitiveTypeSpec = Tuple[CommonReturnTypes] + +InputTypeSpec = Union[IntSpec, FloatSpec, StringSpec, ChoiceSpec, NonPrimitiveTypeSpec] + + +class HiddenSpec(TypedDict, total=True): + prompt: Literal["PROMPT"] + extra_pnginfo: Literal["EXTRA_PNGINFO"] + + +class InputTypes(TypedDict, total=True): required: Dict[str, InputTypeSpec] - optional: Optional[Dict[str, InputTypeSpec]] - hidden: Optional[Dict[str, InputTypeSpec]] + optional: NotRequired[Dict[str, InputTypeSpec]] + hidden: NotRequired[HiddenSpec] ValidateInputsMethod = Optional[Callable[..., Union[bool, str]]] +class FunctionReturnsUIVariables(TypedDict): + ui: dict + + +class SaveNodeResult(TypedDict, total=True): + abs_path: NotRequired[str] + filename: str + subfolder: str + type: Literal["output", "input", "temp"] + + +class UIImagesImagesResult(TypedDict, total=True): + images: List[SaveNodeResult] + + +class UIImagesResult(TypedDict, total=True): + ui: UIImagesImagesResult + + +class UILatentsLatentsResult(TypedDict, total=True): + latents: List[SaveNodeResult] + + +class UILatentsResult(TypedDict, total=True): + ui: UILatentsLatentsResult + + +ValidatedNodeResult = Union[Tuple, UIImagesResult, UILatentsResult, FunctionReturnsUIVariables] + + class CustomNode(Protocol): @classmethod def INPUT_TYPES(cls) -> InputTypes: ... @@ -52,9 +98,9 @@ class CustomNode(Protocol): # 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]]] + RETURN_TYPES: ClassVar[Tuple[CommonReturnTypes, ...]] + RETURN_NAMES: Optional[ClassVar[Tuple[str, ...]]] + OUTPUT_IS_LIST: Optional[ClassVar[Tuple[bool, ...]]] INPUT_IS_LIST: Optional[ClassVar[bool]] FUNCTION: ClassVar[str] CATEGORY: ClassVar[str]