migrate load and save images nodes to v3 schema (rebased)

comfy_api/v3/io.py

@@ -20,6 +20,8 @@ from comfy.clip_vision import ClipVisionModel
 from comfy.clip_vision import Output as ClipVisionOutput_
 from comfy_api.input import VideoInput
 from comfy.hooks import HookGroup, HookKeyframeGroup
+import folder_paths
+import os
 # from comfy_extras.nodes_images import SVG as SVG_ # NOTE: needs to be moved before can be imported due to circular reference
 
 
@@ -207,8 +209,11 @@ class WidgetInputV3(InputV3):
         })
     
     def get_io_type_V1(self):
+        if isinstance(self, Combo.Input):
+            return self.as_value_type_v1()
         return self.widgetType if self.widgetType is not None else super().get_io_type_V1()
 
+
 class OutputV3(IO_V3):
     def __init__(self, id: str, display_name: str=None, tooltip: str=None,
                  is_output_list=False):
@@ -372,7 +377,7 @@ class String(ComfyTypeIO):
     class Input(WidgetInputV3):
         '''String input.'''
         def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
-                    multiline=False, placeholder: str=None, default: int=None,
+                    multiline=False, placeholder: str=None, default: str=None,
                     socketless: bool=None, force_input: bool=None):
             super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, self.io_type, force_input)
             self.multiline = multiline
@@ -389,11 +394,11 @@ class String(ComfyTypeIO):
 class Combo(ComfyType):
     Type = str
     class Input(WidgetInputV3):
-        '''Combo input (dropdown).'''
+        """Combo input (dropdown)."""
         Type = str
         def __init__(self, id: str, options: list[str]=None, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
                     default: str=None, control_after_generate: bool=None,
-                    image_upload: bool=None, image_folder: FolderType=None,
+                    image_upload: bool=None, image_folder: FolderType=None, content_types: list[Literal["image", "video", "audio", "model"]]=None,
                     remote: RemoteOptions=None,
                     socketless: bool=None):
             super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, self.io_type)
@@ -402,6 +407,7 @@ class Combo(ComfyType):
             self.control_after_generate = control_after_generate
             self.image_upload = image_upload
             self.image_folder = image_folder
+            self.content_types = content_types
             self.remote = remote
             self.default: str
 
@@ -412,9 +418,23 @@ class Combo(ComfyType):
                 "control_after_generate": self.control_after_generate,
                 "image_upload": self.image_upload,
                 "image_folder": self.image_folder.value if self.image_folder else None,
+                "content_types": self.content_types if self.content_types else None,
                 "remote": self.remote.as_dict() if self.remote else None,
             })
 
+        def as_value_type_v1(self):
+            if getattr(self, "image_folder"):
+                if self.image_folder == FolderType.input:
+                    target_dir = folder_paths.get_input_directory()
+                elif self.image_folder == FolderType.output:
+                    target_dir = folder_paths.get_output_directory()
+                else:
+                    target_dir = folder_paths.get_temp_directory()
+                files = [f for f in os.listdir(target_dir) if os.path.isfile(os.path.join(target_dir, f))]
+                if self.content_types is None:
+                    return files
+                return sorted(folder_paths.filter_files_content_types(files, self.content_types))
+
 
 @comfytype(io_type="COMBO")
 class MultiCombo(ComfyType):

nodes.py

@@ -26,7 +26,7 @@ import comfy.sd
 import comfy.utils
 import comfy.controlnet
 from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict, FileLocator
-from comfy_api.v3.io import ComfyNodeV3
+from comfy_api.v3 import io
 
 import comfy.clip_vision
 
@@ -1550,36 +1550,36 @@ class KSamplerAdvanced:
             disable_noise = True
         return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
 
-class SaveImage:
+
-    def __init__(self):
+class SaveImage(io.ComfyNodeV3):
-        self.output_dir = folder_paths.get_output_directory()
+    @classmethod
-        self.type = "output"
+    def DEFINE_SCHEMA(cls):
-        self.prefix_append = ""
+        return io.SchemaV3(
-        self.compress_level = 4
+            node_id="SaveImage",
+            display_name="Save Image",
+            description="Saves the input images to your ComfyUI output directory.",
+            category="image",
+            inputs=[
+                io.Image.Input(
+                    "images",
+                    display_name="images",
+                    tooltip="The images to save.",
+                ),
+                io.String.Input(
+                    "filename_prefix",
+                    default="ComfyUI",
+                    tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes.",
+                ),
+            ],
+            hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
+            is_output_node=True,
+        )
 
     @classmethod
-    def INPUT_TYPES(s):
+    def execute(cls, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
-        return {
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
-            "required": {
+            filename_prefix, folder_paths.get_output_directory(), images[0].shape[1], images[0].shape[0]
-                "images": ("IMAGE", {"tooltip": "The images to save."}),
+        )
-                "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
-            },
-            "hidden": {
-                "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
-            },
-        }
-
-    RETURN_TYPES = ()
-    FUNCTION = "save_images"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "image"
-    DESCRIPTION = "Saves the input images to your ComfyUI output directory."
-
-    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
-        filename_prefix += self.prefix_append
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
         results = list()
         for (batch_number, image) in enumerate(images):
             i = 255. * image.cpu().numpy()
@@ -1595,16 +1595,17 @@ class SaveImage:
 
             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=self.compress_level)
+            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4)
             results.append({
                 "filename": file,
                 "subfolder": subfolder,
-                "type": self.type
+                "type": "output",
             })
             counter += 1
 
         return { "ui": { "images": results } }
 
+
 class PreviewImage(SaveImage):
     def __init__(self):
         self.output_dir = folder_paths.get_temp_directory()
@@ -1619,24 +1620,36 @@ class PreviewImage(SaveImage):
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                 }
 
-class LoadImage:
+
+class LoadImage(io.ComfyNodeV3):
     @classmethod
-    def INPUT_TYPES(s):
+    def DEFINE_SCHEMA(cls):
-        input_dir = folder_paths.get_input_directory()
+        return io.SchemaV3(
-        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
+            node_id="LoadImage",
-        files = folder_paths.filter_files_content_types(files, ["image"])
+            display_name="Load Image",
-        return {"required":
+            category="image",
-                    {"image": (sorted(files), {"image_upload": True})},
+            inputs=[
-                }
+                io.Combo.Input(
+                    "image",
+                    display_name="image",
+                    image_upload=True,
+                    image_folder=io.FolderType.input,
+                    content_types=["image"],
+                ),
+            ],
+            outputs=[
+                io.Image.Output(
+                    "IMAGE",
+                ),
+                io.Mask.Output(
+                    "MASK",
+                ),
+            ],
+        )
 
-    CATEGORY = "image"
+    @classmethod
-
+    def execute(cls, image) -> io.NodeOutput:
-    RETURN_TYPES = ("IMAGE", "MASK")
+        img = node_helpers.pillow(Image.open, folder_paths.get_annotated_filepath(image))
-    FUNCTION = "load_image"
-    def load_image(self, image):
-        image_path = folder_paths.get_annotated_filepath(image)
-
-        img = node_helpers.pillow(Image.open, image_path)
 
         output_images = []
         output_masks = []
@@ -1678,7 +1691,7 @@ class LoadImage:
             output_image = output_images[0]
             output_mask = output_masks[0]
 
-        return (output_image, output_mask)
+        return io.NodeOutput(output_image, output_mask)
 
     @classmethod
     def IS_CHANGED(s, image):
@@ -1695,6 +1708,7 @@ class LoadImage:
 
         return True
 
+
 class LoadImageMask:
     _color_channels = ["alpha", "red", "green", "blue"]
     @classmethod
@@ -2162,7 +2176,7 @@ def load_custom_node(module_path: str, ignore=set(), module_parent="custom_nodes
         # V3 node definition
         elif getattr(module, "NODES_LIST", None) is not None:
             for node_cls in module.NODES_LIST:
-                node_cls: ComfyNodeV3
+                node_cls: io.ComfyNodeV3
                 schema = node_cls.GET_SCHEMA()
                 if schema.node_id not in ignore:
                     NODE_CLASS_MAPPINGS[schema.node_id] = node_cls