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

Rename to sampler selection method

Browse Source
This commit is contained in:
nkchocoai
2024-02-17 23:08:39 +09:00
parent a72209e266
commit 92ff5c907f
7 changed files with 144 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/defs/combo.py
+1 -1
View File
@@ -1 +1 @@
SAMPLER_FIND_METHOD = ["Farthest", "Nearest", "By node ID"]
SAMPLER_SELECTION_METHOD = ["Farthest", "Nearest", "By node ID"]
py/nodes/node.py
+8 -6
View File
@@ -14,7 +14,7 @@ from ..capture import Capture
from .. import hook
from ..trace import Trace
from ..defs.combo import SAMPLER_FIND_METHOD
from ..defs.combo import SAMPLER_SELECTION_METHOD
# refer. https://github.com/comfyanonymous/ComfyUI/blob/38b7ac6e269e6ecc5bdd6fefdfb2fb1185b09c9d/nodes.py#L1411
@@ -31,7 +31,7 @@ class SaveImageWithMetaData(BaseNode):
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
"sampler_find_method": (SAMPLER_FIND_METHOD,),
"sampler_selection_method": (SAMPLER_SELECTION_METHOD,),
"sampler_find_node_id": (
"INT",
{"default": 0, "min": 0, "max": 999999999, "step": 1},
@@ -49,7 +49,7 @@ class SaveImageWithMetaData(BaseNode):
self,
images,
filename_prefix="ComfyUI",
sampler_find_method=SAMPLER_FIND_METHOD[0],
sampler_selection_method=SAMPLER_SELECTION_METHOD[0],
sampler_find_node_id=0,
prompt=None,
extra_pnginfo=None,
@@ -68,7 +68,9 @@ class SaveImageWithMetaData(BaseNode):
i = 255.0 * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
pnginfo_dict = self.gen_pnginfo(sampler_find_method, sampler_find_node_id)
pnginfo_dict = self.gen_pnginfo(
sampler_selection_method, sampler_find_node_id
)
if len(images) >= 2:
pnginfo_dict["Batch index"] = index
pnginfo_dict["Batch size"] = len(images)
@@ -100,7 +102,7 @@ class SaveImageWithMetaData(BaseNode):
return {"ui": {"images": results}}
@classmethod
def gen_pnginfo(cls, sampler_find_method, sampler_find_node_id):
def gen_pnginfo(cls, sampler_selection_method, sampler_find_node_id):
# get all node inputs
inputs = Capture.get_inputs()
@@ -109,7 +111,7 @@ class SaveImageWithMetaData(BaseNode):
hook.current_node_id, hook.current_prompt
)
sampler_node_id = Trace.find_sampler_node_id(
trace_tree_from_this_node, sampler_find_method, sampler_find_node_id
trace_tree_from_this_node, sampler_selection_method, sampler_find_node_id
)
# get inputs before sampler node
py/trace.py
+4 -4
View File
@@ -3,7 +3,7 @@ import json
from collections import deque
from .defs.samplers import SAMPLERS
from .defs.combo import SAMPLER_FIND_METHOD
from .defs.combo import SAMPLER_SELECTION_METHOD
class Trace:
@@ -25,8 +25,8 @@ class Trace:
return trace_tree
@classmethod
def find_sampler_node_id(cls, trace_tree, sampler_find_method, node_id):
if sampler_find_method == SAMPLER_FIND_METHOD[2]:
def find_sampler_node_id(cls, trace_tree, sampler_selection_method, node_id):
if sampler_selection_method == SAMPLER_SELECTION_METHOD[2]:
node_id = str(node_id)
_, class_type = trace_tree.get(node_id, (-1, None))
if class_type in SAMPLERS.keys():
@@ -36,7 +36,7 @@ class Trace:
sorted_by_distance_trace_tree = sorted(
[(k, v[0], v[1]) for k, v in trace_tree.items()],
key=lambda x: x[1],
reverse=(sampler_find_method == SAMPLER_FIND_METHOD[0]),
reverse=(sampler_selection_method == SAMPLER_SELECTION_METHOD[0]),
)
for nid, _, class_type in sorted_by_distance_trace_tree:
if class_type in SAMPLERS.keys():