From 0dd77f33ea67725c6915534b8d4e2a8c4e389896 Mon Sep 17 00:00:00 2001 From: yunxiang wei <121092793+yiwangsimple@users.noreply.github.com> Date: Fri, 26 Jul 2024 22:28:33 +0800 Subject: [PATCH] Add files via upload --- __init__.py | 31 +++---- nodes/SD3LongCaptioner_v2.py | 93 +++++++++++++++++++++ nodes/dwimage2.py | 96 +++++++++++++++++++++ nodes/file_based_chat.py | 158 +++++++++++++++++++++++++++++++++++ nodes/groqchat.py | 75 +++++++++++++++++ nodes/moonshot_chat_nodes.py | 120 ++++++++++++++++++++++++++ 6 files changed, 558 insertions(+), 15 deletions(-) create mode 100644 nodes/SD3LongCaptioner_v2.py create mode 100644 nodes/dwimage2.py create mode 100644 nodes/file_based_chat.py create mode 100644 nodes/groqchat.py create mode 100644 nodes/moonshot_chat_nodes.py diff --git a/__init__.py b/__init__.py index 4bce7ef..620c02f 100644 --- a/__init__.py +++ b/__init__.py @@ -1,29 +1,30 @@ -from .groqchat import NODE_CLASS_MAPPINGS as GROQ_CLASS_MAPPINGS -from .groqchat import NODE_DISPLAY_NAME_MAPPINGS as GROQ_DISPLAY_MAPPINGS -from .moonshot_chat_nodes import NODE_CLASS_MAPPINGS as MOONSHOT_CLASS_MAPPINGS -from .moonshot_chat_nodes import NODE_DISPLAY_NAME_MAPPINGS as MOONSHOT_DISPLAY_MAPPINGS -from .SD3LongCaptioner_v2 import NODE_CLASS_MAPPINGS as SD3_V2_CLASS_MAPPINGS -from .SD3LongCaptioner_v2 import NODE_DISPLAY_NAME_MAPPINGS as SD3_V2_DISPLAY_MAPPINGS -from .file_based_chat import NODE_CLASS_MAPPINGS as FILE_CHAT_CLASS_MAPPINGS -from .file_based_chat import NODE_DISPLAY_NAME_MAPPINGS as FILE_CHAT_DISPLAY_MAPPINGS -from .dwimage2 import NODE_CLASS_MAPPINGS as MOONDREAM2_V2_NODE_CLASS_MAPPINGS -from .dwimage2 import NODE_DISPLAY_NAME_MAPPINGS as MOONDREAM2_V2_NODE_DISPLAY_NAME_MAPPINGS - +from .nodes.groqchat import NODE_CLASS_MAPPINGS as GROQ_CLASS_MAPPINGS +from .nodes.groqchat import NODE_DISPLAY_NAME_MAPPINGS as GROQ_DISPLAY_MAPPINGS +from .nodes.moonshot_chat_nodes import NODE_CLASS_MAPPINGS as MOONSHOT_CLASS_MAPPINGS +from .nodes.moonshot_chat_nodes import NODE_DISPLAY_NAME_MAPPINGS as MOONSHOT_DISPLAY_MAPPINGS +from .nodes.SD3LongCaptioner_v2 import NODE_CLASS_MAPPINGS as SD3_V2_CLASS_MAPPINGS +from .nodes.SD3LongCaptioner_v2 import NODE_DISPLAY_NAME_MAPPINGS as SD3_V2_DISPLAY_MAPPINGS +from .nodes.file_based_chat import NODE_CLASS_MAPPINGS as FILE_CHAT_CLASS_MAPPINGS +from .nodes.file_based_chat import NODE_DISPLAY_NAME_MAPPINGS as FILE_CHAT_DISPLAY_MAPPINGS +from .nodes.dwimage2 import NODE_CLASS_MAPPINGS as MOONDREAM2_V2_NODE_CLASS_MAPPINGS +from .nodes.dwimage2 import NODE_DISPLAY_NAME_MAPPINGS as MOONDREAM2_V2_NODE_DISPLAY_NAME_MAPPINGS NODE_CLASS_MAPPINGS = { **GROQ_CLASS_MAPPINGS, **MOONSHOT_CLASS_MAPPINGS, **SD3_V2_CLASS_MAPPINGS, - **FILE_CHAT_CLASS_MAPPINGS - **MOONDREAM2_V2_NODE_CLASS_MAPPINGS + **FILE_CHAT_CLASS_MAPPINGS, + **MOONDREAM2_V2_NODE_CLASS_MAPPINGS, + } NODE_DISPLAY_NAME_MAPPINGS = { **GROQ_DISPLAY_MAPPINGS, **MOONSHOT_DISPLAY_MAPPINGS, **SD3_V2_DISPLAY_MAPPINGS, - **FILE_CHAT_DISPLAY_MAPPINGS - **MOONDREAM2_V2_NODE_DISPLAY_NAME_MAPPINGS + **FILE_CHAT_DISPLAY_MAPPINGS, + **MOONDREAM2_V2_NODE_DISPLAY_NAME_MAPPINGS, + } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/nodes/SD3LongCaptioner_v2.py b/nodes/SD3LongCaptioner_v2.py new file mode 100644 index 0000000..e06939f --- /dev/null +++ b/nodes/SD3LongCaptioner_v2.py @@ -0,0 +1,93 @@ +import os +from pathlib import Path +import torch +from PIL import Image +from torchvision.transforms import ToPILImage +from transformers import AutoProcessor, AutoModelForVision2Seq +import folder_paths +from huggingface_hub import snapshot_download +import gc + +# 定义模型文件存储目录 +files_for_sd3_long_captioner_v2 = Path(os.path.join(folder_paths.models_dir, "LLavacheckpoints", "files_for_sd3_long_captioner_v2")) +files_for_sd3_long_captioner_v2.mkdir(parents=True, exist_ok=True) + +class SD3LongCaptionerV2: + def __init__(self): + self.model_id = "gokaygokay/sd3-long-captioner-v2" + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.model = None + self.processor = None + + def load_model(self): + if self.model is None or self.processor is None: + self.model_path = snapshot_download(self.model_id, + local_dir=files_for_sd3_long_captioner_v2, + force_download=False, + local_files_only=False, + local_dir_use_symlinks="auto") + self.model = AutoModelForVision2Seq.from_pretrained(self.model_path).to(self.device).eval() + self.processor = AutoProcessor.from_pretrained(self.model_path) + + def clear_memory(self): + # 清理内存 + if self.model is not None: + del self.model + self.model = None + if self.processor is not None: + del self.processor + self.processor = None + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"multiline": False, "default": "Describe in detail what's in this image."}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("caption",) + FUNCTION = "generate_caption" + CATEGORY = "🌙DW/SD3LongCaptionerV2" + + def generate_caption(self, image, prompt): + self.load_model() # 确保模型已加载 + + # 将图像张量转换为PIL图像 + pil_image = ToPILImage()(image[0].permute(2, 0, 1)) + + # 处理图像和提示 + model_inputs = self.processor(text=prompt, images=pil_image, return_tensors="pt").to(self.device) + + # 生成描述 + with torch.inference_mode(): + generation = self.model.generate( + **model_inputs, + repetition_penalty=1.05, + max_new_tokens=512, + do_sample=False + ) + + # 解码生成的文本 + decoded = self.processor.decode(generation[0], skip_special_tokens=True) + + # 移除开头的提示词（如果存在） + decoded = decoded.replace(prompt, "", 1).strip() + + # 清理内存 + self.clear_memory() + + return (decoded,) + +NODE_CLASS_MAPPINGS = { + "SD3LongCaptionerV2": SD3LongCaptionerV2 +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SD3LongCaptionerV2": "SD3 Long Captioner V2" +} diff --git a/nodes/dwimage2.py b/nodes/dwimage2.py new file mode 100644 index 0000000..16ac3d8 --- /dev/null +++ b/nodes/dwimage2.py @@ -0,0 +1,96 @@ +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from PIL import Image +from pathlib import Path +import folder_paths +import os +import shutil +import re +from torchvision.transforms import ToPILImage + +# 定义模型存储目录 +models_dir = Path(folder_paths.base_path) / "models" +llava_checkpoints_dir = models_dir / "LLavacheckpoints" +files_for_moondream2 = llava_checkpoints_dir / "files_for_moondream2" +files_for_moondream2.mkdir(parents=True, exist_ok=True) + +class Moondream2Predictor: + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super(Moondream2Predictor, cls).__new__(cls) + cls._instance.model = None + cls._instance.tokenizer = None + return cls._instance + + def load_model(self): + if self.model is None: + self.device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"使用设备: {self.device}") + print("加载模型中...") + self.model = AutoModelForCausalLM.from_pretrained(files_for_moondream2, trust_remote_code=True).to(self.device) + self.tokenizer = AutoTokenizer.from_pretrained(files_for_moondream2) + print("模型加载成功") + + def generate_predictions(self, image_path, question): + self.load_model() # 确保模型已加载 + try: + image_input = Image.open(image_path).convert("RGB") + enc_image = self.model.encode_image(image_input) + generated_text = self.model.answer_question(enc_image, question, self.tokenizer) + return generated_text + finally: + if os.path.exists(image_path): + os.remove(image_path) + + def clear_memory(self): + if self.model: + del self.model + del self.tokenizer + self.model = None + self.tokenizer = None + torch.cuda.empty_cache() + print("模型已卸载，内存已清理") + +class Moondream2model: + def __init__(self): + self.predictor = Moondream2Predictor() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "text_input": ("STRING", {"multiline": True, "default": ""}), + }, + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "moondream2_generate_predictions" + CATEGORY = "🌙DW/dwimage2" + + def moondream2_generate_predictions(self, image, text_input): + try: + pil_image = ToPILImage()(image[0].permute(2, 0, 1)) + temp_path = files_for_moondream2 / "temp_image.png" + pil_image.save(temp_path) + + response = self.predictor.generate_predictions(temp_path, text_input) + response = ' '.join(response.strip().split()) + + # 使用正则表达式移除开头的特定短语 + response = re.sub(r'^(The image contains|The image contains|This picture contains|This picture contains|The image contains|The image contains|This image contains|This image contains)\s*', '', response, flags=re.IGNORECASE) + + return (response,) + finally: + # 确保在函数结束时卸载模型 + self.predictor.clear_memory() + +NODE_CLASS_MAPPINGS = { + "dwimage2": Moondream2model +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "dwimage2": "DW Image2 Chat" +} \ No newline at end of file diff --git a/nodes/file_based_chat.py b/nodes/file_based_chat.py new file mode 100644 index 0000000..8ba1ab3 --- /dev/null +++ b/nodes/file_based_chat.py @@ -0,0 +1,158 @@ +import os +import json +from typing import List, Dict, Any, Optional +from pathlib import Path +import httpx +from openai import OpenAI + +class FileBasedChatNode: + def __init__(self): + self.config = self.load_config() + self.client = OpenAI( + base_url="https://api.moonshot.cn/v1", + api_key=self.config.get("MOONSHOT_API_KEY"), + ) + self.file_messages = [] + self.conversation_history = [] + self.cache_tag = None + + def load_config(self): + config_path = Path(__file__).parent / "config.json" + if config_path.exists(): + with open(config_path, "r") as f: + return json.load(f) + else: + raise FileNotFoundError(f"Config file not found: {config_path}") + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "file_paths": ("STRING", {"multiline": True}), + "user_input": ("STRING", {"multiline": True}), + "use_cache": ("BOOLEAN", {"default": True}), + "cache_tag": ("STRING", {"default": "file_cache"}), + "ttl": ("INT", {"default": 300, "min": 60, "max": 3600}) + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "chat" + CATEGORY = "🌙DW/fileBaseChat" + + def upload_files(self, files: List[str]) -> List[Dict[str, Any]]: + messages = [] + for file in files: + try: + with open(file, 'rb') as f: + response = httpx.post( + f"{self.client.base_url}/files", + headers={"Authorization": f"Bearer {self.client.api_key}"}, + files={"file": f}, + data={"purpose": "assistants"} + ) + response.raise_for_status() + file_id = response.json()['id'] + + # 获取文件内容 + content_response = httpx.get( + f"{self.client.base_url}/files/{file_id}/content", + headers={"Authorization": f"Bearer {self.client.api_key}"} + ) + content_response.raise_for_status() + file_content = content_response.text + + messages.append({ + "role": "system", + "content": file_content, + }) + except Exception as e: + print(f"Error uploading file {file}: {e}") + continue + return messages + + def chat(self, file_paths, user_input, use_cache, cache_tag, ttl): + if not self.file_messages: + files = file_paths.split('\n') + self.file_messages = self.upload_files(files=files) + + messages = [ + *self.file_messages, + { + "role": "system", + "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。" + "你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，" + "种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。", + }, + *self.conversation_history, + { + "role": "user", + "content": user_input, + } + ] + + try: + if use_cache: + cache_response = self.get_cache(cache_tag) + if cache_response: + return (cache_response,) + + completion = self.client.chat.completions.create( + model="moonshot-v1-8k", + messages=messages, + ) + + assistant_response = completion.choices[0].message.content + self.conversation_history.extend([ + {"role": "user", "content": user_input}, + {"role": "assistant", "content": assistant_response} + ]) + + if use_cache: + self.set_cache(cache_tag, assistant_response, ttl) + + return (assistant_response,) + + except Exception as e: + error_message = f"An error occurred: {str(e)}" + print(error_message) + return (error_message,) + + def get_cache(self, cache_tag): + try: + response = httpx.get( + f"{self.client.base_url}/caching", + headers={"Authorization": f"Bearer {self.client.api_key}"}, + params={"tags": cache_tag} + ) + response.raise_for_status() + cache_data = response.json() + if cache_data and cache_data.get('data'): + return cache_data['data'][0]['content'] + except Exception as e: + print(f"Error getting cache: {e}") + return None + + def set_cache(self, cache_tag, content, ttl): + try: + response = httpx.post( + f"{self.client.base_url}/caching", + headers={"Authorization": f"Bearer {self.client.api_key}"}, + json={ + "tags": [cache_tag], + "content": content, + "ttl": ttl + } + ) + response.raise_for_status() + except Exception as e: + print(f"Error setting cache: {e}") + +# 在 ComfyUI 的节点定义文件中注册这个新节点 +NODE_CLASS_MAPPINGS = { + "FileBasedChatNode": FileBasedChatNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FileBasedChatNode": "File-based Chat" +} \ No newline at end of file diff --git a/nodes/groqchat.py b/nodes/groqchat.py new file mode 100644 index 0000000..dda012e --- /dev/null +++ b/nodes/groqchat.py @@ -0,0 +1,75 @@ +import os +import json +from groq import Groq + +class GroqChatNode: + def __init__(self): + self.client = None + self.load_api_key() + + def load_api_key(self): + config_path = os.path.join(os.path.dirname(__file__), 'config.json') + try: + with open(config_path, 'r') as config_file: + config = json.load(config_file) + api_key = config.get('GROQ_API_KEY') + if api_key: + self.client = Groq(api_key=api_key) + else: + print("Error: GROQ_API_KEY not found in config.json") + except FileNotFoundError: + print(f"Error: config.json not found at {config_path}") + except json.JSONDecodeError: + print(f"Error: Invalid JSON in config.json at {config_path}") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": (["gemma-7b-it", "llama3-70b-8192", "mixtral-8x7b-32768", "llama3-8b-8192", "llama-3.1-8b-versatile", "llama-3.1-70b-versatile"],), + "prompt": ("STRING", {"multiline": True}), + "max_tokens": ("INT", {"default": 1000, "min": 1, "max": 32768}), + "temperature": ("FLOAT", {"default": 0.7, "min": 0, "max": 2, "step": 0.1}), + "top_p": ("FLOAT", {"default": 1.0, "min": 0, "max": 1, "step": 0.1}), + }, + "optional": { + "system_message": ("STRING", {"multiline": True}), + "presence_penalty": ("FLOAT", {"default": 0, "min": -2, "max": 2, "step": 0.1}), + "frequency_penalty": ("FLOAT", {"default": 0, "min": -2, "max": 2, "step": 0.1}), + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "generate_text" + CATEGORY = "🌙DW/groqchat" + + def generate_text(self, model, prompt, max_tokens, temperature, top_p, system_message="", presence_penalty=0, frequency_penalty=0): + if not self.client: + return ("Error: GROQ_API_KEY not set or invalid. Please check your config.json file.",) + + messages = [] + if system_message: + messages.append({"role": "system", "content": system_message}) + messages.append({"role": "user", "content": prompt}) + + try: + chat_completion = self.client.chat.completions.create( + model=model, + messages=messages, + max_tokens=max_tokens, + temperature=temperature, + top_p=top_p, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty + ) + return (chat_completion.choices[0].message.content,) + except Exception as e: + return (f"Error: {str(e)}",) + +NODE_CLASS_MAPPINGS = { + "GroqChatNode": GroqChatNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "GroqChatNode": "Groq Chat" +} \ No newline at end of file diff --git a/nodes/moonshot_chat_nodes.py b/nodes/moonshot_chat_nodes.py new file mode 100644 index 0000000..679c6bb --- /dev/null +++ b/nodes/moonshot_chat_nodes.py @@ -0,0 +1,120 @@ +import os +import json +from openai import OpenAI + +def get_moonshot_api_key(): + try: + config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'config.json') + with open(config_path, 'r') as f: + config = json.load(f) + api_key = config["MOONSHOT_API_KEY"] + except: + print("出错啦 Error: API key is required") + return "" + return api_key + +class MoonshotChatBaseNode: + def __init__(self): + self.client = None + self.api_key = get_moonshot_api_key() + if self.api_key: + self.client = OpenAI(api_key=self.api_key, base_url="https://api.moonshot.cn/v1") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"multiline": True}), + "model": (["moonshot-v1-8k", "moonshot-v1-32k", "moonshot-v1-128k"],), + "temperature": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 2.0, "step": 0.1}), + "max_tokens": ("INT", {"default": 1000, "min": 1, "max": 128000}), + }, + "optional": { + "system_message": ("STRING", {"multiline": True}), + } + } + + RETURN_TYPES = ("STRING",) + +class MoonshotSingleChatNode(MoonshotChatBaseNode): + FUNCTION = "generate_single_response" + RETURN_NAMES = ("response",) + + def generate_single_response(self, prompt, model, temperature, max_tokens, system_message=""): + if not self.client: + return ("Error: MOONSHOT_API_KEY not set or invalid. Please check your config.json file.",) + + messages = [] + if system_message: + messages.append({"role": "system", "content": system_message}) + messages.append({"role": "user", "content": prompt}) + + try: + completion = self.client.chat.completions.create( + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + ) + return (completion.choices[0].message.content,) + except Exception as e: + return (f"Error: {str(e)}",) + +class MoonshotMultiChatNode(MoonshotChatBaseNode): + FUNCTION = "generate_chat" + RETURN_NAMES = ("chat_history",) + CATEGORY = "🌙DW/moonshotChat" + + def __init__(self): + super().__init__() + self.conversation_history = [] + + @classmethod + def INPUT_TYPES(cls): + input_types = super().INPUT_TYPES() + input_types["optional"]["reset_conversation"] = ("BOOLEAN", {"default": False}) + return input_types + + def generate_chat(self, prompt, model, temperature, max_tokens, system_message="", reset_conversation=False): + if not self.client: + return ("Error: MOONSHOT_API_KEY not set or invalid. Please check your config.json file.",) + + if reset_conversation: + self.conversation_history = [] + + if not self.conversation_history and system_message: + self.conversation_history.append({"role": "system", "content": system_message}) + + self.conversation_history.append({"role": "user", "content": prompt}) + + try: + completion = self.client.chat.completions.create( + model=model, + messages=self.conversation_history, + temperature=temperature, + max_tokens=max_tokens, + ) + response = completion.choices[0].message.content + self.conversation_history.append({"role": "assistant", "content": response}) + chat_history = self.format_chat_history() + return (chat_history,) + except Exception as e: + return (f"Error: {str(e)}",) + + def format_chat_history(self): + formatted_history = [] + for message in self.conversation_history: + formatted_message = f"{message['role']}: {message['content']}" + formatted_history.append(formatted_message) + formatted_history.append("-" * 40) # 添加分隔线 + return "\n".join(formatted_history) + +NODE_CLASS_MAPPINGS = { + "MoonshotSingleChatNode": MoonshotSingleChatNode, + "MoonshotMultiChatNode": MoonshotMultiChatNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MoonshotSingleChatNode": "🌙Moonshot Single Chat", + "MoonshotMultiChatNode": "🌙Moonshot Multi Chat" +} \ No newline at end of file