86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
import re
|
|
import json
|
|
import os
|
|
import random
|
|
|
|
|
|
class PromptFromTemplate:
|
|
PROMPT_LISTS_DIR = os.path.join(
|
|
os.path.dirname(os.path.abspath(__file__)), "prompt-lists"
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"template": (
|
|
"STRING",
|
|
{"default": "", "multiline": True, "dynamicPrompts": True},
|
|
),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
}
|
|
}
|
|
|
|
def __init__(self):
|
|
self.all_lists = {}
|
|
for list_name in self.get_all_available_lists():
|
|
self.all_lists[list_name] = self.load_list(list_name)
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "generate_prompt_from_template"
|
|
CATEGORY = "Prompter"
|
|
|
|
def generate_prompt_from_template(self, template, seed=0):
|
|
random.seed(seed)
|
|
|
|
def replace_match(match):
|
|
list_name = match.group(1)
|
|
list_name_parts = list_name.split(",")
|
|
list_name = list_name_parts[0].strip()
|
|
item_count = int(list_name_parts[1]) if len(list_name_parts) == 2 else 1
|
|
item_count = min(item_count, 50)
|
|
|
|
if list_name == "random":
|
|
random_list_names = [self.get_random_list() for _ in range(item_count)]
|
|
random_items = [
|
|
self.get_random_items_from_list(list_name, 1)[0]
|
|
for list_name in random_list_names
|
|
]
|
|
return ", ".join(random_items)
|
|
|
|
if list_name not in self.all_lists:
|
|
return f"[{list_name}]"
|
|
|
|
return ", ".join(self.get_random_items_from_list(list_name, item_count))
|
|
|
|
prompt = re.sub(r"\[(.*?)\]", replace_match, template)
|
|
print(prompt)
|
|
return (prompt,)
|
|
|
|
def get_all_available_lists(self):
|
|
lists_path = os.path.join(self.PROMPT_LISTS_DIR, "lists.json")
|
|
with open(lists_path, "r", encoding="utf-8") as file:
|
|
lists = json.load(file)
|
|
return lists
|
|
|
|
def get_random_list(self):
|
|
return random.choice(list(self.all_lists.keys()))
|
|
|
|
def get_random_items_from_list(self, list_name, item_count):
|
|
return random.sample(self.all_lists[list_name], item_count)
|
|
|
|
def load_list(self, list_name):
|
|
hyphenated_list = re.sub(r"([a-z])([A-Z])", r"\1-\2", list_name).lower()
|
|
category, name = hyphenated_list.split(".")
|
|
lists_path = os.path.join(self.PROMPT_LISTS_DIR, f"lists/{category}/{name}.yml")
|
|
with open(lists_path, "r", encoding="utf-8") as file:
|
|
content = file.read()
|
|
list_data = content.split("---")[2].strip().split("\n")
|
|
|
|
return list_data
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Prompt from template 🪴": PromptFromTemplate,
|
|
}
|