161 lines
6.2 KiB
Python
161 lines
6.2 KiB
Python
from ..mz_prompt_utils import Utils
|
|
from ..mz_llama_cpp import *
|
|
from ..mz_llama_core_nodes import *
|
|
from ..mz_prompts import *
|
|
|
|
def base_query_beautify_prompt_text(args_dict):
|
|
model_file = args_dict.get("llama_cpp_model", "")
|
|
text = args_dict.get("text", "")
|
|
style_presets = args_dict.get("style_presets", "")
|
|
options = args_dict.get("llama_cpp_options", {})
|
|
keep_device = args_dict.get("keep_device", False)
|
|
seed = args_dict.get("seed", -1)
|
|
options["seed"] = seed
|
|
|
|
|
|
customize_instruct = args_dict.get("customize_instruct", None)
|
|
Utils.print_log(
|
|
f"customize_instruct: {customize_instruct}")
|
|
try:
|
|
schema = None
|
|
if customize_instruct is None:
|
|
schema = get_schema_obj(
|
|
keys_type={
|
|
"description": get_schema_base_type("string"),
|
|
"long_prompt": get_schema_base_type("string"),
|
|
"main_color_word": get_schema_base_type("string"),
|
|
"camera_angle_word": get_schema_base_type("string"),
|
|
"style_words": get_schema_array("string"),
|
|
"subject_words": get_schema_array("string"),
|
|
"light_words": get_schema_array("string"),
|
|
"environment_words": get_schema_array("string"),
|
|
},
|
|
required=[
|
|
"description",
|
|
"long_prompt",
|
|
"main_color_word",
|
|
"camera_angle_word",
|
|
"style_words",
|
|
"subject_words",
|
|
"light_words",
|
|
"environment_words",
|
|
]
|
|
)
|
|
|
|
question = f"IDEA: {style_presets},{text}"
|
|
if style_presets == "none":
|
|
question = f"IDEA: {text}"
|
|
|
|
system_prompt = mz_prompts.Beautify_Prompt + mz_prompts.Long_prompt + "\n"
|
|
|
|
else:
|
|
|
|
system_prompt = customize_instruct.get("system", "")
|
|
question = customize_instruct.get("instruct", "%text%")
|
|
|
|
system_prompt = system_prompt.replace("%text%", text)
|
|
question = question.replace("%text%", text)
|
|
|
|
Utils.print_log(f"system_prompt: {system_prompt}")
|
|
Utils.print_log(f"question: {question}")
|
|
|
|
if schema is not None:
|
|
response_json = llama_cpp_simple_interrogator_to_json(
|
|
model_file=model_file,
|
|
system=system_prompt,
|
|
question=question,
|
|
schema=schema,
|
|
options=options,
|
|
)
|
|
Utils.print_log(f"response_json: {response_json}")
|
|
|
|
response = json.loads(response_json)
|
|
full_responses = []
|
|
|
|
if response["description"] != "":
|
|
full_responses.append(f"({response['description']})")
|
|
if response["long_prompt"] != "":
|
|
full_responses.append(f"({response['long_prompt']})")
|
|
if response["main_color_word"] != "":
|
|
full_responses.append(f"({response['main_color_word']})")
|
|
if response["camera_angle_word"] != "":
|
|
full_responses.append(f"({response['camera_angle_word']})")
|
|
|
|
response["style_words"] = [
|
|
x for x in response["style_words"] if x != ""]
|
|
if len(response["style_words"]) > 0:
|
|
full_responses.append(
|
|
f"({', '.join(response['style_words'])})")
|
|
|
|
response["subject_words"] = [
|
|
x for x in response["subject_words"] if x != ""]
|
|
if len(response["subject_words"]) > 0:
|
|
full_responses.append(
|
|
f"({', '.join(response['subject_words'])})")
|
|
|
|
response["light_words"] = [
|
|
x for x in response["light_words"] if x != ""]
|
|
if len(response["light_words"]) > 0:
|
|
full_responses.append(
|
|
f"({', '.join(response['light_words'])})")
|
|
|
|
response["environment_words"] = [
|
|
x for x in response["environment_words"] if x != ""]
|
|
if len(response["environment_words"]) > 0:
|
|
full_responses.append(
|
|
f"({', '.join(response['environment_words'])})")
|
|
|
|
full_response = ", ".join(full_responses)
|
|
else:
|
|
full_response = llama_cpp_simple_interrogator(
|
|
model_file=model_file,
|
|
system=system_prompt,
|
|
question=question,
|
|
options=options,
|
|
)
|
|
|
|
start_str = customize_instruct.get("start_str", "")
|
|
if start_str != "" and full_response.find(start_str) != -1:
|
|
full_response_list = full_response.split(start_str)
|
|
# 删除第一个元素
|
|
full_response_list.pop(0)
|
|
full_response = start_str.join(full_response_list)
|
|
|
|
end_str = customize_instruct.get("end_str", "")
|
|
if end_str != "" and full_response.find(end_str) != -1:
|
|
full_response_list = full_response.split(end_str)
|
|
# 删除最后一个元素
|
|
full_response_list.pop()
|
|
full_response = end_str.join(full_response_list)
|
|
|
|
if keep_device is False:
|
|
freed_gpu_memory(model_file=model_file)
|
|
|
|
# 去除换行
|
|
while full_response.find("\n") != -1:
|
|
full_response = full_response.replace("\n", " ")
|
|
|
|
# 句号换成逗号
|
|
while full_response.find(".") != -1:
|
|
full_response = full_response.replace(".", ",")
|
|
|
|
# 去除多余逗号
|
|
while full_response.find(",,") != -1:
|
|
full_response = full_response.replace(",,", ",")
|
|
while full_response.find(", ,") != -1:
|
|
full_response = full_response.replace(", ,", ",")
|
|
|
|
full_response = Utils.prompt_zh_to_en(full_response)
|
|
|
|
style_presets_prompt_text = style_presets_prompt.get(style_presets, "")
|
|
|
|
if style_presets_prompt_text != "":
|
|
full_response = f"{style_presets_prompt_text}, {full_response}"
|
|
|
|
return full_response
|
|
|
|
except Exception as e:
|
|
freed_gpu_memory(model_file=model_file)
|
|
# mz_utils.Utils.print_log(f"Error in auto_prompt_text: {e}")
|
|
raise e
|