ComfyUI/comfy_extras/nodes_metadata.py

import os
import numpy as np
import folder_paths
import av
from fractions import Fraction
import torch
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

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.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.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.Input("metadata", display_name="metadata")
            ],
            outputs=[
                io.Metadata.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.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 SaveWEBMCustomMetadata(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="SaveWebmCustomMetadata",
            category="image/video",
            display_name="Save Webm With Custom Metadata",
            inputs=[
                io.Image.Input("images", tooltip="The images to save."),
                io.String.Input("filename_prefix", default="ComfyUI"),
                io.Combo.Input("codec", ["vp9", "av1"]),
                io.Float.Input("fps", default=24.0, min=0.01, max=1000.0, step=0.01),
                io.Float.Input("crf", default=32.0, min=0.0, max=63.0, step=1.0, tooltip="Higher crf means lower quality with a smaller file size, lower crf means higher quality higher filesize."),
                io.Metadata.Input("metadata", display_name="metadata", tooltip="Metadata to save with the image.", optional=True)
            ],
            outputs=[],
            is_experimental=True,
            is_output_node=True
        )

    @classmethod
    def execute(self, images, codec, fps, filename_prefix, crf, 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])

        file = f"{filename}_{counter:05}_.webm"
        container = av.open(os.path.join(full_output_folder, file), mode="w")

        if metadata is not None:
            for x in metadata:
                container.metadata[x] = json.dumps(metadata[x])

        codec_map = {"vp9": "libvpx-vp9", "av1": "libsvtav1"}
        stream = container.add_stream(codec_map[codec], rate=Fraction(round(fps * 1000), 1000))
        stream.width = images.shape[-2]
        stream.height = images.shape[-3]
        stream.pix_fmt = "yuv420p10le" if codec == "av1" else "yuv420p"
        stream.bit_rate = 0
        stream.options = {'crf': str(crf)}
        if codec == "av1":
            stream.options["preset"] = "6"

        for frame in images:
            frame = av.VideoFrame.from_ndarray(torch.clamp(frame[..., :3] * 255, min=0, max=255).to(device=torch.device("cpu"), dtype=torch.uint8).numpy(), format="rgb24")
            for packet in stream.encode(frame):
                container.mux(packet)
        container.mux(stream.encode())
        container.close()

        results = [{
            "filename": file,
            "subfolder": subfolder,
            "type": "output"
        }]

        return io.NodeOutput(ui = {"images": results, "animated": (True,)})


class MetadataExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [
            GetWorkflowMetadata,
            EmptyMetadata,
            AddMetadataValue,
            SaveImageCustomMetadata,
            SaveWEBMCustomMetadata
        ]

async def comfy_entrypoint() -> MetadataExtension:
    return MetadataExtension()