diff --git a/ChatGLM4_API_Node.py b/ChatGLM4_API_Node.py index 2120d13..19340a5 100644 --- a/ChatGLM4_API_Node.py +++ b/ChatGLM4_API_Node.py @@ -3,22 +3,24 @@ import json import requests import time import jwt -import re import tempfile import random from PIL import Image import base64 import cv2 +from io import BytesIO +import numpy as np +import torch -p = os.path.dirname(os.path.realpath(__file__)) +p = os.path.dirname(os.path.realpath(__file__)) -def get_zpai_api_key(): +def get_zpai_api_key(): try: - config_path = os.path.join(p, 'config.json') + config_path = os.path.join(p, 'config.json') with open(config_path, 'r') as f: config = json.load(f) - api_key = config["Zpai_API_KEY"] + api_key = config["Zpai_API_KEY"] except Exception as e: raise Exception("Error: API key is required", e) return api_key @@ -44,71 +46,95 @@ def generate_token(apikey_s: str, exp_second: int): ) -class ZhipuaiApi_Txt: +class ZhipuaiApi_Txt: def __init__(self): self.api_key = get_zpai_api_key() if self.api_key is not None: - client_key = self.api_key + client_key = self.api_key @classmethod def INPUT_TYPES(cls): return { "required": { "prompt": ("STRING", {"multiline": True, "default": "30 words describe a girl walking on the Moon."}), - "model_name": (["glm-4", "glm-3-turbo"],), + "model_name": (["glm-4", "glm-3-turbo", "cogview-3"],), "max_tokens": ("INT", {"default": 1024, "min": 128, "max": 8192, "step": 128, "display": "slider"}), - "temperature": ("FLOAT", {"default": 0.95, "min": 0.01, "max": 0.99, "step": 0.01, "round": False, "display": "slider"}), + "temperature": ("FLOAT", {"default": 0.95, "min": 0.01, "max": 0.99, "step": 0.01, "round": False, + "display": "slider"}), "output_language": (["English", "Original_language"],), }, "optional": {"prompt_tran2english_only": ("BOOLEAN", {"default": False},), } } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("text",) - FUNCTION = "zhipuai_txt_api" - CATEGORY = "ChatGlm_Api" + RETURN_TYPES = ("STRING", "IMAGE") + RETURN_NAMES = ("text", "image") + FUNCTION = "zhipuai_txt_api" + CATEGORY = "ChatGlm_Api" - - def zhipuai_txt_api(self, prompt, model_name, max_tokens, temperature, output_language, prompt_tran2english_only): + def zhipuai_txt_api(self, prompt, model_name, max_tokens, temperature,output_language, + prompt_tran2english_only): if not self.api_key: raise ValueError("API key is required") if prompt == None: raise ValueError("need prompt") else: - if prompt_tran2english_only: - prompt_txt = '以上内容翻译为成英文' + if model_name == "cogview-3": + url = "https://open.bigmodel.cn/api/paas/v4/images/generations" + exp_seconds = int(9) + client_key = str(get_zpai_api_key()) + token = generate_token(apikey_s=client_key, exp_second=exp_seconds) + header_txt = { + 'Authorization': f'Bearer {token}', + 'Content-Type': 'application/json' + } + data_txt = {"model": f"{model_name}", "prompt": f"{prompt}",} + data_txt = json.dumps(data_txt) + response = requests.post(url=url, headers=header_txt, data=data_txt) + txt_content = response.json() + # print(txt_content) + prompt_txt = txt_content.get("data", 'Default Value') + url = prompt_txt[0]["url"] + response_url = requests.get(url) + image_file = BytesIO(response_url.content) + image = Image.open(image_file) + image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + return (url, image) else: - if output_language == "English": - prompt_txt = 'Answer my questions or requests in English' + if prompt_tran2english_only: + prompt_txt = '翻译以下内容为英文' else: - prompt_txt = '用我使用的语言来回答我的问题或要求' - prompt = ''.join([prompt, prompt_txt]) - url = "https://open.bigmodel.cn/api/paas/v4/chat/completions" - exp_seconds = int(9) - client_key = str(get_zpai_api_key()) - token = generate_token(apikey_s=client_key, exp_second=exp_seconds) - header_txt = { - 'Authorization': f'Bearer {token}', - 'Content-Type': 'application/json' - } - data_txt = {"model": f"{model_name}",'max_tokens': f'{max_tokens}','temperature': f'{temperature}', "messages": [{"role": "user", "content": f"{prompt}"}]} - data_txt = json.dumps(data_txt) - response = requests.post(url=url, headers=header_txt, data=data_txt) - txt_content = response.json() - prompt_txt = txt_content.get("choices", 'Default Value')[0] - prompt_txt = prompt_txt.get("message", 'Default Value') - prompt_txt = prompt_txt.get("content", 'Default Value') - prompt = re.sub(r'\n+', '', prompt_txt) - return (prompt,) + if output_language == "English": + prompt_txt = 'Only respond to the results of the following questions in English' + else: + prompt_txt = '只回复以下问题的结果，用我使用的语言' + prompt = ':'.join([prompt_txt, prompt]) + url = "https://open.bigmodel.cn/api/paas/v4/chat/completions" + exp_seconds = int(9) + client_key = str(get_zpai_api_key()) + token = generate_token(apikey_s=client_key, exp_second=exp_seconds) + header_txt = { + 'Authorization': f'Bearer {token}', + 'Content-Type': 'application/json' + } + data_txt = {"model": f"{model_name}", 'max_tokens': f'{max_tokens}', 'temperature': f'{temperature}', + "messages": [{"role": "user", "content": f"{prompt}"}]} + data_txt = json.dumps(data_txt) + response = requests.post(url=url, headers=header_txt, data=data_txt) + txt_content = response.json() + # print(txt_content) + prompt_txt = txt_content["choices"][0]["message"]["content"] + # prompt_txt = prompt_txt.strip('"') + prompt = str(prompt_txt).replace("\n", "") + return (prompt,None,) -class ZhipuaiApi_Img: +class ZhipuaiApi_Img: def __init__(self): self.api_key = get_zpai_api_key() if self.api_key is not None: - client_key = self.api_key + client_key = self.api_key @classmethod def INPUT_TYPES(cls): @@ -117,7 +143,8 @@ class ZhipuaiApi_Img: "prompt": ("STRING", {"default": "Describe this image", "multiline": True}), "image": ("IMAGE",), "max_tokens": ("INT", {"default": 1024, "min": 128, "max": 8192, "step": 128, "display": "slider"}), - "temperature": ("FLOAT", {"default": 0.8, "min": 0.01, "max": 0.99, "step": 0.01, "round": False, "display": "slider"}), + "temperature": ( + "FLOAT", {"default": 0.8, "min": 0.01, "max": 0.99, "step": 0.01, "round": False, "display": "slider"}), "output_language": (["English", "Original_language"],), } } @@ -142,10 +169,10 @@ class ZhipuaiApi_Img: raise ValueError("Needs img") else: if output_language == "English": - prompt_img = 'Answer my questions or requests in English' + prompt_img = 'Only respond to the results of the following questions in English' else: - prompt_img = '用我使用的语言来回答我的问题或要求' - prompt = ''.join([prompt, prompt_img]) + prompt_img = '只回复以下问题的结果，用我使用的语言' + prompt = ':'.join([prompt_img,prompt]) url = "https://open.bigmodel.cn/api/paas/v4/chat/completions" @@ -175,29 +202,78 @@ class ZhipuaiApi_Img: 'Content-Type': 'application/json' } - data_img = {"model": "glm-4v", 'max_tokens': f'{max_tokens}','temperature': f'{temperature}',"messages": [{"role": "user","content": [{"type": "text", "text": f"{prompt}"},{"type": "image_url","image_url": {"url": f"{img_base64}"}}]}]} + data_img = {"model": "glm-4v", 'max_tokens': f'{max_tokens}', 'temperature': f'{temperature}', "messages": [ + {"role": "user", "content": [{"type": "text", "text": f"{prompt}"}, + {"type": "image_url", "image_url": {"url": f"{img_base64}"}}]}]} data_img = json.dumps(data_img) - response = requests.post(url=url, headers=header_img, data=data_img) - img_content = response.json() - - content_img = img_content.get("choices", 'Default Value')[0] - content_img = content_img.get("message", 'Default Value') - content_img = content_img.get("content", 'Default Value') - prompt = re.sub(r'\n+', '', content_img) - + content_img = img_content["choices"][0]["message"]["content"] + # print(content_img) + # content_img = str(content_img).strip('"') + prompt = content_img.replace("\n", "") os.remove(img_path) + return (prompt,) +class ZhipuaiApi_Character: - return (prompt,) + def __init__(self): + self.api_key = get_zpai_api_key() + if self.api_key is not None: + client_key = self.api_key + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "bot_name": ("STRING", {"default": "苏梦远"}), + "user_name": ("STRING", {"default": "陆星辰"}), + "user_info": ("STRING", {"default": "我是陆星辰，是一个男性，是一位知名编剧，苏梦远是我在剧中创作的虚拟二次元角色。" + "她是一位白发傲娇猫娘，并视我为好朋友。", "multiline": True}), + "bot_info": ("STRING", { + "default": "苏梦远，本名苏喵喵，是一位虚拟的二次元白发傲娇猫娘，她爱吃小鱼干，喜欢温暖的地方，喜欢打滚，讨厌游泳。", + "multiline": True}), + "assistant_prompt": ( + "STRING", {"default": "陆星辰，我希望下一个剧本，我能变身美少女战士，喵。", "multiline": True}), + "user_prompt": ("STRING", {"default": "苏梦远，你本来就是美少女，但不是战士而已。", "multiline": True}) + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + FUNCTION = "zhipuai_character_api" + CATEGORY = "ChatGlm_Api" + + def zhipuai_character_api(self, bot_name,user_name,user_info,bot_info,assistant_prompt, user_prompt): + if not self.api_key: + raise ValueError("API key is required") + url = "https://open.bigmodel.cn/api/paas/v3/model-api/charglm-3/invoke" + exp_seconds = int(9) + client_key = str(get_zpai_api_key()) + token = generate_token(apikey_s=client_key, exp_second=exp_seconds) + header_txt = { + 'Authorization': f'Bearer {token}', + 'Content-Type': 'application/json' + } + data_char = {"model": "charglm-3", "meta":{ "bot_name": f"{bot_name}","user_name": f"{user_name}","user_info": f"{user_info}", "bot_info": f"{bot_info}"}, + "prompt": [{"role": "assistant", "content": f"{assistant_prompt}"}, {"role": "user", "content": f"{user_prompt}"}]} + data_char = json.dumps(data_char) + response = requests.post(url=url, headers=header_txt, data=data_char) + char_content = response.json() + # print(char_content) + prompt_char = char_content['data']["choices"][0]["content"] + assistant_char = str(prompt_char).strip('"') + assistant_char = assistant_char.replace("\\n", "") + return (assistant_char,) NODE_CLASS_MAPPINGS = { "ZhipuaiApi_Txt": ZhipuaiApi_Txt, - "ZhipuaiApi_img": ZhipuaiApi_Img + "ZhipuaiApi_img": ZhipuaiApi_Img, + "ZhipuaiApi_Character": ZhipuaiApi_Character } NODE_DISPLAY_NAME_MAPPINGS = { "ZhipuaiApi_Txt": "ZhipuaiApi_Txt", - "ZhipuaiApi_Img": "ZhipuaiApi_Img" + "ZhipuaiApi_Img": "ZhipuaiApi_Img", + "ZhipuaiApi_Character": "ZhipuaiApi_Character" }