Add files via upload

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.

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

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
+# --- NSFW Example (Use Responsibly) ---
+nsfw_body:
+  - large breasts
+  - medium breasts
+  - small breasts
+  - nude
+  - naked
+  - explicit