95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
import yaml
|
|
import os
|
|
import folder_paths
|
|
from collections import OrderedDict
|
|
|
|
# カスタムYAMLローダーとダンパーを定義
|
|
class OrderedLoader(yaml.SafeLoader):
|
|
pass
|
|
|
|
class OrderedDumper(yaml.SafeDumper):
|
|
pass
|
|
|
|
def dict_constructor(loader, node):
|
|
return OrderedDict(loader.construct_pairs(node))
|
|
|
|
OrderedLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, dict_constructor)
|
|
|
|
def dict_representer(dumper, data):
|
|
return dumper.represent_mapping(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, data.items())
|
|
|
|
OrderedDumper.add_representer(OrderedDict, dict_representer)
|
|
|
|
class ComfyUI_PromptList:
|
|
def __init__(self):
|
|
self.yaml_path = os.path.join(folder_paths.base_path, "custom_nodes", "ComfyUI-PromptList", "prompts.yaml")
|
|
self.data = self.load_yaml()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"selection": (cls.get_prompt_list(),),
|
|
"Prompt Name": ("STRING", {"multiline": True}),
|
|
"Positive Prompt": ("STRING", {"multiline": True}),
|
|
"Negative Prompt": ("STRING", {"multiline": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("positive", "negative")
|
|
FUNCTION = "process"
|
|
CATEGORY = "prompt"
|
|
|
|
def load_yaml(self):
|
|
if os.path.exists(self.yaml_path):
|
|
with open(self.yaml_path, 'r', encoding='utf-8') as file:
|
|
return yaml.load(file, Loader=OrderedLoader) or OrderedDict()
|
|
return OrderedDict()
|
|
|
|
def save_yaml(self):
|
|
with open(self.yaml_path, 'w', encoding='utf-8') as file:
|
|
yaml.dump(self.data, file, Dumper=OrderedDumper, allow_unicode=True)
|
|
|
|
@classmethod
|
|
def get_prompt_list(cls):
|
|
instance = cls()
|
|
return list(instance.data.keys())
|
|
|
|
def process(self, selection, **kwargs):
|
|
prompt_name = kwargs.get("Prompt Name", "")
|
|
positive_prompt = kwargs.get("Positive Prompt", "")
|
|
negative_prompt = kwargs.get("Negative Prompt", "")
|
|
|
|
if prompt_name and positive_prompt and negative_prompt:
|
|
# Update or add new prompt while preserving order
|
|
if prompt_name in self.data:
|
|
self.data[prompt_name]["positive"] = positive_prompt
|
|
self.data[prompt_name]["negative"] = negative_prompt
|
|
else:
|
|
self.data[prompt_name] = OrderedDict([
|
|
("positive", positive_prompt),
|
|
("negative", negative_prompt)
|
|
])
|
|
self.save_yaml()
|
|
return (positive_prompt, negative_prompt)
|
|
elif selection:
|
|
# Return existing prompt from YAML
|
|
prompt = self.data.get(selection, {})
|
|
return (prompt.get("positive", ""), prompt.get("negative", ""))
|
|
|
|
# Default return if no selection or custom input
|
|
return ("", "")
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return float("NaN")
|
|
|
|
# NODE_CLASS_MAPPINGS と NODE_DISPLAY_NAME_MAPPINGS の定義
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ComfyUI-PromptList": ComfyUI_PromptList
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ComfyUI-PromptList": "Prompt List"
|
|
} |