515 lines
17 KiB
Python
515 lines
17 KiB
Python
import os
|
||
import json
|
||
import requests
|
||
|
||
import base64
|
||
from io import BytesIO
|
||
import numpy as np
|
||
from server import PromptServer
|
||
from PIL import Image
|
||
|
||
ALL_CODES_LANGS = ['af', 'sq', 'am', 'ar', 'hy', 'as', 'ay', 'az', 'bm', 'eu', 'be', 'bn', 'bho', 'bs', 'bg', 'ca', 'ceb', 'ny', 'zh-CN', 'zh-TW', 'co', 'hr', 'cs', 'da', 'dv', 'doi', 'nl', 'en', 'eo', 'et', 'ee', 'tl', 'fi', 'fr', 'fy', 'gl', 'ka', 'de', 'el', 'gn', 'gu', 'ht', 'ha', 'haw', 'iw', 'hi', 'hmn', 'hu', 'is', 'ig', 'ilo', 'id', 'ga', 'it', 'ja', 'jw', 'kn', 'kk', 'km', 'rw', 'gom', 'ko', 'kri', 'ku', 'ckb', 'ky', 'lo', 'la', 'lv', 'ln', 'lt', 'lg', 'lb', 'mk', 'mai', 'mg', 'ms', 'ml', 'mt', 'mi', 'mr', 'mni-Mtei', 'lus', 'mn', 'my', 'ne', 'no', 'or', 'om', 'ps', 'fa', 'pl', 'pt', 'pa', 'qu', 'ro', 'ru', 'sm', 'sa', 'gd', 'nso', 'sr', 'st', 'sn', 'sd', 'si', 'sk', 'sl', 'so', 'es', 'su', 'sw', 'sv', 'tg', 'ta', 'tt', 'te', 'th', 'ti', 'ts', 'tr', 'tk', 'ak', 'uk', 'ur', 'ug', 'uz', 'vi', 'cy', 'xh', 'yi', 'yo', 'zu']
|
||
|
||
|
||
ENDPOINT_URL = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
|
||
|
||
# Language models: https://docs.bigmodel.cn/api-reference/%E6%A8%A1%E5%9E%8B-api/%E5%AF%B9%E8%AF%9D%E8%A1%A5%E5%85%A8#%E6%96%87%E6%9C%AC%E6%A8%A1%E5%9E%8B
|
||
LIST_LANGUAGE_MODELS = [
|
||
# GLM-4
|
||
"glm-4-plus",
|
||
"glm-4-air-250414",
|
||
"glm-4-airx",
|
||
"glm-4-flashx",
|
||
"glm-4-flashx-250414",
|
||
# GLM-4.5
|
||
"glm-4.5",
|
||
"glm-4.5-air",
|
||
"glm-4.5-x",
|
||
"glm-4.5-airx",
|
||
"glm-4.5-flash",
|
||
# GLM-4.6
|
||
"glm-4.6",
|
||
# GLM-4.7
|
||
"glm-4.7",
|
||
# GLM-Z1
|
||
"glm-z1-air",
|
||
"glm-z1-airx",
|
||
"glm-z1-flash",
|
||
"glm-z1-flashx",
|
||
]
|
||
|
||
# Multimodal models: https://docs.bigmodel.cn/api-reference/%E6%A8%A1%E5%9E%8B-api/%E5%AF%B9%E8%AF%9D%E8%A1%A5%E5%85%A8#%E8%A7%86%E8%A7%89%E6%A8%A1%E5%9E%8B
|
||
LIST_MULTIMODAL_MODELS = [
|
||
# --- GLM-4v
|
||
"glm-4v-flash",
|
||
"glm-4v",
|
||
"glm-4v-plus-0111",
|
||
# --- GLM-4.1v
|
||
"glm-4.1v-thinking-flashx",
|
||
"glm-4.1v-thinking-flash",
|
||
# --- GLM-4.5v
|
||
"glm-4.5v",
|
||
# --- GLM-4.6v
|
||
"glm-4.6v",
|
||
"glm-4.6v-flash",
|
||
"glm-4.6v-flashx",
|
||
# --- other
|
||
"autoglm-phone",
|
||
]
|
||
|
||
# Multimodal models: https://docs.bigmodel.cn/api-reference/%E6%A8%A1%E5%9E%8B-api/%E5%AF%B9%E8%AF%9D%E8%A1%A5%E5%85%A8#%E8%A7%86%E8%A7%89%E6%A8%A1%E5%9E%8B
|
||
LIST_IMAGE_GENERATION_MODELS = [
|
||
"cogview-4-250304",
|
||
"cogview-4",
|
||
"cogview-3-flash"
|
||
]
|
||
|
||
# CogView: https://docs.bigmodel.cn/api-reference/%E6%A8%A1%E5%9E%8B-api/%E5%9B%BE%E5%83%8F%E7%94%9F%E6%88%90
|
||
LIST_IMAGE_GENERATION_MODELS = [
|
||
"cogview-4-250304",
|
||
"cogview-4",
|
||
"cogview-3-flash"
|
||
]
|
||
|
||
# CogVideo: https://docs.bigmodel.cn/api-reference/%E6%A8%A1%E5%9E%8B-api/%E7%94%9F%E6%88%90%E8%A7%86%E9%A2%91%E5%BC%82%E6%AD%A5
|
||
LIST_VIDEO_GENERATION_MODELS = [
|
||
"cogvideox-3",
|
||
"cogvideox3-flash",
|
||
"cogvideox-2",
|
||
"cogvideox-flash",
|
||
]
|
||
|
||
def getConfigData():
|
||
# Directory node and config file
|
||
dir_node = os.path.dirname(__file__)
|
||
config_path = os.path.join(os.path.abspath(dir_node), "config.json")
|
||
config = {
|
||
"__comment": "Register on the site https://bigmodel.cn and get a key and add it to the field ZHIPUAI_API_KEY. Change default translate languages 'from' and 'to' you use",
|
||
"from_translate": "ru",
|
||
"to_translate": "en",
|
||
"default_language_model": "glm-4.5-flash",
|
||
"default_multimodal_model": "glm-4.6v-flash",
|
||
"ZHIPUAI_API_KEY": "your_api_key"
|
||
}
|
||
|
||
# Load config.js file
|
||
if not os.path.exists(config_path):
|
||
print("[ChatGLMNode] File config.js file not found! Create default config.json...")
|
||
with open(config_path, "w", encoding="utf-8") as f:
|
||
json.dump(config, f, ensure_ascii=False, indent=4)
|
||
return config
|
||
else:
|
||
with open(config_path, "r", encoding="utf-8") as f:
|
||
config = json.load(f)
|
||
return config
|
||
# =====
|
||
|
||
def checkPropValue(obj, key, not_include = []):
|
||
checkVal = lambda v: v is None or v.strip() == "" or v in not_include
|
||
|
||
prop_val = obj.get(key)
|
||
|
||
if checkVal(prop_val):
|
||
obj.update(getConfigData())
|
||
return True if checkVal(obj.get(key)) else False
|
||
|
||
else:
|
||
return False
|
||
|
||
|
||
CONFIG = getConfigData()
|
||
|
||
def createRequest(payload):
|
||
global CONFIG
|
||
|
||
if checkPropValue(CONFIG, "ZHIPUAI_API_KEY", ["your_api_key"]):
|
||
raise ValueError("ZHIPUAI_API_KEY value is empty or missing")
|
||
|
||
ZHIPUAI_API_KEY = CONFIG.get("ZHIPUAI_API_KEY")
|
||
|
||
# Headers
|
||
headers = {
|
||
"Authorization": f"Bearer {ZHIPUAI_API_KEY}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
try:
|
||
response = requests.post(ENDPOINT_URL, headers=headers, json=payload)
|
||
response.raise_for_status()
|
||
|
||
if response.status_code == 200:
|
||
json_data = response.json()
|
||
response_text = json_data.get("choices")[0]["message"]["content"].strip()
|
||
|
||
return response_text
|
||
|
||
except requests.HTTPError as e:
|
||
print(f"Error request ChatGLM: {response.status_code}, {response.text}")
|
||
raise e
|
||
except Exception as e:
|
||
print(f"Error ChatGLM: {e}")
|
||
raise e
|
||
|
||
|
||
def translate(prompt, srcTrans, toTrans, model, max_tokens, temperature, top_p):
|
||
# Check prompt exist
|
||
if prompt is None or prompt.strip() == "":
|
||
return ""
|
||
|
||
# Create body request
|
||
payload = {
|
||
"model": model,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": f"Translate from {srcTrans} to {toTrans} and return only the translated text: {prompt}",
|
||
},
|
||
],
|
||
"max_tokens": round(max_tokens, 2),
|
||
"temperature": round(temperature, 2),
|
||
"top_p": round(top_p, 2),
|
||
}
|
||
|
||
response_translate_text = createRequest(payload)
|
||
|
||
return response_translate_text
|
||
|
||
|
||
class ChatGLM4TranslateCLIPTextEncodeNode:
|
||
@classmethod
|
||
def INPUT_TYPES(self):
|
||
from_lng = CONFIG.get("from_translate") if CONFIG.get("from_translate") in ALL_CODES_LANGS else "ru"
|
||
to_lng = CONFIG.get("to_translate") if CONFIG.get("to_translate") in ALL_CODES_LANGS else "en"
|
||
return {
|
||
"required": {
|
||
"from_translate": (
|
||
ALL_CODES_LANGS,
|
||
{"default": from_lng, "tooltip": "Translation from"},
|
||
),
|
||
"to_translate": (
|
||
ALL_CODES_LANGS,
|
||
{"default": to_lng, "tooltip": "Translation to"},
|
||
),
|
||
"model": (
|
||
LIST_LANGUAGE_MODELS,
|
||
{
|
||
"default": CONFIG.get("default_language_model", "glm-4.5-flash"),
|
||
"tooltip": "The model code to be called. Models with text 'flash' should be free!",
|
||
},
|
||
),
|
||
"max_tokens": (
|
||
"INT",
|
||
{
|
||
"default": 1024,
|
||
"tooltip": "The maximum number of tokens for model output, maximum output is 4095, default value is 1024.",
|
||
},
|
||
),
|
||
"temperature": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.95,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "Sampling temperature, controls the randomness of the output, must be a positive number within the range: [0.0, 1.0], default value is 0.95.",
|
||
},
|
||
),
|
||
"top_p": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.7,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "Another method of temperature sampling, value range is: [0.0, 1.0], default value is 0.7.",
|
||
},
|
||
),
|
||
"text": ("STRING", {"multiline": True, "placeholder": "Input text"}),
|
||
"clip": ("CLIP",),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = (
|
||
"CONDITIONING",
|
||
"STRING",
|
||
)
|
||
FUNCTION = "chatglm_translate_text"
|
||
DESCRIPTION = (
|
||
"This is a node that translates the prompt into another language using ChatGLM."
|
||
)
|
||
CATEGORY = "AlekPet Nodes/conditioning"
|
||
|
||
def chatglm_translate_text(
|
||
self,
|
||
from_translate,
|
||
to_translate,
|
||
model,
|
||
max_tokens,
|
||
temperature,
|
||
top_p,
|
||
text,
|
||
clip,
|
||
):
|
||
|
||
text = translate(
|
||
text, from_translate, to_translate, model, max_tokens, temperature, top_p
|
||
)
|
||
tokens = clip.tokenize(text)
|
||
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
|
||
return ([[cond, {"pooled_output": pooled}]], text)
|
||
|
||
|
||
class ChatGLM4TranslateTextNode(ChatGLM4TranslateCLIPTextEncodeNode):
|
||
@classmethod
|
||
def INPUT_TYPES(self):
|
||
return_types = super().INPUT_TYPES()
|
||
del return_types["required"]["clip"]
|
||
return return_types
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("text",)
|
||
FUNCTION = "chatglm_translate_text"
|
||
|
||
CATEGORY = "AlekPet Nodes/text"
|
||
|
||
def chatglm_translate_text(
|
||
self, from_translate, to_translate, model, max_tokens, temperature, top_p, text
|
||
):
|
||
|
||
text = translate(
|
||
text, from_translate, to_translate, model, max_tokens, temperature, top_p
|
||
)
|
||
|
||
return (text,)
|
||
|
||
|
||
# ChatGLM Instruct Node
|
||
class ChatGLM4InstructNode:
|
||
@classmethod
|
||
def INPUT_TYPES(self):
|
||
return {
|
||
"required": {
|
||
"model": (
|
||
LIST_LANGUAGE_MODELS,
|
||
{
|
||
"default": CONFIG.get("default_language_model", "glm-4.5-flash"),
|
||
"tooltip": "The model code to be called. Models with text 'flash' should be free!",
|
||
},
|
||
),
|
||
"max_tokens": (
|
||
"INT",
|
||
{
|
||
"default": 1024,
|
||
"tooltip": "The maximum number of tokens for model output, maximum output is 4095, default value is 1024.",
|
||
},
|
||
),
|
||
"temperature": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.95,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "Sampling temperature, controls the randomness of the output, must be a positive number within the range: [0.0, 1.0], default value is 0.95.",
|
||
},
|
||
),
|
||
"top_p": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.7,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "Another method of temperature sampling, value range is: [0.0, 1.0], default value is 0.7.",
|
||
},
|
||
),
|
||
"instruct": (
|
||
"STRING",
|
||
{
|
||
"multiline": True,
|
||
"placeholder": "Input instruct text",
|
||
"default": "Generate details text, without quotation marks or the word 'prompt' on english: {query}",
|
||
"tooltip": "Enter the instruction for the neural network to execute and indicate where to insert the query text {query}",
|
||
},
|
||
),
|
||
"query": (
|
||
"STRING",
|
||
{
|
||
"multiline": True,
|
||
"placeholder": "Enter the query text for the instruction",
|
||
"tooltip": "Query field",
|
||
},
|
||
),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
FUNCTION = "chatglm_instruct"
|
||
|
||
CATEGORY = "AlekPet Nodes/Instruct"
|
||
|
||
def chatglm_instruct(self, model, max_tokens, temperature, top_p, instruct, query):
|
||
|
||
if instruct is None or instruct.strip() == "":
|
||
raise ValueError("Instruct text is empty!")
|
||
|
||
if query is None or query.strip() == "":
|
||
raise ValueError("Query text is empty!")
|
||
|
||
instruct = instruct.replace("{query}", query)
|
||
|
||
# Create body request
|
||
payload = {
|
||
"model": model,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": instruct,
|
||
},
|
||
],
|
||
"max_tokens": round(max_tokens, 2),
|
||
"temperature": round(temperature, 2),
|
||
"top_p": round(top_p, 2),
|
||
}
|
||
|
||
answer = createRequest(payload)
|
||
|
||
return (answer,)
|
||
|
||
|
||
# ChatGLM Instruct Media Node
|
||
def toBase64ImgUrl(img):
|
||
bytesIO = BytesIO()
|
||
img.save(bytesIO, format="PNG")
|
||
img_types = bytesIO.getvalue()
|
||
img_base64 = base64.b64encode(img_types)
|
||
return f"data:image/png;base64,{img_base64.decode('utf-8')}"
|
||
|
||
|
||
class ChatGLM4InstructMediaNode:
|
||
@classmethod
|
||
def INPUT_TYPES(self):
|
||
return {
|
||
"optional": {
|
||
"image": ("IMAGE",),
|
||
# "video": ("STRING", {"forceInput": True, "default": ""}),
|
||
},
|
||
"required": {
|
||
"model": (
|
||
LIST_MULTIMODAL_MODELS,
|
||
{
|
||
"default": CONFIG.get("default_multimodal_model", "glm-4.6v-flash"),
|
||
"tooltip": "The model code to be called. Models with text 'flash' should be free!",
|
||
},
|
||
),
|
||
"max_tokens": (
|
||
"INT",
|
||
{
|
||
"default": 1024,
|
||
"tooltip": "The maximum number of tokens for model output, maximum output is 4095, default value is 1024.",
|
||
},
|
||
),
|
||
"temperature": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.8,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "Sampling temperature, controls the randomness of the output, must be a positive number within the range: [0.0, 1.0], default value is 0.95.",
|
||
},
|
||
),
|
||
"top_p": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.6,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "Another method of temperature sampling, value range is: [0.0, 1.0], default value is 0.7.",
|
||
},
|
||
),
|
||
"instruct": (
|
||
"STRING",
|
||
{
|
||
"multiline": True,
|
||
"placeholder": "Input instruct text",
|
||
"default": "What is shown in the picture?",
|
||
"tooltip": "Enter the instruction for the neural network",
|
||
},
|
||
),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
FUNCTION = "chatglm_instruct_media"
|
||
|
||
CATEGORY = "AlekPet Nodes/Instruct"
|
||
|
||
def chatglm_instruct_media(
|
||
self, model, max_tokens, temperature, top_p, instruct, image=None, video=""
|
||
):
|
||
|
||
if instruct is None or instruct.strip() == "":
|
||
raise ValueError("Instruct text is empty!")
|
||
|
||
# video = video.strip()
|
||
|
||
# if image is None and (video is None and video == ""):
|
||
# raise ValueError("Image or Video path is empty!")
|
||
|
||
if image is not None:
|
||
if video != "":
|
||
raise ValueError("You cannot use both an image and a video at the same time!")
|
||
|
||
answer = ""
|
||
payload = {}
|
||
if image is not None:
|
||
img = 255.0 * image.cpu().numpy()
|
||
img = np.squeeze(img)
|
||
img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
|
||
img = toBase64ImgUrl(img)
|
||
|
||
# Create body request for image
|
||
payload = {
|
||
"model": model,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "image_url", "image_url": {"url": img}},
|
||
{"type": "text", "text": instruct},
|
||
],
|
||
}
|
||
],
|
||
"max_tokens": round(max_tokens, 2),
|
||
"temperature": round(temperature, 2),
|
||
"top_p": round(top_p, 2),
|
||
}
|
||
|
||
# if video:
|
||
# # Create body request for video
|
||
# address = PromptServer.instance.address
|
||
# port = PromptServer.instance.port
|
||
# url_video = f"http://{address}:{port}/view?filename={video}&type=input&subfolder="
|
||
|
||
# payload = {
|
||
# "model": model,
|
||
# "messages": [
|
||
# {
|
||
# "role": "user",
|
||
# "content": [
|
||
# {"type": "video_url", "video_url": {"url": url_video}},
|
||
# {"type": "text", "text": instruct},
|
||
# ],
|
||
# }
|
||
# ],
|
||
# "max_tokens": round(max_tokens, 2),
|
||
# "temperature": round(temperature, 2),
|
||
# "top_p": round(top_p, 2),
|
||
# }
|
||
|
||
answer = createRequest(payload)
|
||
|
||
return (answer,)
|