Add input capture hook

.gitignore

@@ -0,0 +1 @@
+__pycache__

__init__.py

@@ -0,0 +1,11 @@
+from .py.nodes.node import SaveImageWithMetaData
+
+NODE_CLASS_MAPPINGS = {
+    "SaveImageWithMetaData": SaveImageWithMetaData,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SaveImageWithMetaData": "Save Image With Metadata",
+}
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

py/__init__.py

@@ -0,0 +1,96 @@
+import functools
+
+from .defs.hook_definition import hooks, samplers
+from .defs.meta import MetaField
+import execution
+
+inputs = {}
+current_node_id = -1
+current_prompt_id = -1
+current_prompt = {}
+positive_nodes = {}
+negative_nodes = {}
+
+
+def prefix_function(function, prefunction):
+    @functools.wraps(function)
+    def run(*args, **kwargs):
+        prefunction(*args, **kwargs)
+        return function(*args, **kwargs)
+
+    return run
+
+
+def pre_execute(self, prompt, prompt_id, extra_data, execute_outputs):
+    global current_prompt
+    global current_prompt_id
+    global positive_nodes
+    global negative_nodes
+
+    current_prompt = prompt
+    current_prompt_id = prompt_id
+    inputs[current_prompt_id] = {}
+    print("pre_execute")
+    print(current_prompt_id)
+    positive_nodes, negative_nodes = get_positives_and_negatives(current_prompt)
+
+
+def pre_get_input_data(*args, **kwargs):
+    global current_node_id
+
+    inputs = args[0]
+    class_def = args[1]
+    unique_id = args[2]
+    current_node_id = unique_id
+    for node_class, metas in hooks.items():
+        if class_def == node_class:
+            print(inputs, class_def, unique_id)
+
+
+def pre_map_node_over_list(*args, **kwargs):
+    global current_node_id
+    global positive_nodes
+    global negative_nodes
+
+    obj, input_data_all, func = args
+    for node_class, metas in hooks.items():
+        if isinstance(obj, node_class):
+            for meta, field_name in metas.items():
+                if meta == MetaField.POSITIVE_PROMPT:
+                    if current_node_id not in positive_nodes.values():
+                        continue
+                elif meta == MetaField.NEGATIVE_PROMPT:
+                    if current_node_id not in negative_nodes.values():
+                        continue
+
+                if meta not in inputs[current_prompt_id]:
+                    inputs[current_prompt_id][meta] = []
+                value = input_data_all.get(field_name)
+                if value is not None:
+                    v = value
+                    if isinstance(value, list) and len(value) > 0:
+                        v = value[0]
+                    inputs[current_prompt_id][meta].append((current_node_id, v))
+            print(inputs[current_prompt_id])
+
+
+execution.PromptExecutor.execute = prefix_function(
+    execution.PromptExecutor.execute, pre_execute
+)
+execution.get_input_data = prefix_function(execution.get_input_data, pre_get_input_data)
+execution.map_node_over_list = prefix_function(
+    execution.map_node_over_list, pre_map_node_over_list
+)
+
+
+def get_positives_and_negatives(prompt):
+    positive_node_id_list = {}
+    negative_node_id_list = {}
+    for node_id, node in prompt.items():
+        for sampler_type, field_map in samplers.items():
+            if node["class_type"] == sampler_type.__name__:
+                positive_node_id = node["inputs"][field_map["positive"]][0]
+                negative_node_id = node["inputs"][field_map["negative"]][0]
+                positive_node_id_list[node_id] = positive_node_id
+                negative_node_id_list[node_id] = negative_node_id
+    return positive_node_id_list, negative_node_id_list

py/defs/hook_definition.py

@@ -0,0 +1,23 @@
+from nodes import CLIPTextEncode, KSampler
+from .meta import MetaField
+
+samplers = {
+    KSampler: {
+        "positive": "positive",
+        "negative": "negative",
+    },
+}
+
+hooks = {
+    CLIPTextEncode: {
+        MetaField.POSITIVE_PROMPT: "text",
+        MetaField.NEGATIVE_PROMPT: "text",
+    },
+    KSampler: {
+        MetaField.SEED: "seed",
+        MetaField.STEPS: "steps",
+        MetaField.CFG: "cfg",
+        MetaField.SAMPLER_NAME: "sampler_name",
+        MetaField.SCHEDULER: "scheduler",
+    },
+}

py/defs/meta.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class MetaField(Enum):
+    MODEL_NAME = 0
+    POSITIVE_PROMPT = 10
+    NEGATIVE_PROMPT = 11
+    SEED = 20
+    STEPS = 21
+    CFG = 22
+    SAMPLER_NAME = 23
+    SCHEDULER = 24
+    IMAGE_WIDTH = 30
+    IMAGE_HEIGHT = 31

py/nodes/base.py

@@ -0,0 +1,2 @@
+class BaseNode:
+    CATEGORY = "SaveImage"

py/nodes/node.py

@@ -0,0 +1,74 @@
+import os
+import json
+
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
+import numpy as np
+
+import folder_paths
+from comfy.cli_args import args
+
+from .base import BaseNode
+
+
+# refer. https://github.com/comfyanonymous/ComfyUI/blob/38b7ac6e269e6ecc5bdd6fefdfb2fb1185b09c9d/nodes.py#L1411
+class SaveImageWithMetaData(BaseNode):
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+        self.compress_level = 4
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "filename_prefix": ("STRING", {"default": "ComfyUI"}),
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "save_images"
+
+    OUTPUT_NODE = True
+
+    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]
+        )
+        results = list()
+        for image in images:
+            i = 255.0 * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            metadata = None
+            if not args.disable_metadata:
+                metadata = PngInfo()
+                if prompt is not None:
+                    metadata.add_text("prompt", json.dumps(prompt))
+                if extra_pnginfo is not None:
+                    for x in extra_pnginfo:
+                        metadata.add_text(x, json.dumps(extra_pnginfo[x]))
+
+            file = f"{filename}_{counter:05}_.png"
+            # img.save(
+            #     os.path.join(full_output_folder, file),
+            #     pnginfo=metadata,
+            #     compress_level=self.compress_level,
+            # )
+            results.append(
+                {"filename": file, "subfolder": subfolder, "type": self.type}
+            )
+            counter += 1
+
+        return {"ui": {"images": results}}