node to create wildcard label prompt from string list (#536)

* node to create wildcard label prompt from string list

* use `is` for checking on None value

* add StringListToString node

* remove `list` from name of wildcard creation node

---------

Co-authored-by: Jonas Krauss <jonas.krauss@stockpulse.de>

__init__.py

@@ -291,6 +291,8 @@ NODE_CLASS_MAPPINGS = {
     "ImpactNeg": ImpactNeg,
     "ImpactConditionalStopIteration": ImpactConditionalStopIteration,
     "ImpactStringSelector": StringSelector,
+    "StringListToString": StringListToString,
+    "WildcardPromptFromString": WildcardPromptFromString,
 
     "RemoveNoiseMask": RemoveNoiseMask,
 
@@ -413,6 +415,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ImpactMakeImageList": "Make Image List",
     "ImpactMakeImageBatch": "Make Image Batch",
     "ImpactStringSelector": "String Selector",
+    "StringListToString": "String List to String",
+    "WildcardPromptFromString": "Wildcard Prompt from String",
     "ImpactIsNotEmptySEGS": "SEGS isn't Empty",
     "SetDefaultImageForSEGS": "Set Default Image for SEGS",
     "RemoveImageFromSEGS": "Remove Image from SEGS",

modules/impact/util_nodes.py

@@ -5,6 +5,7 @@ import torch
 import comfy
 import sys
 import nodes
+import re
 
 
 class GeneralSwitch:
@@ -489,3 +490,91 @@ class StringSelector:
                 selected = lines[select % len(lines)]
 
         return (selected, )
+
+
+class StringListToString:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "join_with": ("STRING", {"default": "\\n"}),
+                "string_list": ("STRING", {"forceInput": True}),
+            }
+        }
+
+    INPUT_IS_LIST = True
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "doit"
+
+    CATEGORY = "ImpactPack/Util"
+
+    def doit(self, join_with, string_list):
+        # convert \\n to newline character
+        if join_with[0] == "\\n":
+            join_with[0] = "\n"
+
+        joined_text = join_with[0].join(string_list)
+
+        return (joined_text,)
+
+
+class WildcardPromptFromString:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "string": ("STRING", {"forceInput": True}),
+                "delimiter": ("STRING", {"multiline": False, "default": "\\n" }),
+                "prefix_all": ("STRING", {"multiline": False}),
+                "postfix_all": ("STRING", {"multiline": False}),
+                "restrict_to_tags": ("STRING", {"multiline": False}),
+                "exclude_tags": ("STRING", {"multiline": False})
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "STRING",)
+    RETURN_NAMES = ("wildcard", "segs_labels",)
+    FUNCTION = "doit"
+
+    CATEGORY = "ImpactPack/Util"
+
+    def doit(self, string, delimiter, prefix_all, postfix_all, restrict_to_tags, exclude_tags):
+        # convert \\n to newline character
+        if delimiter == "\\n":
+            delimiter = "\n"
+
+        # some sanity checks and normalization for later processing
+        if prefix_all is None: prefix_all = ""
+        if postfix_all is None: postfix_all = ""
+        if restrict_to_tags is None: restrict_to_tags = ""
+        if exclude_tags is None: exclude_tags = ""
+        if not isinstance(restrict_to_tags, list):
+            restrict_to_tags = restrict_to_tags.split(", ")
+        if not isinstance(exclude_tags, list):
+            exclude_tags = exclude_tags.split(", ")
+
+        # build the wildcard prompt per list entry
+        output = ["[LAB]"]
+        labels = []
+        for x in string.split(delimiter):
+            label = str(len(labels) + 1)
+            labels.append(label)
+            x = x.split(", ")
+            # restrict to tags
+            if restrict_to_tags != "":
+                x = list(set(x) & set(restrict_to_tags))
+            # remove tags
+            if exclude_tags != "":
+                x = list(set(x) - set(exclude_tags))
+            # next row: <LABEL> <PREFIX> <TAGS> <POSTFIX>
+            prompt_for_seg = f'[{label}] {prefix_all} {", ".join(x)} {postfix_all}'.strip()
+            output.append(prompt_for_seg)
+        output = "\n".join(output)
+
+        # clean string: fixup double spaces, commas etc.
+        output = re.sub(r' ,', ',', output)
+        output = re.sub(r'  +', ' ', output)
+        output = re.sub(r',,+', ',', output)
+        output = re.sub(r'\n, ', '\n', output)
+
+        return (output, ", ".join(labels),)