Skip to content

Commit

Permalink
node to create wildcard label prompt from string list (#536)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
jkrauss82 and Jonas Krauss authored May 1, 2024
1 parent 80e5958 commit b67669d
Show file tree
Hide file tree
Showing 2 changed files with 93 additions and 0 deletions.
4 changes: 4 additions & 0 deletions __init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,8 @@ def wildcard_load():
"ImpactNeg": ImpactNeg,
"ImpactConditionalStopIteration": ImpactConditionalStopIteration,
"ImpactStringSelector": StringSelector,
"StringListToString": StringListToString,
"WildcardPromptFromString": WildcardPromptFromString,

"RemoveNoiseMask": RemoveNoiseMask,

Expand Down Expand Up @@ -413,6 +415,8 @@ def wildcard_load():
"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",
Expand Down
89 changes: 89 additions & 0 deletions modules/impact/util_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import comfy
import sys
import nodes
import re


class GeneralSwitch:
Expand Down Expand Up @@ -489,3 +490,91 @@ def doit(self, strings, multiline, select):
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),)

0 comments on commit b67669d

Please sign in to comment.