import os import folder_paths import random # Default values for Raffle node DEFAULT_FILTER_OUT_TAGS = """monochrome, greyscale, cross-section, cervix, cervical_penetration, uterus, x-ray, internal_cumshot, lactation, forced_lactation, male_lactation, projectile_lactation, lactation_through_clothes, breast_milk, male_focus, male_penetrated, interracial, dark-skinned_male, gaping, extreme_gaping, prolapse, anal_prolapse, fisting, anal_fisting, anal, anal_only, after_anal, anal_fluid, anal_object_insertion, butt_plug, jewel_butt_plug, anal_beads, vibrator_in_anus, anus, cum_in_ass, spread_anus, spread_ass, anus_peek, puckered_anus, dark_anus, presenting_anus, spreading_own_anus, anus_cutout, censored_anus, spread_anus_under_clothes, colored_anus, pubic_hair, female_pubic_hair, pubic_hair_peek, colored_pubic_hair, sparse_pubic_hair, mismatched_pubic_hair, excessive_pubic_hair, shaped_pubic_hair, condom, used_condom, condom_wrapper, condom_in_mouth, holding_condom, condom_on_penis, multiple_condoms, condom_packet_strip, pointless_condom, condom_belt, condom_box, used_condom_on_penis, condom_left_inside, colored_condom, okamoto_condoms, condom_wrapper_in_clothes, condom_thigh_strap, buying_condoms, broken_condom, used_condom_in_clothes, """ DEFAULT_EXCLUDE_TAGLISTS = "comic, 4koma, multiple_girls, multiple_boys, multiple_views, reference_sheet, 2girls, 3girls, 4girls, 5girls, 6+girls, 2boys, 3boys, 4boys, 5boys, 6+boys, gangbang, threesome, mmf_threesome, ffm_threesome, group_sex, cooperative_fellatio, cooperative_paizuri, double_handjob, surrounded_by_penises, furry, obese, yaoi, yuri, otoko_no_ko, strap-on, futa_with_female, futa_without_pussy, implied_futanari, futanari, diaper, fart, pee, peeing, pee_puddle, pee_stain, peeing_self, golden_shower, scat, guro, ero_guro, intestines, vore, horse_penis" DEFAULT_EXCLUDE_CATEGORIES = "artist, character_name, copyright, meta, clothes_and_accessories, female_physical_descriptors, named_garment_exposure, specific_garment_interactions, speech_and_text, standard_physical_descriptors, metadata_and_attribution, intentional_design_exposure, two_handed_character_items, holding_large_items, content_censorship_methods" DEFAULT_TAGLISTS_MUST_INCLUDE = "1girl" # Critical categories that should be excluded to maintain workflow WARNING_ABOUT_NEW_CATEGORIES = {'artist', 'character_name', 'copyright', 'meta'} # Global list of all available categories - used by both Raffle and TagCategoryStrength ALL_CATEGORIES = [ 'abstract_symbols', 'actions', 'artstyle_technique', 'artist', 'background_objects', 'bodily_fluids', 'camera_angle_perspective', 'camera_focus_subject', 'camera_framing_composition', 'character_count', 'character_name', 'clothes_and_accessories', 'color_scheme', 'content_censorship_methods', 'copyright', 'expressions_and_mental_state', 'female_intimate_anatomy', 'female_physical_descriptors', 'format_and_presentation', 'gaze_direction_and_eye_contact', 'general_clothing_exposure', 'generic_clothing_interactions', 'holding_large_items', 'holding_small_items', 'intentional_design_exposure', 'lighting_and_vfx', 'male_intimate_anatomy', 'male_physical_descriptors', 'meta', 'metadata_and_attribution', 'named_garment_exposure', 'nudity_and_absence_of_clothing', 'one_handed_character_items', 'physical_locations', 'poses', 'publicly_visible_anatomy', 'relationships', 'sex_acts', 'sfw_clothed_anatomy', 'special_backgrounds', 'specific_garment_interactions', 'speech_and_text', 'standard_physical_descriptors', 'thematic_settings', 'two_handed_character_items' ] class Raffle: # Class variable to track if the critical categories warning has been shown _critical_warning_shown = False @classmethod def INPUT_TYPES(s): extension_path = os.path.normpath(os.path.dirname(__file__)) return { "required": { "seed": ("INT", { "default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Seed value used to randomly select a taglist from the filtered pool of available taglists" }), "use_general": ("BOOLEAN", {"default": False, "tooltip": "Enable selection from general.txt which contains 100,000 SFW taglists"}), "use_questionable": ("BOOLEAN", {"default": False, "tooltip": "Enable selection from questionable.txt which contains 100,000 questionable taglists"}), "use_sensitive": ("BOOLEAN", {"default": True, "tooltip": "Enable selection from sensitive.txt which contains 100,000 sensitive taglists"}), "use_explicit": ("BOOLEAN", {"default": True, "tooltip": "Enable selection from explicit.txt which contains 100,000 explicit taglists"}), "taglists_must_include": ("STRING", { "multiline": True, "default": DEFAULT_TAGLISTS_MUST_INCLUDE, "tooltip": " Only selects taglists that contain ALL of these tags. WARNING: Each tag added here severely reduces the available pool of taglists. Check the 'Debug info' output to see how many taglists remain available." }), "filter_out_tags": ("STRING", { "multiline": True, "default": DEFAULT_FILTER_OUT_TAGS, "tooltip": " Additional tags to filter out from the final output. Use this to exclude more tags without needing to modify your main negative prompt." }), "exclude_taglists_containing": ("STRING", { "multiline": True, "default": DEFAULT_EXCLUDE_TAGLISTS, "tooltip": " If ANY of these tags appear in the taglist, the entire taglist is removed from the pool of available taglists. Use with caution as each tag listed here can significantly reduce options. For removing individual tags without reducing the pool, use 'filter_out_tags' instead." }), "exclude_tag_categories": ("STRING", { "multiline": True, "default": DEFAULT_EXCLUDE_CATEGORIES, "tooltip": " Exclude entire categories of tags from the final output. Each category contains related tags (e.g., 'poses' contains all pose-related tags). View the complete category list and their tags in the 'Debug info' output. Separate multiple categories with commas." }) }, "optional": { "negative_prompt": ("STRING", { "multiline": True, "forceInput": True, "default": "", "tooltip": " Removes specific tags from the final output without affecting taglist selection. Tags listed here will be filtered out after a taglist is chosen, making this safer to use than 'exclude_taglists_containing'." }) } } CATEGORY = "Raffle" RETURN_TYPES = ("STRING", "STRING", "STRING") RETURN_NAMES = ("Raffled output", "Unfiltered", "Debug info") OUTPUT_TOOLTIPS = ( "The final list of tags selected by the raffle and filtered according to your settings, ready for use", "The complete original taglist that was randomly selected before any filtering was applied", "Information about the selection process, including the size of the available pool of taglists after applying your filters" ) OUTPUT_NODE = True FUNCTION = "process_tags" # Add class variable to cache tag lists _tag_cache = {} def _load_taglist(self, filename, taglists_must_include_tags=None, exclude_tags=None, seed=0): """Load a line from a file, finding taglists that match required tags""" extension_path = os.path.normpath(os.path.dirname(__file__)) filepath = os.path.join(extension_path, "lists", filename) try: valid_taglists = [] # Pre-compute sets for faster lookups - no need to normalize these exclude_tags_set = set(exclude_tags) if exclude_tags else None taglists_must_include_set = set(taglists_must_include_tags) if taglists_must_include_tags else None # First pass: find taglists that match our criteria with open(filepath, 'r', encoding='utf-8') as f: for taglist_num, taglist in enumerate(f): taglist = taglist.strip() if not taglist: continue # Split on comma since file is already normalized taglist_tags = frozenset(tag.strip() for tag in taglist.split(',')) # Check for excluded tags first (using set intersection for speed) if exclude_tags_set and not taglist_tags.isdisjoint(exclude_tags_set): continue # If we have required tags, check if they're all in this taglist if taglists_must_include_set: if taglists_must_include_set.issubset(taglist_tags): valid_taglists.append((filename, taglist)) else: valid_taglists.append((filename, taglist)) return valid_taglists except Exception as e: raise return [] def normalize_tags(self, tag_string): """ Normalize a string of tags to a consistent format: - Convert spaces to underscores within tags - Handle both comma-space and comma separators - Handle newlines as separators - Remove duplicate separators and spaces - Strip whitespace Examples: "red hair, blue hair" -> ["red_hair", "blue_hair"] "red_hair,blue hair" -> ["red_hair", "blue_hair"] "red hair\nblue_hair" -> ["red_hair", "blue_hair"] "red hair,\nblue_hair" -> ["red_hair", "blue_hair"] "red hair,,blue_hair" -> ["red_hair", "blue_hair"] "red hair,blue hair" -> ["red_hair", "blue_hair"] """ # Replace newlines with commas tag_string = tag_string.replace('\r\n', '\n') # Normalize line endings tag_string = tag_string.replace('\n', ',') # Remove multiple consecutive spaces while ' ' in tag_string: tag_string = tag_string.replace(' ', ' ') # Remove multiple consecutive commas while ',,' in tag_string: tag_string = tag_string.replace(',,', ',') # Then split on commas (handling both ", " and "," cases) tags = tag_string.replace(', ', ',').split(',') # Then normalize each tag return [ tag.strip().replace(' ', '_') for tag in tags if tag.strip() ] def process_tags(self, exclude_taglists_containing, taglists_must_include, seed, filter_out_tags="", use_general=True, use_questionable=False, use_sensitive=False, use_explicit=False, exclude_tag_categories="", negative_prompt=""): # Add directory existence check extension_path = os.path.normpath(os.path.dirname(__file__)) lists_path = os.path.join(extension_path, "lists") if not os.path.exists(lists_path): raise ValueError(f"Lists directory not found at {lists_path}") # Check for categorized tags file categorized_tags_file_path = os.path.join(lists_path, "categorized_tags.txt") if not os.path.exists(categorized_tags_file_path): raise ValueError(f"Categorized tags file not found at {categorized_tags_file_path}") # Use the global categories list all_categories = ALL_CATEGORIES # Process excluded categories excluded_categories = [] if exclude_tag_categories.strip(): excluded_categories = self.normalize_tags(exclude_tag_categories) # Check if all excluded categories are valid invalid_categories = [c for c in excluded_categories if c not in all_categories] if invalid_categories: error_msg = (f"Error: Invalid category names: {', '.join(invalid_categories)}. " f"Please check the Debug info output for a complete list of valid categories. " f"Category names may have changed in a new version.") raise ValueError(error_msg) # Check if critical categories are excluded excluded_categories_set = set(excluded_categories) # Only show warning if NONE of the critical categories are excluded critical_categories_excluded = WARNING_ABOUT_NEW_CATEGORIES & excluded_categories_set if not critical_categories_excluded and not Raffle._critical_warning_shown: # Mark that we've shown the warning Raffle._critical_warning_shown = True warning_msg = ( "There's been an update that will potentially break your Raffle generations.\n\n" "To maintain your existing workflow, you will need to manually add some categories " "to the section of the raffle node (bottom section).\n\n" f"Add the following categories to be excluded: artist, character_name, copyright, meta" ) raise ValueError(warning_msg) # Set up categories dictionary - enable all categories except excluded ones categories = {category: (category not in excluded_categories) for category in all_categories} # Create the categories debug output categories_debug = "-- List of Categories --\n" + "\n".join(all_categories) # Load categorized tags with error handling allowed_tags = [] # Changed from set to list to preserve order try: with open(categorized_tags_file_path, 'r', encoding='utf-8') as f: content = f.read() for line in content.splitlines(): line = line.strip() if not line: continue parts = line.split('] ', 1) if len(parts) != 2: continue category = parts[0][1:] # Remove the leading [ tag = parts[1] if categories.get(category): allowed_tags.append(tag) # Add the complete tag except Exception as e: raise # Parse exclude and include lists excluded_tags = set(self.normalize_tags(exclude_taglists_containing)) included_tags = set(self.normalize_tags(taglists_must_include)) # Collect all valid taglists from all enabled files all_valid_taglists = [] if use_general: all_valid_taglists.extend(self._load_taglist("taglists-general.txt", included_tags, excluded_tags, seed)) if use_questionable: all_valid_taglists.extend(self._load_taglist("taglists-questionable.txt", included_tags, excluded_tags, seed)) if use_sensitive: all_valid_taglists.extend(self._load_taglist("taglists-sensitive.txt", included_tags, excluded_tags, seed)) if use_explicit: all_valid_taglists.extend(self._load_taglist("taglists-explicit.txt", included_tags, excluded_tags, seed)) if not all_valid_taglists: raise ValueError("No tags available - no matching taglists found") # Use seed to shuffle and select from all valid taglists rng = random.Random(seed) rng.shuffle(all_valid_taglists) # Take just 1 taglist based on seed selected_index = seed % len(all_valid_taglists) selected_taglist = all_valid_taglists[selected_index] # Extract just the taglist (without filename) _, unfiltered_taglist = selected_taglist # Normalize the unfiltered taglist for consistency in output unfiltered_taglist = ', '.join(self.normalize_tags(unfiltered_taglist)) # Split the taglist into individual tags and normalize them individual_tags = self.normalize_tags(unfiltered_taglist) # Filter tags using allowed_tags and maintain order allowed_tags_set = set(allowed_tags) # For faster lookup filtered_tags = [tag for tag in individual_tags if tag in allowed_tags_set] try: filtered_tags.sort(key=lambda x: allowed_tags.index(x) if x in allowed_tags else len(allowed_tags)) except Exception as e: raise # Remove excluded tags filtered_tags = [tag for tag in filtered_tags if tag not in excluded_tags] # Process negative prompt tags negative_tags = set(self.normalize_tags(negative_prompt)) filtered_tags = [tag for tag in filtered_tags if tag not in negative_tags] # Process negative prompt 2 tags filter_out_tags_set = set(self.normalize_tags(filter_out_tags)) filtered_tags = [tag for tag in filtered_tags if tag not in filter_out_tags_set] debug_info = f"Taglist pool size: {len(all_valid_taglists)}\n\n{categories_debug}" return_values = ( ', '.join(filtered_tags), unfiltered_taglist, debug_info ) return return_values