From 7e4547e4db9409dd599c5fddf24c010d37d77b84 Mon Sep 17 00:00:00 2001 From: Legende <68666984+xLegende@users.noreply.github.com> Date: Sat, 12 Apr 2025 16:22:47 +0200 Subject: [PATCH] Add files via upload --- README.md | 6 +++-- categorized_prompt_formatter.py | 41 +++++++++++++++++++++++++-------- prompt_categories.yaml | 25 +++++++++++++++----- 3 files changed, 54 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 719c12f..f7182ca 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ A custom node for ComfyUI that allows you to filter and restructure text prompts This node takes a standard comma-separated prompt, categorizes each tag according to your definitions, and then rebuilds the prompt based on a template you provide. It's useful for enforcing prompt structure, extracting specific concepts, or reordering elements dynamically. -![Example Node UI (Conceptual - Actual UI may vary slightly)](custom_node.png) +![Example Node UI (Conceptual - Actual UI may vary slightly)](placeholder_image.png) ## Features @@ -13,6 +13,7 @@ This node takes a standard comma-separated prompt, categorizes each tag accordin * **Template-Driven Output:** Structure your output prompt using placeholders like `<|category_name|>`. * **Flexible Delimiters:** Configure input and output delimiters. * **Weight/Emphasis Handling:** Attempts to preserve common syntax like `(tag:1.2)` and `[tag]` while using the base tag for categorization. +* **Underscore/Space Matching:** Optionally treat tags like `red_eyes` and `red eyes` as equivalent during matching (while preserving original format in output). * **Case Sensitivity Option:** Choose whether matching should be case-sensitive or not. * **Unmatched Tag Handling:** Decide whether to discard tags not used in the template, append them to the end, or output them separately. @@ -61,6 +62,7 @@ The node "📝 Categorized Prompt Formatter" should now appear under the "text/f * `category_definition_file`: Enter the name (e.g., `my_categories.yaml`) or full path to your YAML file. The node will search relative to itself, in `ComfyUI/input/`, and absolute paths. * `output_template`: Define your desired output structure using `<|category_name|>` placeholders. Example: `<|quality|>, <|character_num|>, <|eyes|>, <|clothing|>, <|setting|>` * Adjust delimiters, case sensitivity, weight handling, and unmatched tag handling as needed. + * `match_underscores_spaces` (Boolean, default: True): If enabled, the node will try to match input tags against YAML definitions by checking variants with underscores swapped for spaces (e.g., input `red_eyes` can match YAML `red eyes`, and vice-versa). The original format of the input tag is kept for the output string. 5. **Connect Output:** Connect the `formatted_prompt` output to the next node in your workflow (e.g., a CLIP Text Encode node). Use the optional `rejected_prompt` output for debugging or alternative workflows. ## Example @@ -77,4 +79,4 @@ The node "📝 Categorized Prompt Formatter" should now appear under the "text/f * The YAML file path finding tries to be robust, checking relative to the node, the ComfyUI `input` directory, and absolute paths. If it can't find the file, it will print a warning and proceed without categorization. * Weight handling uses regular expressions to parse common patterns. Very complex nested structures might not parse correctly. The base tag extracted is used for lookup. -* If a tag belongs to multiple categories, and the template requests those multiple categories, the tag **will appear multiple times** in the output string by default. +* If a tag belongs to multiple categories, and the template requests those multiple categories, the tag **will appear multiple times** in the output string by default. \ No newline at end of file diff --git a/categorized_prompt_formatter.py b/categorized_prompt_formatter.py index bb92e70..0bae6dc 100644 --- a/categorized_prompt_formatter.py +++ b/categorized_prompt_formatter.py @@ -158,6 +158,7 @@ class CategorizedPromptFormatter: "strip_whitespace": ("BOOLEAN", {"default": True}), "case_sensitive_matching": ("BOOLEAN", {"default": False}), "handle_weights": ("BOOLEAN", {"default": True}), + "match_underscores_spaces": ("BOOLEAN", {"default": True}), "unmatched_tag_handling": (["discard", "append_end", "output_separately"], {"default": "discard"}), } } @@ -206,28 +207,48 @@ class CategorizedPromptFormatter: # --- 2. Parse Input Prompt & Categorize --- categorized_tags = defaultdict(list) all_input_tags_original = set() - processed_input_tags = [] # Store (original_tag, base_tag) + processed_input_tags = [] # Store (original_tag, base_tag for matching) if prompt: raw_tags = prompt.split(input_delimiter) for raw_tag in raw_tags: tag_original = raw_tag.strip() if strip_whitespace else raw_tag - if not tag_original: continue # Skip empty tags resulting from split + if not tag_original: continue all_input_tags_original.add(tag_original) original_parsed, base_parsed = parse_tag(tag_original, handle_weights) - lookup_tag = base_parsed if not case_sensitive_matching else base_parsed.lower() - processed_input_tags.append((original_parsed, lookup_tag)) - + # Prepare the tag used for dictionary lookups (case normalization) + lookup_key_base = base_parsed if case_sensitive_matching else base_parsed.lower() + processed_input_tags.append((original_parsed, lookup_key_base)) # Store original and the key base matched_original_tags = set() - for original_tag, lookup_tag in processed_input_tags: - categories = tag_to_categories_map.get(lookup_tag) - if categories: - for category in categories: + for original_tag, lookup_key_base in processed_input_tags: + found_categories = None + + if match_underscores_spaces: + # Generate variants for lookup + variants_to_check = { + lookup_key_base, + lookup_key_base.replace('_', ' '), + lookup_key_base.replace(' ', '_') + } + # Check each variant against the map + for variant in variants_to_check: + categories = tag_to_categories_map.get(variant) + if categories: + found_categories = categories + break # Found a match, stop checking variants for this tag + else: + # Direct lookup if matching is disabled + found_categories = tag_to_categories_map.get(lookup_key_base) + + # If categories were found (using either direct or variant matching) + if found_categories: + for category in found_categories: + # IMPORTANT: Append the ORIGINAL tag (with weights etc.) categorized_tags[category].append(original_tag) - matched_original_tags.add(original_tag) # Mark as matched if found in *any* category + matched_original_tags.add(original_tag) # Mark the ORIGINAL tag as matched # --- 3. Process Template --- formatted_output_string = output_template diff --git a/prompt_categories.yaml b/prompt_categories.yaml index bf82ead..b61d7aa 100644 --- a/prompt_categories.yaml +++ b/prompt_categories.yaml @@ -6,9 +6,14 @@ quality: - masterpiece - best quality - high quality + - amazing quality - absurdres - highres - ultra detailed + - very aesthetic + - high resolution + - absurd resolution + - ultra-detailed # --- Style Tags --- style: @@ -26,8 +31,14 @@ character_num: - 1girl - 1boy - 2girls + - 2boys + - 3girls + - 3boys + - 4girls + - 4boys - solo - multiple boys + - multiple girls # --- Appearance Details --- eyes: @@ -104,9 +115,11 @@ setting: - simple background - detailed background -# --- NSFW Example (Illustrative - Use Responsibly) --- -# nsfw_body: -# - large breasts -# - small breasts -# - nude -# - explicit \ No newline at end of file +# --- NSFW Example (Use Responsibly) --- +nsfw_body: + - large breasts + - medium breasts + - small breasts + - nude + - naked + - explicit \ No newline at end of file