v3

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()

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 = []