Refactoring

py/capture.py

@@ -0,0 +1,79 @@
+from . import hook
+from .defs.captures import CAPTURE_FIELD_LIST
+from .defs.meta import MetaField
+
+from nodes import NODE_CLASS_MAPPINGS
+from execution import get_input_data
+
+
+class Capture:
+    @classmethod
+    def get_inputs(cls):
+        inputs = {}
+        prompt = hook.current_prompt
+        extra_data = hook.current_extra_data
+        outputs = hook.prompt_executer.outputs
+
+        for node_id, obj in prompt.items():
+            class_type = obj["class_type"]
+            obj_class = NODE_CLASS_MAPPINGS[class_type]
+            node_inputs = prompt[node_id]["inputs"]
+            input_data_all = get_input_data(
+                node_inputs, obj_class, node_id, outputs, prompt, extra_data
+            )
+            for node_class, metas in CAPTURE_FIELD_LIST.items():
+                if class_type == node_class:
+                    for meta, field_data in metas.items():
+                        validate = field_data.get("validate")
+                        if validate is not None and not validate(
+                            node_id, obj, prompt, extra_data, outputs, input_data_all
+                        ):
+                            print("validate failed", meta, node_id, obj)
+                            continue
+
+                        if meta not in inputs:
+                            inputs[meta] = []
+                        field_name = field_data["field_name"]
+                        value = input_data_all.get(field_name)
+                        if value is not None:
+                            format = field_data.get("format")
+                            v = value
+                            if format is not None:
+                                v = format(v)
+                            if isinstance(value, list) and len(value) > 0:
+                                v = value[0]
+                            inputs[meta].append((node_id, v))
+        return inputs
+
+    @classmethod
+    def gen_pnginfo_dict(cls, inputs):
+        pnginfo_dict = {}
+
+        def update_pnginfo_dict(metafield, key):
+            x = inputs.get(metafield, [])
+            if len(x) > 0:
+                pnginfo_dict[key] = x[-1][1]
+
+        update_pnginfo_dict(MetaField.POSITIVE_PROMPT, "")
+        update_pnginfo_dict(MetaField.NEGATIVE_PROMPT, "Negative prompt")
+        update_pnginfo_dict(MetaField.STEPS, "Steps")
+
+        sampler_names = inputs.get(MetaField.SAMPLER_NAME, [])
+        if len(sampler_names) > 0:
+            pnginfo_dict["Sampler"] = sampler_names[0][1]
+
+            schedulers = inputs.get(MetaField.SCHEDULER, [])
+            if len(schedulers) > 0:
+                scheduler = schedulers[0][1]
+                if scheduler != "normal":
+                    pnginfo_dict["Sampler"] += "_" + scheduler
+
+        update_pnginfo_dict(MetaField.CFG, "CFG Scale")
+        update_pnginfo_dict(MetaField.SEED, "Seed")
+
+        image_widths = inputs.get(MetaField.IMAGE_WIDTH, [])
+        image_heights = inputs.get(MetaField.IMAGE_HEIGHT, [])
+        if len(image_widths) > 0 and len(image_heights) > 0:
+            pnginfo_dict["Size"] = f"{image_widths[0][1]}x{image_heights[0][1]}"
+
+        return pnginfo_dict

py/defs/captures.py

@@ -2,7 +2,7 @@ from .meta import MetaField
 from .validators import is_positive_prompt, is_negative_prompt
 
 
-hooks = {
+CAPTURE_FIELD_LIST = {
     "EmptyLatentImage": {
         MetaField.IMAGE_WIDTH: {"field_name": "width"},
         MetaField.IMAGE_HEIGHT: {"field_name": "height"},

py/defs/samplers.py

@@ -1,4 +1,4 @@
-samplers = {
+SAMPLERS = {
     "KSampler": {
         "positive": "positive",
         "negative": "negative",

py/defs/validators.py

@@ -1,10 +1,10 @@
-from .samplers import samplers
+from .samplers import SAMPLERS
 
 
 def is_positive_prompt(node_id, obj, prompt, extra_data, outputs, input_data_all):
     positive_node_id_list = {}
     for nid, node in prompt.items():
-        for sampler_type, field_map in samplers.items():
+        for sampler_type, field_map in SAMPLERS.items():
             if node["class_type"] == sampler_type:
                 positive_node_id = node["inputs"][field_map["positive"]][0]
                 positive_node_id_list[nid] = positive_node_id
@@ -14,7 +14,7 @@ def is_positive_prompt(node_id, obj, prompt, extra_data, outputs, input_data_all
 def is_negative_prompt(node_id, obj, prompt, extra_data, outputs, input_data_all):
     negative_node_id_list = {}
     for nid, node in prompt.items():
-        for sampler_type, field_map in samplers.items():
+        for sampler_type, field_map in SAMPLERS.items():
             if node["class_type"] == sampler_type:
                 negative_node_id = node["inputs"][field_map["negative"]][0]
                 negative_node_id_list[nid] = negative_node_id

py/nodes/node.py

@@ -8,13 +8,8 @@ import folder_paths
 from comfy.cli_args import args
 
 from .base import BaseNode
-from .. import hook
-from ..defs.hooks import hooks
-from ..defs.samplers import samplers
-from ..defs.meta import MetaField
 
-from nodes import NODE_CLASS_MAPPINGS
-from execution import get_input_data
+from ..capture import Capture
 
 
 # refer. https://github.com/comfyanonymous/ComfyUI/blob/38b7ac6e269e6ecc5bdd6fefdfb2fb1185b09c9d/nodes.py#L1411
@@ -57,6 +52,14 @@ class SaveImageWithMetaData(BaseNode):
         for image in images:
             i = 255.0 * image.cpu().numpy()
             img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+
+            inputs = Capture.get_inputs()
+            print("[get_inputs]")
+            print(inputs)
+            pnginfo_dict = Capture.gen_pnginfo_dict(inputs)
+            print("[PNGInfo Dict]")
+            print(pnginfo_dict)
+
             metadata = None
             if not args.disable_metadata:
                 metadata = PngInfo()
@@ -73,84 +76,9 @@ class SaveImageWithMetaData(BaseNode):
             #     pnginfo=metadata,
             #     compress_level=self.compress_level,
             # )
-            inputs = self.get_inputs()
-            print("[get_inputs]")
-            print(inputs)
-            pnginfo_dict = self.gen_pnginfo_dict(inputs)
-            print("[PNGInfo Dict]")
-            print(pnginfo_dict)
             results.append(
                 {"filename": file, "subfolder": subfolder, "type": self.type}
             )
             counter += 1
 
         return {"ui": {"images": results}}
-
-    def get_inputs(self):
-        inputs = {}
-        prompt = hook.current_prompt
-        extra_data = hook.current_extra_data
-        outputs = hook.prompt_executer.outputs
-
-        for node_id, obj in prompt.items():
-            class_type = obj["class_type"]
-            obj_class = NODE_CLASS_MAPPINGS[class_type]
-            node_inputs = prompt[node_id]["inputs"]
-            input_data_all = get_input_data(
-                node_inputs, obj_class, node_id, outputs, prompt, extra_data
-            )
-            for node_class, metas in hooks.items():
-                if class_type == node_class:
-                    for meta, field_data in metas.items():
-                        validate = field_data.get("validate")
-                        if validate is not None and not validate(
-                            node_id, obj, prompt, extra_data, outputs, input_data_all
-                        ):
-                            print("validate failed", meta, node_id, obj)
-                            continue
-
-                        if meta not in inputs:
-                            inputs[meta] = []
-                        field_name = field_data["field_name"]
-                        value = input_data_all.get(field_name)
-                        if value is not None:
-                            format = field_data.get("format")
-                            v = value
-                            if format is not None:
-                                v = format(v)
-                            if isinstance(value, list) and len(value) > 0:
-                                v = value[0]
-                            inputs[meta].append((node_id, v))
-        return inputs
-
-    def gen_pnginfo_dict(self, inputs):
-        pnginfo_dict = {}
-
-        def update_pnginfo_dict(metafield, key):
-            x = inputs.get(metafield, [])
-            if len(x) > 0:
-                pnginfo_dict[key] = x[-1][1]
-
-        update_pnginfo_dict(MetaField.POSITIVE_PROMPT, "")
-        update_pnginfo_dict(MetaField.NEGATIVE_PROMPT, "Negative prompt")
-        update_pnginfo_dict(MetaField.STEPS, "Steps")
-
-        sampler_names = inputs.get(MetaField.SAMPLER_NAME, [])
-        if len(sampler_names) > 0:
-            pnginfo_dict["Sampler"] = sampler_names[0][1]
-
-            schedulers = inputs.get(MetaField.SCHEDULER, [])
-            if len(schedulers) > 0:
-                scheduler = schedulers[0][1]
-                if scheduler != "normal":
-                    pnginfo_dict["Sampler"] += "_" + scheduler
-
-        update_pnginfo_dict(MetaField.CFG, "CFG Scale")
-        update_pnginfo_dict(MetaField.SEED, "Seed")
-
-        image_widths = inputs.get(MetaField.IMAGE_WIDTH, [])
-        image_heights = inputs.get(MetaField.IMAGE_HEIGHT, [])
-        if len(image_widths) > 0 and len(image_heights) > 0:
-            pnginfo_dict["Size"] = f"{image_widths[0][1]}x{image_heights[0][1]}"
-
-        return pnginfo_dict