363 lines
18 KiB
Python
363 lines
18 KiB
Python
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 |