From 3cee9a9db40f816cb3e7b1c482a4305d95bf638c Mon Sep 17 00:00:00 2001 From: nkchocoai Date: Sat, 15 Jun 2024 14:29:18 +0900 Subject: [PATCH] #22 Search for PP and NP recursively. --- py/defs/validators.py | 34 +++++++++++++++++++++++----------- 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/py/defs/validators.py b/py/defs/validators.py index 2d5d151..746abda 100644 --- a/py/defs/validators.py +++ b/py/defs/validators.py @@ -1,21 +1,33 @@ +from collections import deque + 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() + return node_id in _get_node_id_list(prompt, "positive") def is_negative_prompt(node_id, obj, prompt, extra_data, outputs, input_data_all): - negative_node_id_list = {} + return node_id in _get_node_id_list(prompt, "negative") + + +def _get_node_id_list(prompt, field_name): + 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() + # There are nodes between "KSampler" and "CLIP Text Encode" in the SD3 workflow + d = deque() + d.append(node["inputs"][field_map[field_name]][0]) + while len(d) > 0: + nid2 = d.popleft() + class_type = prompt[nid2]["class_type"] + if class_type == "CLIPTextEncode": + node_id_list[nid] = nid2 + break + inputs = prompt[nid2]["inputs"] + for k, v in inputs.items(): + if isinstance(v, list): + d.append(v[0]) + + return node_id_list.values()