From 71511da1b23eb6dffa466592bb18c2179d610e90 Mon Sep 17 00:00:00 2001 From: ZHO-ZHO-ZHO <140084057+ZHO-ZHO-ZHO@users.noreply.github.com> Date: Tue, 30 Jan 2024 17:07:05 +0800 Subject: [PATCH] V1.0 --- QwenVL_API_Node.py | 282 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 282 insertions(+) create mode 100644 QwenVL_API_Node.py diff --git a/QwenVL_API_Node.py b/QwenVL_API_Node.py new file mode 100644 index 0000000..d7c4def --- /dev/null +++ b/QwenVL_API_Node.py @@ -0,0 +1,282 @@ +import os +import io +import json +import requests +import torch +import dashscope +from dashscope import MultiModalConversation +from io import BytesIO +from PIL import Image, ImageChops +from datetime import datetime +import tempfile +import random +import platform +import hashlib + +p = os.path.dirname(os.path.realpath(__file__)) + +def get_qwenvl_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["QWENVL_API_KEY"] + except: + print("出错啦 Error: API key is required") + return "" + return api_key + + +class QWenVL_API_S_Zho: + + def __init__(self): + self.api_key = get_qwenvl_api_key() + if self.api_key is not None: + dashscope.api_key=self.api_key + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"default": "Describe this image", "multiline": True}), + "model_name": (["qwen-vl-plus", "qwen-vl-max"],), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + FUNCTION = "qwen_vl_generation" + + CATEGORY = "Zho模块组/💫QWenVL" + + def tensor_to_image(self, tensor): + # 确保张量是在CPU上 + tensor = tensor.cpu() + + # 将张量数据转换为0-255范围并转换为整数 + # 这里假设张量已经是H x W x C格式 + image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() + + # 创建PIL图像 + image = Image.fromarray(image_np, mode='RGB') + return image + + def qwen_vl_generation(self, image, prompt, model_name, seed): + if not self.api_key: + raise ValueError("API key is required") + + if image == None: + raise ValueError("qwen_vl needs a image") + else: + # 转换图像 + pil_image = self.tensor_to_image(image) + + # 生成临时文件路径 + temp_directory = tempfile.gettempdir() + unique_suffix = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5)) + filename = f"image{unique_suffix}.png" + temp_image_path = os.path.join(temp_directory, filename) + #temp_image_url = f"file://{temp_image_path}" + + # 根据操作系统选择正确的文件URL格式 + if platform.system() == 'Windows': + temp_image_url = f"file://{temp_image_path}" + else: + temp_image_url = f"file:///{temp_image_path}" + + temp_image_url = temp_image_url.replace('\\', '/') + + # 保存图像到临时路径 + pil_image.save(temp_image_path) + + + messages = [ + { + "role": "user", + "content": [ + {"image": temp_image_url}, + {"text": prompt} + ] + } + ] + + #print("temp_image_url:", temp_image_url) + #print("prompt:", prompt) + + torch.manual_seed(seed) + + response = dashscope.MultiModalConversation.call(model=model_name, messages=messages, seed=seed) + #print(response) + + response_json = response + if 'output' in response_json and 'choices' in response_json['output']: + choices = response_json['output']['choices'] + if choices and 'message' in choices[0]: + message_content = choices[0]['message']['content'] + if message_content and 'text' in message_content[0]: + text_output = message_content[0]['text'] + #print(text_output) + else: + print("No text content found.") + else: + print("No message found in the first choice.") + else: + print("No choices found in the output.") + + os.remove(temp_image_path) + #print("remove : done!" ) + + return (text_output, ) + + +class QWenVL_API_S_Multi_Zho: + + def __init__(self): + self.api_key = get_qwenvl_api_key() + self.messages = [] # 初始化对话历史为空 + self.last_image_hash = None + if self.api_key is not None: + dashscope.api_key=self.api_key + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"default": "Describe this image", "multiline": True}), + "model_name": (["qwen-vl-plus", "qwen-vl-max"],), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + FUNCTION = "qwen_vl_generation" + + CATEGORY = "Zho模块组/💫QWenVL" + + def tensor_to_image(self, tensor): + # 确保张量是在CPU上 + tensor = tensor.cpu() + + # 将张量数据转换为0-255范围并转换为整数 + # 这里假设张量已经是H x W x C格式 + image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() + + # 创建PIL图像 + image = Image.fromarray(image_np, mode='RGB') + return image + + def format_qwchat_history(self): + formatted_history = [] + for message in self.messages: + role = message['role'] + contents = message['content'] + for content in contents: + if 'text' in content: + text = content['text'] + formatted_message = f"{role}: {text}" + formatted_history.append(formatted_message) + formatted_history.append("-" * 40) # 添加分隔线 + return "\n".join(formatted_history) + + def get_image_hash(self, pil_image): + # 将图像转换为字节 + image_bytes = pil_image.tobytes() + # 使用哈希函数计算哈希值 + return hashlib.md5(image_bytes).hexdigest() + + def qwen_vl_generation(self, image, prompt, model_name, seed): + if not self.api_key: + raise ValueError("API key is required") + + if image == None: + raise ValueError("qwen_vl needs a image") + else: + # 转换图像 + pil_image = self.tensor_to_image(image) + + # 在当前文件目录下创建 "qw" 文件夹（如果不存在） + qw_folder = os.path.join(p, "qw") + os.makedirs(qw_folder, exist_ok=True) + + # 获取当前图像的哈希值 + current_image_hash = self.get_image_hash(pil_image) + + # 构建文件名 + local_image_filename = f"image_{current_image_hash}.png" + local_image_path = os.path.join(qw_folder, local_image_filename) + + # 根据操作系统选择正确的文件URL格式 + if platform.system() == 'Windows': + local_image_url = f"file://{local_image_path}" + else: + local_image_url = f"file:///{local_image_path}" + + # 保证路径中的反斜杠被替换为正斜杠 + local_image_url = local_image_url.replace('\\', '/') + + # 如果当前图像与上次的不同 + if current_image_hash != self.last_image_hash: + pil_image.save(local_image_path) + # 更新last_image_hash + self.last_image_hash = current_image_hash + print(f"Image saved as {local_image_filename}") + else: + print("Image not saved as it is identical to the last one.") + + self.messages.append({ + "role": "user", + "content": [ + {"image": local_image_url}, + {"text": prompt} + ] + }) + + #print("local_image_url:", local_image_url) + #print("prompt:", prompt) + + torch.manual_seed(seed) + + response = dashscope.MultiModalConversation.call(model=model_name, messages=self.messages, seed=seed) + #print(response) + + # 更新对话历史 + if response and response.output and response.output.choices: + choice = response.output.choices[0] + if choice and choice.message: + self.messages.append({'role': choice.message.role, 'content': choice.message.content}) + + response_json = response + if 'output' in response_json and 'choices' in response_json['output']: + choices = response_json['output']['choices'] + if choices and 'message' in choices[0]: + message_content = choices[0]['message']['content'] + if message_content and 'text' in message_content[0]: + text_output = message_content[0]['text'] + #print(text_output) + else: + print("No text content found.") + else: + print("No message found in the first choice.") + else: + print("No choices found in the output.") + + + # 获取格式化的对话历史 + chat_history = self.format_qwchat_history() + + return (chat_history, ) + + +NODE_CLASS_MAPPINGS = { + "QWenVL_API_S_Zho": QWenVL_API_S_Zho, + "QWenVL_API_S_Multi_Zho": QWenVL_API_S_Multi_Zho, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "QWenVL_API_S_Zho": "㊙️QWenVL_Zho", + "QWenVL_API_S_Multi_Zho": "㊙️QWenVL_Chat_Zho", +}