V3 Nodes: refactor ComfyNodeV3 class; use of ui.SavedResult; ported SaveAnimatedPNG and SaveAnimatedWEBP nodes

Author: bigcat88

comfy_api/v3/helpers.py

@@ -0,0 +1,16 @@
+from typing import Optional, Callable
+
+
+def first_real_override(cls: type, name: str, *, base: type) -> Optional[Callable]:
+    """Return the *callable* override of `name` visible on `cls`, or None if every
+    implementation up to (and including) `base` is the placeholder defined on `base`.
+    """
+    base_func = getattr(base, name).__func__
+    for c in cls.mro():                       # NodeB, NodeA, ComfyNodeV3, object …
+        if c is base:                         # reached the placeholder – we're done
+            break
+        if name in c.__dict__:                # first class that *defines* the attr
+            func = getattr(c, name).__func__
+            if func is not base_func:         # real override
+                return getattr(cls, name)     # bound to *cls*
+    return None

comfy_api/v3/io.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Any, Literal, TYPE_CHECKING, TypeVar, Callable, Optional, cast, TypedDict
+from typing import Any, Literal, TypeVar, Callable, TypedDict
 from typing_extensions import NotRequired
 from enum import Enum
 from abc import ABC, abstractmethod
@@ -108,7 +108,7 @@ class ComfyType(ABC):
 def comfytype(io_type: str, **kwargs):
     '''
     Decorator to mark nested classes as ComfyType; io_type will be bound to the class.
-    
+
     A ComfyType may have the following attributes:
     - Type = <type hint here>
     - class Input(InputV3): ...
@@ -206,15 +206,15 @@ def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str
         self.socketless = socketless
         self.widgetType = widgetType
         self.force_input = force_input
-    
+
     def as_dict_V1(self):
         return super().as_dict_V1() | prune_dict({
             "default": self.default,
             "socketless": self.socketless,
             "widgetType": self.widgetType,
             "forceInput": self.force_input,
         })
-    
+
     def get_io_type_V1(self):
         return self.widgetType if self.widgetType is not None else super().get_io_type_V1()
 
@@ -289,21 +289,21 @@ def __setattr__(self, key: str, value: Any):
             super().__setattr__(key, value)
         else:
             self.local_state[key] = value
-    
+
     def __setitem__(self, key: str, value: Any):
         self.local_state[key] = value
-    
+
     def __getitem__(self, key: str):
         return self.local_state[key]
-    
+
     def __delitem__(self, key: str):
         del self.local_state[key]
 
 
 @comfytype(io_type="BOOLEAN")
 class Boolean(ComfyTypeIO):
     Type = bool
-    
+
     class Input(WidgetInputV3):
         '''Boolean input.'''
         def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
@@ -313,7 +313,7 @@ def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str
             self.label_on = label_on
             self.label_off = label_off
             self.default: bool
-        
+
         def as_dict_V1(self):
             return super().as_dict_V1() | prune_dict({
                 "label_on": self.label_on,
@@ -385,7 +385,7 @@ def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str
             self.multiline = multiline
             self.placeholder = placeholder
             self.default: str
-        
+
         def as_dict_V1(self):
             return super().as_dict_V1() | prune_dict({
                 "multiline": self.multiline,
@@ -500,7 +500,7 @@ class PooledDict(TypedDict):
         By default, the dimensions are based on total pixel amount, but the first value can be set to "percentage" to use a percentage of the image size instead.
 
         (1024, 1024, 0, 0) would apply conditioning to the top-left 1024x1024 pixels.
-        
+
         ("percentage", 0.5, 0.5, 0, 0) would apply conditioning to the top-left 50% of the image.''' # TODO: verify its actually top-left
         strength: NotRequired[float]
         '''Strength of conditioning. Default strength is 1.0.'''
@@ -755,7 +755,7 @@ def __init__(self, id: str | InputV3, types: list[type[ComfyType] | ComfyType],
                     self.input_override.widgetType = self.input_override.get_io_type_V1()
             super().__init__(id, display_name, optional, tooltip, lazy, extra_dict)
             self._io_types = types
-        
+
         @property
         def io_types(self) -> list[type[InputV3]]:
             '''
@@ -768,14 +768,14 @@ def io_types(self) -> list[type[InputV3]]:
                 else:
                     io_types.append(x)
             return io_types
-        
+
         def get_io_type_V1(self):
             # ensure types are unique and order is preserved
             str_types = [x.io_type for x in self.io_types]
             if self.input_override is not None:
                 str_types.insert(0, self.input_override.get_io_type_V1())
             return ",".join(list(dict.fromkeys(str_types)))
-        
+
         def as_dict_V1(self):
             if self.input_override is not None:
                 return self.input_override.as_dict_V1() | super().as_dict_V1()
@@ -870,7 +870,7 @@ def __init__(self, unique_id: str, prompt: Any,
     def __getattr__(self, key: str):
         '''If hidden variable not found, return None.'''
         return None
-    
+
     @classmethod
     def from_dict(cls, d: dict | None):
         if d is None:
@@ -1088,7 +1088,7 @@ class ComfyNodeV3:
 
     RELATIVE_PYTHON_MODULE = None
     SCHEMA = None
-    
+
     # filled in during execution
     state: NodeState = None
     resources: Resources = None
@@ -1097,28 +1097,24 @@ class ComfyNodeV3:
     @classmethod
     @abstractmethod
     def DEFINE_SCHEMA(cls) -> SchemaV3:
-        """
-        Override this function with one that returns a SchemaV3 instance.
-        """
-        return None
-    DEFINE_SCHEMA = None
+        """Override this function with one that returns a SchemaV3 instance."""
+        raise NotImplementedError
 
     @classmethod
     @abstractmethod
     def execute(cls, **kwargs) -> NodeOutput:
-        pass
-    execute = None
+        raise NotImplementedError
 
     @classmethod
     def validate_inputs(cls, **kwargs) -> bool:
-        """Optionally, define this function to validate inputs; equivalnet to V1's VALIDATE_INPUTS."""
-        pass
-    validate_inputs = None
+        """Optionally, define this function to validate inputs; equivalent to V1's VALIDATE_INPUTS."""
+        raise NotImplementedError
 
     @classmethod
     def fingerprint_inputs(cls, **kwargs) -> Any:
         """Optionally, define this function to fingerprint inputs; equivalent to V1's IS_CHANGED."""
-        pass
+        raise NotImplementedError
+
     fingerprint_inputs = None
 
     @classmethod
@@ -1135,8 +1131,8 @@ def check_lazy_status(cls, **kwargs) -> list[str]:
 
         Comfy Docs: https://docs.comfy.org/custom-nodes/backend/lazy_evaluation#defining-check-lazy-status
         """
-        need = [name for name in kwargs if kwargs[name] is None]
-        return need
+        return [name for name in kwargs if kwargs[name] is None]
+
     check_lazy_status = None
 
     @classmethod
@@ -1405,7 +1401,7 @@ def result(self):
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "NodeOutput":
         args = ()
-        ui = None   
+        ui = None
         expand = None
         if "result" in data:
             result = data["result"]

comfy_api/v3/resources.py

@@ -32,7 +32,7 @@ def __str__(self):
 class Resources(ABC):
     def __init__(self):
         ...
-    
+
     @abstractmethod
     def get(self, key: ResourceKey, default: Any=...) -> Any:
         pass

comfy_api/v3/ui.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
-from abc import ABC, abstractmethod
 
-from comfy_api.v3.io import Image, Mask, FolderType, _UIOutput, ComfyNodeV3
+from comfy_api.v3.io import Image, FolderType, _UIOutput, ComfyNodeV3
 # used for image preview
 from comfy.cli_args import args
 import folder_paths
@@ -13,33 +12,33 @@
 import numpy as np
 
 
-class SavedResult:
+class SavedResult(dict):
     def __init__(self, filename: str, subfolder: str, type: FolderType):
-        self.filename = filename
-        self.subfolder = subfolder
-        self.type = type
-    
-    def as_dict(self):
-        return {
-            "filename": self.filename,
-            "subfolder": self.subfolder,
-            "type": self.type
-        }
+        super().__init__(filename=filename, subfolder=subfolder,type=type.value)
+
+    @property
+    def filename(self) -> str:
+        return self["filename"]
+
+    @property
+    def subfolder(self) -> str:
+        return self["subfolder"]
+
+    @property
+    def type(self) -> FolderType:
+        return FolderType(self["type"])
+
 
 class PreviewImage(_UIOutput):
     def __init__(self, image: Image.Type, animated: bool=False, cls: ComfyNodeV3=None, **kwargs):
         output_dir = folder_paths.get_temp_directory()
-        type = "temp"
         prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
-        compress_level = 1
-        filename_prefix = "ComfyUI"
+        filename_prefix = "ComfyUI" + prefix_append
 
-        filename_prefix += prefix_append
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, output_dir, image[0].shape[1], image[0].shape[0])
         results = list()
         for (batch_number, image) in enumerate(image):
-            i = 255. * image.cpu().numpy()
-            img = PILImage.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            img = PILImage.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8))
             metadata = None
             if not args.disable_metadata and cls is not None:
                 metadata = PngInfo()
@@ -51,17 +50,16 @@ def __init__(self, image: Image.Type, animated: bool=False, cls: ComfyNodeV3=Non
 
             filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
             file = f"{filename_with_batch_num}_{counter:05}_.png"
-            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=compress_level)
-            results.append(SavedResult(file, subfolder, type))
+            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=1)
+            results.append(SavedResult(file, subfolder, FolderType.temp))
             counter += 1
-        
+
         self.values = results
         self.animated = animated
-    
+
     def as_dict(self):
-        values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
         return {
-            "images": values,
+            "images": self.values,
             "animated": (self.animated,)
         }
 
@@ -111,36 +109,29 @@ def __init__(self, mask: PreviewMask.Type, animated: bool=False, cls: ComfyNodeV
 #         comfy.utils.save_torch_file(output, file, metadata=metadata)
 
 #         self.values = values
-    
+
 #     def as_dict(self):
-#         values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
 #         return {
-#             "latents": values,
+#             "latents": self.values,
 #         }
 
 class PreviewAudio(_UIOutput):
     def __init__(self, values: list[SavedResult | dict], **kwargs):
         self.values = values
-    
+
     def as_dict(self):
-        values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
-        return {
-            "audio": values,
-        }
+        return {"audio": self.values}
 
 class PreviewUI3D(_UIOutput):
     def __init__(self, values: list[SavedResult | dict], **kwargs):
         self.values = values
-    
+
     def as_dict(self):
-        values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values]
-        return {
-            "3d": values,
-        }
+        return {"3d": self.values}
 
 class PreviewText(_UIOutput):
     def __init__(self, value: str, **kwargs):
         self.value = value
-    
+
     def as_dict(self):
         return {"text": (self.value,)}

comfy_extras/nodes_v1_test.py

@@ -20,7 +20,7 @@ def INPUT_TYPES(cls):
                 "mask": (IO.MASK,),
             }
         }
-    
+
     RETURN_TYPES = (IO.INT, IO.IMAGE)
     RETURN_NAMES = ("INT", "img🖼️")
     OUTPUT_TOOLTIPS = (None, "This is an image")

comfy_extras/nodes_v3_test.py

@@ -1,11 +1,11 @@
 import torch
 import time
 from comfy_api.v3 import io, ui, resources
-import logging
+import logging  # noqa
 import folder_paths
 import comfy.utils
 import comfy.sd
-from typing import Any
+
 
 @io.comfytype(io_type="XYZ")
 class XYZ:
@@ -88,11 +88,11 @@ def execute(cls, image: io.Image.Type, some_int: int, combo: io.Combo.Type, comb
         expected_int = 123
         if "thing" not in cls.state:
             cls.state["thing"] = "hahaha"
-            yyy = cls.state["thing"]
+            yyy = cls.state["thing"]    # noqa
             del cls.state["thing"]
         if cls.state.get_value("int2") is None:
             cls.state.set_value("int2", 123)
-            zzz = cls.state.get_value("int2")
+            zzz = cls.state.get_value("int2")   # noqa
             cls.state.pop("int2")
         if cls.state.my_int is None:
             cls.state.my_int = expected_int
@@ -175,7 +175,7 @@ def DEFINE_SCHEMA(cls):
                 io.Image.Output(),
             ],
         )
-    
+
     @classmethod
     def validate_inputs(cls, nmock, nmock2):
         return True