修改优化，新增提示词保存，自定义提示词处理
This commit is contained in:
DreamFly
+11
@@ -9,6 +9,9 @@ from .nodes.prompt_portrait import PortraitSkinPrompt
|
||||
from .nodes.prompt_portrait import PortraitFashionPrompt
|
||||
from .nodes.prompt_portrait import PortraitPosePrompt
|
||||
from .nodes.prompt_portrait import PortraitCosmeticPrompt
|
||||
from .nodes.prompt_combine import CombinePrompt, MultiCombinePrompt
|
||||
from .nodes.prompt_custom import CustomPrompt
|
||||
from .nodes.prompt_save import SavePrompt
|
||||
|
||||
|
||||
|
||||
@@ -24,6 +27,10 @@ NODE_CLASS_MAPPINGS = {
|
||||
"PortraitFashionPrompt": PortraitFashionPrompt,
|
||||
"PortraitPosePrompt": PortraitPosePrompt,
|
||||
"PortraitCosmeticPrompt": PortraitCosmeticPrompt,
|
||||
"CombinePrompt": CombinePrompt,
|
||||
"MultiCombinePrompt": MultiCombinePrompt,
|
||||
"CustomPrompt": CustomPrompt,
|
||||
"SavePrompt": SavePrompt,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -38,6 +45,10 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PortraitFashionPrompt": "Portrait Fashion Prompt",
|
||||
"PortraitPosePrompt": "Portrait Pose Prompt",
|
||||
"PortraitCosmeticPrompt": "Portrait Cosmetic Prompt",
|
||||
"CombinePrompt": "Combine Prompt",
|
||||
"MultiCombinePrompt": "Multi Combine Prompt",
|
||||
"CustomPrompt": "Custom Prompt",
|
||||
"SavePrompt": "Save Prompt",
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
|
||||
|
||||
# Combine Prompt
|
||||
class CombinePrompt:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"text_A": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"text_B": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"text_C": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
FUNCTION = "combine"
|
||||
CATEGORY = "PromptWrapper"
|
||||
|
||||
def combine(self, text_A="", text_B="", text_C=""):
|
||||
|
||||
texts = [text_A, text_B, text_C]
|
||||
prompts = []
|
||||
for text in texts:
|
||||
if text != "" and text != "\n" and text != "\r\n":
|
||||
prompts.append(text)
|
||||
|
||||
|
||||
output_prompt = ""
|
||||
if len(prompts) > 0:
|
||||
output_prompt = ",".join(prompts)
|
||||
|
||||
return (output_prompt, )
|
||||
|
||||
|
||||
|
||||
# Multi Combine Prompt
|
||||
class MultiCombinePrompt:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"text_A": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"text_B": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"text_C": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"text_D": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"text_E": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"text_F": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
FUNCTION = "combine"
|
||||
CATEGORY = "PromptWrapper"
|
||||
|
||||
def combine(self, text_A="", text_B="", text_C="", text_D="", text_E="", text_F=""):
|
||||
|
||||
texts = [text_A, text_B, text_C, text_D, text_E, text_F]
|
||||
prompts = []
|
||||
for text in texts:
|
||||
if text != "" and text != "\n" and text != "\r\n":
|
||||
prompts.append(text)
|
||||
|
||||
|
||||
output_prompt = ""
|
||||
if len(prompts) > 0:
|
||||
output_prompt = ",".join(prompts)
|
||||
|
||||
return (output_prompt, )
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
import os
|
||||
from ..modules.build_prompt import BuildPrompt
|
||||
from ..modules.build_prompt import buildOptionList, buildClassifyDatas
|
||||
from ..modules.build_prompt import EMPTY_OPTION, RANDOM_OPTION, DEFAULT_OPTION
|
||||
from ..modules.config import config
|
||||
|
||||
|
||||
|
||||
def _ReadClassifyList():
|
||||
# 获取Script的目录路径
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__)) # Script directory
|
||||
custom_dir = os.path.join(script_dir, "..", "user")
|
||||
filenames = os.listdir(custom_dir)
|
||||
|
||||
classify_list = []
|
||||
for filename in filenames:
|
||||
if filename.endswith("_datas.jsonl") and os.path.isfile(os.path.join(custom_dir, filename)):
|
||||
data_name = filename.replace("_datas.jsonl", "")
|
||||
data_name = data_name.replace("_", "")
|
||||
classify_list.append(data_name)
|
||||
|
||||
return classify_list
|
||||
|
||||
|
||||
|
||||
# Custom Prompt
|
||||
class CustomPrompt:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
classify_list = _ReadClassifyList()
|
||||
classify_list = buildOptionList(classify_list, True, False)
|
||||
custom_data_list = buildOptionList([], True, True)
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"classify": (classify_list, {
|
||||
"default": DEFAULT_OPTION
|
||||
}),
|
||||
"prompt": (custom_data_list, {
|
||||
"default": DEFAULT_OPTION
|
||||
}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||||
"preset_prompt": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"input_prompt": ("STRING", {"forceInput": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
FUNCTION = "prompt_generate"
|
||||
CATEGORY = "PromptWrapper"
|
||||
|
||||
def prompt_generate(self, classify, prompt="", seed=0, preset_prompt="", input_prompt=""):
|
||||
|
||||
prompt_words = []
|
||||
|
||||
if input_prompt != "":
|
||||
prompt_words.append(input_prompt)
|
||||
|
||||
if preset_prompt != "":
|
||||
prompt_words.append(preset_prompt)
|
||||
|
||||
if classify != EMPTY_OPTION:
|
||||
|
||||
if prompt == RANDOM_OPTION:
|
||||
# 获取Script的目录路径
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__)) # Script directory
|
||||
classify_datas = buildClassifyDatas(classify)
|
||||
custom_filename = f"{classify_datas}.jsonl"
|
||||
custom_data_filename = os.path.join(script_dir, "..", "user", custom_filename)
|
||||
|
||||
build = BuildPrompt(custom_data_filename)
|
||||
prompt = build.generate_prompt(seed)
|
||||
if prompt != "":
|
||||
prompt_words.append(prompt)
|
||||
|
||||
elif prompt != EMPTY_OPTION:
|
||||
prompt_words.append(prompt)
|
||||
|
||||
output_prompt = ""
|
||||
if len(prompt_words) > 0:
|
||||
output_prompt = ", ".join(prompt_words)
|
||||
|
||||
return (output_prompt, )
|
||||
|
||||
|
||||
|
||||
@@ -1047,7 +1047,7 @@ class PortraitPosePrompt:
|
||||
}),
|
||||
# 胸部动作权重
|
||||
"chest_action_weight": ("FLOAT", {
|
||||
"default": 0,
|
||||
"default": 1,
|
||||
"min": 0,
|
||||
"max": max_float_value,
|
||||
"round": 0.01,
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
import os
|
||||
import json
|
||||
|
||||
|
||||
# Save Prompt
|
||||
class SavePrompt:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {
|
||||
"forceInput": True
|
||||
}),
|
||||
"filename": ("STRING", {
|
||||
"default": "custom"
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
RETURN_NAMES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "save"
|
||||
CATEGORY = "PromptWrapper"
|
||||
|
||||
|
||||
def save(self, text="", filename=""):
|
||||
|
||||
text = text.strip()
|
||||
if text != "" and text != "\n" and text != "\r\n":
|
||||
# 获取Script的目录路径
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
save_dir = os.path.join(script_dir, "..", "user")
|
||||
|
||||
if not os.path.exists(save_dir):
|
||||
os.makedirs(save_dir)
|
||||
|
||||
save_path = os.path.join(save_dir, f"{filename}_datas.jsonl")
|
||||
with open(save_path, "a+", encoding="utf-8") as fp:
|
||||
data = { "text": text }
|
||||
text_data = json.dumps(data, ensure_ascii=False)
|
||||
size = fp.tell()
|
||||
if size == 0:
|
||||
fp.write(text_data)
|
||||
else:
|
||||
fp.write(f"\n{text_data}")
|
||||
|
||||
|
||||
return {"ui": {"text": text}}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user