Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

node to create wildcard label prompt from string list #536

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions __init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,8 @@ def wildcard_load():
"ImpactNeg": ImpactNeg,
"ImpactConditionalStopIteration": ImpactConditionalStopIteration,
"ImpactStringSelector": StringSelector,
"StringListToString": StringListToString,
"WildcardPromptFromString": WildcardPromptFromString,

"RemoveNoiseMask": RemoveNoiseMask,

Expand Down Expand Up @@ -433,6 +435,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 @@ -485,3 +486,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),)