Add files via upload

__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"]

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"
+}

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"
+}

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"
+}

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"
+}

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"
+}