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

This commit is contained in:
bigcat88 2025-07-08 13:18:27 +03:00
parent 5f91e2905a
commit 36770c1658
No known key found for this signature in database
GPG Key ID: 1F0BF0EC3CF22721
2 changed files with 87 additions and 53 deletions

View File

@ -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):
@ -969,7 +989,7 @@ class SchemaV3:
issues.append(f"Ids must be unique between inputs and outputs, but {intersection} are not.")
if len(issues) > 0:
raise ValueError("\n".join(issues))
def finalize(self):
"""Add hidden based on selected schema options."""
# if is an api_node, will need key-related hidden

112
nodes.py
View File

@ -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):
self.output_dir = folder_paths.get_output_directory()
self.type = "output"
self.prefix_append = ""
self.compress_level = 4
class SaveImage(io.ComfyNodeV3):
@classmethod
def DEFINE_SCHEMA(cls):
return io.SchemaV3(
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):
return {
"required": {
"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])
def execute(cls, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
filename_prefix, folder_paths.get_output_directory(), 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):
input_dir = folder_paths.get_input_directory()
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
files = folder_paths.filter_files_content_types(files, ["image"])
return {"required":
{"image": (sorted(files), {"image_upload": True})},
}
def DEFINE_SCHEMA(cls):
return io.SchemaV3(
node_id="LoadImage",
display_name="Load Image",
category="image",
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"
RETURN_TYPES = ("IMAGE", "MASK")
FUNCTION = "load_image"
def load_image(self, image):
image_path = folder_paths.get_annotated_filepath(image)
img = node_helpers.pillow(Image.open, image_path)
@classmethod
def execute(cls, image) -> io.NodeOutput:
img = node_helpers.pillow(Image.open, folder_paths.get_annotated_filepath(image))
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