'update'

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)

__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)

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}")

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,

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

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.