Add list sampler node
This commit is contained in:
@@ -4,11 +4,35 @@ import os
|
||||
import random
|
||||
|
||||
|
||||
class PromptFromTemplate:
|
||||
class BasePrompt:
|
||||
PROMPT_LISTS_DIR = os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)), "prompt-lists"
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
@classmethod
|
||||
def get_all_available_lists(cls):
|
||||
lists_path = os.path.join(cls.PROMPT_LISTS_DIR, "lists.json")
|
||||
with open(lists_path, "r", encoding="utf-8") as file:
|
||||
lists = json.load(file)
|
||||
return lists
|
||||
|
||||
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
|
||||
|
||||
|
||||
class PromptFromTemplate(BasePrompt):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
@@ -21,11 +45,6 @@ class PromptFromTemplate:
|
||||
}
|
||||
}
|
||||
|
||||
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"
|
||||
@@ -57,29 +76,66 @@ class PromptFromTemplate:
|
||||
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
|
||||
class PromptListSampler(BasePrompt):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"list_name": (cls.get_all_available_lists(),),
|
||||
"mode": (["random", "sequential"], {"default": "sequential"}),
|
||||
"number_of_items": ("INT", {"default": 1, "min": 1}),
|
||||
"join_with": ("STRING", {"default": ", "}),
|
||||
"index": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xFFFFFFFFFFFFFFFF,
|
||||
"control_after_generate": True,
|
||||
},
|
||||
),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
FUNCTION = "get_item_from_list"
|
||||
CATEGORY = "Prompter"
|
||||
|
||||
def get_item_from_list(
|
||||
self,
|
||||
list_name,
|
||||
index=0,
|
||||
number_of_items=1,
|
||||
mode="sequential",
|
||||
join_with=", ",
|
||||
seed=0,
|
||||
):
|
||||
random.seed(seed)
|
||||
|
||||
if list_name not in self.all_lists:
|
||||
return f"[{list_name}]"
|
||||
|
||||
if mode == "random":
|
||||
items = random.sample(
|
||||
self.all_lists[list_name],
|
||||
min(number_of_items, len(self.all_lists[list_name])),
|
||||
)
|
||||
else:
|
||||
index = index % len(self.all_lists[list_name])
|
||||
items = self.all_lists[list_name][index : index + number_of_items]
|
||||
|
||||
return (join_with.join(items),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Prompt from template 🪴": PromptFromTemplate,
|
||||
"List sampler 🪴": PromptListSampler,
|
||||
}
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "comfyui-prompter-fofrai"
|
||||
description = "A prompt helper for ComfyUI, based on prompter.fofr.ai"
|
||||
version = "1.0.0"
|
||||
version = "1.0.1"
|
||||
license = "LICENSE"
|
||||
|
||||
[project.urls]
|
||||
|
||||
Reference in New Issue
Block a user