diff --git a/__init__.py b/__init__.py index 9fa6c91..86479cf 100644 --- a/__init__.py +++ b/__init__.py @@ -6,12 +6,13 @@ from .execution_delay import ExecutionDelayNode from .file_loader import WildcardFileLoaderNode from .randomizer import MultilineRandomChanceNode from .text_concatenator import TextConcatenator4Node, TextConcatenator8Node -from .text_tag_replacer import TextRandomReplacerNode +from .text_tag_replacer import TextRandomReplacerNode, TextMultiLineRandomReplacerNode NODE_CLASS_MAPPINGS = { "MultilineRandomChanceNode": MultilineRandomChanceNode, "ConsoleLogNode": ConsoleLogNode, "TextRandomReplacerNode": TextRandomReplacerNode, + "TextMultiLineRandomReplacerNode": TextMultiLineRandomReplacerNode, "WildcardFileLoaderNode": WildcardFileLoaderNode, "ExecutionDelayNode": ExecutionDelayNode, "TextConcatenator4Node": TextConcatenator4Node, @@ -22,6 +23,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "MultilineRandomChanceNode": "🎲 Weighted Multiline Randomizer", "ConsoleLogNode": "🖥️ Log to Console", "TextRandomReplacerNode": "🏷️ Text Random Replacer", + "TextMultiLineRandomReplacerNode": "🏷️ Text Random Replacer (MultiLine)", "WildcardFileLoaderNode": "📂 Wildcard File Loader", "ExecutionDelayNode": "⏳ Execution Delay", "TextConcatenator4Node": "🧩 Text Concatenator (4-Way)", diff --git a/text_tag_replacer.py b/text_tag_replacer.py index 4ac9d8c..c264a00 100644 --- a/text_tag_replacer.py +++ b/text_tag_replacer.py @@ -4,6 +4,101 @@ from .randomizer import MultilineRandomChanceNode class TextRandomReplacerNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ( + "STRING", + { + "multiline": False, + "default": "#TAG" + # "default": "A {professional photo|2d sketch} of a #CHARACTER in a #LOCATION, cinematic lighting.", + }, + ), + #"target_tag": ("STRING", {"default": "#CHARACTER"}), + "target_tag": ("STRING", {"default": "#TAG"}), + "replacement_options": ( + "STRING", + { + "multiline": True, + "default": "" + # "default": "ancient elven {wizard|mage}\ncyberpunk {hacker|netrunner} :: 2", + }, + ), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("completed_prompt",) + FUNCTION = "replace_tag_with_random" + CATEGORY = "TextRandomizer/Text" + + def _process_single_string( + self, single_text, target_tag, chosen_replacement, randomizer + ): + """Helper to swap the tag and expand any curly brackets in a single prompt string.""" + if target_tag and target_tag in single_text: + text_with_tag_swapped = single_text.replace(target_tag, chosen_replacement) + else: + text_with_tag_swapped = single_text + + bracket_pattern = re.compile(r"\{([^{}]+)\}") + final_prompt = text_with_tag_swapped + + loop_guard = 0 + while bracket_pattern.search(final_prompt) and loop_guard < 10: + final_prompt = bracket_pattern.sub( + randomizer._replace_brackets, final_prompt + ) + loop_guard += 1 + + return final_prompt + + def replace_tag_with_random(self, text, target_tag, replacement_options, seed): + randomizer = MultilineRandomChanceNode() + + # 1. Standardise the replacement options into a flat list of pure strings + if isinstance(replacement_options, list): + options_list = replacement_options + else: + # FIX: We unpack index [0] to extract the pure string from the randomizer tuple + chosen_tuple = randomizer.process_random_text(replacement_options, seed) + options_list = [chosen_tuple[0]] + + # 2. Standardise the input template text into a list + if isinstance(text, list): + text_list = text + else: + text_list = [text] + + final_batch_output = [] + + # 3. Zip or loop through the text templates and your batch of options + max_batch = max(len(text_list), len(options_list)) + + for i in range(max_batch): + single_text_item = text_list[i % len(text_list)] + chosen_replacement = options_list[i % len(options_list)] + + processed_item = self._process_single_string( + single_text_item, target_tag, chosen_replacement, randomizer + ) + final_batch_output.append(processed_item) + + # 4. Clean up the output format for ComfyUI pipeline compliance + if ( + len(final_batch_output) == 1 + and not isinstance(text, list) + and not isinstance(replacement_options, list) + ): + return (final_batch_output[0],) + else: + return (final_batch_output,) + + +class TextMultiLineRandomReplacerNode: @classmethod def INPUT_TYPES(cls): return {