#22 Search for PP and NP recursively.

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()