diff --git a/__init__.py b/__init__.py index e5a4f74..e1d6708 100644 --- a/__init__.py +++ b/__init__.py @@ -5,6 +5,7 @@ from .py.node.const import * from .py.node.format import * from .py.node.preset import * from .py.node.weight import * +from .py.node.replace import * from .py.server import * from .py.preset import PresetManager @@ -19,6 +20,7 @@ NODE_CLASS_MAPPINGS = { "PromptUtilitiesConstStringMultiLine": PromptUtilitiesConstStringMultiLine, "PromptUtilitiesPromptWeight": PromptUtilitiesPromptWeight, "PromptUtilitiesRoundPromptWeight": PromptUtilitiesRoundPromptWeight, + "PromptUtilitiesReplaceOrInsertTag": PromptUtilitiesReplaceOrInsertTag, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -32,6 +34,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PromptUtilitiesConstStringMultiLine": "Const String (multi line)", "PromptUtilitiesPromptWeight": "Prompt Weight", "PromptUtilitiesRoundPromptWeight": "Round Prompt Weight", + "PromptUtilitiesReplaceOrInsertTag": "Replace or Insert Tag", } WEB_DIRECTORY = "./js" diff --git a/py/node/replace.py b/py/node/replace.py new file mode 100644 index 0000000..ae98ae9 --- /dev/null +++ b/py/node/replace.py @@ -0,0 +1,57 @@ +from .base import BaseNode + +import re +from comfy.sd1_clip import token_weights + + +class PromptUtilitiesReplaceOrInsertTag(BaseNode): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "text": ( + "STRING", + {"multiline": True, "default": "", "forceInput": True}, + ), + "pattern": ("STRING", {"default": "", "multiline": False}), + "value": ("STRING", {"default": "", "multiline": False}), + "mode": (["replace", "insert"],), + "inherit_weight": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = ("STRING",) + + FUNCTION = "replace_tag" + + def replace_tag(self, text, pattern, value, mode, inherit_weight): + text = text.replace("\(", "{[{[") + text = text.replace("\)", "}]}]") + weights = token_weights(text, 1.0) + tags = [] + for t, weight in weights: + for tag in t.split(","): + tag = tag.strip() + if tag: + tags.append((tag, weight)) + tags2 = [] + for i, v in enumerate(tags): + if re.match(pattern, v[0]): + w = 1.0 + if inherit_weight: + w = v[1] + if mode == "insert": + tags2.append(v) + tags2.append((value, w)) + else: + tags2.append(v) + result = [] + for tag, weight in tags2: + if weight == 1.0: + result.append(tag) + else: + result.append(f"({tag}:{weight})") + s = ", ".join(result) + s = s.replace("{[{[", "\(") + s = s.replace("}]}]", "\)") + return (s,)