diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ed8ebf5 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e81f8bb --- /dev/null +++ b/__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"] diff --git a/py/__init__.py b/py/__init__.py new file mode 100644 index 0000000..07724b1 --- /dev/null +++ b/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 diff --git a/py/defs/__init__.py b/py/defs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/py/defs/hook_definition.py b/py/defs/hook_definition.py new file mode 100644 index 0000000..e6125a1 --- /dev/null +++ b/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", + }, +} diff --git a/py/defs/meta.py b/py/defs/meta.py new file mode 100644 index 0000000..d63ca4d --- /dev/null +++ b/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 diff --git a/py/nodes/__init__.py b/py/nodes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/py/nodes/base.py b/py/nodes/base.py new file mode 100644 index 0000000..a82188a --- /dev/null +++ b/py/nodes/base.py @@ -0,0 +1,2 @@ +class BaseNode: + CATEGORY = "SaveImage" diff --git a/py/nodes/node.py b/py/nodes/node.py new file mode 100644 index 0000000..3c0df76 --- /dev/null +++ b/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}}