diff --git a/__init__.py b/__init__.py
index b8dcde3..494238a 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,37 +1,44 @@
# 从各节点文件导入映射
-from .ollama_node import (
+from .node.ollama_node import (
NODE_CLASS_MAPPINGS as OLLAMA_NODE_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as OLLAMA_DISPLAY_NAMES
)
-from .online_api import (
+from .node.online_api import (
NODE_CLASS_MAPPINGS as LLM_ONLINE_API,
NODE_DISPLAY_NAME_MAPPINGS as LLM_ONLINE_API_DISPLAY_NAMES
)
-from .string_utils import (
+from .node.string_utils import (
NODE_CLASS_MAPPINGS as STRING_UTILS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as STRING_UTILS_DISPLAY_NAMES
)
-from .qiniu_uploader import (
+from .node.qiniu_uploader import (
NODE_CLASS_MAPPINGS as QINIU_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as QINIU_DISPLAY_NAMES
)
+from .node.img_utils import (
+ NODE_CLASS_MAPPINGS as IMG_UTIL_MAPPINGS,
+ NODE_DISPLAY_NAME_MAPPINGS as IMG_UTIL_DISPLAY_NAMES
+)
+
# 合并所有映射
NODE_CLASS_MAPPINGS = {
**OLLAMA_NODE_MAPPINGS,
**LLM_ONLINE_API,
**STRING_UTILS_MAPPINGS,
- **QINIU_MAPPINGS
+ **QINIU_MAPPINGS,
+ **IMG_UTIL_MAPPINGS
}
NODE_DISPLAY_NAME_MAPPINGS = {
**OLLAMA_DISPLAY_NAMES,
**LLM_ONLINE_API_DISPLAY_NAMES,
**STRING_UTILS_DISPLAY_NAMES,
- **QINIU_DISPLAY_NAMES
+ **QINIU_DISPLAY_NAMES,
+ **IMG_UTIL_DISPLAY_NAMES
}
# 定义公开接口
diff --git a/node/img_utils.py b/node/img_utils.py
new file mode 100644
index 0000000..60d1d69
--- /dev/null
+++ b/node/img_utils.py
@@ -0,0 +1,54 @@
+from PIL import Image, ImageOps
+from PIL.PngImagePlugin import PngInfo
+import requests
+import torch
+import numpy as np
+
+def pil2tensor(images: Image.Image | list[Image.Image]) -> torch.Tensor:
+ """Converts a PIL Image or a list of PIL Images to a tensor."""
+
+ def single_pil2tensor(image: Image.Image) -> torch.Tensor:
+ np_image = np.array(image).astype(np.float32) / 255.0
+ if np_image.ndim == 2: # Grayscale
+ return torch.from_numpy(np_image).unsqueeze(0) # (1, H, W)
+ else: # RGB or RGBA
+ return torch.from_numpy(np_image).unsqueeze(0) # (1, H, W, C)
+
+ if isinstance(images, Image.Image):
+ return single_pil2tensor(images)
+ else:
+ return torch.cat([single_pil2tensor(img) for img in images], dim=0)
+
+class LoadImageFromUrl:
+ """Load an image from the given URL"""
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "url": (
+ "STRING",
+ {
+ "default": "https://upload.wikimedia.org/wikipedia/commons/thumb/a/a7/Example.jpg/800px-Example.jpg"
+ },
+ ),
+ }
+ }
+
+ RETURN_TYPES = ("IMAGE",)
+ FUNCTION = "load"
+ CATEGORY = "云服务/加载网络图片"
+
+ def load(self, url):
+ # get the image from the url
+ image = Image.open(requests.get(url, stream=True).raw)
+ image = ImageOps.exif_transpose(image)
+ return (pil2tensor(image),)
+
+NODE_CLASS_MAPPINGS = {
+ "LoadImageFromUrl": LoadImageFromUrl
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "LoadImageFromUrl": "☁️ 加载图片"
+}
\ No newline at end of file
diff --git a/ollama_node.py b/node/ollama_node.py
similarity index 97%
rename from ollama_node.py
rename to node/ollama_node.py
index 3545155..d54583b 100644
--- a/ollama_node.py
+++ b/node/ollama_node.py
@@ -1,231 +1,231 @@
-import requests
-import json
-import logging
-import re
-import os
-import threading
-from threading import Lock
-from typing import Optional, List, Union, Dict, Any
-
-# ----------------------------
-# 日志系统配置
-# ----------------------------
-desktop_path = os.path.join(os.path.expanduser("~"), "Desktop")
-log_path = os.path.join(desktop_path, "ollama_comfyui.log")
-
-# 强制删除旧日志
-try:
- if os.path.exists(log_path):
- os.remove(log_path)
-except Exception as e:
- print(f"旧日志删除失败: {str(e)}")
-
-# 日志初始化
-logging.basicConfig(
- filename=log_path,
- level=logging.DEBUG,
- format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
- filemode='w'
-)
-logger = logging.getLogger(__name__)
-logger.addHandler(logging.StreamHandler()) # 控制台输出
-
-class ComfyUI_LLM_Ollama:
- """
- Ollama LLM集成节点
-
- 特性：
- - 支持流式响应生成
- - 上下文记忆管理
- - 多线程安全访问
- - 动态模型列表加载
- - 详细的日志记录
- """
- # 在类定义顶部添加
- WEB_DIRECTORY = "./js"
-
- # 类级共享状态
- _conn_lock = Lock()
- _api_lock = Lock()
- _connection_checked = False
- _connection_status = False
- _available_models = ["llama3", "deepseek-r1:7b"] # 默认值
-
- # 类级日志器 ✅ 修正点1
- logger = logging.getLogger("ComfyUI-Ollama")
-
- @classmethod
- def INPUT_TYPES(cls):
- """动态生成输入配置"""
- # 在类加载时尝试获取可用模型
- if not cls._connection_checked:
- cls._check_connection()
-
- return {
- "required": {
- "prompt": ("STRING", {
- "multiline": True,
- "default": "请用简洁的语言回答...",
- "dynamicPrompts": False,
- "lineCount": 4 # ✅ 初始显示4行高度
- }),
- "model": (cls._available_models, {"default": "deepseek-r1:7b"}),
- "temperature": ("FLOAT", {
- "default": 0.7,
- "min": 0.0,
- "max": 2.0,
- "step": 0.1,
- "display": "slider"
- }),
- "max_tokens": ("INT", {
- "default": 1024,
- "min": 1,
- "max": 4096,
- "step": 64,
- "display": "number"
- }),
- "hide_thoughts": ("BOOLEAN", {"default": False}),
- },
- "optional": {
- "context": ("STRING", {"forceInput": True}),
- "system_message": ("STRING", {
- "multiline": True,
- "default": "你是有帮助的AI助手",
- "lazy": True,
- "lineCount": 3 # ✅ 初始显示3行高度
- }),
- "stop_sequences": ("STRING", {
- "default": "",
- "lazy": True
- }),
- }
- }
-
- RETURN_TYPES = ("STRING", "STRING")
- RETURN_NAMES = ("response", "context")
- FUNCTION = "generate"
- CATEGORY = "LLM"
- OUTPUT_NODE = False
-
- def __init__(self):
- self.ollama_url = "http://localhost:11434"
- self.headers = {"Content-Type": "application/json"}
- self.timeout = 120
- self._init_logging() # ✅ 修正点2：实例初始化日志配置
-
- def _init_logging(self):
- """初始化日志系统"""
- self.logger = logging.getLogger("ComfyUI-Ollama")
- self.logger.propagate = False
- if not self.logger.handlers:
- handler = logging.StreamHandler()
- formatter = logging.Formatter('[%(name)s] %(levelname)s - %(message)s')
- handler.setFormatter(formatter)
- self.logger.addHandler(handler)
- self.logger.setLevel(logging.INFO)
-
- @classmethod
- def _check_connection(cls):
- """类方法使用类级日志器 ✅ 修正点3"""
- with cls._conn_lock:
- if not cls._connection_checked:
- try:
- # 使用类级日志器
- cls.logger.info(f"🛠 正在检查Ollama连接...")
-
- response = requests.get(
- "http://localhost:11434/api/tags",
- timeout=10
- )
- if response.status_code == 200:
- models = response.json().get("models", [])
- cls._available_models = [m["name"] for m in models]
- cls._connection_status = True
- cls.logger.info(f"✅ 可用模型: {cls._available_models}")
- else:
- cls.logger.warning(f"连接失败，状态码：{response.status_code}")
- except Exception as e:
- cls.logger.error(f"连接异常：{str(e)}")
- finally:
- cls._connection_checked = True
-
- def _build_payload(self, **kwargs):
- """构造API请求负载"""
- stop_sequences = [s.strip() for s in kwargs['stop_sequences'].split(',')] if kwargs['stop_sequences'] else []
-
- return {
- "model": kwargs['model'],
- "prompt": kwargs['prompt'],
- "system": kwargs.get('system_message', '你是有帮助的AI助手'),
- "context": self._parse_context(kwargs.get('context')),
- "options": {
- "temperature": kwargs['temperature'],
- "num_predict": kwargs['max_tokens'],
- "stop": stop_sequences,
- }
- }
-
- def _parse_context(self, context: Optional[str]) -> List:
- """解析上下文数据"""
- try:
- return json.loads(context) if context else []
- except json.JSONDecodeError:
- self.logger.warning("上下文解析失败，使用空上下文")
- return []
-
- def _clean_response(self, text: str, hide_thoughts: bool) -> str:
- """响应后处理"""
- if hide_thoughts:
- return re.sub(r'.*?', '', text, flags=re.DOTALL).strip()
- return text.strip()
-
- def generate(self, **kwargs):
- """主执行方法"""
- if not self._connection_status:
- return ("Ollama服务不可用，请检查", "")
-
- with self._api_lock:
- try:
- payload = self._build_payload(**kwargs)
- self.logger.debug(f"请求参数：{json.dumps(payload, indent=2)}")
-
- response_text = ""
- context = []
-
- with requests.post(
- f"{self.ollama_url}/api/generate",
- json=payload,
- headers=self.headers,
- stream=True,
- timeout=self.timeout
- ) as response:
- response.raise_for_status()
-
- for line in response.iter_lines():
- if line:
- data = json.loads(line.decode('utf-8'))
- response_text += data.get("response", "")
- if data.get("done"):
- context = data.get("context", [])
-
- cleaned_response = self._clean_response(response_text, kwargs['hide_thoughts'])
- self.logger.info(f"📥 响应长度: {len(cleaned_response)}字符")
- return (cleaned_response, json.dumps(context))
-
- except requests.RequestException as e:
- error_msg = f"API请求失败: {str(e)}"
- self.logger.error(error_msg)
- return (error_msg, "")
- except Exception as e:
- error_msg = f"处理错误: {str(e)}"
- self.logger.exception(error_msg)
- return (error_msg, "")
-
-# 节点注册
-NODE_CLASS_MAPPINGS = {
- "ComfyUI_LLM_Ollama": ComfyUI_LLM_Ollama
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
- "ComfyUI_LLM_Ollama": "🤖 Ollama LLM"
+import requests
+import json
+import logging
+import re
+import os
+import threading
+from threading import Lock
+from typing import Optional, List, Union, Dict, Any
+
+# ----------------------------
+# 日志系统配置
+# ----------------------------
+desktop_path = os.path.join(os.path.expanduser("~"), "Desktop")
+log_path = os.path.join(desktop_path, "ollama_comfyui.log")
+
+# 强制删除旧日志
+try:
+ if os.path.exists(log_path):
+ os.remove(log_path)
+except Exception as e:
+ print(f"旧日志删除失败: {str(e)}")
+
+# 日志初始化
+logging.basicConfig(
+ filename=log_path,
+ level=logging.DEBUG,
+ format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+ filemode='w'
+)
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.StreamHandler()) # 控制台输出
+
+class ComfyUI_LLM_Ollama:
+ """
+ Ollama LLM集成节点
+
+ 特性：
+ - 支持流式响应生成
+ - 上下文记忆管理
+ - 多线程安全访问
+ - 动态模型列表加载
+ - 详细的日志记录
+ """
+ # 在类定义顶部添加
+ WEB_DIRECTORY = "./js"
+
+ # 类级共享状态
+ _conn_lock = Lock()
+ _api_lock = Lock()
+ _connection_checked = False
+ _connection_status = False
+ _available_models = ["llama3", "deepseek-r1:7b"] # 默认值
+
+ # 类级日志器 ✅ 修正点1
+ logger = logging.getLogger("ComfyUI-Ollama")
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ """动态生成输入配置"""
+ # 在类加载时尝试获取可用模型
+ if not cls._connection_checked:
+ cls._check_connection()
+
+ return {
+ "required": {
+ "prompt": ("STRING", {
+ "multiline": True,
+ "default": "请用简洁的语言回答...",
+ "dynamicPrompts": False,
+ "lineCount": 4 # ✅ 初始显示4行高度
+ }),
+ "model": (cls._available_models, {"default": "deepseek-r1:7b"}),
+ "temperature": ("FLOAT", {
+ "default": 0.7,
+ "min": 0.0,
+ "max": 2.0,
+ "step": 0.1,
+ "display": "slider"
+ }),
+ "max_tokens": ("INT", {
+ "default": 1024,
+ "min": 1,
+ "max": 4096,
+ "step": 64,
+ "display": "number"
+ }),
+ "hide_thoughts": ("BOOLEAN", {"default": False}),
+ },
+ "optional": {
+ "context": ("STRING", {"forceInput": True}),
+ "system_message": ("STRING", {
+ "multiline": True,
+ "default": "你是有帮助的AI助手",
+ "lazy": True,
+ "lineCount": 3 # ✅ 初始显示3行高度
+ }),
+ "stop_sequences": ("STRING", {
+ "default": "",
+ "lazy": True
+ }),
+ }
+ }
+
+ RETURN_TYPES = ("STRING", "STRING")
+ RETURN_NAMES = ("response", "context")
+ FUNCTION = "generate"
+ CATEGORY = "LLM"
+ OUTPUT_NODE = False
+
+ def __init__(self):
+ self.ollama_url = "http://localhost:11434"
+ self.headers = {"Content-Type": "application/json"}
+ self.timeout = 120
+ self._init_logging() # ✅ 修正点2：实例初始化日志配置
+
+ def _init_logging(self):
+ """初始化日志系统"""
+ self.logger = logging.getLogger("ComfyUI-Ollama")
+ self.logger.propagate = False
+ if not self.logger.handlers:
+ handler = logging.StreamHandler()
+ formatter = logging.Formatter('[%(name)s] %(levelname)s - %(message)s')
+ handler.setFormatter(formatter)
+ self.logger.addHandler(handler)
+ self.logger.setLevel(logging.INFO)
+
+ @classmethod
+ def _check_connection(cls):
+ """类方法使用类级日志器 ✅ 修正点3"""
+ with cls._conn_lock:
+ if not cls._connection_checked:
+ try:
+ # 使用类级日志器
+ cls.logger.info(f"🛠 正在检查Ollama连接...")
+
+ response = requests.get(
+ "http://localhost:11434/api/tags",
+ timeout=10
+ )
+ if response.status_code == 200:
+ models = response.json().get("models", [])
+ cls._available_models = [m["name"] for m in models]
+ cls._connection_status = True
+ cls.logger.info(f"✅ 可用模型: {cls._available_models}")
+ else:
+ cls.logger.warning(f"连接失败，状态码：{response.status_code}")
+ except Exception as e:
+ cls.logger.error(f"连接异常：{str(e)}")
+ finally:
+ cls._connection_checked = True
+
+ def _build_payload(self, **kwargs):
+ """构造API请求负载"""
+ stop_sequences = [s.strip() for s in kwargs['stop_sequences'].split(',')] if kwargs['stop_sequences'] else []
+
+ return {
+ "model": kwargs['model'],
+ "prompt": kwargs['prompt'],
+ "system": kwargs.get('system_message', '你是有帮助的AI助手'),
+ "context": self._parse_context(kwargs.get('context')),
+ "options": {
+ "temperature": kwargs['temperature'],
+ "num_predict": kwargs['max_tokens'],
+ "stop": stop_sequences,
+ }
+ }
+
+ def _parse_context(self, context: Optional[str]) -> List:
+ """解析上下文数据"""
+ try:
+ return json.loads(context) if context else []
+ except json.JSONDecodeError:
+ self.logger.warning("上下文解析失败，使用空上下文")
+ return []
+
+ def _clean_response(self, text: str, hide_thoughts: bool) -> str:
+ """响应后处理"""
+ if hide_thoughts:
+ return re.sub(r'.*?', '', text, flags=re.DOTALL).strip()
+ return text.strip()
+
+ def generate(self, **kwargs):
+ """主执行方法"""
+ if not self._connection_status:
+ return ("Ollama服务不可用，请检查", "")
+
+ with self._api_lock:
+ try:
+ payload = self._build_payload(**kwargs)
+ self.logger.debug(f"请求参数：{json.dumps(payload, indent=2)}")
+
+ response_text = ""
+ context = []
+
+ with requests.post(
+ f"{self.ollama_url}/api/generate",
+ json=payload,
+ headers=self.headers,
+ stream=True,
+ timeout=self.timeout
+ ) as response:
+ response.raise_for_status()
+
+ for line in response.iter_lines():
+ if line:
+ data = json.loads(line.decode('utf-8'))
+ response_text += data.get("response", "")
+ if data.get("done"):
+ context = data.get("context", [])
+
+ cleaned_response = self._clean_response(response_text, kwargs['hide_thoughts'])
+ self.logger.info(f"📥 响应长度: {len(cleaned_response)}字符")
+ return (cleaned_response, json.dumps(context))
+
+ except requests.RequestException as e:
+ error_msg = f"API请求失败: {str(e)}"
+ self.logger.error(error_msg)
+ return (error_msg, "")
+ except Exception as e:
+ error_msg = f"处理错误: {str(e)}"
+ self.logger.exception(error_msg)
+ return (error_msg, "")
+
+# 节点注册
+NODE_CLASS_MAPPINGS = {
+ "ComfyUI_LLM_Ollama": ComfyUI_LLM_Ollama
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "ComfyUI_LLM_Ollama": "🤖 Ollama LLM"
}
\ No newline at end of file
diff --git a/online_api.py b/node/online_api.py
similarity index 100%
rename from online_api.py
rename to node/online_api.py
diff --git a/qiniu_uploader.py b/node/qiniu_uploader.py
similarity index 100%
rename from qiniu_uploader.py
rename to node/qiniu_uploader.py
diff --git a/string_utils.py b/node/string_utils.py
similarity index 100%
rename from string_utils.py
rename to node/string_utils.py
diff --git a/qiniu_config.json b/qiniu_config.json
new file mode 100644
index 0000000..33ff87f
--- /dev/null
+++ b/qiniu_config.json
@@ -0,0 +1,8 @@
+{
+ "access_key": "GRvkqEByqR6JZ6tSrIYPQUmxVnwuJMcQLWM4fQOk",
+ "secret_key": "4ryqDs__yUkx4grSVQpMg_ZaLbRrR6N55VcQZCrD",
+ "bucket_name": "dyaigc",
+ "input_dir": "input",
+ "output_dir": "output",
+ "domain": "swqqsa5wv.hb-bkt.clouddn.com"
+}
\ No newline at end of file
diff --git a/qiniu_config.json.example b/qiniu_config.json.example
new file mode 100644
index 0000000..b1c2a7e
--- /dev/null
+++ b/qiniu_config.json.example
@@ -0,0 +1,8 @@
+{
+ "access_key": "",
+ "secret_key": "",
+ "bucket_name": "dyaigc",
+ "input_dir": "input",
+ "output_dir": "output",
+ "domain": "swqqsa5wv.hb-bkt.clouddn.com"
+}
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 4da8c99..e0b9ddd 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,5 @@ qiniu
pillow
numpy
openai
+requests
+torch