From bca2f53ecd4d24f4faf76022669e13dc7500a99d Mon Sep 17 00:00:00 2001 From: nkchocoai Date: Sat, 27 Jan 2024 16:06:49 +0900 Subject: [PATCH] Add PromptUtilitiesLoadPresetAdvanced node --- __init__.py | 4 +- advanced_presets/sample/animagineXL.json | 6 +++ advanced_presets/sample/clothes.json | 14 ++++++ advanced_presets/sample/loras.json | 32 +++++++++++++ py/node/preset.py | 59 ++++++++++++++++++++++-- py/preset.py | 46 ++++++++++++++---- 6 files changed, 147 insertions(+), 14 deletions(-) create mode 100644 advanced_presets/sample/animagineXL.json create mode 100644 advanced_presets/sample/clothes.json create mode 100644 advanced_presets/sample/loras.json diff --git a/__init__.py b/__init__.py index 404b61a..3b9da71 100644 --- a/__init__.py +++ b/__init__.py @@ -7,6 +7,7 @@ NODE_CLASS_MAPPINGS = { "PromptUtilitiesFormatString": PromptUtilitiesFormatString, "PromptUtilitiesJoinStringList": PromptUtilitiesJoinStringList, "PromptUtilitiesLoadPreset": PromptUtilitiesLoadPreset, + "PromptUtilitiesLoadPresetAdvanced": PromptUtilitiesLoadPresetAdvanced, "PromptUtilitiesConstString": PromptUtilitiesConstString, "PromptUtilitiesConstStringMultiLine": PromptUtilitiesConstStringMultiLine, } @@ -15,8 +16,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PromptUtilitiesFormatString": "Format String", "PromptUtilitiesJoinStringList": "Join String List", "PromptUtilitiesLoadPreset": "Load Preset", + "PromptUtilitiesLoadPresetAdvanced": "Load Preset (Advanced)", "PromptUtilitiesConstString": "Const String", - "PromptUtilitiesConstStringMultiLine": "Const String(multi line)", + "PromptUtilitiesConstStringMultiLine": "Const String (multi line)", } WEB_DIRECTORY = "./js" diff --git a/advanced_presets/sample/animagineXL.json b/advanced_presets/sample/animagineXL.json new file mode 100644 index 0000000..22b8313 --- /dev/null +++ b/advanced_presets/sample/animagineXL.json @@ -0,0 +1,6 @@ +{ + "postfix": { + "positive_prompt": "masterpiece, best quality", + "negative_prompt": "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name" + } +} \ No newline at end of file diff --git a/advanced_presets/sample/clothes.json b/advanced_presets/sample/clothes.json new file mode 100644 index 0000000..847d48a --- /dev/null +++ b/advanced_presets/sample/clothes.json @@ -0,0 +1,14 @@ +{ + "maid": { + "positive_prompt": "maid, maid apron, maid headdress" + }, + "school uniform": { + "positive_prompt": "school uniform, serafuku, pleated skirt, sailor collar, neckerchief" + }, + "playboy bunny": { + "positive_prompt": "playboy bunny, strapless leotard, rabbit ears, detached collar" + }, + "bikini": { + "positive_prompt": "bikini, swimsuit" + } +} \ No newline at end of file diff --git a/advanced_presets/sample/loras.json b/advanced_presets/sample/loras.json new file mode 100644 index 0000000..5e386d4 --- /dev/null +++ b/advanced_presets/sample/loras.json @@ -0,0 +1,32 @@ +{ + "lora_simple": { + "positive_prompt": "hello, world", + "lora": { + "lora_name": "foo.safetensors", + "weight": 1.0 + } + }, + "lora_advanced": { + "positive_prompt": "positive, prompt", + "negative_prompt": "negative, prompt", + "lora": { + "lora_name": "bar.safetensors", + "strength_model": 0.5, + "strength_clip": 0.75 + } + }, + "lora_stack": { + "positive_prompt": "foo, bar, baz", + "loras": [ + { + "lora_name": "foo.safetensors", + "weight": 1.0 + }, + { + "lora_name": "bar.safetensors", + "strength_model": 0.5, + "strength_clip": 0.75 + } + ] + } +} \ No newline at end of file diff --git a/py/node/preset.py b/py/node/preset.py index f6dc46d..4ddebf9 100644 --- a/py/node/preset.py +++ b/py/node/preset.py @@ -1,5 +1,5 @@ from .base import BaseNode -from ..preset import PresetManager +from ..preset import PresetManager, PresetManagerAdvanced class PromptUtilitiesLoadPreset(BaseNode): @@ -15,5 +15,58 @@ class PromptUtilitiesLoadPreset(BaseNode): FUNCTION = "load_preset" def load_preset(self, preset): - prompt = PresetManager.get_prompt(preset) - return (prompt,) \ No newline at end of file + prompt = PresetManager.get_preset(preset) + return (prompt,) + + +class PromptUtilitiesLoadPresetAdvanced(BaseNode): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "preset": (list(PresetManagerAdvanced.get_presets().keys()), ), + } + } + + RETURN_TYPES = ("STRING","STRING","STRING","FLOAT","FLOAT","LORA_STACK",) + RETURN_NAMES = ("positive prompt","negative prompt","lora name","strength model","strength clip","lora stack",) + FUNCTION = "load_preset" + + def load_preset(self, preset): + preset = PresetManagerAdvanced.get_preset(preset) + positive_prompt = preset.get("positive_prompt","") + negative_prompt = preset.get("negative_prompt","") + + lora = preset.get("lora", None) + lora_name, strength_model, sterngth_clip = self.load_lora(lora) + + loras = preset.get("loras", None) + lora_stack = self.load_loras(loras) + + return (positive_prompt, negative_prompt, lora_name, strength_model, sterngth_clip, lora_stack) + + def load_lora(self, lora): + if lora is None: + return "None", 0, 0 + + lora_name = lora.get("lora_name", "None") + if "weight" in lora: + strength_model = lora.get("weight", 0) + sterngth_clip = lora.get("weight", 0) + else: + strength_model = lora.get("strength_model", 0) + sterngth_clip = lora.get("strength_clip", 0) + + return lora_name, strength_model, sterngth_clip + + + def load_loras(self, loras): + lora_stack = [] + if loras is None: + return lora_stack + + for lora in loras: + lora_name, strength_model, sterngth_clip = self.load_lora(lora) + loras.append((lora_name, strength_model, sterngth_clip)) + + return lora_stack \ No newline at end of file diff --git a/py/preset.py b/py/preset.py index 7a04cee..274b6f8 100644 --- a/py/preset.py +++ b/py/preset.py @@ -1,13 +1,19 @@ import csv +import json import os import folder_paths -class PresetManager: +class PresetManagerBase: _presets = None custom_nodes_dir = folder_paths.get_folder_paths("custom_nodes")[0] - presets_dir = os.path.join(custom_nodes_dir, "ComfyUI-PromptUtilities", "presets") + + file_extensions = [] + + @classmethod + def get_presets_dir(cls): + return os.path.join(cls.custom_nodes_dir, "ComfyUI-PromptUtilities", cls.presets_dir) @classmethod def get_presets(cls): @@ -16,23 +22,43 @@ class PresetManager: return cls._presets @classmethod - def get_prompt(cls, key): + def get_preset(cls, key): presets = cls.get_presets() print(key, presets, presets[key]) return presets[key] + @classmethod + def get_preset_filename_list(cls): + files, _ = folder_paths.recursive_search(cls.get_presets_dir(), excluded_dir_names=[".git"]) + + return folder_paths.filter_files_extensions(files, cls.file_extensions) + @classmethod def load_presets(cls): cls._presets = dict() preset_filename_list = cls.get_preset_filename_list() for preset_filename in preset_filename_list: - with open(os.path.join(cls.presets_dir, preset_filename),"r") as f: - reader = csv.DictReader(f) - for row in reader: - cls._presets[f"{preset_filename[:-4]} : {row['name']}"] = row['prompt'] + with open(os.path.join(cls.get_presets_dir(), preset_filename),"r") as f: + cls.load_file(f, preset_filename) + + +class PresetManager(PresetManagerBase): + presets_dir = "presets" + file_extensions = ['.csv'] @classmethod - def get_preset_filename_list(cls): - files, _ = folder_paths.recursive_search(cls.presets_dir, excluded_dir_names=[".git"]) + def load_file(cls, f, preset_filename): + reader = csv.DictReader(f) + for row in reader: + cls._presets[f"{preset_filename[:-4]} : {row['name']}"] = row['prompt'] - return folder_paths.filter_files_extensions(files, ['.csv']) + +class PresetManagerAdvanced(PresetManagerBase): + presets_dir = "advanced_presets" + file_extensions = ['.json'] + + @classmethod + def load_file(cls, f, preset_filename): + data = json.load(f) + for k,v in data.items(): + cls._presets[f"{preset_filename[:-5]} : {k}"] = v