80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
from zhipuai import ZhipuAI
|
|
import json
|
|
import os
|
|
|
|
p = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
|
|
def get_api_key():
|
|
try:
|
|
config_path = os.path.join(p, 'config.json')
|
|
with open(config_path, 'r') as f:
|
|
config = json.load(f)
|
|
api_key = config["CHATGLM_API_KEY"]
|
|
except:
|
|
print("Error: API key is required")
|
|
return ""
|
|
return api_key
|
|
|
|
|
|
class GlmPromptNode:
|
|
|
|
def __init__(self):
|
|
api_key = get_api_key()
|
|
if api_key is not None:
|
|
self.client = ZhipuAI(api_key=api_key)
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True}),
|
|
# 可选列表
|
|
"model": (["glm-4-flash", "glm-4", "glm-4-flashx", "glm-4-plus", "glm-4v-plus", "glm-4-0520",
|
|
"glm-4-air", "glm-4-airx"],
|
|
{"default": "glm-4-flash"}
|
|
),
|
|
"max_tokens": ("INT", {"default": 1024, "min": 128, "max": 4096}),
|
|
},
|
|
"optional": {
|
|
"system": ("STRING", {"multiline": True}),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "SimpleNode"
|
|
OUTPUT_NODE = True
|
|
RETURN_NAMES = ("clip-L", "clip-T5")
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
FUNCTION = "PROCESSOR"
|
|
|
|
DESCRIPTION = """
|
|
generate prompt using chatglm api
|
|
"""
|
|
|
|
def PROCESSOR(self, prompt, model, max_tokens, system):
|
|
if not prompt:
|
|
return ("", "")
|
|
if not hasattr(self, "client"):
|
|
#raise "请先在Config中设置CHATGLM_API_KEY"
|
|
raise "please set CHATGLM_API_KEY in Config.json"
|
|
response = self.client.chat.completions.create(
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
messages=[
|
|
{"role": "system", "content": system or """你是Flux绘画模型的提示词构建师，通过精心构建的提示词，帮助用户实现写真级照片提示词。你的能力有:
|
|
- 理解用户需求，提炼关键信息；
|
|
- 构建详细且准确的提示词；
|
|
- 调整提示词以优化出图效果。
|
|
|
|
结果应给出英文纯文本格式的json，含有以下字段：theme，图片的主题。clip-L，图片的风格描述，如写实照片，素描，高清等等。clip-T5：图片的内容，对人物和周边环境应该详细描述。
|
|
json应该为有效的json，纯文本格式，不需要markdown格式"""},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
)
|
|
result = response.choices[0].message.content.strip('```').strip('json').strip('\\n').replace('\"', '"')
|
|
print(result)
|
|
result = json.loads(result)
|
|
print(result)
|
|
return (result["clip-L"], result["clip-T5"])
|