commit aa10c1fb99efe588fd666419b9146f4781b6e7fe Author: dotaevo-linux-do Date: Tue Nov 11 14:18:37 2025 +0800 new file: README.md new file: __init__.py new file: image2prompt_node.py new file: requirements.txt diff --git a/README.md b/README.md new file mode 100644 index 0000000..ca7929d --- /dev/null +++ b/README.md @@ -0,0 +1,80 @@ +# ComfyUI NudeJS Image2Prompt 节点 + +这是一个 ComfyUI 自定义节点,可以将图片转换为文本提示。 + +## 功能特性 + +- **图片预处理**:自动调整图片尺寸(最大宽度 1000px,最大高度 800px) +- **Base64 编码**:将图片转换为 base64 格式 +- **API 调用**:通过 POST 请求调用外部 API +- **智能重试**:支持配置重试次数,应对网络波动 +- **双模式输出**:支持中性文本和提示词文本两种输出模式 + +## 安装方法 + +### 方法一:复制到 ComfyUI 目录 + +1. 将整个 `ComfyUI-NudeJS-Image2Prompt` 文件夹复制到 ComfyUI 的 `custom_nodes` 目录中: + ``` + ComfyUI/custom_nodes/ComfyUI-NudeJS-Image2Prompt/ + ``` + +2. 重启 ComfyUI + +### 方法二:符号链接 + +在 Linux/macOS 系统中: +```bash +ln -s /path/to/ComfyUI-NudeJS-Image2Prompt /path/to/ComfyUI/custom_nodes/ +``` + +## 使用方法 + +在 ComfyUI 工作流中添加 "NudeJS/Image2Prompt" 节点: + +### 输入参数 + +- **images** (必需): 输入图片 +- **APIUrl** (必需): API 接口地址,例如:`https://api.example.com` +- **APIKey** (必需): API 鉴权密钥 +- **RetryCount** (可选): 请求失败后的重试次数,默认值:3 +- **Neutral** (可选): 是否使用中性文本模式,默认值:true + +### 输出 + +- **text**: 转换后的文本提示 + +## 工作流程示例 + +1. 添加图片加载节点(如 Load Image、Checkpoint Loader 等) +2. 添加 "Image2Prompt" 节点 +3. 将图片连接到 Image2Prompt 节点的 "images" 输入 +4. 配置 APIUrl 和 APIKey +5. 运行工作流 + +## 技术实现 + +- **图片处理**:使用 Pillow 进行图片缩放和格式转换 +- **网络请求**:使用 requests 库,支持重试和超时 +- **Base64 编码**:使用标准库 base64 进行编码 + +## 错误处理 + +- 网络请求失败时会自动重试(最多 RetryCount 次) +- 失败时会在控制台输出错误信息 +- 失败时返回 `None` + +## 依赖 + +- ComfyUI +- Python 3.7+ +- requests +- Pillow (PIL) + +## 许可证 + +MIT License + +## 贡献 + +欢迎提交 Issue 和 Pull Request! diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..cd9a7a9 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .image2prompt_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/image2prompt_node.py b/image2prompt_node.py new file mode 100644 index 0000000..f75eeaa --- /dev/null +++ b/image2prompt_node.py @@ -0,0 +1,247 @@ +import base64 +import io +from typing import Dict, Any, List, Tuple, Union + +import torch +from PIL import Image +import requests +from requests.adapters import HTTPAdapter +from requests.packages.urllib3.util.retry import Retry + + +class Image2Prompt: + """ + ComfyUI 自定义节点:将图片转换为文本提示 + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + return { + "required": { + "images": ("IMAGE",), + "APIUrl": ("STRING", {"default": "", "multiline": False}), + "APIKey": ("STRING", {"default": "", "multiline": False}), + "RetryCount": ("INT", {"default": 3, "min": 0, "max": 100}), + "Neutral": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + FUNCTION = "execute" + CATEGORY = "NudeJS" + + def __init__(self): + self.session = requests.Session() + # 配置重试策略 + retry_strategy = Retry( + total=3, + backoff_factor=1, + status_forcelist=[429, 500, 502, 503, 504], + ) + adapter = HTTPAdapter(max_retries=retry_strategy) + self.session.mount("http://", adapter) + self.session.mount("https://", adapter) + + def execute( + self, + images: torch.Tensor, + APIUrl: str, + APIKey: str, + RetryCount: int, + Neutral: bool, + ) -> Tuple[Union[str, None]]: + """ + 执行图片转提示词的主逻辑 + + Args: + images: 输入的图片张量 + APIUrl: API 接口地址 + APIKey: API 密钥 + RetryCount: 重试次数 + Neutral: 是否使用中性文本 + + Returns: + 转换后的文本提示 + """ + if len(images) == 0: + return (None,) + + # 获取第一张图片(ComfyUI 通常批量处理,这里取第一张) + img_tensor = images[0] + + # 转换为 PIL Image + pil_image = self._tensor_to_pil(img_tensor) + + # 检查并缩放图片 + resized_image = self._resize_image(pil_image) + + # 转换为 base64 + img_base64 = self._image_to_base64(resized_image) + + # 发送请求获取提示词 + prompt_text = self._call_api( + APIUrl=APIUrl, + APIKey=APIKey, + imgBase64=img_base64, + retry_count=RetryCount, + neutral=Neutral, + ) + + return (prompt_text,) + + def _tensor_to_pil(self, img_tensor: torch.Tensor) -> Image.Image: + """ + 将 ComfyUI 的张量转换为 PIL Image + + Args: + img_tensor: ComfyUI 中的图片张量 (B, H, W, C) 或 (H, W, C) + + Returns: + PIL Image + """ + # 确保张量在 CPU 上 + if img_tensor.is_cuda: + img_tensor = img_tensor.cpu() + + # 处理不同的张量形状 + if len(img_tensor.shape) == 4: + # 批量图片,取第一张 + img_tensor = img_tensor[0] + elif len(img_tensor.shape) == 3: + # 单张图片 + pass + else: + raise ValueError(f"Unexpected tensor shape: {img_tensor.shape}") + + # 转换为 PIL Image + # ComfyUI 的图片张量通常值范围是 [0, 1],形状是 [H, W, C] + if img_tensor.shape[-1] in [3, 4]: # RGB 或 RGBA + array = (img_tensor.numpy() * 255).astype("uint8") + pil_image = Image.fromarray(array, mode="RGB") + else: + raise ValueError(f"Unexpected number of channels: {img_tensor.shape[-1]}") + + return pil_image + + def _resize_image(self, image: Image.Image) -> Image.Image: + """ + 缩放图片,确保宽度不超过 1000,高度不超过 800 + + Args: + image: 输入的 PIL 图片 + + Returns: + 缩放后的图片 + """ + width, height = image.size + new_width = width + new_height = height + + # 检查宽度 + if width > 1000: + new_width = 1000 + new_height = int((height * 1000) / width) + + # 检查高度 + if new_height > 800: + new_height = 800 + new_width = int((new_width * 800) / new_height) + + # 如果需要缩放,则调整大小 + if new_width != width or new_height != height: + image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) + + return image + + def _image_to_base64(self, image: Image.Image) -> str: + """ + 将 PIL 图片转换为 base64 字符串(不包含 data:image/png;base64, 前缀) + + Args: + image: 输入的 PIL 图片 + + Returns: + base64 编码的字符串 + """ + # 转换为 RGB 模式(如果需要) + if image.mode != "RGB": + image = image.convert("RGB") + + # 保存到内存中的字节流 + buffer = io.BytesIO() + image.save(buffer, format="PNG") + img_bytes = buffer.getvalue() + + # 编码为 base64 + img_base64 = base64.b64encode(img_bytes).decode("utf-8") + + return img_base64 + + def _call_api( + self, + APIUrl: str, + APIKey: str, + imgBase64: str, + retry_count: int, + neutral: bool, + ) -> Union[str, None]: + """ + 调用 API 接口获取提示词 + + Args: + APIUrl: API 地址 + APIKey: API 密钥 + imgBase64: base64 编码的图片 + retry_count: 最大重试次数 + neutral: 是否使用中性文本 + + Returns: + 提示词文本,失败返回 None + """ + url = f"{APIUrl}/generate" + headers = {"authorization": APIKey} + data = {"base64": imgBase64} + + max_retries = max(0, retry_count) + + for attempt in range(max_retries + 1): + try: + response = self.session.post(url, json=data, headers=headers, timeout=30) + response.raise_for_status() + + result = response.json() + + if result.get("code") == 200: + if neutral: + return result.get("neutralText") + else: + return result.get("promptText") + else: + # API 返回错误码,记录日志 + print(f"API request failed with code: {result.get('code')}, message: {result.get('message', 'unknown error')}") + + except requests.exceptions.RequestException as e: + print(f"Attempt {attempt + 1} failed: {str(e)}") + + except Exception as e: + print(f"Unexpected error on attempt {attempt + 1}: {str(e)}") + + # 如果不是最后一次尝试,等待一下再重试 + if attempt < max_retries: + import time + time.sleep(1) + + # 所有重试都失败了 + print(f"Failed after {max_retries} retries") + return None + + +# 注册节点 +NODE_CLASS_MAPPINGS = { + "Image2Prompt": Image2Prompt, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Image2Prompt": "Image2Prompt", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..4b26e68 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +requests>=2.25.0 +Pillow>=8.0.0