diff --git a/py/__init__.py b/py/__init__.py index 07724b1..ca14bf6 100644 --- a/py/__init__.py +++ b/py/__init__.py @@ -1,17 +1,10 @@ import functools -from .defs.hook_definition import hooks, samplers -from .defs.meta import MetaField +from .hook import pre_execute import execution -inputs = {} -current_node_id = -1 -current_prompt_id = -1 -current_prompt = {} -positive_nodes = {} -negative_nodes = {} - +# refer. https://stackoverflow.com/a/35758398 def prefix_function(function, prefunction): @functools.wraps(function) def run(*args, **kwargs): @@ -21,76 +14,6 @@ def prefix_function(function, prefunction): 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/hook_definition.py b/py/defs/hook_definition.py deleted file mode 100644 index e6125a1..0000000 --- a/py/defs/hook_definition.py +++ /dev/null @@ -1,23 +0,0 @@ -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/hooks.py b/py/defs/hooks.py new file mode 100644 index 0000000..9374f10 --- /dev/null +++ b/py/defs/hooks.py @@ -0,0 +1,27 @@ +from .meta import MetaField +from .validators import is_positive_prompt, is_negative_prompt + + +hooks = { + "EmptyLatentImage": { + MetaField.IMAGE_WIDTH: {"field_name": "width"}, + MetaField.IMAGE_HEIGHT: {"field_name": "height"}, + }, + "CLIPTextEncode": { + MetaField.POSITIVE_PROMPT: { + "field_name": "text", + "validate": is_positive_prompt, + }, + MetaField.NEGATIVE_PROMPT: { + "field_name": "text", + "validate": is_negative_prompt, + }, + }, + "KSampler": { + MetaField.SEED: {"field_name": "seed"}, + MetaField.STEPS: {"field_name": "steps"}, + MetaField.CFG: {"field_name": "cfg"}, + MetaField.SAMPLER_NAME: {"field_name": "sampler_name"}, + MetaField.SCHEDULER: {"field_name": "scheduler"}, + }, +} diff --git a/py/defs/samplers.py b/py/defs/samplers.py new file mode 100644 index 0000000..34cb447 --- /dev/null +++ b/py/defs/samplers.py @@ -0,0 +1,6 @@ +samplers = { + "KSampler": { + "positive": "positive", + "negative": "negative", + }, +} diff --git a/py/defs/validators.py b/py/defs/validators.py new file mode 100644 index 0000000..f48558f --- /dev/null +++ b/py/defs/validators.py @@ -0,0 +1,21 @@ +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(): + if node["class_type"] == sampler_type: + positive_node_id = node["inputs"][field_map["positive"]][0] + positive_node_id_list[nid] = positive_node_id + return node_id in positive_node_id_list.values() + + +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(): + if node["class_type"] == sampler_type: + negative_node_id = node["inputs"][field_map["negative"]][0] + negative_node_id_list[nid] = negative_node_id + return node_id in negative_node_id_list.values() diff --git a/py/hook.py b/py/hook.py new file mode 100644 index 0000000..d45d78a --- /dev/null +++ b/py/hook.py @@ -0,0 +1,13 @@ +current_prompt = {} +current_extra_data = {} +prompt_executer = None + + +def pre_execute(self, prompt, prompt_id, extra_data, execute_outputs): + global current_prompt + global current_extra_data + global prompt_executer + + current_prompt = prompt + current_extra_data = extra_data + prompt_executer = self diff --git a/py/nodes/node.py b/py/nodes/node.py index 3c0df76..4ec4fa4 100644 --- a/py/nodes/node.py +++ b/py/nodes/node.py @@ -1,4 +1,3 @@ -import os import json from PIL import Image @@ -9,6 +8,13 @@ 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 # refer. https://github.com/comfyanonymous/ComfyUI/blob/38b7ac6e269e6ecc5bdd6fefdfb2fb1185b09c9d/nodes.py#L1411 @@ -37,7 +43,6 @@ class SaveImageWithMetaData(BaseNode): def save_images( self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None ): - filename_prefix += self.prefix_append ( full_output_folder, filename, @@ -48,12 +53,14 @@ class SaveImageWithMetaData(BaseNode): filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0] ) results = list() + print("Save Image With Metadata") 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() + # metadata.add_text("parameters", comment) if prompt is not None: metadata.add_text("prompt", json.dumps(prompt)) if extra_pnginfo is not None: @@ -66,9 +73,84 @@ 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