Add metadata to image previews, add a finalize function on SchemaV3 to automatically add hidden values that are required by certain toggles on node definition

comfy_api/v3/io.py

@@ -968,6 +968,26 @@ class SchemaV3:
         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
+        if self.is_api_node:
+            if self.hidden is None:
+                self.hidden = []
+            if Hidden.auth_token_comfy_org not in self.hidden:
+                self.hidden.append(Hidden.auth_token_comfy_org)
+            if Hidden.api_key_comfy_org not in self.hidden:
+                self.hidden.append(Hidden.api_key_comfy_org)
+        # if is an output_node, will need prompt and extra_pnginfo
+        if self.is_output_node:
+            if self.hidden is None:
+                self.hidden = []
+            if Hidden.prompt not in self.hidden:
+                self.hidden.append(Hidden.prompt)
+            if Hidden.extra_pnginfo not in self.hidden:
+                self.hidden.append(Hidden.extra_pnginfo)
+
+
 class Serializer:
     def __init_subclass__(cls, io_type: str, **kwargs):
         cls.io_type = io_type
@@ -1144,7 +1164,7 @@ class ComfyNodeV3:
 
     @classmethod
     def INPUT_TYPES(cls, include_hidden=True, return_schema=False) -> dict[str, dict] | tuple[dict[str, dict], SchemaV3]:
-        schema = cls.DEFINE_SCHEMA()
+        schema = cls.FINALIZE_SCHEMA()
         # for V1, make inputs be a dict with potential keys {required, optional, hidden}
         input = {
             "required": {}
@@ -1165,9 +1185,17 @@ class ComfyNodeV3:
         return input
 
     @classmethod
-    def GET_SCHEMA(cls) -> SchemaV3:
+    def FINALIZE_SCHEMA(cls):
-        cls.VALIDATE_CLASS()
+        """Call DEFINE_SCHEMA and finalize it."""
         schema = cls.DEFINE_SCHEMA()
+        schema.finalize()
+        return schema
+
+    @classmethod
+    def GET_SCHEMA(cls) -> SchemaV3:
+        """Validate node class, finalize schema, validate schema, and set expected class properties."""
+        cls.VALIDATE_CLASS()
+        schema = cls.FINALIZE_SCHEMA()
         schema.validate()
         if cls._DESCRIPTION is None:
             cls._DESCRIPTION = schema.description

comfy_api/v3/ui.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
 from abc import ABC, abstractmethod
 
-from comfy_api.v3.io import Image, Mask, FolderType, _UIOutput
+from comfy_api.v3.io import Image, Mask, FolderType, _UIOutput, ComfyNodeV3
 # used for image preview
+from comfy.cli_args import args
 import folder_paths
 import random
 from PIL import Image as PILImage
+from PIL.PngImagePlugin import PngInfo
 import os
+import json
 import numpy as np
 
 
@@ -24,7 +27,7 @@ class SavedResult:
         }
 
 class PreviewImage(_UIOutput):
-    def __init__(self, image: Image.Type, animated: bool=False, **kwargs):
+    def __init__(self, image: Image.Type, animated: bool=False, node: ComfyNodeV3=None, **kwargs):
         output_dir = folder_paths.get_temp_directory()
         type = "temp"
         prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
@@ -38,13 +41,13 @@ class PreviewImage(_UIOutput):
             i = 255. * image.cpu().numpy()
             img = PILImage.fromarray(np.clip(i, 0, 255).astype(np.uint8))
             metadata = None
-            # if not args.disable_metadata:
+            if not args.disable_metadata and node is not None:
-            #     metadata = PngInfo()
+                metadata = PngInfo()
-            #     if prompt is not None:
+                if node.hidden.prompt is not None:
-            #         metadata.add_text("prompt", json.dumps(prompt))
+                    metadata.add_text("prompt", json.dumps(node.hidden.prompt))
-            #     if extra_pnginfo is not None:
+                if node.hidden.extra_pnginfo is not None:
-            #         for x in extra_pnginfo:
+                    for x in node.hidden.extra_pnginfo:
-            #             metadata.add_text(x, json.dumps(extra_pnginfo[x]))
+                        metadata.add_text(x, json.dumps(node.hidden.extra_pnginfo[x]))
 
             filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
             file = f"{filename_with_batch_num}_{counter:05}_.png"
@@ -63,9 +66,9 @@ class PreviewImage(_UIOutput):
         }
 
 class PreviewMask(PreviewImage):
-    def __init__(self, mask: PreviewMask.Type, animated: bool=False, **kwargs):
+    def __init__(self, mask: PreviewMask.Type, animated: bool=False, node: ComfyNodeV3=None, **kwargs):
         preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
-        super().__init__(preview, animated, **kwargs)
+        super().__init__(preview, animated, node, **kwargs)
 
 # class UILatent(_UIOutput):
 #     def __init__(self, values: list[SavedResult | dict], **kwargs):