From 52eb31658d05479cf0aab22a39d8d16aec54af1e Mon Sep 17 00:00:00 2001 From: wailovet Date: Mon, 13 May 2024 00:59:25 +0800 Subject: [PATCH] 'update' --- README.md | 3 + __init__.py | 54 ++++++++++++++- mz_llama_cpp.py | 6 +- mz_llava.py | 6 +- mz_openaiapi.py | 171 ++++++++++++++++++++++++++++++++++++++++++++++++ mz_prompts.py | 2 +- 6 files changed, 234 insertions(+), 8 deletions(-) create mode 100644 mz_openaiapi.py diff --git a/README.md b/README.md index 3997a2b..981803e 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ Use llama.cpp to assist in generating some nodes related to prompt words, including beautifying prompt words and image recognition similar to clip-interrogator ## Recent changes +* [2024-05-13] 新增OpenAI API节点 (add OpenAI API node) * [2024-04-30] 支持自定义指令 (Support for custom instructions) * [2024-04-30] 添加llava-v1.6-vicuna-13b (add llava-v1.6-vicuna-13b) * [2024-04-30] 添加翻译 @@ -17,6 +18,8 @@ Use llama.cpp to assist in generating some nodes related to prompt words, includ 2. Restart ComfyUI. ## Nodes ++ CLIPTextEncode (OpenAI API) + + CLIPTextEncode (Phi-3) ![image](https://github.com/MinusZoneAI/ComfyUI-Prompt-MZ/assets/5035199/c4b97aeb-23c0-4cf1-a6a5-d259fdf83f6e) diff --git a/__init__.py b/__init__.py index 11ec40d..ee28d13 100644 --- a/__init__.py +++ b/__init__.py @@ -82,7 +82,7 @@ class MZ_CustomizeInstruct: return { "required": { "system": ("STRING", {"multiline": True, "default": mz_prompts.Long_prompt}), - "instruct": ("STRING", {"multiline": True, "default": "Short: {text}"}), + "instruct": ("STRING", {"multiline": True, "default": "Short: %text%"}), "start_str": ("STRING", {"default": "Long: "}), "end_str": ("STRING", {"default": ""}), }, @@ -96,9 +96,11 @@ class MZ_CustomizeInstruct: def create(self, **kwargs): return (kwargs,) + NODE_CLASS_MAPPINGS["MZ_CustomizeInstruct"] = MZ_CustomizeInstruct NODE_DISPLAY_NAME_MAPPINGS["MZ_CustomizeInstruct"] = f"{AUTHOR_NAME} - CustomizeInstruct" + def getCommonCLIPTextEncodeInput(): style_presets = mz_llama_cpp.get_style_presets() CommonCLIPTextEncodeInput = { @@ -115,6 +117,7 @@ def getCommonCLIPTextEncodeInput(): "clip": ("CLIP", ), "llama_cpp_options": ("LLamaCPPOptions", ), "customize_instruct": ("CustomizeInstruct", ), + # "customize_json_schema": ("STRING", ), } } @@ -420,5 +423,54 @@ class MZ_LLamaCPPInterrogator: # NODE_DISPLAY_NAME_MAPPINGS["MZ_LLamaCPPInterrogator"] = f"{AUTHOR_NAME} - LLamaCPP simple interrogator" +class MZ_OpenAIApiCLIPTextEncode: + @classmethod + def INPUT_TYPES(s): + importlib.reload(mz_llama_cpp) + + result = { + "required": { + "base_url": ("STRING", {"default": ""}), + "api_key": ("STRING", {"default": "", "placeholder": ""}), + "model_name": ("STRING", {"default": "gpt-3.5-turbo-1106"}), + }, + "optional": { + }, + } + + common_input = getCommonCLIPTextEncodeInput() + for key in common_input["required"]: + result["required"][key] = common_input["required"][key] + for key in common_input["optional"]: + if key != "llama_cpp_options": + result["optional"][key] = common_input["optional"][key] + + return result + RETURN_TYPES = ("STRING", "CONDITIONING",) + RETURN_NAMES = ("text", "conditioning",) + FUNCTION = "encode" + CATEGORY = CATEGORY_NAME + + def encode(self, **kwargs): + import mz_openaiapi + + importlib.reload(mz_openaiapi) + + text = mz_openaiapi.query_beautify_prompt_text(kwargs) + conditionings = None + clip = kwargs.get("clip", None) + if clip is not None: + tokens = clip.tokenize(text) + cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True,) + conditionings = [[cond, {"pooled_output": pooled}]] + + return (text, conditionings) + + +NODE_CLASS_MAPPINGS["MZ_OpenAIApiCLIPTextEncode"] = MZ_OpenAIApiCLIPTextEncode +NODE_DISPLAY_NAME_MAPPINGS[ + "MZ_OpenAIApiCLIPTextEncode"] = f"{AUTHOR_NAME} - CLIPTextEncode(OpenAIApi)" + + import mz_gen_translate mz_gen_translate.gen_translate(NODE_DISPLAY_NAME_MAPPINGS, NODE_CLASS_MAPPINGS) diff --git a/mz_llama_cpp.py b/mz_llama_cpp.py index b1ed531..9242447 100644 --- a/mz_llama_cpp.py +++ b/mz_llama_cpp.py @@ -473,10 +473,10 @@ def base_query_beautify_prompt_text(args_dict): else: system_prompt = customize_instruct.get("system", "") - question = customize_instruct.get("instruct", "{text}") + question = customize_instruct.get("instruct", "%text%") - system_prompt = system_prompt.replace("{text}", text) - question = question.replace("{text}", text) + system_prompt = system_prompt.replace("%text%", text) + question = question.replace("%text%", text) mz_prompt_utils.Utils.print_log(f"system_prompt: {system_prompt}") mz_prompt_utils.Utils.print_log(f"question: {question}") diff --git a/mz_llava.py b/mz_llava.py index 6953557..5987b92 100644 --- a/mz_llava.py +++ b/mz_llava.py @@ -147,9 +147,9 @@ def base_image_interrogator(args_dict): else: system_prompt = customize_instruct.get("system", "") - question = customize_instruct.get("instruct", "{text}") - system_prompt = system_prompt.replace("{text}", "") - question = question.replace("{text}", "") + question = customize_instruct.get("instruct", "%text%") + system_prompt = system_prompt.replace("%text%", "") + question = question.replace("%text%", "") full_response = mz_llama_cpp.llava_cpp_simple_interrogator( model_file=model_file, diff --git a/mz_openaiapi.py b/mz_openaiapi.py new file mode 100644 index 0000000..7f58d34 --- /dev/null +++ b/mz_openaiapi.py @@ -0,0 +1,171 @@ +import os +import sys +import json +import subprocess + +import mz_prompt_utils +import mz_llama_cpp +import mz_prompts + + +def query_beautify_prompt_text(args_dict): + try: + from openai import OpenAI + import openai + except ImportError: + subprocess.check_call( + [sys.executable, "-m", "pip", "install", "openai"]) + from openai import OpenAI + import openai + + api_key = args_dict.get("api_key", None) + base_url = args_dict.get("base_url", None) + + text = args_dict.get("text", "") + style_presets = args_dict.get("style_presets", "") + + if api_key is None: + raise ValueError("api_key is required") + + client = OpenAI( + api_key=api_key, + default_headers={"x-foo": "true"} + ) + + if base_url is not None: + client.base_url = base_url + + model_name = args_dict.get("model_name", "gpt-3.5-turbo") + + options = args_dict.get("options", {}) + + customize_instruct = args_dict.get("customize_instruct", None) + mz_prompt_utils.Utils.print_log( + f"customize_instruct: {customize_instruct}") + + + schema = None + if customize_instruct is None: + schema = mz_llama_cpp.get_schema_obj( + keys_type={ + "description": mz_llama_cpp.get_schema_base_type("string"), + "long_prompt": mz_llama_cpp.get_schema_base_type("string"), + "main_color_word": mz_llama_cpp.get_schema_base_type("string"), + "camera_angle_word": mz_llama_cpp.get_schema_base_type("string"), + "style_words": mz_llama_cpp.get_schema_array("string"), + "subject_words": mz_llama_cpp.get_schema_array("string"), + "light_words": mz_llama_cpp.get_schema_array("string"), + "environment_words": mz_llama_cpp.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) + + mz_prompt_utils.Utils.print_log(f"system_prompt: {system_prompt}") + mz_prompt_utils.Utils.print_log(f"question: {question}") + # print(f"system_prompt: {system_prompt}") + # print(f"question: {question}") + + output = None + if schema is not None: + + output = client.chat.completions.create( + model=model_name, + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": question}, + ], + tools=[{ + "type": "function", + "function": { + "name": "beautify_prompt_text", + "description": "required Beautify Prompt Text", + "parameters": schema, + } + }], + tool_choice={"type": "function", + "function": {"name": "beautify_prompt_text"}}, + ) + tool_calls = output.choices[0].message.tool_calls + + + functions_args = {} + for tool_call in tool_calls: + function_name = tool_call.function.name + function_args = json.loads(tool_call.function.arguments) + functions_args[function_name] = function_args + beautify_prompt_text_result = functions_args.get( + "beautify_prompt_text", {}) + + results = [] + for key, value in beautify_prompt_text_result.items(): + if type(value) == list: + value = [item for item in value if item != ""] + value = [mz_prompt_utils.Utils.prompt_zh_to_en(item) + for item in value] + if len(value) == 0: + continue + item_str = ", ".join(value) + results.append(f"({item_str})") + else: + if value == "": + continue + value = mz_prompt_utils.Utils.prompt_zh_to_en(value) + results.append(f"({value})") + + full_response = ", ".join(results) + + else: + output = client.chat.completions.create( + model=model_name, + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": question}, + ], + ) + + full_response = output.choices[0].message.content + + mz_prompt_utils.Utils.print_log( + f"OPENAI_OUTPUT: \n{output.model_dump_json()}") + # print(output.model_dump_json()) + + + # 去除换行 + 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(", ,", ",") + style_presets_prompt_text = mz_llama_cpp.style_presets_prompt.get( + style_presets, "") + if style_presets_prompt_text != "": + full_response = f"{style_presets_prompt_text}, {full_response}" + return full_response diff --git a/mz_prompts.py b/mz_prompts.py index 121147e..1a01b57 100644 --- a/mz_prompts.py +++ b/mz_prompts.py @@ -12,7 +12,7 @@ Below is a list of prompts that can be used to generate images with Stable Diffu I want you to write me a list of detailed prompts exactly about the idea written after IDEA. Follow the structure of the example prompts. This means a very short description of the scene, followed by modifiers divided by commas to alter the mood, style, lighting, and more. """ -Long_prompt = "Long prompt version should consist of 3 to 5 sentences. Long prompt version must sepcify the color, shape, texture or spatial relation of the included objects. DO NOT generate sentences that describe any atmosphere!!!" +Long_prompt = "Long prompt version should consist of 3 to 5 sentences. Long prompt version must sepcify the color, shape, texture or spatial relation of the included objects. DO NOT generate sentences that describe any atmosphere!!! The language of reply is English only!!!" Standardize_Prompt = """ Extract the content about Stable Diffusion style from the following input and combine it into a json array. Note that the output will be directly used in the program.