diff --git a/comfy_extras/nodes_metadata.py b/comfy_extras/nodes_metadata.py new file mode 100644 index 000000000..0906099eb --- /dev/null +++ b/comfy_extras/nodes_metadata.py @@ -0,0 +1,128 @@ +import os +import numpy as np +import folder_paths +from PIL import Image +from PIL.PngImagePlugin import PngInfo +from comfy.cli_args import args +import json +from typing_extensions import override +from comfy_api.latest import ComfyExtension, io + +io_Metadata_Type = io.Custom(io_type="Metadata") + +class GetWorkflowMetadata(io.ComfyNode): + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="GetWorkflowMetadata", + category="image", + display_name="Get Workflow Metadata", + description="Gets the workflow metadata from the current workflow.", + inputs=[], + hidden=[ + io.Hidden.extra_pnginfo + ], + outputs=[ + io_Metadata_Type.Output("metadata", display_name="workflow") + ], + ) + @classmethod + def execute(self, extra_pnginfo): + metadata = PngInfo() + if not args.disable_metadata: + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + return io.NodeOutput(metadata) + +class EmptyMetadata(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="EmptyMetadata", + category="image", + display_name="Empty Metadata", + description="Create a blank / empty metadata to add upon.", + inputs=[], + outputs=[ + io_Metadata_Type.Output("metadata", display_name="metadata") + ], + ) + @classmethod + def execute(self): + metadata = PngInfo() + return io.NodeOutput(metadata) + +class AddMetadataValue(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="AddMetadataValue", + category="image", + display_name="Add Metadata Value", + description="Add an arbitrary value to the metadata.", + inputs=[ + io.String.Input("key", tooltip="Key to save the value at."), + io.String.Input("value", tooltip="What to add to the metadata."), + io_Metadata_Type.Input("metadata", display_name="metadata") + ], + outputs=[ + io_Metadata_Type.Output("modified_metadata", display_name="metadata") + ], + ) + @classmethod + def execute(self, key, value, metadata): + metadata.add_text(key, json.dumps(value)) + return io.NodeOutput(metadata) + + + +class SaveImageCustomMetadata(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SaveImageCustomMetadata", + category="image", + display_name="Save Image With Custom Metadata", + description="Saves the input images to your ComfyUI output directory with custom metadata.", + inputs=[ + io.Image.Input("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."), + io_Metadata_Type.Input("metadata", display_name="metadata", tooltip="Metadata to save with the image.", optional=True) + ], + outputs=[], + is_output_node=True, + ) + @classmethod + def execute(self, images, filename_prefix, metadata=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() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + 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) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": "output" + }) + counter += 1 + + return io.NodeOutput(ui={ "images": results }) + + +class MetadataExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + GetWorkflowMetadata, + EmptyMetadata, + AddMetadataValue, + SaveImageCustomMetadata + ] + +async def comfy_entrypoint() -> MetadataExtension: + return MetadataExtension() diff --git a/nodes.py b/nodes.py index 9bedbcaca..ebc507a9b 100644 --- a/nodes.py +++ b/nodes.py @@ -2320,6 +2320,7 @@ async def init_builtin_extra_nodes(): "nodes_camera_trajectory.py", "nodes_edit_model.py", "nodes_tcfg.py", + "nodes_metadata.py", ] import_failed = []