90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
import random
|
|
|
|
class InlineWildcardGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"variable_name": ("STRING", {"default": "my_var", "multiline": False}),
|
|
"wildcard_text": ("STRING", {"default": "option1, option2, option3", "multiline": True}),
|
|
"delimiter": ("STRING", {"default": ",", "multiline": False}),
|
|
"mix_count": ("INT", {"default": 1, "min": 1, "max": 50}),
|
|
"random_mix": ("BOOLEAN", {"default": False}),
|
|
"join_separator": ("STRING", {"default": ", ", "multiline": False}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
"optional": {
|
|
"prev_vars": ("WILDCARD_BUNDLE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("WILDCARD_BUNDLE",)
|
|
RETURN_NAMES = ("wildcard_bundle",)
|
|
FUNCTION = "generate_variable"
|
|
CATEGORY = "utils/wildcards"
|
|
|
|
def generate_variable(self, variable_name, wildcard_text, delimiter, mix_count, random_mix, join_separator, seed, prev_vars=None):
|
|
if prev_vars is None:
|
|
variables = {}
|
|
else:
|
|
variables = prev_vars.copy()
|
|
|
|
processed_text = wildcard_text.replace("\n", delimiter)
|
|
options = [x.strip() for x in processed_text.split(delimiter) if x.strip()]
|
|
|
|
if not options:
|
|
final_string = ""
|
|
else:
|
|
rng = random.Random(seed)
|
|
|
|
if random_mix:
|
|
count_to_pick = min(len(options), rng.randint(1, mix_count))
|
|
else:
|
|
count_to_pick = min(len(options), mix_count)
|
|
|
|
selected_items = rng.sample(options, count_to_pick)
|
|
|
|
final_string = join_separator.join(selected_items)
|
|
|
|
clean_key = variable_name.strip()
|
|
|
|
variables[clean_key] = final_string
|
|
|
|
return (variables,)
|
|
|
|
|
|
class WildcardPromptBuilder:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"default": "A photo of a <my_var>", "multiline": True}),
|
|
"wildcard_bundle": ("WILDCARD_BUNDLE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("final_prompt",)
|
|
FUNCTION = "build_prompt"
|
|
CATEGORY = "utils/wildcards"
|
|
|
|
def build_prompt(self, text, wildcard_bundle):
|
|
result_text = text
|
|
|
|
for key, value in wildcard_bundle.items():
|
|
placeholder = f"<{key}>"
|
|
if placeholder in result_text:
|
|
result_text = result_text.replace(placeholder, value)
|
|
|
|
return (result_text,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"InlineWildcardGenerator": InlineWildcardGenerator,
|
|
"WildcardPromptBuilder": WildcardPromptBuilder
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"InlineWildcardGenerator": "Define Inline Wildcard",
|
|
"WildcardPromptBuilder": "Build Prompt from Wildcards"
|
|
}
|