529 lines
21 KiB
Python
529 lines
21 KiB
Python
"""
|
|
Random Wildcard Loader Node for ComfyUI
|
|
Loads random wildcard files and combines their contents with user text.
|
|
Supports inline wildcard expansion using __wildcard__ syntax.
|
|
Supports $prompt keyword for Mad Libs style random wildcard insertion.
|
|
"""
|
|
|
|
import os
|
|
import random
|
|
import re
|
|
import folder_paths
|
|
|
|
class RandomWildcardLoader:
|
|
"""
|
|
A ComfyUI node that loads random wildcard files from the wildcards directory
|
|
and combines their contents with user-provided text.
|
|
Supports inline wildcard expansion in the prompt field using __wildcard__ syntax.
|
|
Supports $prompt keyword for Mad Libs style random wildcard insertion.
|
|
"""
|
|
|
|
# Get the base path for ComfyUI
|
|
BASE_PATH = folder_paths.base_path
|
|
WILDCARDS_DIR = os.path.join(BASE_PATH, "wildcards")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"num_files": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"prompt": ("STRING", {
|
|
"default": "",
|
|
"multiline": True,
|
|
"placeholder": "Enter your prompt here.\nUse __wildcard__ for specific wildcards: A __colors__ cat\nUse $prompt for random wildcards: The dog plays with a $prompt"
|
|
}),
|
|
"loaded_contents": ("STRING", {
|
|
"default": "",
|
|
"multiline": True,
|
|
"placeholder": "Randomly selected wildcard content will appear here..."
|
|
}),
|
|
"loaded_wildcards": ("STRING", {
|
|
"default": "",
|
|
"multiline": True,
|
|
"placeholder": "Names of loaded wildcards will appear here..."
|
|
}),
|
|
"refresh": ("BOOLEAN", {
|
|
"default": False,
|
|
"label_on": "Refresh",
|
|
"label_off": "Cached"
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("combined_text",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "load_wildcards"
|
|
CATEGORY = "text/wildcards"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, num_files, seed, prompt, loaded_contents, loaded_wildcards, refresh):
|
|
"""Force refresh when refresh is toggled or seed changes."""
|
|
if refresh:
|
|
return float("NaN")
|
|
return seed
|
|
|
|
def get_wildcard_files(self):
|
|
"""Get all .txt files from the wildcards directory recursively."""
|
|
wildcard_files = []
|
|
|
|
if not os.path.exists(self.WILDCARDS_DIR):
|
|
os.makedirs(self.WILDCARDS_DIR, exist_ok=True)
|
|
return wildcard_files
|
|
|
|
for root, dirs, files in os.walk(self.WILDCARDS_DIR):
|
|
for file in files:
|
|
if file.endswith('.txt'):
|
|
full_path = os.path.join(root, file)
|
|
# Get relative path from wildcards dir for display
|
|
rel_path = os.path.relpath(full_path, self.WILDCARDS_DIR)
|
|
wildcard_files.append({
|
|
'path': full_path,
|
|
'name': rel_path,
|
|
'basename': os.path.splitext(file)[0]
|
|
})
|
|
|
|
return wildcard_files
|
|
|
|
def read_wildcard_file(self, file_path):
|
|
"""Read contents of a wildcard file and return a random line."""
|
|
try:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
lines = [line.strip() for line in f.readlines() if line.strip()]
|
|
if lines:
|
|
return random.choice(lines)
|
|
return ""
|
|
except Exception as e:
|
|
print(f"Error reading wildcard file {file_path}: {e}")
|
|
return ""
|
|
|
|
def get_wildcard_content(self, wildcard_name):
|
|
"""Get random content from a specific wildcard file by name."""
|
|
# Handle nested paths (e.g., "folder/subfolder/name")
|
|
wildcard_path = wildcard_name.replace('/', os.sep).replace('\\', os.sep)
|
|
|
|
# Try with .txt extension
|
|
file_path = os.path.join(self.WILDCARDS_DIR, wildcard_path + '.txt')
|
|
if os.path.exists(file_path):
|
|
return self.read_wildcard_file(file_path)
|
|
|
|
# Try without extension (maybe user included it)
|
|
file_path = os.path.join(self.WILDCARDS_DIR, wildcard_path)
|
|
if os.path.exists(file_path):
|
|
return self.read_wildcard_file(file_path)
|
|
|
|
# Wildcard not found, return the original syntax
|
|
return f"__{wildcard_name}__"
|
|
|
|
def get_random_wildcard_content(self):
|
|
"""
|
|
Get random content from a randomly selected wildcard file.
|
|
Returns tuple of (content, wildcard_name) or (None, None) if no files available.
|
|
"""
|
|
wildcard_files = self.get_wildcard_files()
|
|
if not wildcard_files:
|
|
return None, None
|
|
|
|
# Pick a random wildcard file
|
|
selected_file = random.choice(wildcard_files)
|
|
content = self.read_wildcard_file(selected_file['path'])
|
|
|
|
return content, selected_file['name']
|
|
|
|
def format_wildcard_display_name(self, name):
|
|
"""
|
|
Format wildcard name for display: folder/__Filename__ format.
|
|
Example: 'random/Weapon.txt' -> 'random/__Weapon__'
|
|
Example: 'colors.txt' -> '__colors__'
|
|
"""
|
|
# Remove .txt extension and normalize separators
|
|
name_clean = name.replace('.txt', '').replace('\\', '/')
|
|
|
|
# Split into path and filename
|
|
if '/' in name_clean:
|
|
path_part, file_part = name_clean.rsplit('/', 1)
|
|
return f"{path_part}/__{file_part}__"
|
|
else:
|
|
return f"__{name_clean}__"
|
|
|
|
def expand_random_prompts(self, text, prefix="", suffix=""):
|
|
"""
|
|
Expand all $prompt keywords in text with random wildcard content.
|
|
Mad Libs style - each $prompt gets a random result from a random wildcard file.
|
|
Optionally wraps content with prefix and suffix.
|
|
Returns expanded text and lists of wildcards/contents that were used.
|
|
"""
|
|
expanded_wildcards = []
|
|
expanded_contents = []
|
|
|
|
# Find all $prompt occurrences
|
|
# Pattern matches $prompt that's not preceded by another $ (to allow escaping)
|
|
pattern = r'(?<!\$)\$prompt'
|
|
|
|
def replace_prompt(match):
|
|
content, wildcard_name = self.get_random_wildcard_content()
|
|
if content and wildcard_name:
|
|
# Format wildcard name: remove .txt extension and use path without extension
|
|
name_without_ext = wildcard_name.replace('.txt', '').replace('\\', '/')
|
|
expanded_wildcards.append(name_without_ext)
|
|
expanded_contents.append(content)
|
|
# Apply prefix and suffix
|
|
return f"{prefix}{content}{suffix}"
|
|
return match.group(0) # Return original if no wildcards available
|
|
|
|
expanded_text = re.sub(pattern, replace_prompt, text)
|
|
|
|
return expanded_text, expanded_wildcards, expanded_contents
|
|
|
|
def expand_wildcards(self, text, rng, prefix="", suffix=""):
|
|
"""
|
|
Expand all __wildcard__ patterns in text.
|
|
Optionally wraps content with prefix and suffix.
|
|
Returns expanded text and list of wildcards that were expanded.
|
|
"""
|
|
pattern = r'__([a-zA-Z0-9_/\\-]+)__'
|
|
expanded_wildcards = []
|
|
expanded_contents = []
|
|
|
|
def replace_wildcard(match):
|
|
wildcard_name = match.group(1)
|
|
content = self.get_wildcard_content(wildcard_name)
|
|
# Only record if wildcard was found (not returned as original syntax)
|
|
if content != f"__{wildcard_name}__":
|
|
expanded_wildcards.append(wildcard_name)
|
|
expanded_contents.append(content)
|
|
# Apply prefix and suffix
|
|
return f"{prefix}{content}{suffix}"
|
|
return content
|
|
|
|
# Set random state for consistent expansion
|
|
expanded_text = re.sub(pattern, replace_wildcard, text)
|
|
|
|
return expanded_text, expanded_wildcards, expanded_contents
|
|
|
|
def load_wildcards(self, num_files, seed, prompt, loaded_contents, loaded_wildcards, refresh):
|
|
"""Main function to load wildcards and combine with user text."""
|
|
|
|
# Initialize random with seed
|
|
random.seed(seed)
|
|
|
|
# Get all available wildcard files
|
|
wildcard_files = self.get_wildcard_files()
|
|
|
|
# Track all wildcards and contents (from both prompt and random selection)
|
|
all_wildcard_names = []
|
|
all_wildcard_contents = []
|
|
|
|
# First, expand $prompt keywords (Mad Libs style random wildcards)
|
|
expanded_prompt, random_prompt_wildcards, random_prompt_contents = self.expand_random_prompts(prompt)
|
|
|
|
# Add $prompt wildcards to the tracking lists
|
|
all_wildcard_names.extend(random_prompt_wildcards)
|
|
all_wildcard_contents.extend(random_prompt_contents)
|
|
|
|
# Then, expand __wildcard__ syntax in the prompt field
|
|
expanded_prompt, prompt_wildcards, prompt_contents = self.expand_wildcards(expanded_prompt, random)
|
|
|
|
# Add __wildcard__ expansions to the tracking lists
|
|
all_wildcard_names.extend(prompt_wildcards)
|
|
all_wildcard_contents.extend(prompt_contents)
|
|
|
|
# Select and load random wildcard files
|
|
random_contents = []
|
|
|
|
if wildcard_files and num_files > 0:
|
|
# Limit num_files to available files
|
|
num_to_select = min(num_files, len(wildcard_files))
|
|
|
|
# Randomly select files
|
|
selected_files = random.sample(wildcard_files, num_to_select)
|
|
|
|
for file_info in selected_files:
|
|
content = self.read_wildcard_file(file_info['path'])
|
|
if content:
|
|
random_contents.append(content)
|
|
all_wildcard_names.append(file_info['name'])
|
|
all_wildcard_contents.append(content)
|
|
|
|
# Build the combined text output
|
|
combined_parts = []
|
|
|
|
# Add the expanded prompt if not empty
|
|
if expanded_prompt.strip():
|
|
combined_parts.append(expanded_prompt.strip())
|
|
|
|
# Add random wildcard contents
|
|
if random_contents:
|
|
combined_parts.extend(random_contents)
|
|
|
|
combined_text = "\n".join(combined_parts) if combined_parts else ""
|
|
|
|
# Format display strings
|
|
# Wildcards display: format as folder/__name__, comma-separated
|
|
wildcards_display = ", ".join([self.format_wildcard_display_name(name) for name in all_wildcard_names]) if all_wildcard_names else ""
|
|
# Contents display: show the combined output text
|
|
contents_display = combined_text
|
|
|
|
# Return result with UI updates
|
|
return {
|
|
"ui": {
|
|
"loaded_wildcards": [wildcards_display],
|
|
"loaded_contents": [contents_display]
|
|
},
|
|
"result": (combined_text,)
|
|
}
|
|
|
|
|
|
class RandomWildcardLoaderAdvanced(RandomWildcardLoader):
|
|
"""
|
|
Advanced version of the Random Wildcard Loader with additional options.
|
|
Includes subfolder filtering, prefix/suffix, and nested folder control.
|
|
"""
|
|
|
|
INCLUDE_NESTED_OPTIONS = ["All Levels", "Top Level Only"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"num_files": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"prompt": ("STRING", {
|
|
"default": "",
|
|
"multiline": True,
|
|
"placeholder": "Enter your prompt here.\nUse __wildcard__ for specific wildcards: A __colors__ cat\nUse $prompt for random wildcards: The dog plays with a $prompt"
|
|
}),
|
|
"loaded_contents": ("STRING", {
|
|
"default": "",
|
|
"multiline": True,
|
|
"placeholder": "Randomly selected wildcard content will appear here..."
|
|
}),
|
|
"loaded_wildcards": ("STRING", {
|
|
"default": "",
|
|
"multiline": True,
|
|
"placeholder": "Names of loaded wildcards will appear here..."
|
|
}),
|
|
"separator": ("STRING", {
|
|
"default": ", ",
|
|
"multiline": False,
|
|
"placeholder": "Separator between wildcards"
|
|
}),
|
|
"subfolder": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"placeholder": "Subfolder to load from (leave empty for all)"
|
|
}),
|
|
"prefix": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"placeholder": "Prefix for each wildcard"
|
|
}),
|
|
"suffix": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"placeholder": "Suffix for each wildcard"
|
|
}),
|
|
"include_nested": (["All Levels", "Top Level Only"], {
|
|
"default": "All Levels"
|
|
}),
|
|
"same_file": (["Different Files", "Same File"], {
|
|
"default": "Different Files"
|
|
}),
|
|
"refresh": ("BOOLEAN", {
|
|
"default": False,
|
|
"label_on": "Refresh",
|
|
"label_off": "Cached"
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("combined_text",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "load_wildcards_advanced"
|
|
CATEGORY = "text/wildcards"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, num_files, seed, prompt, loaded_contents, loaded_wildcards, separator, subfolder, prefix, suffix, include_nested, same_file, refresh):
|
|
"""Force refresh when refresh is toggled or seed changes."""
|
|
if refresh:
|
|
return float("NaN")
|
|
return seed
|
|
|
|
def get_wildcard_files_filtered(self, subfolder="", include_nested=True):
|
|
"""Get wildcard files with optional subfolder filtering."""
|
|
wildcard_files = []
|
|
|
|
# Determine the search directory
|
|
if subfolder and subfolder.strip():
|
|
search_dir = os.path.join(self.WILDCARDS_DIR, subfolder.strip())
|
|
else:
|
|
search_dir = self.WILDCARDS_DIR
|
|
|
|
if not os.path.exists(search_dir):
|
|
return wildcard_files
|
|
|
|
if include_nested:
|
|
# Recursive search
|
|
for root, dirs, files in os.walk(search_dir):
|
|
for file in files:
|
|
if file.endswith('.txt'):
|
|
full_path = os.path.join(root, file)
|
|
rel_path = os.path.relpath(full_path, self.WILDCARDS_DIR)
|
|
wildcard_files.append({
|
|
'path': full_path,
|
|
'name': rel_path,
|
|
'basename': os.path.splitext(file)[0]
|
|
})
|
|
else:
|
|
# Top level only
|
|
for file in os.listdir(search_dir):
|
|
if file.endswith('.txt'):
|
|
full_path = os.path.join(search_dir, file)
|
|
if os.path.isfile(full_path):
|
|
rel_path = os.path.relpath(full_path, self.WILDCARDS_DIR)
|
|
wildcard_files.append({
|
|
'path': full_path,
|
|
'name': rel_path,
|
|
'basename': os.path.splitext(file)[0]
|
|
})
|
|
|
|
return wildcard_files
|
|
|
|
def load_wildcards_advanced(self, num_files, seed, prompt, loaded_contents, loaded_wildcards, separator, subfolder, prefix, suffix, include_nested, same_file, refresh):
|
|
"""Main function to load wildcards with advanced options."""
|
|
|
|
# Initialize random with seed
|
|
random.seed(seed)
|
|
|
|
# Determine if we should include nested folders
|
|
include_nested_bool = (include_nested == "All Levels")
|
|
|
|
# Determine if all selections should come from same file
|
|
same_file_bool = (same_file == "Same File")
|
|
|
|
# Get wildcard files with filtering
|
|
wildcard_files = self.get_wildcard_files_filtered(subfolder, include_nested_bool)
|
|
|
|
# Track all wildcards and contents (from both prompt and random selection)
|
|
all_wildcard_names = []
|
|
all_wildcard_contents = []
|
|
|
|
# First, expand $prompt keywords (Mad Libs style random wildcards) - no prefix/suffix yet
|
|
expanded_prompt, random_prompt_wildcards, random_prompt_contents = self.expand_random_prompts(prompt)
|
|
|
|
# Add $prompt wildcards to the tracking lists
|
|
all_wildcard_names.extend(random_prompt_wildcards)
|
|
all_wildcard_contents.extend(random_prompt_contents)
|
|
|
|
# Then, expand __wildcard__ syntax in the prompt field - no prefix/suffix yet
|
|
expanded_prompt, prompt_wildcards, prompt_contents = self.expand_wildcards(expanded_prompt, random)
|
|
|
|
# Add __wildcard__ expansions to the tracking lists
|
|
all_wildcard_names.extend(prompt_wildcards)
|
|
all_wildcard_contents.extend(prompt_contents)
|
|
|
|
# Select and load random wildcard files
|
|
random_contents = []
|
|
|
|
if wildcard_files and num_files > 0:
|
|
if same_file_bool:
|
|
# All selections from the same randomly chosen file
|
|
selected_file = random.choice(wildcard_files)
|
|
for _ in range(num_files):
|
|
content = self.read_wildcard_file(selected_file['path'])
|
|
if content:
|
|
random_contents.append(content)
|
|
all_wildcard_names.append(selected_file['name'])
|
|
all_wildcard_contents.append(content)
|
|
else:
|
|
# Selections from different files (original behavior)
|
|
# Limit num_files to available files
|
|
num_to_select = min(num_files, len(wildcard_files))
|
|
|
|
# Randomly select files
|
|
selected_files = random.sample(wildcard_files, num_to_select)
|
|
|
|
for file_info in selected_files:
|
|
content = self.read_wildcard_file(file_info['path'])
|
|
if content:
|
|
# Store raw content (no prefix/suffix applied individually)
|
|
random_contents.append(content)
|
|
all_wildcard_names.append(file_info['name'])
|
|
all_wildcard_contents.append(content)
|
|
|
|
# Build the combined text output
|
|
combined_parts = []
|
|
|
|
# Add the expanded prompt if not empty
|
|
if expanded_prompt.strip():
|
|
combined_parts.append(expanded_prompt.strip())
|
|
|
|
# Add random wildcard contents
|
|
if random_contents:
|
|
combined_parts.extend(random_contents)
|
|
|
|
# Join all parts with separator
|
|
inner_text = separator.join(combined_parts) if combined_parts else ""
|
|
|
|
# Apply prefix and suffix to the ENTIRE combined output
|
|
# Prefix gets a space added at the end, suffix gets ", " added at the beginning
|
|
if inner_text:
|
|
prefix_output = f"{prefix} " if prefix else ""
|
|
suffix_output = f", {suffix}" if suffix else ""
|
|
combined_text = f"{prefix_output}{inner_text}{suffix_output}"
|
|
else:
|
|
combined_text = ""
|
|
|
|
# Format display strings
|
|
# Wildcards display: format as folder/__name__, comma-separated
|
|
wildcards_display = ", ".join([self.format_wildcard_display_name(name) for name in all_wildcard_names]) if all_wildcard_names else ""
|
|
# Contents display: show the combined output text
|
|
contents_display = combined_text
|
|
|
|
# Return result with UI updates
|
|
return {
|
|
"ui": {
|
|
"loaded_wildcards": [wildcards_display],
|
|
"loaded_contents": [contents_display]
|
|
},
|
|
"result": (combined_text,)
|
|
}
|
|
|
|
|
|
# Node registration for ComfyUI
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RandomWildcardLoader": RandomWildcardLoader,
|
|
"RandomWildcardLoaderAdvanced": RandomWildcardLoaderAdvanced
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RandomWildcardLoader": "Random Wildcard Loader",
|
|
"RandomWildcardLoaderAdvanced": "Random Wildcard Loader (Advanced)"
|
|
}
|