diff --git a/__init__.py b/__init__.py index c1a2e2f..ccb2feb 100644 --- a/__init__.py +++ b/__init__.py @@ -57,15 +57,18 @@ SELECTOR_CLASS_MAPPINGS = { TEXT_CLASS_MAPPINGS = { "Sage_SetText": Sage_SetText, - "Sage_SaveText": Sage_SaveText, + "Sage_TextSwitch": Sage_TextSwitch, "SageSetWildcardText": SageSetWildcardText, + "Sage_TextSubstitution": Sage_TextSubstitution, "Sage_JoinText": Sage_JoinText, "Sage_TripleJoinText": Sage_TripleJoinText, "Sage_CleanText": Sage_CleanText, + "Sage_TextSelectLine": Sage_TextSelectLine, "Sage_TextRandomLine": Sage_TextRandomLine, "Sage_TextWeight": Sage_TextWeight, "Sage_ViewAnything": Sage_ViewAnything, #"Sage_ViewNotes": Sage_ViewNotes, + "Sage_SaveText": Sage_SaveText, "Sage_PonyPrefix": Sage_PonyPrefix, "Sage_PonyStyle": Sage_PonyStyle, "Sage_HiDreamE1_Instruction": Sage_HiDreamE1_Instruction @@ -199,15 +202,18 @@ SELECTOR_NAME_MAPPINGS = { TEXT_NAME_MAPPINGS = { "Sage_SetText": "Set Text", - "Sage_SaveText": "Save Text", + "Sage_TextSwitch": "Text Switch", "SageSetWildcardText": "Set Text (Wildcards)", + "Sage_TextSubstitution": "Text Substitution", "Sage_JoinText": "Join Text", "Sage_TripleJoinText": "Join Text x3", "Sage_CleanText": "Clean Text", + "Sage_TextSelectLine": "Select Line from Text", "Sage_TextRandomLine": "Random Line from Text", "Sage_TextWeight": "Text Weight", "Sage_ViewAnything": "View Any Node as Text", #"Sage_ViewNotes": "View Notes", + "Sage_SaveText": "Save Text", "Sage_PonyPrefix": "Add Pony v6 Prefixes", "Sage_PonyStyle": "Add Pony Style", "Sage_HiDreamE1_Instruction": "HiDreamE1 Instruction" diff --git a/nodes/text.py b/nodes/text.py index 281a6cb..fda7cd4 100644 --- a/nodes/text.py +++ b/nodes/text.py @@ -81,6 +81,26 @@ class Sage_SetText(ComfyNodeABC): def pass_str(self, str, prefix=None, suffix=None) -> tuple[str]: return (f"{prefix or ''}{str}{suffix or ''}",) +class Sage_TextSwitch(ComfyNodeABC): + @classmethod + def INPUT_TYPES(cls) -> InputTypeDict: + return { + "required": { + "str": (IO.STRING, {"forceInput": True, "multiline": True}), + "active": (IO.BOOLEAN, {"defaultInput": False, "default": True}) + } + } + + RETURN_TYPES = (IO.STRING,) + RETURN_NAMES = ("str",) + + FUNCTION = "text_switch" + + CATEGORY = "Sage Utils/text" + DESCRIPTION = "Passes the text if active is true, otherwise passes an empty string." + + def text_switch(self, str: str, active: bool) -> tuple[str]: + return (str if active else "",) class Sage_SaveText(ComfyNodeABC): @classmethod @@ -248,6 +268,73 @@ class Sage_TextRandomLine(ComfyNodeABC): random.seed(seed) random_line = lines[random.randint(0, len(lines) - 1)].strip() # Use seed to select a line return (random_line,) + +class Sage_TextSelectLine(ComfyNodeABC): + @classmethod + def INPUT_TYPES(cls) -> InputTypeDict: + return { + "required": { + "text": (IO.STRING, {"defaultInput": True, "multiline": True}), + "line_number": (IO.INT, {"defaultInput": False, "default": 0, "min": 0, "max": 10000, "step": 1}) + } + } + + RETURN_TYPES = (IO.STRING,) + RETURN_NAMES = ("selected_line",) + + FUNCTION = "select_line" + + CATEGORY = "Sage Utils/text" + DESCRIPTION = "Selects a specific line from the given text based on the line number. Line numbers start from 0. If the line number is out of range, it will select the first or last line as appropriate." + + def select_line(self, text: str, line_number: int) -> tuple[str]: + lines = text.splitlines() + if line_number < 0: + line_number = 0 + elif line_number >= len(lines): + line_number = len(lines) - 1 + + return (lines[line_number].strip(),) + +# This node takes a text box (with a prefix and suffix) and three optional strings, and substitutes the strings for the +# placeholders in the text box. +class Sage_TextSubstitution(ComfyNodeABC): + @classmethod + def INPUT_TYPES(cls) -> InputTypeDict: + return { + "required": { + "text": (IO.STRING, {"defaultInput": True, "multiline": True}) + }, + "optional": { + "prefix": (IO.STRING, {"defaultInput": True, "multiline": True}), + "suffix": (IO.STRING, {"defaultInput": True, "multiline": True}), + "str1": (IO.STRING, {"defaultInput": True, "multiline": True}), + "str2": (IO.STRING, {"defaultInput": True, "multiline": True}), + "str3": (IO.STRING, {"defaultInput": True, "multiline": True}), + } + } + + RETURN_TYPES = (IO.STRING,) + RETURN_NAMES = ("result",) + + FUNCTION = "substitute_text" + CATEGORY = "Sage Utils/text" + DESCRIPTION = "Substitutes the placeholders in the text with the provided strings. The placeholders are {str1}, {str2}, and {str3}. The prefix and suffix are added to the final result." + + def substitute_text(self, text: str, prefix: str = "", suffix: str = "", str1: str = "", str2: str = "", str3: str = "") -> tuple[str]: + """ + Substitutes the placeholders in the text with the provided strings. + The placeholders are {str1}, {str2}, and {str3}. + The prefix and suffix are added to the final result. + """ + # Replace placeholders with actual strings + result = text.format(str1=str1, str2=str2, str3=str3) + + # Add prefix and suffix + result = f"{prefix}{result}{suffix}" + + return (result,) + class Sage_PonyPrefix(ComfyNodeABC): @classmethod def INPUT_TYPES(cls) -> InputTypeDict: @@ -276,7 +363,8 @@ class Sage_TextWeight(ComfyNodeABC): return { "required": { "text": (IO.STRING, {"defaultInput": True, "multiline": True}), - "weight": (IO.FLOAT, {"defaultInput": False, "default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}) + "weight": (IO.FLOAT, {"defaultInput": False, "default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}), + "separator": (IO.STRING, {"defaultInput": False, "default": ', '}) } } @@ -288,8 +376,8 @@ class Sage_TextWeight(ComfyNodeABC): CATEGORY = "Sage Utils/text" DESCRIPTION = "Applies a weight to a text string." - def weight_text(self, text: str, weight: float) -> tuple[str]: - return (f"({text}:{weight:.2f}),",) + def weight_text(self, text: str, weight: float, separator: str) -> tuple[str]: + return (f"({text}:{weight:.2f}){separator}",) class Sage_HiDreamE1_Instruction(ComfyNodeABC): @classmethod