nodes/nkchocoai-ComfyUI-SaveImageWithMetaData
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Get PNGInfo

Browse Source
This commit is contained in:
nkchocoai
2024-02-15 21:48:25 +09:00
parent 4d6de8ce09
commit 75664ff31c
7 changed files with 153 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/__init__.py
+2 -79
View File
@@ -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
py/defs/hook_definition.py
-23
View File
@@ -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",
},
}
py/defs/hooks.py
+27
View File
@@ -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"},
},
}
py/defs/samplers.py
+6
View File
@@ -0,0 +1,6 @@
samplers = {
"KSampler": {
"positive": "positive",
"negative": "negative",
},
}
py/defs/validators.py
+21
View File
@@ -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()
py/hook.py
+13
View File
@@ -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
py/nodes/node.py
+84 -2
View File
@@ -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