raffle.py

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": "<taglists_must_include> 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": "<filter_out_tags> 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": "<exclude_taglists_containing> 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_tag_categories> 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": "<negative_prompt> 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 <exclude_tag_categories> 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