commit 9ab209b2b791f558fa9f8e28b7a48525686dd9e2 Author: o1key <951565127@qq.com> Date: Fri Apr 3 16:18:45 2026 +0800 feat: sync latest local version as authoritative codebase Complete rewrite/sync of comfyui_o1key custom nodes. Treat this commit as the current canonical version. Co-Authored-By: Claude Sonnet 4.5 diff --git a/.cursorrules b/.cursorrules new file mode 100644 index 0000000..d6e3250 --- /dev/null +++ b/.cursorrules @@ -0,0 +1,756 @@ +# Comfyui_o1key 开发指南 + +## 对话原则 +始终使用中文进行对话。 + +## 编码规范 ⚠️ 重要 + +### 文件编码要求 +- **所有文本文件必须使用 UTF-8 编码（无 BOM）** +- **行结束符使用 LF（Unix 风格），Windows 批处理文件除外（CRLF）** +- 项目已配置 `.gitattributes` 和 `.editorconfig` 来自动处理编码 + +### 编辑器配置 +确保编辑器设置： +- 文件编码：UTF-8（无 BOM） +- 行结束符：LF +- 自动插入文件末尾空行：开启 + +## Git 提交规范 + +### Commit Message 规范 +- **所有 commit message 必须使用英文**，避免中文编码问题 +- 使用 Conventional Commits 格式：`: ` + +### 常用类型 +- `feat`: 新增功能 +- `fix`: 修复问题 +- `docs`: 文档更新 +- `refactor`: 代码重构 +- `style`: 代码格式调整 +- `test`: 测试相关 +- `chore`: 构建/工具配置 + +### 示例 +```bash +git commit -m "feat: add new model support" +git commit -m "fix: resolve image encoding issue" +git commit -m "docs: update README installation guide" +``` + +## 配置文件管理 + +### 基本原则 + +`.config` 文件包含敏感信息（API 密钥），已添加到 `.gitignore` 中，**不会被提交到版本控制**。 + +### 配置方式 + +用户通过以下方式创建本地配置： + +1. **快捷脚本**（推荐） + - Windows: 双击 `设置API密钥(win).bat` + - Linux/Mac: 运行 `./设置API密钥(mac).sh` + - 脚本会自动创建 `.config` 文件 + +2. **手动创建** + - 参考 `.config.example` 模板 + - 在插件根目录创建 `.config` 文件 + - 填写 API 密钥 + +3. **环境变量** + - 设置 `O1KEY_API_KEY` 环境变量 + - 无需创建配置文件 + +### 注意事项 + +- `.config` 文件仅存在于本地，不会被 Git 追踪 +- 开发者无需担心意外提交密钥的问题 +- 提交代码时会自动忽略 `.config` 文件 + +## 项目概述 + +这是一个 ComfyUI 自定义节点插件，通过 api.o1key.com 调用 AI 模型进行图像生成。 + +### 技术栈 +- Python 3.7+ +- ComfyUI 框架 +- aiohttp (异步 HTTP) +- Pillow (图像处理) +- PyTorch (张量处理) + +--- + +## 目录结构 + +``` +Comfyui_o1key/ +├── __init__.py # 节点注册入口 +├── models_config.py # 模型配置中心 ⭐ 管理所有支持的模型 +├── version.txt # 版本号文件 +├── update.bat # Windows 自动更新脚本 +├── update.sh # Linux/Mac 自动更新脚本 +├── nodes/ # 节点模块 +│ ├── __init__.py +│ ├── nano_banana_pro.py # NanoBananaPro 节点 +│ └── batch_nano_banana_pro.py # 批量节点 +├── utils/ # 工具模块 +│ ├── __init__.py +│ ├── image_utils.py # 图像转换工具 +│ ├── config.py # 配置管理 +│ └── update_checker.py # 更新检查器 +├── clients/ # API 客户端 +│ ├── __init__.py +│ ├── base_client.py # 客户端基类 +│ └── gemini_client.py # Gemini API 客户端 +├── .config.example # 配置文件模板 +├── requirements.txt # 依赖包 +└── README.md # 用户文档 +├── 设置API密钥(win).bat # Windows 配置脚本 +└── 设置API密钥(mac).sh # Mac/Linux 配置脚本 + +注：.config 文件在本地自动创建，不提交到版本控制 +``` + +--- + +## 模型管理系统 + +### 概述 + +所有 Nano Banana Pro 支持的模型都在 `models_config.py` 中统一管理。要添加新模型或临时关闭某个模型，只需编辑这个文件即可。 + +### 模型配置文件 (models_config.py) + +#### 配置结构 + +```python +GEMINI_MODELS = [ + { + "id": "gemini-3-pro-image-preview-url", + "description": "URL 模式,根据分辨率自动选择端点 (1K/2K/4K)", + "enabled": True, + "endpoint_type": "dynamic", + "endpoint": None # 动态端点，由代码根据分辨率选择 + }, + { + "id": "gemini-3-pro-image-preview", + "description": "标准模式,固定端点", + "enabled": True, + "endpoint_type": "standard", + "endpoint": "/v1beta/models/gemini-3-pro-image-preview:generateContent" + }, + # 更多模型... +] +``` + +#### 字段说明 + +| 字段 | 类型 | 必需 | 说明 | +|------|------|------|------| +| `id` | string | 是 | 模型标识符，用于 API 调用 | +| `description` | string | 是 | 模型描述，说明特点和适用场景 | +| `enabled` | boolean | 是 | 是否启用该模型（false 则在节点中隐藏） | +| `endpoint_type` | string | 是 | 端点类型："dynamic", "standard", "flatfee" | +| `endpoint` | string | 是 | API 端点路径（动态端点设为 None） | + +#### 端点类型说明 + +- **dynamic**: 根据分辨率动态选择端点（如 gemini-3-pro-image-preview-url） +- **standard**: 使用固定端点（如 gemini-3-pro-image-preview） +- **flatfee**: 固定费用模式端点（如 gemini-3-pro-image-preview-flatfee） + +### 常见操作 + +#### 1. 添加新模型 + +在 `GEMINI_MODELS` 列表末尾添加新模型： + +```python +GEMINI_MODELS = [ + # ... 现有模型 ... + { + "id": "gemini-新模型名称", + "description": "新模型的描述和特点", + "enabled": True, + "endpoint_type": "standard", # 根据实际情况选择 + "endpoint": "/v1beta/models/gemini-新模型名称:generateContent" # 配置端点 + } +] +``` + +**注意**： +- **固定端点模型**：直接在 `endpoint` 字段填写完整的端点路径即可，无需修改代码 +- **动态端点模型**：如果模型需要根据分辨率动态选择端点，设置 `endpoint_type: "dynamic"` 和 `endpoint: None`，并在 `gemini_client.py` 的 `get_endpoint()` 方法中添加对应逻辑 + +#### 2. 临时关闭模型 + +将模型的 `enabled` 字段设为 `False`： + +```python +{ + "id": "gemini-3-pro-image-preview-url", + "description": "URL 模式", + "enabled": False, # 临时关闭 + "endpoint_type": "dynamic" +} +``` + +关闭后，该模型将不会出现在 ComfyUI 节点的下拉列表中。 + +#### 3. 重新启用模型 + +将 `enabled` 改回 `True`： + +```python +{ + "id": "gemini-3-pro-image-preview-url", + "enabled": True, # 重新启用 + # ... +} +``` + +#### 4. 修改模型描述 + +直接编辑 `description` 字段： + +```python +{ + "id": "gemini-3-pro-image-preview", + "description": "标准模式,固定端点,适用于常规图像生成", # 更新描述 + # ... +} +``` + +### 工具函数 + +`models_config.py` 提供了一些工具函数，可在代码中使用： + +```python +from ..models_config import ( + get_enabled_models, # 获取启用的模型列表 + get_all_models, # 获取所有模型（包括禁用的） + get_model_config, # 获取指定模型的完整配置 + is_model_enabled, # 检查模型是否启用 + get_model_description, # 获取模型描述 + get_endpoint_type, # 获取端点类型 + get_model_endpoint # 获取模型端点 +) + +# 示例：获取启用的模型 +enabled = get_enabled_models() +# ['gemini-3-pro-image-preview-url', 'gemini-3-pro-image-preview', ...] + +# 示例：获取模型配置 +config = get_model_config("gemini-3-pro-image-preview-url") +# {'id': '...', 'description': '...', 'enabled': True, 'endpoint_type': 'dynamic', 'endpoint': None} + +# 示例：获取模型端点 +endpoint = get_model_endpoint("gemini-3-pro-image-preview") +# '/v1beta/models/gemini-3-pro-image-preview:generateContent' +``` + +### 节点集成 + +所有使用模型列表的节点都会自动从 `models_config.py` 加载： + +```python +from ..models_config import get_enabled_models + +class NanoBananaPro: + @classmethod + def INPUT_TYPES(cls): + # 自动从配置加载启用的模型 + enabled_models = get_enabled_models() + + return { + "required": { + "模型": (enabled_models, { + "default": enabled_models[0] + }), + # ... + } + } +``` + +### 配置验证 + +`models_config.py` 在加载时会自动验证配置： + +- 检查每个模型是否有必需字段（id, description, enabled, endpoint_type, endpoint） +- 检查 `endpoint_type` 是否合法（dynamic, standard, flatfee） +- 检查非动态端点模型必须配置有效的 `endpoint` +- 检查端点格式是否正确（应以 `/v1beta/models/` 开头） +- 确保至少有一个模型是启用的 + +如果配置不合法，会在终端打印警告信息。 + +### 最佳实践 + +1. **添加新模型前**： + - 确认模型使用 Gemini 原生接口格式 + - 确认端点规则（dynamic/standard/flatfee） + - 编写清晰的描述说明 + +2. **临时测试**： + - 关闭其他模型，只启用测试模型 + - 验证功能后再重新启用其他模型 + +3. **版本控制**： + - `models_config.py` 应纳入版本控制 + - 重大模型变更应记录在 `CHANGELOG.md` 中 + +4. **文档更新**： + - 添加新模型后，更新 `README.md` 中的模型列表 + - 如有特殊使用说明，添加到文档中 + +--- + +## 开发新节点流程 + +### 1. 创建节点文件 + +在 `nodes/` 目录下创建新的 Python 文件： + +```python +# nodes/my_new_node.py + +from typing import Optional, Tuple +import torch + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor +from ..clients.gemini_client import GeminiAPIClient + + +class MyNewNode: + """节点描述""" + + def __init__(self): + self.client = None + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"default": "", "multiline": True}), + # 更多参数... + }, + "optional": { + "images": ("IMAGE",) + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "execute" + CATEGORY = "image/generation" + + def execute(self, prompt: str, images: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor]: + # 实现逻辑 + pass +``` + +### 2. 注册节点 + +在 `nodes/__init__.py` 中添加导出： + +```python +from .my_new_node import MyNewNode +__all__ = ['NanoBananaPro', 'MyNewNode'] +``` + +在根 `__init__.py` 中注册： + +```python +from .nodes import NanoBananaPro, MyNewNode + +NODE_CLASS_MAPPINGS = { + "NanoBananaPro": NanoBananaPro, + "MyNewNode": MyNewNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NanoBananaPro": "Nano Banana Pro", + "MyNewNode": "My New Node" +} +``` + +### 3. 更新 CHANGELOG.md + +记录新增功能。 + +--- + +## ComfyUI 节点规范 + +### INPUT_TYPES 参数类型 + +| 类型 | 格式 | 示例 | +|------|------|------| +| 字符串 | `("STRING", {...})` | `("STRING", {"default": "", "multiline": True})` | +| 整数 | `("INT", {...})` | `("INT", {"default": 1, "min": 1, "max": 100})` | +| 浮点数 | `("FLOAT", {...})` | `("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1})` | +| 下拉选项 | `([...], {...})` | `(["option1", "option2"], {"default": "option1"})` | +| 图像 | `("IMAGE",)` | 放在 optional 中 | + +### 返回值规范 + +```python +RETURN_TYPES = ("IMAGE", "MASK", "STRING") # 类型元组 +RETURN_NAMES = ("images", "mask", "text") # 名称元组 +``` + +### 必须的类属性 + +```python +FUNCTION = "execute" # 执行函数名 +CATEGORY = "image/generation" # 节点分类路径 +``` + +--- + +## 工具模块使用 + +### 图像转换 (utils/image_utils.py) + +```python +from ..utils.image_utils import tensor_to_pil, pil_to_tensor + +# ComfyUI Tensor → PIL Image 列表 +pil_images = tensor_to_pil(tensor) # tensor: [B, H, W, C], range [0, 1] + +# PIL Image 列表 → ComfyUI Tensor +tensor = pil_to_tensor(pil_images) # 返回 [B, H, W, C], range [0, 1] + +# PIL → Base64 +from ..utils.image_utils import encode_image_to_base64 +b64_str = encode_image_to_base64(pil_image) + +# Base64 → PIL +from ..utils.image_utils import decode_base64_to_pil +pil_image = decode_base64_to_pil(b64_str) +``` + +### 配置管理 (utils/config.py) + +```python +from ..utils.config import get_api_key, get_api_key_or_raise, load_config, get_api_base_url + +# 获取 API 密钥（返回 None 如果未找到） +api_key = get_api_key("O1KEY_API_KEY") + +# 获取 API 密钥（抛出异常如果未找到） +api_key = get_api_key_or_raise("O1KEY_API_KEY") + +# 获取 API 基础 URL（统一配置） +base_url = get_api_base_url() # 默认: https://vip.o1key.com + +# 加载完整配置 +config = load_config() +``` + +### API 基础 URL 配置 + +所有 API 客户端都使用统一的基础 URL 配置，默认为 `https://vip.o1key.com`。 + +#### 配置优先级 + +1. **环境变量** `O1KEY_API_BASE_URL`（优先级最高） +2. **.config 文件**中的 `O1KEY_API_BASE_URL` 配置项 +3. **默认值** `https://vip.o1key.com`（在 `utils/config.py` 中定义） + +#### 修改 API 地址 + +**方法 1：修改默认值（影响所有用户）** + +编辑 `utils/config.py`： + +```python +# 修改此常量 +DEFAULT_API_BASE_URL = "https://your-api-domain.com" +``` + +**方法 2：使用环境变量（推荐，不影响代码）** + +在系统环境变量中设置： +```bash +# Windows +set O1KEY_API_BASE_URL=https://your-api-domain.com + +# Linux/Mac +export O1KEY_API_BASE_URL=https://your-api-domain.com +``` + +**方法 3：在 .config 文件中配置** + +在插件根目录的 `.config` 文件中添加： +``` +O1KEY_API_BASE_URL=https://your-api-domain.com +``` + +#### 使用示例 + +所有客户端会自动使用统一配置： + +```python +from ..utils.config import get_api_base_url + +# 获取当前配置的 API 地址 +base_url = get_api_base_url() +print(f"当前 API 地址: {base_url}") +``` + +--- + +## API 客户端使用 + +### 使用 GeminiAPIClient + +```python +from ..clients.gemini_client import GeminiAPIClient + +# 初始化（自动读取配置） +client = GeminiAPIClient() + +# 同步生成（用于 ComfyUI 节点） +images = client.generate_sync( + prompt="描述文字", + model="gemini-3-pro-image-preview-url", + resolution="2K", + aspect_ratio="1:1", + batch_size=1, + images=None, # 可选：输入图像列表 + progress_callback=None +) +``` + +### 创建新的 API 客户端 + +继承 `BaseAPIClient` 并实现抽象方法： + +```python +from ..clients.base_client import BaseAPIClient + +class MyAPIClient(BaseAPIClient): + def __init__(self): + super().__init__( + base_url="https://api.example.com", + api_key=get_api_key_or_raise("MY_API_KEY"), + max_request_size=20 * 1024 * 1024 + ) + + def get_endpoint(self, **kwargs) -> str: + return "/v1/generate" + + def build_request_body(self, **kwargs) -> dict: + return {"prompt": kwargs.get("prompt", "")} + + def parse_response(self, response: dict) -> Any: + return response.get("result") +``` + +--- + +## API 端点说明 + +### Gemini 模型端点 + +**gemini-3-pro-image-preview-url** (根据分辨率动态选择): +- 1K: `/v1beta/models/gemini-3-pro-image-preview-url:generateContent` +- 2K: `/v1beta/models/gemini-3-pro-image-preview-2k-url:generateContent` +- 4K: `/v1beta/models/gemini-3-pro-image-preview-4k-url:generateContent` + +**gemini-3-pro-image-preview** (固定端点): +- `/v1beta/models/gemini-3-pro-image-preview:generateContent` + +**gemini-3-pro-image-preview-flatfee** (固定端点): +- `/v1beta/models/gemini-3-pro-image-preview-flatfee:generateContent` + +### 请求格式 + +```json +{ + "contents": [{ + "role": "user", + "parts": [ + {"text": "提示词"}, + {"inline_data": {"mime_type": "image/png", "data": "base64..."}} + ] + }], + "generationConfig": { + "responseModalities": ["TEXT", "IMAGE"], + "imageConfig": { + "aspectRatio": "1:1", + "imageSize": "2K" + } + } +} +``` + +--- + +## 代码规范 + +### 命名约定 + +- 类名：PascalCase（如 `NanoBananaPro`） +- 函数/方法：snake_case（如 `tensor_to_pil`） +- 常量：UPPER_CASE（如 `API_BASE_URL`） +- 私有方法：前缀下划线（如 `_load_config`） + +### 类型注解 + +所有公开函数必须有类型注解： + +```python +def function_name(param1: str, param2: Optional[int] = None) -> List[Image.Image]: + pass +``` + +### 文档字符串 + +使用 Google 风格的 docstring： + +```python +def function_name(param1: str, param2: int) -> bool: + """ + 函数简短描述 + + Args: + param1: 参数1说明 + param2: 参数2说明 + + Returns: + 返回值说明 + + Raises: + ValueError: 异常情况说明 + + Example: + >>> result = function_name("test", 42) + >>> print(result) + True + """ + pass +``` + +### 错误处理 + +```python +try: + # 业务逻辑 + pass +except ValueError as e: + # 用户输入错误 + print(f"节点名: 输入错误 - {str(e)}") + raise +except RuntimeError as e: + # API 或网络错误 + print(f"节点名: API 错误 - {str(e)}") + raise +except Exception as e: + # 未知错误 + print(f"节点名: 未知错误 - {str(e)}") + raise +``` + +--- + +## 限制与约束 + +| 限制项 | 值 | 说明 | +|--------|-----|------| +| 请求体大小 | 20MB | 超过会报错 | +| 输入图像数量 | 14张 | 图生图模式限制 | +| 批次大小 | 1-1000 | 并发生成数量 | +| 支持的分辨率 | 1K/2K/4K | API 限制 | + +--- + +## 测试检查清单 + +新节点开发完成后，验证以下场景： + +- [ ] 文生图基础功能 +- [ ] 图生图功能（如支持） +- [ ] 不同分辨率（1K/2K/4K） +- [ ] 不同宽高比 +- [ ] 批量生成 +- [ ] 错误处理（无 API 密钥、网络错误等） +- [ ] 边界条件（最大图像数、最大批次） + +--- + +## 更新日志 + +修改代码后，更新 `CHANGELOG.md` 记录变更。 + +格式： +```markdown +## [版本号] - 日期 + +### Added +- 新增功能 + +### Changed +- 变更内容 + +### Fixed +- 修复问题 +``` + +--- + +## 版本发布流程 + +### 1. 准备发布 + +发布新版本前确认以下事项： + +- [ ] 所有功能测试通过 +- [ ] 更新 `CHANGELOG.md`（记录本次变更） +- [ ] 更新 `version.txt`（更新版本号） +- [ ] 更新 `README.md`（如有新功能需要说明） + +### 2. 版本号规范 + +遵循语义化版本 (Semantic Versioning)： + +- **主版本号** (Major): 重大架构变更、不兼容的 API 修改 +- **次版本号** (Minor): 新增功能、向后兼容 +- **修订号** (Patch): Bug 修复、小改进 + +示例：`v1.10.2` → Major.Minor.Patch + +### 3. 发布步骤 + +```bash +# 1. 更新版本号 +echo "v1.11.0" > version.txt + +# 2. 提交变更 +git add . +git commit -m "Release v1.11.0: 添加新功能描述" + +# 3. 创建标签 +git tag v1.11.0 + +# 4. 推送到远程 +git push origin main --tags +``` + +### 4. 用户更新 + +用户运行更新脚本即可获取最新版本： + +- **Windows**: 双击 `update.bat` +- **Linux/Mac**: 运行 `./update.sh` + +更新脚本会自动： +- 检查远程更新 +- 备份配置文件 +- 拉取最新代码 +- 更新依赖包 +- 显示更新日志 + +--- \ No newline at end of file diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..e5a183e --- /dev/null +++ b/.editorconfig @@ -0,0 +1,38 @@ +# EditorConfig 配置文件 +# https://editorconfig.org + +root = true + +# 默认配置 +[*] +charset = utf-8 +end_of_line = lf +insert_final_newline = true +trim_trailing_whitespace = true +indent_style = space +indent_size = 4 + +# Python 文件 +[*.py] +indent_size = 4 + +# Shell 脚本 +[*.sh] +indent_size = 4 + +# Windows 批处理文件 +[*.{bat,cmd}] +end_of_line = crlf +indent_size = 4 + +# Markdown 文件 +[*.md] +trim_trailing_whitespace = false + +# YAML 文件 +[*.{yml,yaml}] +indent_size = 2 + +# JSON 文件 +[*.json] +indent_size = 2 diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..84d1262 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,31 @@ +# 默认自动处理行结束符 +* text=auto + +# Python 文件使用 LF +*.py text eol=lf + +# Shell 脚本使用 LF +*.sh text eol=lf + +# Windows 批处理文件使用 CRLF +*.bat text eol=crlf +*.cmd text eol=crlf + +# 配置文件使用 LF +.config text eol=lf +.config.* text eol=lf + +# Markdown 文档使用 LF +*.md text eol=lf + +# 二进制文件 +*.png binary +*.jpg binary +*.jpeg binary +*.gif binary +*.ico binary +*.mov binary +*.mp4 binary +*.mp3 binary +*.zip binary +*.psd binary diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4b2985f --- /dev/null +++ b/.gitignore @@ -0,0 +1,27 @@ +# .config 文件包含敏感信息，不提交到版本控制 +# 用户可通过 setup_api_key.bat 自动创建本地配置 +.config + +# Python 缓存 +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python + +# 环境 +.env +.venv +env/ +venv/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..95db019 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,601 @@ +# Changelog + +本项目的所有重要变更都将记录在此文件中。 + +格式基于 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.0.0/)。 + +--- + +## [Unreleased] + +### Added ✨ +- **快捷配置脚本** + - 新增 `设置API密钥(win).bat` - Windows 一键配置工具 + - 新增 `设置API密钥(mac).sh` - Mac/Linux 一键配置工具 + - 自动生成 `.config` 配置文件 + - 交互式提示引导用户输入 API 密钥 + - 自动检测并提示覆盖已存在的配置文件 + - 彩色输出和友好的用户提示信息 +- **配置模板文件** + - 新增 `.config.example` 作为配置文件示例 + +### Changed +- **502 错误提示优化** (`clients/base_client.py`) + - 当 API 返回 502 时，弹框显示友好文案：「糟糕！请求到上游时遇到超时或过载！别担心，过会儿再次点击运行即可！」 + - 在 `request_async` 与 `request_get_async` 中均增加 502 专用分支 +- **配置管理策略** + - `.config` 文件现在完全忽略提交（添加到 `.gitignore`） + - 简化配置流程，用户通过快捷脚本自动创建本地配置 + - 移除配置文件安全检查机制（不再需要） +- **README 文档** + - 更新配置章节，添加快捷脚本使用说明 + - 调整配置方法优先级：快捷脚本 > 环境变量 > 手动配置 + - 简化安全提示说明 + +### Removed +- **安全检查工具**（不再需要） + - 删除 `check_config_safety.py` 配置安全检查脚本 + - 删除 `.git-hooks-install.bat` Git Hook 安装脚本 + - 彻底杜绝配置文件泄密风险 + +--- + +## [1.9.1] - 2026-02-05 + +### Security 🔐 +- **错误信息伪装 - 极致隐蔽性** + - 未配置情况下的错误信息伪装成商业授权提示 + - 前端弹框显示：`未授权！` + - 控制台日志显示：`请联系作者授权后方可使用！` + - 完全隐藏真实的配置方式（环境变量） + - 让外部用户以为这是需要购买授权的商业插件 + +### Changed +- **配置模块** (`utils/config.py`) + - 将未找到 API 密钥的错误消息改为简洁的"未授权！" + - 移除所有暴露配置方式的信息 + +- **节点异常处理** (所有节点) + - `NanoBananaPro` - 添加授权错误检测，特殊处理授权失败日志 + - `BatchNanoBananaPro` - 添加授权错误检测，特殊处理授权失败日志 + - `GoogleGemini` - 添加授权错误检测，特殊处理授权失败日志 + - 其他 ValueError 类型仍正常显示详细错误信息 + +### Benefits +- ✅ **极致隐蔽** - 完全看不出需要配置环境变量 +- ✅ **伪装商业** - 外部用户以为需要付费授权 +- ✅ **信息安全** - 不泄露任何配置相关信息 +- ✅ **功能完整** - 其他错误类型仍正常提示 +- ✅ **向后兼容** - 已配置用户完全不受影响 + +--- + +## [1.9.0] - 2026-02-05 + +### Security 🔐 +- **配置方式优化 - 提升隐蔽性和安全性** + - 移除所有配置相关文件（`setup_config.ps1`, `.config`, `.config.example`） + - 改为使用系统环境变量存储 API 密钥 + - 避免在项目中留下配置痕迹，提高隐蔽性 + - 配置方式对外部用户完全不可见 + +### Changed +- **配置管理模块重构** (`utils/config.py`) + - 调整读取优先级：环境变量优先 > .config 文件（向后兼容） + - 简化错误提示：仅提示设置环境变量，不再提及配置文件 + - 更新模块说明：从"处理 .config 文件"改为"处理环境变量" + +- **文档更新** + - `README.md` - 简化配置说明，仅保留环境变量设置方法 + - `批量节点使用指南.md` - 更新常见问题中的配置说明 + - `.gitignore` - 注释配置文件规则（已弃用） + +### Benefits +- ✅ **高隐蔽性** - 项目中无任何配置相关文件 +- ✅ **高安全性** - 敏感信息存储在系统级别，不在项目目录 +- ✅ **简化维护** - 一行命令创建/更新/删除配置 +- ✅ **多项目共享** - 环境变量可被其他项目复用 +- ✅ **向后兼容** - 仍支持从 .config 文件读取（如果存在） + +--- + +## [1.8.0] - 2026-02-04 + +### Added +- **集成 ComfyUI 原生进度条** 🎉 + - `NanoBananaPro` 节点现在支持 UI 绿色进度条显示 + - `BatchNanoBananaPro` 节点现在支持 UI 绿色进度条显示 + - 使用 `comfy.utils.ProgressBar` 实现实时进度更新 + - 进度条在节点上方显示，从 0% 平滑更新到 100% + - 兼容性检查：如果 ProgressBar 不可用，自动降级到终端进度显示 + +### Improved +- **用户体验提升** + - 生图过程中可视化进度反馈更直观 + - 节点运行时自动显示绿色边框（ComfyUI 原生） + - 保留详细的终端进度日志，方便调试 + - 批量提示词模式下进度条总数自动调整 + +### Technical +- 在 `nodes/nano_banana_pro.py` 中集成 ProgressBar + - 创建进度条实例：`ProgressBar(生图数量)` + - 在 `progress_callback` 中调用 `pbar.update(1)` + - 批量提示词模式重新创建进度条以匹配实际总数 +- 在 `nodes/batch_nano_banana_pro.py` 中集成 ProgressBar + - 将 `pbar` 参数传递给 `_process_batch_async` 方法 + - 每完成一个任务立即更新进度条 + - 支持大批量任务的实时进度显示 +- 添加 `PROGRESS_BAR_AVAILABLE` 标志进行兼容性检测 + +### Impact +- ✅ 所有图像生成节点现在都有 UI 进度条 +- ✅ 后续新增的视频生成节点可直接复用此实现 +- ✅ 不影响现有功能，完全向后兼容 + +--- + +## [1.7.0] - 2026-02-03 + +### Added +- **新增 Google Gemini 节点** (`GoogleGemini`) + - 用于调用 Gemini 3 Flash 模型进行多模态文本生成 + - **输入支持**： + - 提示词（必填）：用户提示词，支持多行 + - 系统指令（可选）：系统级指令，引导模型行为 + - 思考深度：不思考（默认）/ 高 + - 图片（可选）：支持 ComfyUI IMAGE 类型输入 + - 视频（可选）：支持 ComfyUI VIDEO 类型输入 + - **输出**：文本内容（STRING 类型） + - **支持的视频格式**：mp4, mpeg, mov, avi, flv, webm, wmv, 3gpp + - **端点映射**： + - 不思考 → `/v1beta/models/gemini-3-flash-preview-nothinking:generateContent` + - 高 → `/v1beta/models/gemini-3-flash-preview-high:generateContent` + +- **新增 Gemini Flash API 客户端** (`clients/gemini_flash_client.py`) + - 继承 `BaseAPIClient` 基类 + - 支持系统指令配置 + - 支持图片和视频的 base64 编码发送 + - 智能超时设置（视频请求 5 分钟，其他 3 分钟） + +### Technical +- 新增 `GeminiFlashClient` 类处理 Gemini Flash 模型 API 调用 +- 新增 `GoogleGemini` 节点类实现多模态文本生成 +- 支持自动检测视频 MIME 类型 +- 视频文件大小限制 20MB + +--- + +## [1.6.2] - 2026-02-02 + +### Added +- **模型端点配置化** (`models_config.py`) + - 在模型配置中新增 `endpoint` 字段，集中管理每个模型的 API 端点 + - 新增 `get_model_endpoint()` 工具函数，用于获取模型端点 + - 添加新模型时只需在配置文件中填写端点，无需修改代码 + +### Changed +- **简化端点获取逻辑** (`clients/gemini_client.py`) + - 重构 `get_endpoint()` 方法，从配置文件读取端点而非硬编码 + - 保留动态端点模型（gemini-3-pro-image-preview-url）的特殊处理逻辑 + - 其他模型自动从 `models_config.py` 读取端点配置 + +### Improved +- **增强配置验证** (`models_config.py`) + - 验证非动态端点模型必须配置有效的 `endpoint` + - 验证端点格式是否正确（应以 `/v1beta/models/` 开头） + - 更新必需字段列表，包含 `endpoint` 字段 + +- **更新开发文档** (`.cursorrules`) + - 更新模型配置示例，说明 `endpoint` 字段的使用方法 + - 更新添加新模型的指南，强调配置端点的方式 + - 更新工具函数列表，添加 `get_model_endpoint()` 说明 + - 更新配置验证规则说明 + +### Benefits +- 降低维护成本：添加新模型只需修改配置文件 +- 提高可读性：端点集中管理，一目了然 +- 减少错误：配置验证确保端点格式正确 +- 保持灵活性：动态端点模型仍可使用代码逻辑 + +--- + +## [1.6.1] - 2026-02-02 + +### Fixed +- **修复 gemini-3-pro-image-preview-flatfee 模型 504 错误** + - 暂时禁用 `gemini-3-pro-image-preview-flatfee` 模型（端点返回 504 Gateway Timeout） + - 更新模型描述标注"暂时不可用-504错误" + - 建议用户使用其他可用模型（如 gemini-3-pro-image-preview 或 gemini-3-pro-image-preview-url） + +### Improved +- **增强 HTTP 错误处理** (`clients/base_client.py`) + - 新增针对 504 Gateway Timeout 的友好错误提示 + - 说明原因：服务器响应超时或端点暂时不可用 + - 提供解决建议：尝试其他模型、稍后重试、降低分辨率等 + - 新增针对 503 Service Unavailable 的错误提示 + - 说明原因：模型服务过载或维护中 + - 提供解决建议：稍后重试或尝试其他模型 + - 新增针对 429 Too Many Requests 的错误提示 + - 说明原因：API 配额用尽或请求过于频繁 + - 提供解决建议：等待后重试或检查配额 + - 新增针对 404 Not Found 的错误提示 + - 说明原因：端点路径错误或模型不存在 + - 提供解决建议：检查模型名称或使用其他模型 + - 统一错误信息格式：错误类型 + 原因 + 建议 + - 同时优化 POST 和 GET 请求的错误处理 + +### Technical +- 在 `request_async()` 和 `request_get_async()` 中添加状态码判断逻辑 +- 提供更详细的错误诊断信息，帮助用户快速定位和解决问题 + +--- + +## [1.6.0] - 2026-02-02 + +### Added +- **模型管理系统** (`models_config.py`) + - 创建集中式模型配置文件，所有 Nano Banana Pro 支持的模型统一管理 + - 支持快速添加新模型、临时关闭或启用模型 + - 模型配置包含：模型ID、描述、启用状态、端点类型 + - 提供工具函数： + - `get_enabled_models()` - 获取启用的模型列表 + - `get_all_models()` - 获取所有模型（包括禁用的） + - `get_model_config()` - 获取指定模型的完整配置 + - `is_model_enabled()` - 检查模型是否启用 + - `get_model_description()` - 获取模型描述 + - `get_endpoint_type()` - 获取端点类型 + - 自动配置验证，确保配置完整性和合法性 + +### Changed +- **NanoBananaPro 节点重构** + - 移除硬编码的 `MODELS` 列表 + - 改为从 `models_config.py` 动态加载模型列表 + - 节点在 ComfyUI 中显示的模型列表自动同步配置文件 + +- **BatchNanoBananaPro 节点重构** + - 移除硬编码的 `MODELS` 列表 + - 改为从 `models_config.py` 动态加载模型列表 + - 保持与 NanoBananaPro 节点的模型列表一致性 + +### Improved +- **开发指南更新** (`.cursorrules`) + - 新增"模型管理系统"章节 + - 详细说明模型配置文件结构和字段含义 + - 提供添加新模型、关闭/启用模型、修改描述等常见操作指南 + - 新增工具函数使用示例和节点集成说明 + - 补充配置验证机制和最佳实践建议 + +- **目录结构更新** + - 在开发指南中添加 `models_config.py` 文件说明 + - 标记为模型配置中心 ⭐ + +### Benefits +- ✅ **集中管理** - 所有模型定义在一个文件，易于维护 +- ✅ **易于扩展** - 添加新模型只需在配置文件中添加一个字典 +- ✅ **快速开关** - 修改 `enabled` 字段即可临时关闭或启用模型 +- ✅ **文档化** - 每个模型都有 `description` 说明特点和适用场景 +- ✅ **类型安全** - Python 文件支持代码提示和类型检查 +- ✅ **自动同步** - 所有节点自动使用最新的模型配置 + +--- + +## [1.5.7] - 2026-02-02 + +### Performance 🚀 +- **极致性能优化：彻底解决异步阻塞问题** + - 将 `parse_response()` 改为 `parse_response_async()`，实现完全异步的图片下载 + - 使用 `aiohttp` 替代同步的 `requests.get()` 下载图片 + - **问题**：之前在异步事件循环中使用同步 HTTP 请求会阻塞整个事件循环 + - **影响**：虽然 API 请求是并发的，但图片下载变成了串行操作 + - **效果**：现在图片下载也是完全并发的，真正实现端到端的异步性能 + +- **性能提升幅度**： + - 使用 `gemini-3-pro-image-preview-url` 模型时提升最明显 + - 批量生成 4 张图时，从"串行下载 4 张"变为"并发下载 4 张" + - 预计性能提升 2-4 倍（取决于网络延迟和图片大小） + +### Changed +- `GeminiAPIClient.parse_response()` → `parse_response_async()` + - 新增 `session` 参数，用于复用 aiohttp 会话 + - 支持并发下载多个图片 URL + - 自动管理 session 生命周期 +- `GeminiAPIClient.generate_single_async()` 现在调用异步解析方法 +- 移除 `requests` 依赖，统一使用 `aiohttp` + +### Impact +- 所有节点自动受益： + - ✅ `NanoBananaPro` - 批量生成速度显著提升 + - ✅ `BatchNanoBananaPro` - 大批量任务性能大幅改善 + - ✅ 所有使用 `gemini-3-pro-image-preview-url` 模型的场景 + +--- + +## [1.5.6] - 2026-02-02 + +### Improved +- **优化批量生成实时进度显示** (`Nano Banana Pro`) + - 使用 `asyncio.as_completed` 替代 `asyncio.gather`，实现真正的实时进度 + - 每完成一个请求立即显示进度，而非等待所有请求完成后批量显示 + - 进度信息包含成功/失败状态： + - 成功：`✓ [1/4] 第 1 张生成成功` + - 失败：`✗ [2/4] 生成失败 - 错误原因` + - 最终统计显示成功和失败数量 + +### Fixed +- **修复批量生成失败信息丢失问题** + - 之前：失败的请求被静默忽略，用户不知道哪些请求失败 + - 现在：每个失败的请求都会显示错误原因，方便排查问题 + +### Technical +- 更新 `generate_batch_async()` 和 `generate_multi_prompts_async()` 方法 +- 进度回调签名变更：`(current, total)` → `(current, total, success, error_msg)` +- 错误信息自动截取第一行，避免过长输出 + +--- + +## [1.5.5] - 2026-02-02 + +### Added +- **增强 API 错误处理机制** + - 新增 `candidatesTokenCount = 0` 检测（最高优先级） + - 自动检测内容审核拒绝情况 + - 提供明确的拒绝原因和改进建议 + - 新增 `finishReason` 异常检测（次优先级） + - 支持检测 `PROHIBITED_CONTENT`（违禁内容） + - 支持检测 `SAFETY`（安全过滤器） + - 支持检测 `RECITATION`（版权问题） + - 支持检测 `MAX_TOKENS`（Token 超限） + - 针对每种错误类型提供具体的解决建议 + - 新增 API 文本响应拒绝检测 + - 当 API 返回文本而非图像时，自动提取拒绝说明 + - 直接展示 API 的拒绝理由给用户 + +### Improved +- **优化错误提示格式** (`clients/gemini_client.py`) + - 统一错误信息格式：错误类型 + 原因 + 建议 + - 所有错误以 `RuntimeError` 抛出，便于节点层面捕获 + - 提供清晰的多行格式化错误信息 + - 包含针对性的操作建议，帮助用户快速解决问题 + +### Technical +- 在 `GeminiAPIClient.parse_response()` 方法中实现三层错误检测 +- 错误检测按优先级顺序执行，确保最重要的问题优先报告 +- 保持向后兼容，不影响正常的图像生成流程 + +--- + +## [1.5.4] - 2026-02-01 + +### Added +- **新增余额查询功能** + - 在每次图像生成请求完成后自动查询并显示用户余额 + - 支持查询 API 名称和当前可用余额 + - 余额格式：`当前余额：$XX.XX | API：xxx` + - 查询失败时显示警告信息，不影响主流程 + - 适用于 `NanoBananaPro` 和 `BatchNanoBananaPro` 节点 + +### Changed +- **扩展 API 客户端功能** (`clients/base_client.py`) + - 新增 `request_get_async()` 方法支持 GET 请求 + - 扩展 `get_headers()` 方法支持 Bearer Token 认证 + - 兼容现有的 x-goog-api-key 认证方式 + +- **增强 Gemini 客户端** (`clients/gemini_client.py`) + - 新增 `query_balance_async()` 异步查询余额方法 + - 新增 `query_balance_sync()` 同步查询余额方法（用于节点） + - 新增 `format_balance_info()` 格式化余额信息方法 + - 余额转换公式：实际显示 = total_available / 500000 + +--- + +## [1.5.3] - 2026-02-01 + +### Changed +- **重大更新新手使用指南** (`GUIDE.md`) + - 新增批量提示词功能详细说明（节点详细说明部分） + - 新增场景 5：批量提示词生成（多提示词并发） + - 新增场景 6：批量提示词 + 图生图（共享参考图） + - 新增场景 7：批量提示词精准匹配规则详解 + - 单提示词模式与批量提示词模式对比 + - 详细的行为示例（纯文生图、图生图、多参考图） + - 常见误区与正确用法对照 + - 使用决策树帮助用户选择合适的节点和模式 + - 新增节点功能对比表和提示词模式对比表 + - 新增 Q6-Q7：批量提示词相关常见问题 + - 新增技巧 7：批量提示词最佳实践（4个子技巧） + - 更新场景编号（原场景 6-8 → 新场景 8） + - 修正批量提示词的描述（所有提示词共享输入图像，而非 1:1 匹配） + - 强化对 `---` 分隔符格式要求的说明 + +--- + +## [1.5.2] - 2026-02-01 + +### Added +- **新增新手使用指南** (`GUIDE.md`) + - 详细的安装配置步骤 + - 三个核心节点的完整文档和参数说明 + - 六种常见使用场景的工作流示例 + - 常见问题解答和解决方案 + - 六个进阶使用技巧 + - 面向初次使用插件的用户，提供从零到一的完整指导 + +### Changed +- **更新开发规范** (`.cursorrules`) + - 新增"新手指南维护规则"章节 + - 规定每次新增或变更节点时必须同步更新 `GUIDE.md` + - 提供节点文档和使用场景的标准模板 + - 明确文档维护的五大原则（用户视角、实用性、完整性、同步性、可读性） + +--- + +## [1.5.1] - 2026-02-01 + +### Fixed +- **修复批量 Nano Banana Pro 节点事件循环冲突** + - 修复 `RuntimeError: Cannot run the event loop while another loop is running` 错误 + - 使用 `ThreadPoolExecutor` 在独立线程中运行异步事件循环 + - 避免与 ComfyUI 主事件循环冲突 + - 确保批量处理任务稳定执行 + +--- + +## [1.5.0] - 2026-02-01 + +### Changed +- **批量 Nano Banana Pro 节点重构** (`BatchNanoBananaPro`) + - 参数重命名： + - `手动参考图` → `加载参考图` + - `预览图像` → `输出图像` + - `配对模式` → `图片配对模式` + - 配对模式选项值重命名： + - `1:1索引配对` → `1:1` + - `笛卡尔积` → `1*N` + - 参数顺序调整：`文件夹2-4` 移到 `文件夹1` 下方（从 optional 移至 required） + - 固定并发控制：移除 `最大并发数` 参数，默认自动管理（每批最多 100 并发） + - 固定生成数量：移除 `每组生成数量` 参数，每组配对固定生成 1 张图 + +### Added +- **批量 Nano Banana Pro 节点新增像素缩放功能** + - 新增 `像素缩放` 参数（BOOLEAN，默认 True） + - 新增 `分辨率像素` 参数（FLOAT，默认 1.0，范围 0.1-100.0） + - 支持对文件夹加载的图片和手动参考图进行缩放 + - 使用 Lanczos 重采样算法保持图片质量 + - 缩放在发送 API 前应用 + +### Removed +- **批量 Nano Banana Pro 节点移除处理报告功能** + - 移除返回值中的 `处理报告` 字符串输出 + - 简化返回类型为单一 `IMAGE` 输出 + - 统计信息仍通过控制台打印输出 + +### Improved +- **输出图像优化** + - `输出图像` 现在返回所有生成的图片（而非仅预览最后几张） + - 提供完整的批处理结果输出 + +--- + +## [1.4.0] - 2026-02-01 + +### Added +- **批量 Nano Banana Pro 节点** (`BatchNanoBananaPro`) + - 支持从 1-4 个文件夹批量加载图片 + - 两种配对模式：1:1 索引配对 / 笛卡尔积 + - 支持手动参考图输入（可与文件夹图片混合使用） + - 智能命名保存（保留原始文件名 + 自动后缀） + - 并发控制（默认最大 100，超过自动分批） + - 完整的处理报告输出 + - 预览最后生成的图片 + +- **文件处理工具模块** (`utils/file_utils.py`) + - `load_images_from_folder()` - 从文件夹加载图片 + - `pair_images_indexed()` - 1:1 索引配对 + - `pair_images_cartesian()` - 笛卡尔积配对 + - `generate_output_filename()` - 智能输出文件名生成 + - `save_image()` - 保存图片到指定路径 + +### Technical +- 新增 `ImageInfo` 命名元组，携带图片元数据 +- 支持 jpg/jpeg/png/webp/bmp/gif 图片格式 +- 文件名按字母顺序排序，确保配对顺序一致 + +--- + +## [1.3.0] - 2026-02-01 + +### Added +- **批量提示词功能（Nano Banana Pro）** + - 支持使用单行 `---` 分隔符同时提交多个不同提示词 + - 所有提示词并发生成，提高效率 + - 每个提示词可生成指定数量的图像（提示词数量 × 生图数量） + - 示例：3个提示词 × 2张/提示词 = 6张图 + - 触发条件：`---` 必须单独占据一行 + - 自动过滤空提示词 + +### Changed +- 优化 Nano Banana Pro 节点日志输出 + - 批量提示词模式显示 "X个提示词 × Y张/提示词 = Z张图" + - 单提示词模式保持原有输出格式 + +### Technical +- 新增 `parse_batch_prompts()` 工具函数（utils/image_utils.py） +- 新增 `generate_multi_prompts_async()` 方法（clients/gemini_client.py） +- 新增 `generate_multi_prompts_sync()` 方法（clients/gemini_client.py） + +--- + +## [1.2.1] - 2026-02-01 + +### Changed +- **加载批次图像（Nano Banana Pro）** 节点优化 + - 节点显示名称改为 "加载批次图像（Nano Banana Pro）" + - 移除 "参考图数量" 参数，改为自动检测所有有效输入图像数量 + - "开启像素缩放" 改名为 "像素缩放"，类型改为 BOOLEAN 开关（默认开启） + - "目标像素数_百万" 改名为 "分辨率像素" + +### Improved +- 简化用户操作流程，无需手动设置图像数量 +- 支持灵活的图像输入方式（1-14 张任意数量） + +--- + +## [1.2.0] - 2026-02-01 + +### Added +- **加载批次图像** 节点 (`BatchImageLoader`) + - 支持动态输入数量（1-14 张图像） + - 可选的像素缩放功能（保持纵横比） + - 使用 Lanczos 重采样算法进行高质量缩放 + - 支持设置目标像素数（0.1-100 百万像素） + - 可与原生"加载图像"节点连接使用 + - 输出批次张量供其他节点使用 + +--- + +## [1.1.1] - 2026-02-01 + +### Changed +- 将 Nano Banana Pro 节点的随机种子参数改为 ComfyUI 原生格式 + - `随机种子` → `seed` (参数名符合 ComfyUI 标准) + - 移除 `-1` 自动随机逻辑 + - 种子默认值改为 `0`，取值范围 `[0, 2^64-1]` + +### Removed +- 移除 `control_after_generate` 参数（简化节点参数） + +--- + +## [1.1.0] - 2026-02-01 + +### Changed +- 重构项目结构，模块化设计 + - 新增 `nodes/` 目录存放节点实现 + - 新增 `utils/` 目录存放工具函数 + - 新增 `clients/` 目录存放 API 客户端 +- 封装图像转换工具到 `utils/image_utils.py` +- 封装配置管理到 `utils/config.py` +- 重构 API 客户端，拆分为基类和具体实现 +- 精简文档结构，整合为 README.md + CHANGELOG.md + .cursorrules + +### Added +- `BaseAPIClient` 抽象基类，支持快速开发新 API 客户端 +- `.cursorrules` 开发指导文档，统一开发规范 +- `get_api_key_or_raise()` 函数，简化密钥获取逻辑 + +--- + +## [1.0.0] - 2026-02-01 + +### Added +- 初始版本发布 +- Nano Banana Pro 节点 + - 文生图功能 + - 图生图功能（最多 14 张输入图像） + - 批量并发生成（最多 1000 张） + - 多分辨率支持（1K / 2K / 4K） + - 10 种宽高比选项 + - 可控随机种子 +- 通过 api.o1key.com 调用 Gemini 3 Pro 模型 +- 支持 .config 文件和环境变量配置 API 密钥 +- 完善的错误处理和日志输出 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..e33e38b --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Comfyui_o1key + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..9eaa3f1 --- /dev/null +++ b/README.md @@ -0,0 +1,267 @@ +# Comfyui_o1key + +通过 `api.o1key.com` 调用 AI 模型的 ComfyUI 自定义节点集合。 + +## 功能特性 + +- 🎨 文生图 / 图生图 +- 🔄 批量并发生成（最多 1000 张） +- 📐 10 种宽高比 +- 🎯 3 种分辨率（1K / 2K / 4K） +- 🌱 可控随机种子 + +--- + +## 📦 安装 + +### 方法一：通过 ComfyUI Manager（推荐） + +1. 在 ComfyUI 中打开 Manager +2. 搜索 `Comfyui_o1key` +3. 点击安装 +4. 重启 ComfyUI + +### 方法二：手动安装 + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/lizhongyi1209/comfyui_o1key.git +cd comfyui_o1key +pip install -r requirements.txt +``` + +然后重启 ComfyUI。 + +### 国内用户安装（GitHub 拉取慢或失败时） + +使用 Gitee 镜像安装与更新，避免网络问题： + +```bash +cd ComfyUI/custom_nodes +git clone https://gitee.com/resonLzy/comfyui_o1key.git +cd comfyui_o1key +pip install -r requirements.txt +``` + +自动更新脚本（见下方「更新插件」）已改为从 Gitee 拉取，国内用户可直接使用。 + +--- + +## ⚙️ 配置 + +### 获取 API 密钥 + +1. 访问 [vip.o1key.com](https://vip.o1key.com) +2. 注册并获取 API 密钥 + +### 配置方式 + +#### 配置 API 密钥（必需） + +**方法一：快捷脚本配置（最简单）⭐** + +我们提供了一键配置脚本，自动创建配置文件： + +**Windows 用户：** +双击运行 `设置API密钥(win).bat`，按提示输入 API 密钥即可。 + +**Linux/Mac 用户：** +```bash +# 添加执行权限（仅首次需要） +chmod +x 设置API密钥(mac).sh + +# 运行配置脚本 +./设置API密钥(mac).sh +``` + +按提示输入 API 密钥，配置完成后重启 ComfyUI。 + +**方法二：环境变量（推荐）** + +**Windows 用户：** +1. 右键 "此电脑" → 属性 → 高级系统设置 → 环境变量 +2. 在"用户变量"中新建： + - 变量名：`O1KEY_API_KEY` + - 变量值：你的 API 密钥 +3. 重启 ComfyUI + +**Linux/Mac 用户：** + +在 `~/.bashrc` 或 `~/.zshrc` 中添加： +```bash +export O1KEY_API_KEY="你的API密钥" +``` + +然后执行 `source ~/.bashrc` 并重启 ComfyUI。 + +**方法三：手动创建配置文件** + +在插件目录下创建 `.config` 文件（参考 `.config.example`）： +``` +O1KEY_API_KEY=你的API密钥 +``` + +> **⚠️ 安全提示** +> +> `.config` 文件包含敏感信息，已添加到 `.gitignore` 中，不会被提交到版本控制。 +> 请妥善保管你的 API 密钥，不要分享给他人。 + +#### 配置 API 地址（可选） + +默认使用 `https://vip.o1key.com`，通常无需修改。 + +如需自定义 API 地址，可通过以下方式： + +1. **环境变量**（推荐）： + ```bash + # Windows + set O1KEY_API_BASE_URL=https://your-api-domain.com + + # Linux/Mac + export O1KEY_API_BASE_URL=https://your-api-domain.com + ``` + +2. **配置文件**：在 `.config` 中添加： + ``` + O1KEY_API_BASE_URL=https://your-api-domain.com + ``` + +3. **修改默认值**：编辑 `utils/config.py` 中的 `DEFAULT_API_BASE_URL` 常量 + +--- + +## 🔄 更新插件 + +自动更新脚本**已改为从国内镜像（Gitee）拉取**，国内用户无需科学上网即可更新。 + +### 方法一：自动更新（推荐）⭐ + +**Windows 用户：** +1. 进入插件目录：`ComfyUI\custom_nodes\comfyui_o1key` +2. 双击运行 `自动更新插件(win).bat` +3. 等待更新完成 +4. 重启 ComfyUI + +**Linux/Mac 用户：** +```bash +cd ComfyUI/custom_nodes/comfyui_o1key +chmod +x "自动更新插件(mac).sh" # 首次运行需要添加执行权限 +./"自动更新插件(mac).sh" +``` + +### 方法二：手动更新 + +从 Gitee 镜像拉取（国内推荐）： +```bash +cd ComfyUI/custom_nodes/comfyui_o1key +git remote get-url gitee &>/dev/null || git remote add gitee https://gitee.com/resonLzy/comfyui_o1key.git +git pull gitee main +pip install -r requirements.txt --upgrade +``` + +从 GitHub 拉取： +```bash +cd ComfyUI/custom_nodes/comfyui_o1key +git pull origin main +pip install -r requirements.txt --upgrade +``` + +**💡 提示：** +- 自动更新脚本会自动备份和恢复你的 `.config` 配置文件 +- 更新会保留环境变量中配置的 API 密钥 +- 更新检查在每次启动 ComfyUI 时自动进行（不会影响性能） +- 如果发现新版本，终端会显示更新提示 + +--- + +## 📚 节点说明 + +### Nano Banana Pro + +高性能图像生成节点，支持文生图和图生图。 + +**参数：** +- **提示词**：描述你想生成的图像 +- **模型**：选择使用的 AI 模型 +- **分辨率**：1K / 2K / 4K +- **宽高比**：1:1, 16:9, 9:16, 4:3, 3:4, 21:9, 9:21, 3:2, 2:3, 16:10 +- **批次大小**：单次生成的图像数量（1-1000） +- **随机种子**：控制生成的随机性（-1 为随机） +- **输入图像**（可选）：用于图生图模式 + +### Batch Nano Banana Pro + +批量并发生成节点，适合大量图像生成。 + +### Google Gemini + +Google Gemini 模型节点，支持更多模型选择。 + +--- + +## 📝 更新日志 + +查看 [CHANGELOG.md](./CHANGELOG.md) 了解详细的版本更新记录。 + +--- + +## 📄 许可证 + +本项目采用 Apache License 2.0 许可证。 + +--- + +## 🤝 贡献 + +欢迎提交 Issue 和 Pull Request！ + +--- + +## ⚠️ 开发者注意事项 + +### 维护者：发布流程与镜像同步 + +代码**先提交并推送到 GitHub**，再**同步到 Gitee 镜像**，国内用户通过 Gitee 拉取以解决网络问题。 + +**首次配置**（仅需一次）： +```bash +git remote add gitee https://gitee.com/resonLzy/comfyui_o1key.git +``` + +**每次发布**： +```bash +git push origin main # 先更新 GitHub +git push gitee main # 再同步到 Gitee 镜像 +``` + +### 文件编码要求 + +**所有文本文件必须使用 UTF-8 编码（无 BOM）！** + +如果你在 GitHub 上看到中文乱码，说明文件编码有问题。请使用以下方法修复： + +**Windows 用户：** + +```powershell +.\fix_encoding.ps1 +``` + +**Linux/Mac 用户：** + +```bash +chmod +x fix_encoding.sh +./fix_encoding.sh +``` + +详细说明请查看 [编码修复指南.md](./编码修复指南.md) + +--- + +## 📮 联系方式 + +- GitHub: [@lizhongyi1209](https://github.com/lizhongyi1209) +- 项目地址: https://github.com/lizhongyi1209/comfyui_o1key + +--- + +**当前版本：v1.10.1** diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..cfd87ad --- /dev/null +++ b/__init__.py @@ -0,0 +1,108 @@ +""" +Comfyui_o1key - ComfyUI 自定义节点集合 +通过 api.o1key.com 调用 AI 模型进行图像生成和文本生成 + +项目结构: +├── nodes/ # 节点实现 +├── utils/ # 工具模块 +├── clients/ # API 客户端 +└── __init__.py # 节点注册入口 +""" + +# 检查更新（仅在启动时检查一次） +try: + from .utils.update_checker import check_for_updates, notify_update_available + + if check_for_updates(): + notify_update_available() +except Exception: + # 静默失败，不影响插件加载 + pass + +import ssl + +from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, SaveCleanImage, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, QuanNengShengTu, BatchQuanNengShengTu, AspectRatioPreset, MultiResPreview, BatchImagesO1key + +# 报错弹框友好文案（不修改原节点代码，仅在外层统一处理） +_MSG_TIMEOUT = "API 请求超时，请稍后重试或检查网络。" +_MSG_SSL_NETWORK = ( + "本地网络不太稳定!解决方案如下:\n" + "1. 重启程序再试试看 (优先)\n" + "2. 调整一下网络环境,如wifi或宽带等\n" + "3. 切换VPN节点,或更换代理模式\n" + "4. 关掉杀毒软件或防火墙\n" + "5. 关掉浏览器VPN插件,避免冲突" +) + +def _wrap_generate_for_error_display(cls, attr="generate"): + original = getattr(cls, attr, None) + if original is None: + return + def wrapped(self, *args, **kwargs): + try: + return original(self, *args, **kwargs) + except TimeoutError as e: + msg = (str(e) or "").strip() + if not msg: + msg = _MSG_TIMEOUT + raise TimeoutError(msg) from None + except (ssl.SSLError, OSError) as e: + err_str = str(e) + if "DECRYPTION_FAILED_OR_BAD_RECORD_MAC" in err_str or "decryption failed or bad record mac" in err_str.lower(): + raise RuntimeError(_MSG_SSL_NETWORK) from None + raise + setattr(cls, attr, wrapped) + +_wrap_generate_for_error_display(NanoBananaPro) +_wrap_generate_for_error_display(BatchNanoBananaPro) +_wrap_generate_for_error_display(QuanNengShengTu) +_wrap_generate_for_error_display(BatchQuanNengShengTu, "process_batch") + +# ComfyUI 节点注册 +NODE_CLASS_MAPPINGS = { + "NanoBananaPro": NanoBananaPro, + "BatchNanoBananaPro": BatchNanoBananaPro, + "GoogleGemini": GoogleGemini, + "LoadFile": LoadFile, + "ImageStitchPro": ImageStitchPro, + "SaveCleanImage": SaveCleanImage, + "BatchCleanMetadata": BatchCleanMetadata, + "VideoPreview": VideoPreview, + "GoogleVeo": GoogleVeo, + "FluxImageEdit": FluxImageEdit, + "UniversalLLMChat": UniversalLLMChat, + "KlingVideo": KlingVideo, + "KlingFirstLastFrame": KlingFirstLastFrame, + "KlingMotionControlTest": KlingMotionControlTest, + "QuanNengShengTu": QuanNengShengTu, + "BatchQuanNengShengTu": BatchQuanNengShengTu, + "AspectRatioPreset": AspectRatioPreset, + "MultiResPreview": MultiResPreview, + "BatchImagesO1key": BatchImagesO1key, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NanoBananaPro": "Nano Banana", + "BatchNanoBananaPro": "批量 Nano Banana", + "GoogleGemini": "Google Gemini", + "LoadFile": "加载文件", + "ImageStitchPro": "图像拼接 Pro", + "SaveCleanImage": "保存图像（防AI识别）", + "BatchCleanMetadata": "批量任务（防AI识别）", + "VideoPreview": "视频预览", + "GoogleVeo": "Google Veo - ab", + "FluxImageEdit": "Flux2 图像编辑", + "UniversalLLMChat": "全能LLM对话助手", + "KlingVideo": "自研模型 3.0 视频", + "KlingFirstLastFrame": "自研模型 3.0 首尾帧到视频", + "KlingMotionControlTest": "自研模型 动作控制（测试）", + "QuanNengShengTu": "全能生图", + "BatchQuanNengShengTu": "全能生图（批量）", + "AspectRatioPreset": "图片宽高比预设", + "MultiResPreview": "预览图像（v2）", + "BatchImagesO1key": "加载图像（批量）", +} + +WEB_DIRECTORY = "./web" + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] diff --git a/clients/__init__.py b/clients/__init__.py new file mode 100644 index 0000000..e96f850 --- /dev/null +++ b/clients/__init__.py @@ -0,0 +1,14 @@ +""" +API 客户端模块 +包含与外部 API 通信的客户端实现 +""" + +from .base_client import BaseAPIClient +from .gemini_client import GeminiAPIClient +from .gemini_flash_client import GeminiFlashClient +from .sora_client import SoraClient +from .kling_client import KlingClient +from .veo_client import VeoClient +from .openai_client import OpenAIAPIClient + +__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'OpenAIAPIClient'] diff --git a/clients/base_client.py b/clients/base_client.py new file mode 100644 index 0000000..3830a97 --- /dev/null +++ b/clients/base_client.py @@ -0,0 +1,545 @@ +""" +API 客户端基类 +提供通用的 HTTP 请求、响应解析和错误处理功能 +""" + +import asyncio +import json +import threading +from abc import ABC, abstractmethod +from typing import Any, Callable, Dict, List, Optional + +import aiohttp + + +class BaseAPIClient(ABC): + """ + API 客户端抽象基类 + + 子类需要实现以下方法： + - get_endpoint(): 获取 API 端点 + - build_request_body(): 构建请求体 + - parse_response(): 解析响应 + """ + + def __init__( + self, + base_url: str, + api_key: str, + max_request_size: int = 100 * 1024 * 1024 + ): + """ + 初始化客户端 + + Args: + base_url: API 基础 URL + api_key: API 密钥 + max_request_size: 最大请求体大小（字节），默认 100MB + """ + self.base_url = base_url + self.api_key = api_key + self.max_request_size = max_request_size + + @abstractmethod + def get_endpoint(self, **kwargs) -> str: + """ + 获取 API 端点路径 + + Args: + **kwargs: 额外参数（如模型名、分辨率等） + + Returns: + 端点路径字符串 + """ + pass + + @abstractmethod + def build_request_body(self, **kwargs) -> Dict[str, Any]: + """ + 构建 API 请求体 + + Args: + **kwargs: 请求参数 + + Returns: + 请求体字典 + """ + pass + + @abstractmethod + def parse_response(self, response: Dict[str, Any]) -> Any: + """ + 解析 API 响应 + + Args: + response: API 响应字典 + + Returns: + 解析后的结果 + """ + pass + + def get_headers(self, use_bearer_token: bool = False) -> Dict[str, str]: + """ + 获取请求头 + + Args: + use_bearer_token: 是否使用 Bearer Token 认证（默认为 False） + + Returns: + 请求头字典 + """ + if use_bearer_token: + return { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json" + } + else: + return { + "x-goog-api-key": self.api_key, + "Content-Type": "application/json" + } + + def check_request_size(self, request_body: Dict[str, Any]) -> None: + """ + 检查请求体大小是否超过限制 + + Args: + request_body: 请求体字典 + + Raises: + ValueError: 如果请求体超过限制 + """ + request_json = json.dumps(request_body) + request_size = len(request_json.encode('utf-8')) + + if request_size > self.max_request_size: + raise ValueError( + "请求体积超过100MB限制，请调整分辨率或减少图片数量" + ) + + def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]: + """ + 子类可重写：为指定 HTTP 状态码返回自定义错误文案。 + 若返回 None，则使用基类默认拼接文案。 + + Args: + status_code: HTTP 状态码（如 429、503） + error_message: API 返回的原始错误信息 + + Returns: + 自定义完整错误文案，或 None 表示使用默认 + """ + return None + + async def request_async( + self, + endpoint: str, + request_body: Dict[str, Any], + session: Optional[aiohttp.ClientSession] = None, + use_bearer_token: bool = False, + timeout: Optional[int] = None + ) -> Dict[str, Any]: + """ + 发送异步 HTTP 请求（带详细计时） + + Args: + endpoint: API 端点 + request_body: 请求体 + session: aiohttp 会话（可选） + use_bearer_token: 是否使用 Bearer Token 认证 + timeout: 超时时间（秒）- 已废弃，由服务器端控制 + + Returns: + 响应 JSON + + Raises: + RuntimeError: 请求失败时 + """ + import time + + url = f"{self.base_url}{endpoint}" + headers = self.get_headers(use_bearer_token) + + # 检查请求大小 + self.check_request_size(request_body) + + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + try: + # 连接计时 + connect_start = time.time() + + async with session.post(url, json=request_body, headers=headers) as response: + connect_time = time.time() - connect_start + + if response.status != 200: + error_text = await response.text() + + # 尝试解析 JSON 错误信息，提取关键内容 + error_message = error_text + try: + error_json = json.loads(error_text) + # 尝试从多个常见位置提取错误信息 + if "error" in error_json: + if isinstance(error_json["error"], dict): + error_message = error_json["error"].get("message", error_text) + else: + error_message = str(error_json["error"]) + elif "message" in error_json: + error_message = error_json["message"] + except: + # 如果不是 JSON，使用原始文本 + pass + + # 针对常见错误状态码提供友好提示 + if response.status == 400: + raise RuntimeError( + f"请求参数错误 (400 Bad Request)\n" + f"API 返回错误：{error_message}\n" + f"建议：\n" + f" - 检查 API 密钥是否有效\n" + f" - 确认请求参数格式正确" + ) + elif response.status == 401: + raise RuntimeError( + f"认证失败 (401 Unauthorized)\n" + f"API 返回错误：{error_message}\n" + f"建议：\n" + f" - 检查 API 密钥是否正确\n" + f" - 确认 API 密钥是否过期" + ) + elif response.status == 403: + raise RuntimeError( + f"权限不足 (403 Forbidden)\n" + f"API 返回错误：{error_message}\n" + f"建议：\n" + f" - 检查 API 密钥权限\n" + f" - 确认账户余额充足" + ) + elif response.status == 404: + raise RuntimeError( + f"端点不存在 (404 Not Found)\n" + f"API 返回错误：{error_message}\n" + f"建议：\n" + f" - 检查模型名称是否正确\n" + f" - 使用其他可用模型" + ) + elif response.status == 429: + custom = self.get_http_error_message(429, error_message) + if custom is not None: + raise RuntimeError(custom) + raise RuntimeError( + f"请求频率超限 (429 Too Many Requests)\n" + f"API 返回错误：{error_message}\n" + f"建议：\n" + f" - 等待一段时间后重试\n" + f" - 检查 API 配额是否充足" + ) + elif response.status == 503: + custom = self.get_http_error_message(503, error_message) + if custom is not None: + raise RuntimeError(custom) + raise RuntimeError( + f"服务暂时不可用 (503 Service Unavailable)\n" + f"API 返回错误：{error_message}\n" + f"建议：\n" + f" - 稍后重试\n" + f" - 尝试使用其他模型" + ) + elif response.status == 504: + raise RuntimeError( + f"API 请求超时 (504 Gateway Timeout)\n" + f"API 返回错误：{error_message}\n" + f"建议：\n" + f" - 尝试使用其他模型\n" + f" - 稍后重试\n" + f" - 降低分辨率或减少输入图像数量" + ) + elif response.status == 502: + raise RuntimeError( + "糟糕！请求到上游时遇到超时或过载！别担心，过会儿再次点击运行即可！" + ) + else: + raise RuntimeError( + f"API 请求失败 (状态码: {response.status})\n" + f"API 返回错误：{error_message}" + ) + + # 接收响应体 + wait_start = time.time() + response_data = await response.json() + download_time = time.time() - wait_start + + # 附加计时信息到响应数据（供上层使用） + response_size = len(str(response_data)) + if not isinstance(response_data, dict): + response_data = {"data": response_data} + + # 将计时信息存储在响应的元数据中 + response_data["_timing"] = { + "connect_time": connect_time, + "download_time": download_time, + "response_size": response_size + } + + return response_data + + finally: + if close_session: + await session.close() + + async def request_get_async( + self, + endpoint: str, + session: Optional[aiohttp.ClientSession] = None, + use_bearer_token: bool = True, + timeout: Optional[int] = None + ) -> Dict[str, Any]: + """ + 发送异步 HTTP GET 请求 + + Args: + endpoint: API 端点 + session: aiohttp 会话（可选） + use_bearer_token: 是否使用 Bearer Token 认证（默认为 True） + timeout: 超时时间（秒）- 已废弃，由服务器端控制 + + Returns: + 响应 JSON + + Raises: + RuntimeError: 请求失败时 + """ + url = f"{self.base_url}{endpoint}" + headers = self.get_headers(use_bearer_token) + + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + try: + async with session.get(url, headers=headers) as response: + if response.status != 200: + error_text = await response.text() + + # 尝试解析 JSON 错误信息，提取关键内容 + error_message = error_text + try: + error_json = json.loads(error_text) + # 尝试从多个常见位置提取错误信息 + if "error" in error_json: + if isinstance(error_json["error"], dict): + error_message = error_json["error"].get("message", error_text) + else: + error_message = str(error_json["error"]) + elif "message" in error_json: + error_message = error_json["message"] + except: + # 如果不是 JSON，使用原始文本 + pass + + # 针对常见错误状态码提供友好提示 + if response.status == 400: + raise RuntimeError( + f"请求参数错误 (400 Bad Request)\n" + f"API 返回错误：{error_message}\n" + f"建议：检查请求参数" + ) + elif response.status == 401: + raise RuntimeError( + f"认证失败 (401 Unauthorized)\n" + f"API 返回错误：{error_message}\n" + f"建议：检查 API 密钥" + ) + elif response.status == 429: + custom = self.get_http_error_message(429, error_message) + if custom is not None: + raise RuntimeError(custom) + raise RuntimeError( + f"请求频率超限 (429 Too Many Requests)\n" + f"API 返回错误：{error_message}\n" + f"建议：等待一段时间后重试" + ) + elif response.status == 503: + custom = self.get_http_error_message(503, error_message) + if custom is not None: + raise RuntimeError(custom) + raise RuntimeError( + f"服务暂时不可用 (503 Service Unavailable)\n" + f"API 返回错误：{error_message}\n" + f"建议：稍后重试" + ) + elif response.status == 504: + raise RuntimeError( + f"API 请求超时 (504 Gateway Timeout)\n" + f"API 返回错误：{error_message}\n" + f"建议：稍后重试" + ) + elif response.status == 502: + raise RuntimeError( + "糟糕！请求到上游时遇到超时或过载！别担心，过会儿再次点击运行即可！" + ) + else: + raise RuntimeError( + f"API 请求失败 (状态码: {response.status})\n" + f"API 返回错误：{error_message}" + ) + + return await response.json() + + finally: + if close_session: + await session.close() + + async def batch_request_async( + self, + requests: List[Dict[str, Any]], + progress_callback: Optional[Callable[[int, int], None]] = None + ) -> List[Any]: + """ + 批量并发请求 + + Args: + requests: 请求列表，每个元素包含 endpoint 和 request_body + progress_callback: 进度回调函数 (current, total) + + Returns: + 响应结果列表 + """ + results = [] + completed = 0 + total = len(requests) + + # 创建无限制的连接器 + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + + async with aiohttp.ClientSession(connector=connector) as session: + tasks = [] + + for req in requests: + task = self.request_async( + endpoint=req['endpoint'], + request_body=req['request_body'], + session=session + ) + tasks.append(task) + + # 并发执行 + responses = await asyncio.gather(*tasks, return_exceptions=True) + + for i, resp in enumerate(responses): + if isinstance(resp, Exception): + print(f"⚠️ 第 {i+1} 个请求失败: {str(resp)}") + continue + + try: + parsed = self.parse_response(resp) + results.append(parsed) + completed += 1 + + if progress_callback: + progress_callback(completed, total) + + except Exception as e: + print(f"⚠️ 第 {i+1} 个响应解析失败: {str(e)}") + + return results + + def run_async_in_thread(self, coro) -> Any: + """ + 在独立线程中运行异步代码（用于 ComfyUI 同步接口） + + Args: + coro: 协程对象 + + Returns: + 协程执行结果 + """ + result_container = [] + error_container = [] + + def run_in_thread(): + try: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + try: + result = loop.run_until_complete(coro) + result_container.append(result) + finally: + loop.close() + + except Exception as e: + error_container.append(e) + + thread = threading.Thread(target=run_in_thread) + thread.start() + thread.join() + + if error_container: + raise error_container[0] + + if not result_container: + raise RuntimeError("异步任务未返回结果") + + return result_container[0] + + async def query_balance_async(self) -> Dict[str, Any]: + """ + 异步查询账户余额 + + Returns: + 余额信息字典，包含 name、total_available 等字段 + + Raises: + RuntimeError: 查询失败时 + """ + endpoint = "/api/usage/token" + response = await self.request_get_async(endpoint, use_bearer_token=True) + + if not response.get("code"): + raise RuntimeError("余额查询响应格式错误") + + data = response.get("data", {}) + return data + + def query_balance_sync(self) -> Dict[str, Any]: + """ + 同步查询账户余额（用于 ComfyUI 节点） + + Returns: + 余额信息字典 + + Raises: + RuntimeError: 查询失败时 + """ + coro = self.query_balance_async() + return self.run_async_in_thread(coro) + + def format_balance_info(self, balance_data: Dict[str, Any]) -> str: + """ + 格式化余额信息为展示文本 + + Args: + balance_data: 余额信息字典 + + Returns: + 格式化文本，如 "当前余额：100.00 | API：xxx" + + Example: + >>> data = {"name": "test-api", "total_available": 50000000} + >>> client.format_balance_info(data) + '当前余额：100.00 | API：test-api' + """ + api_name = balance_data.get("name", "未知") + total_available = balance_data.get("total_available", 0) + + # 实际显示余额 = total_available / 500000，单位：美元 + balance_in_dollars = total_available / 500000 + + return f"当前余额：{balance_in_dollars:.2f} | API：{api_name}" diff --git a/clients/flux_edit_client.py b/clients/flux_edit_client.py new file mode 100644 index 0000000..3c41416 --- /dev/null +++ b/clients/flux_edit_client.py @@ -0,0 +1,198 @@ +""" +Flux 图像编辑 API 客户端 +通过 vip.o1key.com 调用 Flux2 图像编辑 + SeedVR2 超分辨率服务 + +工作流程： +1. submit_task → POST /v1/images/edits (multipart/form-data 提交主图+参考图+提示词) +2. poll_result → GET /v1/images/edits/{task_id} (直连容器轮询) +""" + +import base64 +import time +from typing import Optional + +import requests + +from ..utils.config import get_api_key_or_raise, get_api_base_url + + +# 显示名 → 实际请求值的映射 +SIZE_DISPLAY_MAP = { + "2K": "2048", + "4K": "4096", +} + +# 轮询直连容器地址，绕过代理层 +POLL_BASE_URL = "https://xrrh7tn08tfgwa8w-8188.container.x-gpu.com" + + +class FluxEditClient: + """ + Flux 图像编辑客户端 + + 对接 vip.o1key.com 上的 /v1/images/edits 接口， + 将图像编辑+超分辨率任务提交到远程服务器执行。 + """ + + SUBMIT_ENDPOINT = "/v1/images/edits" + STATUS_ENDPOINT = "/v1/images/edits/{task_id}" + + DEFAULT_POLL_INTERVAL = 15 # 秒 + + def __init__(self): + self.api_key = get_api_key_or_raise() + self.base_url = get_api_base_url() + + # ------------------------------------------------------------------ + # 同步方法（供 ComfyUI 节点调用） + # ------------------------------------------------------------------ + + def submit_and_wait( + self, + image_bytes: bytes, + mask_bytes: bytes, + prompt: str, + size: str = "4K", + poll_interval: int = DEFAULT_POLL_INTERVAL, + progress_callback=None, + ) -> bytes: + """ + 提交任务并同步等待结果（阻塞直到完成） + + Args: + image_bytes: 主图二进制数据 + mask_bytes: 参考图二进制数据 + prompt: 编辑提示词 + size: 分辨率显示名 ("2K" 或 "4K") + poll_interval: 轮询间隔（秒） + progress_callback: 进度回调 fn(status_str) + + Returns: + 结果图像的二进制数据 + + Raises: + RuntimeError: 任务失败 + """ + size_value = SIZE_DISPLAY_MAP.get(size, size) + + # 1. 提交任务（走代理） + task_id = self._submit_task_sync(image_bytes, mask_bytes, prompt, size_value) + if progress_callback: + progress_callback(f"任务已提交: {task_id[:8]}...") + + # 2. 轮询等待（直连容器） + return self._poll_result_sync( + task_id, poll_interval, progress_callback + ) + + def _submit_task_sync( + self, + image_bytes: bytes, + mask_bytes: bytes, + prompt: str, + size: str, + ) -> str: + """同步提交任务，返回 task_id""" + url = f"{self.base_url}{self.SUBMIT_ENDPOINT}" + headers = {"Authorization": f"Bearer {self.api_key}"} + + files = { + "image": ("image.jpg", image_bytes, "image/jpeg"), + "mask": ("mask.jpg", mask_bytes, "image/jpeg"), + } + data = { + "prompt": prompt, + "size": size, + "model": "flux2-fp8-dualr", + } + + try: + resp = requests.post(url, files=files, data=data, headers=headers, timeout=60) + except requests.exceptions.Timeout: + raise RuntimeError("提交任务超时，请检查网络连接") + except requests.exceptions.ConnectionError: + raise RuntimeError("无法连接到服务器，请检查网络或服务器地址") + + if resp.status_code != 200: + raise RuntimeError( + f"提交任务失败 (HTTP {resp.status_code})\n" + f"响应: {resp.text[:500]}" + ) + + result = resp.json() + task_id = result.get("id") + if not task_id: + raise RuntimeError(f"服务器返回异常: 未获取到任务ID\n{result}") + + return task_id + + def _poll_result_sync( + self, + task_id: str, + poll_interval: int, + progress_callback=None, + ) -> bytes: + """同步轮询任务状态（直连容器），返回结果图像二进制""" + url = f"{POLL_BASE_URL}{self.STATUS_ENDPOINT.format(task_id=task_id)}" + + start_time = time.time() + last_status = None + + while True: + elapsed = time.time() - start_time + + try: + resp = requests.get(url, timeout=30) + except requests.exceptions.ConnectionError: + raise RuntimeError("轮询时无法连接到服务器，请检查网络") + + if resp.status_code != 200: + raise RuntimeError( + f"查询任务状态失败 (HTTP {resp.status_code})\n" + f"响应: {resp.text[:500]}" + ) + + result = resp.json() + status = result.get("status", "unknown") + + # 状态变化时打印日志 + if status != last_status: + elapsed_str = f"{elapsed:.0f}s" + print(f"Flux Edit: [{elapsed_str}] 任务 {task_id[:8]}... → {status}") + last_status = status + + if progress_callback: + elapsed_str = f"{elapsed:.0f}s" + status_desc = { + "pending": "排队中", + "processing": "处理中", + "generating": "生图中，请耐心等待，预计耗时140s左右", + }.get(status, status) + progress_callback(f"{status_desc} (当前进度：{elapsed_str})") + + if status == "completed": + # 解码 base64 图像 + b64_data = result.get("result") + if not b64_data: + raise RuntimeError("任务完成但未返回图像数据") + return base64.b64decode(b64_data) + + elif status == "failed": + error_msg = result.get("error", "未知错误") + raise RuntimeError( + f"图像编辑任务失败\n" + f"错误: {error_msg}" + ) + + elif status in ("not_found",): + raise RuntimeError( + f"任务未找到: {task_id}\n" + f"可能已被清理或 ID 无效" + ) + + # 继续等待 + time.sleep(poll_interval) + + def query_balance_sync(self) -> dict: + """查询余额（兼容现有节点的 finally 块调用）""" + return {"name": "flux-edit", "total_available": 0} diff --git a/clients/gemini_client.py b/clients/gemini_client.py new file mode 100644 index 0000000..35e4ba0 --- /dev/null +++ b/clients/gemini_client.py @@ -0,0 +1,986 @@ +""" +Gemini API 客户端 +处理与 api.o1key.com 的通信，用于图像生成 +""" + +import re +import time +from io import BytesIO +from typing import Any, Callable, Dict, List, Optional + +import aiohttp +from PIL import Image + +from ..utils.image_utils import encode_image_to_base64, decode_base64_to_pil +from ..utils.config import get_api_key_or_raise, get_api_base_url +from .base_client import BaseAPIClient + + +class GeminiAPIClient(BaseAPIClient): + """ + Gemini API 客户端 + 用于调用 Gemini 3 Pro 模型进行图像生成 + """ + + def __init__(self, api_key: Optional[str] = None): + """ + 初始化客户端 + + Args: + api_key: API 密钥，如果为 None 则从配置文件或环境变量读取 + """ + if api_key is None: + api_key = get_api_key_or_raise("O1KEY_API_KEY") + + super().__init__( + base_url=get_api_base_url(), + api_key=api_key, + max_request_size=100 * 1024 * 1024 + ) + + def get_endpoint(self, model: str = "", resolution: str = "2K", **kwargs) -> str: + """ + 根据模型和分辨率获取 API 端点 + + Args: + model: 模型名称 + resolution: 分辨率（1K, 2K, 4K） + + Returns: + API 端点路径 + """ + from ..models_config import get_model_endpoint + + # 特殊处理：动态端点模型（根据分辨率选择） + if model == "nano-banana-pro-限时特价": + if resolution == "1K": + return "/v1beta/models/nano-banana-pro:generateContent" + elif resolution == "2K": + return "/v1beta/models/nano-banana-pro-2k:generateContent" + elif resolution == "4K": + return "/v1beta/models/nano-banana-pro-4k:generateContent" + else: + return "/v1beta/models/nano-banana-pro-2k:generateContent" + + elif model == "nano-banana-2-官方计费": + if resolution == "512": + return "/v1beta/models/nano-banana-2-0.5k-official:generateContent" + elif resolution == "1K": + return "/v1beta/models/nano-banana-2-1k-official:generateContent" + elif resolution == "2K": + return "/v1beta/models/nano-banana-2-2k-official:generateContent" + elif resolution == "4K": + return "/v1beta/models/nano-banana-2-4k-official:generateContent" + else: + return "/v1beta/models/nano-banana-2-2k-official:generateContent" + + elif model == "nano-banana-pro-官方计费": + if resolution == "1K": + return "/v1beta/models/nano-banana-pro-1k-official:generateContent" + elif resolution == "2K": + return "/v1beta/models/nano-banana-pro-2k-official:generateContent" + elif resolution == "4K": + return "/v1beta/models/nano-banana-pro-4k-official:generateContent" + else: + return "/v1beta/models/nano-banana-pro-2k-official:generateContent" + + elif model == "gemini-3-pro-image-preview-url": + if resolution == "1K": + return "/v1beta/models/gemini-3-pro-image-preview-url:generateContent" + elif resolution == "2K": + return "/v1beta/models/gemini-3-pro-image-preview-2k-url:generateContent" + elif resolution == "4K": + return "/v1beta/models/gemini-3-pro-image-preview-4k-url:generateContent" + else: + return "/v1beta/models/gemini-3-pro-image-preview-2k-url:generateContent" + + # 其他模型：从配置文件读取端点 + endpoint = get_model_endpoint(model) + if endpoint: + return endpoint + + # 兜底：使用标准模式端点 + return "/v1beta/models/gemini-3-pro-image-preview:generateContent" + + def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]: + """Gemini 请求 429/503 时返回图中约定的多行错误框文案。""" + if status_code == 429: + return ( + "莫慌!该模型暂时超出速率限制啦\n" + "解决方案如下(任意一种):\n" + "1.切换当前模型\n" + "2.前往后台,修改令牌分组" + ) + if status_code == 503: + return ( + "警报!谷歌服务器当前过载!\n" + "解决方案如下:\n" + "1.摸会儿鱼吧,我也没办法,谷歌会尽快恢复,嘿嘿~\n" + "2.切换其他模型\n" + "3.前往后台,修改令牌分组" + ) + return None + + def build_request_body( + self, + prompt: str = "", + images: Optional[List[Image.Image]] = None, + aspect_ratio: str = "1:1", + resolution: str = "2K", + enable_grounding: bool = False, + enable_image_search: bool = False, + candidate_count: int = 1, + **kwargs + ) -> Dict[str, Any]: + """ + 构建 API 请求体 + + Args: + prompt: 提示词 + images: 输入图像列表（可选） + aspect_ratio: 宽高比 + resolution: 分辨率 + enable_grounding: 是否启用 Google Search Grounding + enable_image_search: 是否同时启用 Google Image Search（仅 Gemini 3.1 Flash 支持） + candidate_count: 单次请求返回的候选图数量，默认 1 + + Returns: + 请求体字典 + """ + parts = [] + + # 添加文本部分 + parts.append({"text": prompt}) + + # 添加图像部分（如果有） + if images: + for img in images: + img_base64 = encode_image_to_base64(img) + parts.append({ + "inline_data": { + "mime_type": "image/png", + "data": img_base64 + } + }) + + # 构建请求体 + request_body = { + "contents": [ + { + "role": "user", + "parts": parts + } + ], + "generationConfig": { + "candidateCount": candidate_count, + "responseModalities": ["TEXT", "IMAGE"], + "imageConfig": { + "aspectRatio": aspect_ratio, + "imageSize": resolution + } + } + } + + # 添加 Google Search Grounding 工具（如果启用） + # 注意：enable_image_search=True 时会自动隐含 enable_grounding + if enable_grounding or enable_image_search: + if enable_image_search: + # 同时启用网页搜索和图片搜索（仅 nano-banana-2 / gemini-3.1-flash-image-preview 支持） + request_body["tools"] = [ + { + "google_search": { + "searchTypes": { + "webSearch": {}, + "imageSearch": {} + } + } + } + ] + else: + # 仅启用网页搜索（通用） + request_body["tools"] = [{"google_search": {}}] + + return request_body + + def parse_response(self, response: Dict[str, Any]) -> List[Image.Image]: + """ + 同步解析 API 响应（保留以满足抽象基类要求） + + 注意：此方法仅用于兼容基类接口，实际使用请调用 parse_response_async() + + Args: + response: API 响应字典 + + Returns: + 图像列表 + + Raises: + RuntimeError: 此方法不应被直接调用 + """ + raise RuntimeError( + "parse_response() 不应被直接调用。" + "请使用 generate_single_async() 或 generate_batch_async() 等高级方法。" + ) + + async def parse_response_async( + self, + response: Dict[str, Any], + session: Optional[aiohttp.ClientSession] = None + ) -> tuple[List[Image.Image], Dict[str, Any]]: + """ + 异步解析 API 响应，提取生成的图像 + + Args: + response: API 响应字典 + session: aiohttp 会话（用于下载图片） + + Returns: + (图像列表, 格式信息字典) + 格式信息包含: type (base64/url), size, resolution, download_speed (仅URL) + + Raises: + RuntimeError: 解析失败或 API 拒绝时 + """ + + # 初始化格式信息 + format_info = { + "type": None, # "base64" or "url" + "size": 0, + "resolution": None, + "download_speed": None + } + + candidates = response.get("candidates", []) + + # ========== 错误检测（按优先级顺序）========== + + # 1. 检查 candidatesTokenCount（最高优先级） + usage_metadata = response.get("usageMetadata", {}) + candidates_token_count = usage_metadata.get("candidatesTokenCount", -1) + + if candidates_token_count == 0: + error_msg = ( + "Damn!你触发顶级风控啦!还没到生图阶段就被拒了。\n" + "赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~" + ) + raise RuntimeError(error_msg) + + # 2. 检查 finishReason（次优先级） + candidates = response.get("candidates", []) + if candidates: + for candidate in candidates: + finish_reason = candidate.get("finishReason", "") + + if finish_reason and finish_reason != "STOP": + error_msg = ( + "Ohh no! 生图过程触发风控,图片被拒绝生成!\n" + "可能原因如下:\n" + "1.违禁内容\n" + "2.触发安全过滤器\n" + "3.涉及版权问题\n" + "4. Token超限\n" + "赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~" + ) + raise RuntimeError(error_msg) + + # ========== 图像提取 ========== + + images = [] + text_responses = [] # 收集文本响应 + + # 需要关闭 session 的标记 + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + try: + for candidate_idx, candidate in enumerate(candidates): + content = candidate.get("content", {}) + parts = content.get("parts", []) + + for part_idx, part in enumerate(parts): + # 方式1: inline_data 或 inlineData (base64) + # 兼容两种命名方式：蛇形（inline_data）和驼峰（inlineData） + inline_data_key = None + if "inline_data" in part: + inline_data_key = "inline_data" + elif "inlineData" in part: + inline_data_key = "inlineData" + + if inline_data_key: + inline_data = part[inline_data_key] + # 同样兼容 data/mimeType 的命名 + img_data = inline_data.get("data") or inline_data.get("data", "") + + if img_data: + img = decode_base64_to_pil(img_data) + images.append(img) + + # 记录格式信息 + if format_info["type"] is None: + format_info["type"] = "base64" + format_info["size"] = len(img_data) * 3 / 4 # Base64 解码后的字节数 + format_info["resolution"] = f"{img.size[0]}x{img.size[1]}" + + # 方式2: text 中的 URL - 改为异步下载 + elif "text" in part: + text = part["text"] + + # 收集文本响应（用于后续错误检测） + text_responses.append(text) + + # 尝试 markdown 格式: ![alt](url) + url_pattern_md = r'!\[.*?\]\((https?://[^\)]+)\)' + urls = re.findall(url_pattern_md, text) + + # 如果没找到，尝试纯 URL 格式 + if not urls: + url_pattern_plain = r'https?://[^\s<>"{}|\\^`\[\]]+' + urls = re.findall(url_pattern_plain, text) + + if urls: + for url_idx, url in enumerate(urls): + try: + # 使用 aiohttp 异步下载 + download_start = time.time() + async with session.get(url) as img_response: + if img_response.status == 200: + img_data = await img_response.read() + download_time = time.time() - download_start + img_size = len(img_data) + speed = img_size / download_time if download_time > 0 else 0 + + img = Image.open(BytesIO(img_data)) + images.append(img) + + # 记录格式信息（只记录第一张） + if format_info["type"] is None: + format_info["type"] = "url" + format_info["size"] = img_size + format_info["resolution"] = f"{img.size[0]}x{img.size[1]}" + format_info["download_speed"] = speed + except Exception as e: + pass # 静默失败，继续尝试其他URL + + # 方式3: 直接的 URL 字段 - 也改为异步 + elif "imageUrl" in part or "url" in part: + url = part.get("imageUrl") or part.get("url") + try: + download_start = time.time() + async with session.get(url) as img_response: + if img_response.status == 200: + img_data = await img_response.read() + download_time = time.time() - download_start + img_size = len(img_data) + speed = img_size / download_time if download_time > 0 else 0 + + img = Image.open(BytesIO(img_data)) + images.append(img) + + # 记录格式信息 + if format_info["type"] is None: + format_info["type"] = "url" + format_info["size"] = img_size + format_info["resolution"] = f"{img.size[0]}x{img.size[1]}" + format_info["download_speed"] = speed + except Exception as e: + pass # 静默失败 + + except Exception as e: + raise RuntimeError(f"解析 API 响应失败: {str(e)}") + + finally: + if close_session: + await session.close() + + # 3. 检查 API 文本响应拒绝说明 + if not images and text_responses: + # API 返回了文本但没有图片，说明请求被拒绝 + combined_text = "\n".join(text_responses) + error_msg = ( + f"API 拒绝响应\n\n" + f"API 返回说明：\n{combined_text}\n\n" + f"建议：\n" + f" - 根据上述说明调整请求内容\n" + f" - 确保提示词和参考图符合使用规范" + ) + raise RuntimeError(error_msg) + + if not images: + raise RuntimeError("API 响应中未找到生成的图像") + + return images, format_info + + async def generate_single_async( + self, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: Optional[List[Image.Image]] = None, + session=None, + task_index: Optional[int] = None, + total_tasks: Optional[int] = None, + debug: bool = False, + debug_request: bool = False, + enable_grounding: bool = False, + enable_image_search: bool = False, + candidate_count: int = 1 + ) -> tuple[List[Image.Image], Dict[str, Any]]: + """ + 单次异步生成请求（极简单行日志） + + Args: + prompt: 提示词 + model: 模型名称 + resolution: 分辨率 + aspect_ratio: 宽高比 + images: 输入图像列表 + session: aiohttp 会话 + task_index: 任务索引（用于批量任务） + total_tasks: 总任务数（用于批量任务） + debug: 是否打印完整 API 响应 + debug_request: 是否打印发送的请求体（base64 图片数据将被截断） + enable_grounding: 是否启用 Google Search Grounding + enable_image_search: 是否同时启用 Google Image Search + + Returns: + (生成的图像列表, 计时信息字典) + """ + import json + + total_start = time.time() + + # 任务前缀 + task_prefix = f"[{task_index}/{total_tasks}]" if task_index is not None and total_tasks else "" + + # ========== 1. 构建请求 ========== + build_start = time.time() + endpoint = self.get_endpoint(model=model, resolution=resolution) + request_body = self.build_request_body( + prompt=prompt, + images=images, + aspect_ratio=aspect_ratio, + resolution=resolution, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + candidate_count=candidate_count + ) + build_time = time.time() - build_start + + # ========== 调试日志：打印请求体 ========== + if debug_request: + import json as _json + + def _truncate_base64_req(obj, max_len=200): + if isinstance(obj, dict): + return {k: _truncate_base64_req(v, max_len) for k, v in obj.items()} + elif isinstance(obj, list): + return [_truncate_base64_req(item, max_len) for item in obj] + elif isinstance(obj, str) and len(obj) > max_len: + if all(c in 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=' for c in obj[:50]): + return f"" + return obj + return obj + + safe_request = _truncate_base64_req(request_body) + print( + f"\n{'='*60}\n" + f"[请求体日志] 任务 {task_prefix or '?'} 发送请求体：\n" + f"端点: {endpoint}\n" + f"{_json.dumps(safe_request, ensure_ascii=False, indent=2)}\n" + f"{'='*60}\n" + ) + + # 计算请求体大小 + request_size = len(json.dumps(request_body).encode('utf-8')) + if request_size < 1024 * 1024: + size_str = f"{request_size / 1024:.2f}KB" + else: + size_str = f"{request_size / (1024 * 1024):.2f}MB" + + # ========== 2. 发送网络请求 ========== + request_start = time.time() + + try: + response = await self.request_async(endpoint, request_body, session) + except Exception as e: + request_time = time.time() - request_start + error_first_line = str(e).split('\n')[0] + print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → 失败: {error_first_line} ✗") + raise + + request_time = time.time() - request_start + + # ========== 调试日志：打印完整响应 ========== + if debug: + import json as _json + # 构建可安全序列化的响应副本（截断 base64 图片数据避免输出过长） + def _truncate_base64(obj, max_len=200): + if isinstance(obj, dict): + return {k: _truncate_base64(v, max_len) for k, v in obj.items()} + elif isinstance(obj, list): + return [_truncate_base64(item, max_len) for item in obj] + elif isinstance(obj, str) and len(obj) > max_len: + # 判断是否为 base64 图片数据（不含空格/换行的长字符串） + if all(c in 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=' for c in obj[:50]): + return f"" + return obj + return obj + + safe_response = _truncate_base64(response) + print( + f"\n{'='*60}\n" + f"[调试日志] 任务 {task_prefix or '?'} 完整 API 响应：\n" + f"{_json.dumps(safe_response, ensure_ascii=False, indent=2)}\n" + f"{'='*60}\n" + ) + + # ========== 3. 解析响应 ========== + parse_start = time.time() + + try: + result_images, format_info = await self.parse_response_async(response, session) + except Exception as e: + parse_time = time.time() - parse_start + error_first_line = str(e).split('\n')[0] + print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → 解析失败: {error_first_line} ✗") + raise + + parse_time = time.time() - parse_start + + # ========== 4. 格式化输出（单行） ========== + # 格式化图像大小 + img_size = format_info.get("size", 0) + if img_size < 1024 * 1024: + img_size_str = f"{img_size / 1024:.2f}KB" + else: + img_size_str = f"{img_size / (1024 * 1024):.2f}MB" + + # 根据类型构建下载信息 + if format_info.get("type") == "base64": + download_info = f"Base64 {img_size_str} ({parse_time:.1f}s)" + elif format_info.get("type") == "url": + speed = format_info.get("download_speed", 0) + speed_str = f"{speed / (1024 * 1024):.1f}MB/s" + download_info = f"URL {img_size_str} ({parse_time:.1f}s, {speed_str})" + else: + download_info = f"{img_size_str}" + + # 单行输出 + print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → {download_info} ✓") + + # 返回结果和计时信息 + total_time = time.time() - total_start + timing_info = { + "build_time": build_time, + "request_time": request_time, + "parse_time": parse_time, + "total_time": total_time, + "format_type": format_info.get("type", "unknown") + } + + return result_images, timing_info + + async def generate_batch_async( + self, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + batch_size: int, + images: Optional[List[Image.Image]] = None, + progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None, + debug: bool = False, + debug_request: bool = False, + enable_grounding: bool = False, + enable_image_search: bool = False, + candidate_count: int = 1 + ) -> List[Image.Image]: + """ + 批量全并发生成 - 改进版：支持分批处理和内存管理 + + Args: + prompt: 提示词 + model: 模型名称 + resolution: 分辨率 + aspect_ratio: 宽高比 + batch_size: 批次大小 + images: 输入图像列表 + progress_callback: 进度回调，签名为 (completed, total, success, error_msg) + debug: 是否打印完整 API 响应 + debug_request: 是否打印发送的请求体 + enable_grounding: 是否启用 Google Search Grounding + enable_image_search: 是否同时启用 Google Image Search + candidate_count: 单次请求返回的候选图数量 + + Returns: + 生成的图像列表 + """ + import aiohttp + import asyncio + + all_images = [] + completed = 0 + success_count = 0 + fail_count = 0 + first_error = None # 保存第一个错误 + + # 分批处理配置 + max_concurrent = 10 # 最大并发数 + save_batch_size = 10 # 分批保存大小 + + # 计算需要多少批次 + num_batches = (batch_size + max_concurrent - 1) // max_concurrent + + print(f"GeminiClient: 批量生成 {batch_size} 张图片，并发数: {max_concurrent}，分 {num_batches} 批执行") + + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + + async with aiohttp.ClientSession(connector=connector) as session: + # 分批执行 + for batch_idx in range(num_batches): + batch_start = batch_idx * max_concurrent + batch_end = min(batch_start + max_concurrent, batch_size) + batch_size_current = batch_end - batch_start + + if num_batches > 1: + print(f"GeminiClient: 执行第 {batch_idx + 1}/{num_batches} 批 ({batch_start + 1}-{batch_end})...") + + # 创建当前批次的任务 + tasks = [] + for i in range(batch_size_current): + task_index = batch_start + i + task = asyncio.create_task( + self.generate_single_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=images, + session=session, + task_index=task_index + 1, + total_tasks=batch_size, + debug=debug, + debug_request=debug_request, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + candidate_count=candidate_count + ), + name=f"task_{task_index}" + ) + tasks.append(task) + + # 收集当前批次的结果 + batch_images = [] + batch_completed = 0 + + for coro in asyncio.as_completed(tasks): + batch_completed += 1 + completed += 1 + + try: + result_images, timing_info = await coro + if result_images: + # 立即处理生成的图片 + for img in result_images: + batch_images.append(img) + all_images.append(img) + + success_count += 1 + + # 通知进度 + if progress_callback: + progress_callback(completed, batch_size, True, None) + + # 每成功生成一张图片就打印日志 + print(f"GeminiClient: 任务 {completed}/{batch_size} 成功生成图片 ✓") + + except Exception as e: + fail_count += 1 + # 保存第一个错误（用于后续抛出） + if first_error is None: + first_error = e + error_msg = str(e) + + # 传递完整的错误信息（用于排查问题） + if progress_callback: + progress_callback(completed, batch_size, False, error_msg) + + print(f"GeminiClient: 任务 {completed}/{batch_size} 失败 ✗") + + # 当前批次完成后，立即清理内存 + if batch_images: + print(f"GeminiClient: 第 {batch_idx + 1} 批完成，生成 {len(batch_images)} 张图片") + + # 强制垃圾回收，释放内存 + import gc + gc.collect() + + # 短暂暂停，让系统处理内存 + await asyncio.sleep(0.1) + + # 清空当前批次图片引用，帮助垃圾回收 + batch_images = [] + + # 最终结果检查 + if not all_images: + # 如果有保存的原始错误，直接抛出原始错误 + if first_error: + raise first_error + raise RuntimeError(f"批量生成失败，{fail_count} 个请求全部失败") + + print(f"GeminiClient: 批量生成完成，成功 {success_count}/{batch_size}，失败 {fail_count}") + return all_images + + def generate_sync( + self, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + batch_size: int, + images: Optional[List[Image.Image]] = None, + progress_callback: Optional[Callable[[int, int], None]] = None, + debug: bool = False, + debug_request: bool = False, + enable_grounding: bool = False, + enable_image_search: bool = False, + candidate_count: int = 1 + ) -> List[Image.Image]: + """ + 同步生成接口（用于 ComfyUI） + + Args: + prompt: 提示词 + model: 模型名称 + resolution: 分辨率 + aspect_ratio: 宽高比 + batch_size: 批次大小 + images: 输入图像列表 + progress_callback: 进度回调 + debug: 是否打印完整 API 响应 + debug_request: 是否打印发送的请求体 + enable_grounding: 是否启用 Google Search Grounding + enable_image_search: 是否同时启用 Google Image Search + candidate_count: 单次请求返回的候选图数量 + + Returns: + 生成的图像列表 + """ + coro = self.generate_batch_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + batch_size=batch_size, + images=images, + progress_callback=progress_callback, + debug=debug, + debug_request=debug_request, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + candidate_count=candidate_count + ) + + return self.run_async_in_thread(coro) + + async def generate_multi_prompts_async( + self, + prompts: List[str], + model: str, + resolution: str, + aspect_ratio: str, + images_per_prompt: int, + images: Optional[List[Image.Image]] = None, + progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None, + debug: bool = False, + debug_request: bool = False, + enable_grounding: bool = False, + enable_image_search: bool = False + ) -> List[Image.Image]: + """ + 多提示词批量生成 - 改进版：支持分批处理和内存管理 + + 为每个提示词生成指定数量的图像，分批并发执行。 + + Args: + prompts: 提示词列表 + model: 模型名称 + resolution: 分辨率 + aspect_ratio: 宽高比 + images_per_prompt: 每个提示词生成的图像数量 + images: 输入图像列表（所有提示词共享） + progress_callback: 进度回调，签名为 (completed, total, success, error_msg) + debug: 是否打印完整 API 响应 + debug_request: 是否打印发送的请求体 + enable_grounding: 是否启用 Google Search Grounding + enable_image_search: 是否同时启用 Google Image Search + + Returns: + 生成的图像列表（长度 = len(prompts) * images_per_prompt） + """ + import aiohttp + import asyncio + + all_images = [] + completed = 0 + success_count = 0 + fail_count = 0 + first_error = None # 保存第一个错误 + total_tasks = len(prompts) * images_per_prompt + + # 分批处理配置 + max_concurrent = 10 # 最大并发数 + + print(f"GeminiClient: 多提示词批量生成，共 {total_tasks} 个任务，{len(prompts)} 个提示词，每个 {images_per_prompt} 张") + + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + + async with aiohttp.ClientSession(connector=connector) as session: + # 创建所有任务 + tasks = [] + task_idx = 0 + for prompt in prompts: + for _ in range(images_per_prompt): + task = asyncio.create_task( + self.generate_single_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=images, + session=session, + task_index=task_idx + 1, + total_tasks=total_tasks, + debug=debug, + debug_request=debug_request, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search + ), + name=f"task_{task_idx}" + ) + tasks.append(task) + task_idx += 1 + + # 分批处理：每10个任务为一组 + batch_size = max_concurrent + num_batches = (total_tasks + batch_size - 1) // batch_size + + for batch_idx in range(num_batches): + batch_start = batch_idx * batch_size + batch_end = min(batch_start + batch_size, total_tasks) + batch_tasks = tasks[batch_start:batch_end] + + if num_batches > 1: + print(f"GeminiClient: 执行第 {batch_idx + 1}/{num_batches} 批 ({batch_start + 1}-{batch_end})...") + + # 收集当前批次的结果 + batch_images = [] + + for coro in asyncio.as_completed(batch_tasks): + completed += 1 + + try: + result_images, timing_info = await coro + if result_images: + # 立即处理生成的图片 + for img in result_images: + batch_images.append(img) + all_images.append(img) + + success_count += 1 + + # 通知进度 + if progress_callback: + progress_callback(completed, total_tasks, True, None) + + # 每成功生成一张图片就打印日志 + print(f"GeminiClient: 任务 {completed}/{total_tasks} 成功生成图片 ✓") + + except Exception as e: + fail_count += 1 + # 保存第一个错误（用于后续抛出） + if first_error is None: + first_error = e + error_msg = str(e) + + # 传递完整的错误信息（用于排查问题） + if progress_callback: + progress_callback(completed, total_tasks, False, error_msg) + + print(f"GeminiClient: 任务 {completed}/{total_tasks} 失败 ✗") + + # 当前批次完成后，立即清理内存 + if batch_images: + print(f"GeminiClient: 第 {batch_idx + 1} 批完成，生成 {len(batch_images)} 张图片") + + # 强制垃圾回收，释放内存 + import gc + gc.collect() + + # 短暂暂停，让系统处理内存 + await asyncio.sleep(0.1) + + # 清空当前批次图片引用，帮助垃圾回收 + batch_images = [] + + if not all_images: + # 如果有保存的原始错误，直接抛出原始错误 + if first_error: + raise first_error + raise RuntimeError(f"批量生成失败，{fail_count} 个请求全部失败") + + print(f"GeminiClient: 多提示词批量生成完成，成功 {success_count}/{total_tasks}，失败 {fail_count}") + return all_images + + def generate_multi_prompts_sync( + self, + prompts: List[str], + model: str, + resolution: str, + aspect_ratio: str, + images_per_prompt: int, + images: Optional[List[Image.Image]] = None, + progress_callback: Optional[Callable[[int, int], None]] = None, + debug: bool = False, + debug_request: bool = False, + enable_grounding: bool = False, + enable_image_search: bool = False + ) -> List[Image.Image]: + """ + 多提示词批量生成（同步接口，用于 ComfyUI） + + Args: + prompts: 提示词列表 + model: 模型名称 + resolution: 分辨率 + aspect_ratio: 宽高比 + images_per_prompt: 每个提示词生成的图像数量 + images: 输入图像列表 + progress_callback: 进度回调 + debug: 是否打印完整 API 响应 + debug_request: 是否打印发送的请求体 + enable_grounding: 是否启用 Google Search Grounding + enable_image_search: 是否同时启用 Google Image Search + + Returns: + 生成的图像列表 + """ + coro = self.generate_multi_prompts_async( + prompts=prompts, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images_per_prompt=images_per_prompt, + images=images, + progress_callback=progress_callback, + debug=debug, + debug_request=debug_request, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search + ) + + return self.run_async_in_thread(coro) + \ No newline at end of file diff --git a/clients/gemini_flash_client.py b/clients/gemini_flash_client.py new file mode 100644 index 0000000..cfe989b --- /dev/null +++ b/clients/gemini_flash_client.py @@ -0,0 +1,303 @@ +""" +Gemini Flash API 客户端 +用于调用 Gemini 3 Flash 模型进行多模态文本生成 +""" + +import asyncio +from typing import Any, Dict, List, Optional + +import aiohttp + +from ..utils.config import get_api_key_or_raise, get_api_base_url +from ..models_config import ( + get_flash_model_endpoint, + get_enabled_flash_models, + get_flash_model_thinking_level_value, +) +from .base_client import BaseAPIClient + + +class GeminiFlashClient(BaseAPIClient): + """ + Gemini Flash API 客户端 + 用于调用 Gemini 3 Flash 模型进行多模态文本生成 + + 特点： + - 支持图片和视频输入 + - 支持动态思考等级端点（不思考/低/中/高） + """ + + def __init__(self, api_key: Optional[str] = None): + """ + 初始化客户端 + + Args: + api_key: API 密钥，如果为 None 则从配置文件或环境变量读取 + """ + if api_key is None: + api_key = get_api_key_or_raise("O1KEY_API_KEY") + + super().__init__( + base_url=get_api_base_url(), + api_key=api_key, + max_request_size=100 * 1024 * 1024 # 100MB + ) + + def get_endpoint( + self, + model: str = "gemini-3-flash-preview", + **kwargs + ) -> str: + """ + 获取模型的 API 端点 + + Args: + model: 模型名称 + + Returns: + API 端点路径 + """ + endpoint = get_flash_model_endpoint(model) + + if endpoint is None: + # 回退到第一个启用的模型端点 + default_models = get_enabled_flash_models() + if default_models: + endpoint = get_flash_model_endpoint(default_models[0]) + + if endpoint is None: + raise ValueError(f"无法获取模型 '{model}' 的端点") + + return endpoint + + def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]: + """Gemini 请求 429/503 时返回图中约定的多行错误框文案。""" + if status_code == 429: + return ( + "莫慌!该模型暂时超出速率限制啦\n" + "解决方案如下(任意一种):\n" + "1.切换当前模型\n" + "2.前往后台,修改令牌分组" + ) + if status_code == 503: + return ( + "警报!谷歌服务器当前过载!\n" + "解决方案如下:\n" + "1.摸会儿鱼吧,我也没办法,谷歌会尽快恢复,嘿嘿~\n" + "2.切换其他模型\n" + "3.前往后台,修改令牌分组" + ) + return None + + def build_request_body( + self, + prompt: str = "", + model: str = "gemini-3-flash-preview", + thinking_level: str = "不思考", + image_data: Optional[List[Dict[str, str]]] = None, + video_data: Optional[Dict[str, str]] = None, + document_data: Optional[Dict[str, str]] = None, + **kwargs + ) -> Dict[str, Any]: + """ + 构建 API 请求体 + + Args: + prompt: 用户提示词 + model: 模型名称 + thinking_level: 思考等级（不思考/低/中/高）- 通过动态端点控制，不需要在请求体中传递 + image_data: 图片数据列表，每个元素包含 mime_type 和 data + video_data: 视频数据，包含 mime_type 和 data + document_data: 文档数据，包含 mime_type 和 data + + Returns: + 请求体字典 + """ + parts = [] + + # 添加文本部分 + if prompt: + parts.append({"text": prompt}) + + # 添加图片部分（如果有） + if image_data: + for img in image_data: + parts.append({ + "inline_data": { + "mime_type": img["mime_type"], + "data": img["data"] + } + }) + + # 添加视频部分（如果有） + if video_data: + parts.append({ + "inline_data": { + "mime_type": video_data["mime_type"], + "data": video_data["data"] + } + }) + + # 添加文档部分（如果有） + if document_data: + parts.append({ + "inline_data": { + "mime_type": document_data["mime_type"], + "data": document_data["data"] + } + }) + + # 构建请求体 + request_body = { + "contents": [ + { + "parts": parts + } + ] + } + + # 对于支持 thinkingConfig 的固定端点模型（如 gemini-3-pro-preview） + # 通过请求体传递思考等级；动态端点模型（如 gemini-3-flash-preview） + # 通过不同 URL 端点控制，无需此字段 + thinking_level_value = get_flash_model_thinking_level_value(model, thinking_level) + if thinking_level_value is not None: + request_body["generationConfig"] = { + "thinkingConfig": { + "thinkingLevel": thinking_level_value + } + } + + return request_body + + def parse_response(self, response: Dict[str, Any]) -> str: + """ + 解析 API 响应，提取生成的文本 + + Args: + response: API 响应字典 + + Returns: + 生成的文本内容 + + Raises: + RuntimeError: 解析失败或 API 拒绝时 + """ + # 检查 candidatesTokenCount + usage_metadata = response.get("usageMetadata", {}) + candidates_token_count = usage_metadata.get("candidatesTokenCount", -1) + + if candidates_token_count == 0: + raise RuntimeError( + "内容审核拒绝 - candidatesTokenCount = 0\n\n" + "原因：提示词或输入内容包含不适当内容\n" + "建议：检查并调整输入内容" + ) + + # 检查 finishReason + candidates = response.get("candidates", []) + if candidates: + for candidate in candidates: + finish_reason = candidate.get("finishReason", "") + + if finish_reason and finish_reason not in ["STOP", "MAX_TOKENS"]: + reason_messages = { + "PROHIBITED_CONTENT": "违禁内容拒绝", + "SAFETY": "安全过滤器拒绝", + "RECITATION": "版权问题" + } + error_title = reason_messages.get(finish_reason, f"生成异常 ({finish_reason})") + raise RuntimeError(f"{error_title}\n建议：调整输入内容后重试") + + # 提取文本内容 + text_parts = [] + + for candidate in candidates: + content = candidate.get("content", {}) + parts = content.get("parts", []) + + for part in parts: + if "text" in part: + text_parts.append(part["text"]) + + if not text_parts: + raise RuntimeError("API 响应中未找到生成的文本") + + # 合并所有文本部分 + return "\n".join(text_parts) + + async def generate_async( + self, + prompt: str, + model: str = "gemini-3-flash-preview", + thinking_level: str = "不思考", + image_data: Optional[List[Dict[str, str]]] = None, + video_data: Optional[Dict[str, str]] = None, + document_data: Optional[Dict[str, str]] = None, + session: Optional[aiohttp.ClientSession] = None + ) -> str: + """ + 异步生成文本 + + Args: + prompt: 用户提示词 + model: 模型名称 + thinking_level: 思考等级（不思考/低/中/高） + image_data: 图片数据列表 + video_data: 视频数据 + document_data: 文档数据 + session: aiohttp 会话 + + Returns: + 生成的文本内容 + """ + endpoint = self.get_endpoint(model=model) + request_body = self.build_request_body( + prompt=prompt, + model=model, + thinking_level=thinking_level, + image_data=image_data, + video_data=video_data, + document_data=document_data + ) + + response = await self.request_async( + endpoint, + request_body, + session + ) + + return self.parse_response(response) + + def generate_sync( + self, + prompt: str, + model: str = "gemini-3-flash-preview", + thinking_level: str = "不思考", + image_data: Optional[List[Dict[str, str]]] = None, + video_data: Optional[Dict[str, str]] = None, + document_data: Optional[Dict[str, str]] = None + ) -> str: + """ + 同步生成文本（用于 ComfyUI 节点） + + Args: + prompt: 用户提示词 + model: 模型名称 + thinking_level: 思考等级（不思考/低/中/高） + image_data: 图片数据列表 + video_data: 视频数据 + document_data: 文档数据 + + Returns: + 生成的文本内容 + """ + coro = self.generate_async( + prompt=prompt, + model=model, + thinking_level=thinking_level, + image_data=image_data, + video_data=video_data, + document_data=document_data + ) + + return self.run_async_in_thread(coro) diff --git a/clients/kling_client.py b/clients/kling_client.py new file mode 100644 index 0000000..f843311 --- /dev/null +++ b/clients/kling_client.py @@ -0,0 +1,174 @@ +""" +Kling 视频生成 API 客户端 +""" + +import asyncio +import json +import os +from typing import Any, Callable, Dict, Optional + +import aiohttp + +from ..utils.config import get_api_key_or_raise, get_api_base_url + + +class KlingClient: + """Kling 视频生成客户端""" + + ENDPOINTS = { + "image2video": "/kling/v1/videos/image2video", + "text2video": "/kling/v1/videos/text2video", + "motion_control": "/kling/v1/videos/motion-control", + } + + POLL_INITIAL_INTERVAL = 3 + POLL_MAX_INTERVAL = 15 + + def __init__(self): + self.api_key = get_api_key_or_raise() + self.base_url = get_api_base_url() + + def _headers(self) -> Dict[str, str]: + return { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + + # ── 提交任务 ────────────────────────────────────────────────────── + + async def create_video_async( + self, + endpoint_type: str, + body: Dict[str, Any], + session: aiohttp.ClientSession, + ) -> Dict[str, Any]: + url = f"{self.base_url}{self.ENDPOINTS[endpoint_type]}" + + async with session.post(url, json=body, headers=self._headers()) as resp: + text = await resp.text() + if resp.status != 200: + raise RuntimeError(f"提交失败 ({resp.status}): {text}") + return json.loads(text) + + # ── 轮询状态 ────────────────────────────────────────────────────── + + async def poll_status_async( + self, + task_id: str, + endpoint_type: str, + session: aiohttp.ClientSession, + on_progress: Optional[Callable[[int], None]] = None, + ) -> Dict[str, Any]: + url = f"{self.base_url}{self.ENDPOINTS[endpoint_type]}/{task_id}" + interval = self.POLL_INITIAL_INTERVAL + + while True: + async with session.get(url, headers=self._headers()) as resp: + text = await resp.text() + if resp.status != 200: + raise RuntimeError(f"状态查询失败 ({resp.status}): {text}") + result = json.loads(text) + + data = result.get("data", {}) + inner_data = data.get("data", {}) if isinstance(data, dict) else {} + status = ( + data.get("status") or + inner_data.get("task_status") or + result.get("status") or + "" + ) + status = status.lower() if status else "" + + progress_str = data.get("progress", "0%") + try: + progress_pct = int(str(progress_str).replace("%", "").strip()) + except (ValueError, AttributeError): + progress_pct = 0 + + print(f"[视频生成] 生成中 {progress_pct}%") + + if on_progress: + on_progress(progress_pct) + + if status in ("success", "completed", "done", "finished", "succeed"): + return result + elif status in ("failed", "fail"): + error_info = result.get("error", {}) + if isinstance(error_info, dict): + error_msg = error_info.get("message", "未知错误") + else: + error_msg = str(error_info) + raise RuntimeError(f"生成失败：{error_msg}") + + await asyncio.sleep(interval) + interval = min(interval * 1.5, self.POLL_MAX_INTERVAL) + + # ── 下载视频 ────────────────────────────────────────────────────── + + async def download_video_async( + self, + video_url: str, + save_path: str, + session: aiohttp.ClientSession, + ) -> str: + print("[视频生成] 下载视频...") + async with session.get(video_url, allow_redirects=True) as resp: + if resp.status != 200: + raise RuntimeError(f"视频下载失败 ({resp.status})") + os.makedirs(os.path.dirname(save_path), exist_ok=True) + with open(save_path, "wb") as f: + async for chunk in resp.content.iter_chunked(8192): + f.write(chunk) + return save_path + + # ── 异步入口（供节点调用）──────────────────────────────────────── + + async def generate_async( + self, + endpoint_type: str, + body: Dict[str, Any], + save_path: str, + on_stage: Optional[Callable[[str], None]] = None, + on_progress: Optional[Callable[[int], None]] = None, + ) -> str: + """提交 → 轮询 → 下载，返回本地文件路径""" + connector = aiohttp.TCPConnector(force_close=True) + async with aiohttp.ClientSession(connector=connector) as session: + if on_stage: + on_stage("submitting") + + result = await self.create_video_async(endpoint_type, body, session) + # 提交响应结构：result.data.task_id + task_id = result.get("task_id") or result.get("data", {}).get("task_id") + if not task_id: + raise RuntimeError(f"API 未返回任务 ID，响应：{result}") + if on_stage: + on_stage(f"submitted:{task_id}") + + if on_stage: + on_stage("polling") + final = await self.poll_status_async( + task_id, endpoint_type, session, on_progress=on_progress + ) + + # 兼容多种URL路径 + # 响应结构：result.data.result_url 或 result.data.data.task_result.videos[0].url + data = final.get("data", {}) + inner_data = data.get("data", {}) if isinstance(data, dict) else {} + video_url = ( + data.get("result_url") or + final.get("url") or + final.get("video_url") or + (inner_data.get("task_result", {}).get("videos", [{}])[0].get("url") + if inner_data.get("task_result", {}).get("videos") else None) + ) + if not video_url: + raise RuntimeError(f"API 未返回视频 URL，响应：{final}") + + if on_stage: + on_stage("downloading") + path = await self.download_video_async(video_url, save_path, session) + + if on_stage: + on_stage("done") + return path diff --git a/clients/openai_client.py b/clients/openai_client.py new file mode 100644 index 0000000..ba99088 --- /dev/null +++ b/clients/openai_client.py @@ -0,0 +1,775 @@ +""" +OpenAI 兼容 API 客户端 +端点固定为 /v1/chat/completions，模型名放入请求体 model 字段 +""" + +import re +import time +from io import BytesIO +from typing import Any, Callable, Dict, List, Optional + +import aiohttp +from PIL import Image + +from ..utils.image_utils import encode_image_to_base64, decode_base64_to_pil +from ..utils.config import get_api_key_or_raise, get_api_base_url +from .base_client import BaseAPIClient + + +# 固定端点 +_ENDPOINT = "/v1/chat/completions" + + +class OpenAIAPIClient(BaseAPIClient): + """ + OpenAI 兼容格式的图像生成客户端 + + 与 GeminiAPIClient 的主要区别： + - 端点固定为 /v1/chat/completions（不再动态拼模型名到 URL） + - 解析后的模型字符串放入请求体的 model 字段 + - 请求体采用 messages 数组格式，图片以 data URI 内联 + - 顶层追加 modalities 和 image_config 字段 + - 响应解析对应 choices[0].message.content 结构 + """ + + def __init__(self, api_key: Optional[str] = None): + if api_key is None: + api_key = get_api_key_or_raise("O1KEY_API_KEY") + + super().__init__( + base_url=get_api_base_url(), + api_key=api_key, + max_request_size=100 * 1024 * 1024 + ) + + # ------------------------------------------------------------------ # + # 模型名解析 # + # 原 GeminiAPIClient.get_endpoint() 里动态拼 URL 的逻辑 # + # 现在改为：同样的输入 → 返回纯模型名字符串，放进请求体 # + # ------------------------------------------------------------------ # + + def resolve_model_name(self, model: str, resolution: str) -> str: + """ + 将「节点选中的模型 ID + 分辨率」解析为实际请求所用的模型名称。 + + 对应关系与原 GeminiAPIClient.get_endpoint() 完全一致， + 只是把拼在 URL 路径里的模型段提取出来单独返回。 + + Args: + model: 节点下拉框中的模型 ID，如 "nano-banana-pro-限时特价" + resolution: 分辨率字符串，如 "1K" / "2K" / "4K" / "512" + + Returns: + 实际模型名，如 "nano-banana-pro-2k" + """ + # ── 动态端点模型 ────────────────────────────────────────────────── + if model == "nano-banana-pro-限时特价": + if resolution == "1K": + return "nano-banana-pro" + elif resolution == "4K": + return "nano-banana-pro-4k" + else: # 2K（默认） + return "nano-banana-pro-2k" + + elif model == "nano-banana-pro-官方计费": + if resolution == "1K": + return "nano-banana-pro-1k-official" + elif resolution == "4K": + return "nano-banana-pro-4k-official" + else: # 2K（默认） + return "nano-banana-pro-2k-official" + + elif model == "nano-banana-2-官方计费": + if resolution == "512": + return "nano-banana-2-0.5k-official" + elif resolution == "1K": + return "nano-banana-2-1k-official" + elif resolution == "4K": + return "nano-banana-2-4k-official" + else: # 2K（默认） + return "nano-banana-2-2k-official" + + elif model == "gemini-3-pro-image-preview-url": + if resolution == "1K": + return "gemini-3-pro-image-preview-url" + elif resolution == "4K": + return "gemini-3-pro-image-preview-4k-url" + else: # 2K（默认） + return "gemini-3-pro-image-preview-2k-url" + + # ── 固定端点模型：从 models_config 里取端点，提取模型名段 ────────── + from ..models_config import get_model_endpoint + endpoint = get_model_endpoint(model) + if endpoint: + # 端点格式：/v1beta/models/:generateContent + # 提取 部分 + match = re.search(r"/models/([^:]+):", endpoint) + if match: + return match.group(1) + + # ── 兜底：直接用 model ID ────────────────────────────────────────── + return model + + # ------------------------------------------------------------------ # + # BaseAPIClient 抽象方法实现 # + # ------------------------------------------------------------------ # + + def get_endpoint(self, **kwargs) -> str: + """固定返回 /v1/chat/completions，模型信息已移入请求体。""" + return _ENDPOINT + + def build_request_body( + self, + prompt: str = "", + images: Optional[List[Image.Image]] = None, + aspect_ratio: str = "1:1", + resolution: str = "2K", + model: str = "", + **kwargs + ) -> Dict[str, Any]: + """ + 构建 OpenAI /v1/chat/completions 格式请求体。 + + 文生图示例输出： + { + "model": "nano-banana-pro-2k", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "一个中国女子的OOTD"} + ] + } + ], + "modalities": ["image", "text"], + "stream": false, + "extra_body": { + "google": { + "image_config": { + "aspect_ratio": "16:9", + "image_size": "2K" + } + } + } + } + + 图生图时 content 数组追加若干 image_url 块： + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,<...>"} + } + + Args: + prompt: 提示词 + images: 参考图列表（可选，图生图时传入） + aspect_ratio: 宽高比，如 "16:9" + resolution: 分辨率，如 "2K" + model: 已解析好的模型名（由 resolve_model_name 返回） + """ + # ── 构建 content 数组 ───────────────────────────────────────────── + content: List[Dict[str, Any]] = [] + + # 1. 文本部分（始终在最前） + content.append({ + "type": "text", + "text": prompt + }) + + # 2. 图片部分（图生图时追加，每张图一个 image_url block） + if images: + for img in images: + b64 = encode_image_to_base64(img) + content.append({ + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{b64}" + } + }) + + # ── 分辨率映射（节点内部值 → API 所需值） ──────────────────────────── + _resolution_map = {"512": "0.5K", "1K": "1K", "2K": "2K", "4K": "4K"} + api_image_size = _resolution_map.get(resolution, resolution) + + # ── 组装完整请求体 ───────────────────────────────────────────────── + request_body: Dict[str, Any] = { + "model": model, + "messages": [ + { + "role": "user", + "content": content + } + ], + "modalities": ["image", "text"], + "stream": False, + "extra_body": { + "google": { + "image_config": { + "aspect_ratio": aspect_ratio, + "image_size": api_image_size + } + } + } + } + + return request_body + + def parse_response(self, response: Dict[str, Any]) -> List[Image.Image]: + """同步 parse_response，仅为满足抽象基类要求，实际不应被直接调用。""" + raise RuntimeError( + "parse_response() 不应被直接调用。" + "请使用 generate_single_async() 等高级方法。" + ) + + def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]: + """429 / 503 友好文案。""" + if status_code == 429: + return ( + "莫慌!该模型暂时超出速率限制啦\n" + "解决方案如下(任意一种):\n" + "1.切换当前模型\n" + "2.前往后台,修改令牌分组" + ) + if status_code == 503: + return ( + "警报!服务器当前过载!\n" + "解决方案如下:\n" + "1.摸会儿鱼吧,稍后会恢复,嘿嘿~\n" + "2.切换其他模型\n" + "3.前往后台,修改令牌分组" + ) + return None + + # ------------------------------------------------------------------ # + # 响应解析 # + # ------------------------------------------------------------------ # + + async def parse_response_async( + self, + response: Dict[str, Any], + session: Optional[aiohttp.ClientSession] = None + ) -> tuple[List[Image.Image], Dict[str, Any]]: + """ + 异步解析 /v1/chat/completions 格式响应，提取生成的图像。 + + 响应结构（OpenAI 格式）： + { + "choices": [ + { + "message": { + "role": "assistant", + "content": [ + {"type": "text", "text": "..."}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}} + // 或直接 inline_data / inlineData（兼容 Gemini 风格回包） + ] + }, + "finish_reason": "stop" + } + ], + "usage": {...} + } + """ + format_info: Dict[str, Any] = { + "type": None, # "base64" | "url" + "size": 0, + "resolution": None, + "download_speed": None + } + + # ── 错误前置检测 ─────────────────────────────────────────────────── + + # 1. usage.completion_tokens == 0 → 风控拦截（对齐 Gemini 的 candidatesTokenCount==0） + usage = response.get("usage", {}) + completion_tokens = usage.get("completion_tokens", -1) + if completion_tokens == 0: + raise RuntimeError( + "Damn!你触发顶级风控啦!还没到生图阶段就被拒了。\n" + "赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~" + ) + + # 2. finish_reason 不是 "stop" → 安全过滤 / token 超限等 + choices = response.get("choices", []) + if choices: + for choice in choices: + finish_reason = choice.get("finish_reason", "") + if finish_reason and finish_reason != "stop": + raise RuntimeError( + "Ohh no! 生图过程触发风控,图片被拒绝生成!\n" + "可能原因如下:\n" + "1.违禁内容\n" + "2.触发安全过滤器\n" + "3.涉及版权问题\n" + "4. Token超限\n" + "赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~" + ) + + # ── 图像提取 ─────────────────────────────────────────────────────── + images: List[Image.Image] = [] + text_responses: List[str] = [] + + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + try: + for choice in choices: + message = choice.get("message", {}) + + # ── 优先从 message.images 提取（非标准扩展字段） ────────────── + # 部分服务端把图片放在独立的 images 字段，content 同时为 null + msg_images = message.get("images") or [] + for img_part in msg_images: + part_type = img_part.get("type", "") + if part_type == "image_url": + url_obj = img_part.get("image_url", {}) + url = url_obj.get("url", "") + if url.startswith("data:"): + try: + _, b64_data = url.split(",", 1) + img = decode_base64_to_pil(b64_data) + images.append(img) + if format_info["type"] is None: + format_info["type"] = "base64" + format_info["size"] = len(b64_data) * 3 / 4 + format_info["resolution"] = f"{img.size[0]}x{img.size[1]}" + except Exception: + pass + elif url.startswith("http"): + try: + dl_start = time.time() + async with session.get(url) as img_resp: + if img_resp.status == 200: + img_data = await img_resp.read() + dl_time = time.time() - dl_start + speed = len(img_data) / dl_time if dl_time > 0 else 0 + img = Image.open(BytesIO(img_data)) + images.append(img) + if format_info["type"] is None: + format_info["type"] = "url" + format_info["size"] = len(img_data) + format_info["resolution"] = f"{img.size[0]}x{img.size[1]}" + format_info["download_speed"] = speed + except Exception: + pass + + # ── 再从 message.content 提取（标准 OpenAI 格式） ───────────── + # content 为 null 时用空列表兜底，避免 for in None 崩溃 + raw_content = message.get("content") or [] + + # content 可能是字符串（纯文本）或数组（多模态） + if isinstance(raw_content, str): + text_responses.append(raw_content) + continue + + for part in raw_content: + part_type = part.get("type", "") + + # ── 情况 A：OpenAI image_url 格式 ───────────────────── + if part_type == "image_url": + url_obj = part.get("image_url", {}) + url = url_obj.get("url", "") + + if url.startswith("data:"): + # data URI → 直接 base64 解码 + # 格式：data:image/png;base64, + try: + header, b64_data = url.split(",", 1) + img = decode_base64_to_pil(b64_data) + images.append(img) + if format_info["type"] is None: + format_info["type"] = "base64" + format_info["size"] = len(b64_data) * 3 / 4 + format_info["resolution"] = f"{img.size[0]}x{img.size[1]}" + except Exception: + pass + + elif url.startswith("http"): + # 远程 URL → 异步下载 + try: + dl_start = time.time() + async with session.get(url) as img_resp: + if img_resp.status == 200: + img_data = await img_resp.read() + dl_time = time.time() - dl_start + speed = len(img_data) / dl_time if dl_time > 0 else 0 + img = Image.open(BytesIO(img_data)) + images.append(img) + if format_info["type"] is None: + format_info["type"] = "url" + format_info["size"] = len(img_data) + format_info["resolution"] = f"{img.size[0]}x{img.size[1]}" + format_info["download_speed"] = speed + except Exception: + pass + + # ── 情况 B：Gemini 风格 inline_data / inlineData（兼容） ─ + elif part_type in ("inline_data", "inlineData") or \ + "inline_data" in part or "inlineData" in part: + inline_key = "inline_data" if "inline_data" in part else "inlineData" + inline = part.get(inline_key, {}) + b64_data = inline.get("data", "") + if b64_data: + try: + img = decode_base64_to_pil(b64_data) + images.append(img) + if format_info["type"] is None: + format_info["type"] = "base64" + format_info["size"] = len(b64_data) * 3 / 4 + format_info["resolution"] = f"{img.size[0]}x{img.size[1]}" + except Exception: + pass + + # ── 情况 C：text 中嵌套 URL（markdown 或纯链接） ───────── + elif part_type == "text": + text = part.get("text", "") + text_responses.append(text) + + # markdown 图片链接：![alt](url) + urls = re.findall(r'!\[.*?\]\((https?://[^\)]+)\)', text) + if not urls: + urls = re.findall(r'https?://[^\s<>"{}|\\^`\[\]]+', text) + + for url in urls: + try: + dl_start = time.time() + async with session.get(url) as img_resp: + if img_resp.status == 200: + img_data = await img_resp.read() + dl_time = time.time() - dl_start + speed = len(img_data) / dl_time if dl_time > 0 else 0 + img = Image.open(BytesIO(img_data)) + images.append(img) + if format_info["type"] is None: + format_info["type"] = "url" + format_info["size"] = len(img_data) + format_info["resolution"] = f"{img.size[0]}x{img.size[1]}" + format_info["download_speed"] = speed + except Exception: + pass + + except RuntimeError: + raise + except Exception as e: + raise RuntimeError(f"解析 API 响应失败: {str(e)}") + finally: + if close_session: + await session.close() + + # ── 3. 无图像但有文本 → API 拒绝说明 ───────────────────────────── + if not images and text_responses: + combined = "\n".join(text_responses) + raise RuntimeError( + f"API 拒绝响应\n\n" + f"API 返回说明：\n{combined}\n\n" + f"建议：\n" + f" - 根据上述说明调整请求内容\n" + f" - 确保提示词和参考图符合使用规范" + ) + + if not images: + raise RuntimeError("API 响应中未找到生成的图像") + + return images, format_info + + # ------------------------------------------------------------------ # + # 核心生成方法（接口与 GeminiAPIClient 保持一致，节点可无缝切换） # + # ------------------------------------------------------------------ # + + async def generate_single_async( + self, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: Optional[List[Image.Image]] = None, + session: Optional[aiohttp.ClientSession] = None, + task_index: Optional[int] = None, + total_tasks: Optional[int] = None, + debug: bool = False, + debug_request: bool = False, + enable_grounding: bool = False, # 保留签名兼容，OpenAI 格式暂不使用 + enable_image_search: bool = False # 保留签名兼容，OpenAI 格式暂不使用 + ) -> tuple[List[Image.Image], Dict[str, Any]]: + """ + 单次异步生成请求（OpenAI /v1/chat/completions 格式）。 + + Args: + prompt: 提示词 + model: 节点选中的模型 ID（将自动解析为实际模型名） + resolution: 分辨率 + aspect_ratio: 宽高比 + images: 参考图列表（图生图时传入） + session: 复用的 aiohttp 会话 + task_index: 任务序号（批量时用于日志） + total_tasks: 总任务数（批量时用于日志） + debug: 打印完整 API 响应 + debug_request: 打印请求体（base64 自动截断） + + Returns: + (生成的图像列表, 计时信息字典) + """ + import json + + total_start = time.time() + task_prefix = f"[{task_index}/{total_tasks}]" if task_index is not None and total_tasks else "" + + # ── 1. 解析模型名 & 构建请求体 ──────────────────────────────────── + build_start = time.time() + resolved_model = self.resolve_model_name(model, resolution) + endpoint = self.get_endpoint() + + request_body = self.build_request_body( + prompt=prompt, + images=images, + aspect_ratio=aspect_ratio, + resolution=resolution, + model=resolved_model + ) + build_time = time.time() - build_start + + # ── 调试：打印请求体 ─────────────────────────────────────────────── + if debug_request: + import json as _json + def _shorten_b64(obj): + if isinstance(obj, dict): + return {k: _shorten_b64(v) for k, v in obj.items()} + if isinstance(obj, list): + return [_shorten_b64(i) for i in obj] + if isinstance(obj, str): + if obj.startswith("data:"): + header, _, data = obj.partition(",") + return f"{header}," + if len(obj) > 200 and all( + c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=" + for c in obj[:64] + ): + return f"" + return obj + print( + f"\n{'='*60}\n" + f"[请求体日志] 任务 {task_prefix or '?'}\n" + f"端点: {self.base_url}{endpoint}\n" + f"{_json.dumps(_shorten_b64(request_body), ensure_ascii=False, indent=2)}\n" + f"{'='*60}\n" + ) + + # ── 2. 计算请求体大小 ───────────────────────────────────────────── + request_size = len(json.dumps(request_body).encode("utf-8")) + size_str = ( + f"{request_size / 1024:.2f}KB" + if request_size < 1024 * 1024 + else f"{request_size / (1024 * 1024):.2f}MB" + ) + + # ── 3. 发送请求（Bearer Token 认证） ───────────────────────────── + request_start = time.time() + try: + response = await self.request_async( + endpoint, + request_body, + session, + use_bearer_token=True + ) + except Exception as e: + request_time = time.time() - request_start + error_first_line = str(e).split("\n")[0] + print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → 失败: {error_first_line} ✗") + raise + + request_time = time.time() - request_start + + # ── 调试：打印完整响应 ───────────────────────────────────────────── + if debug: + import json as _json + def _shorten_b64(obj): + if isinstance(obj, dict): + return {k: _shorten_b64(v) for k, v in obj.items()} + if isinstance(obj, list): + return [_shorten_b64(i) for i in obj] + if isinstance(obj, str): + if obj.startswith("data:"): + header, _, data = obj.partition(",") + return f"{header}," + if len(obj) > 200 and all( + c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=" + for c in obj[:64] + ): + return f"" + return obj + print( + f"\n{'='*60}\n" + f"[调试日志] 任务 {task_prefix or '?'} 完整 API 响应：\n" + f"{_json.dumps(_shorten_b64(response), ensure_ascii=False, indent=2)}\n" + f"{'='*60}\n" + ) + + # ── 4. 解析响应 ─────────────────────────────────────────────────── + parse_start = time.time() + try: + result_images, format_info = await self.parse_response_async(response, session) + except Exception as e: + parse_time = time.time() - parse_start + error_first_line = str(e).split("\n")[0] + print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → 解析失败: {error_first_line} ✗") + raise + + parse_time = time.time() - parse_start + + # ── 5. 单行日志输出 ─────────────────────────────────────────────── + img_size = format_info.get("size", 0) + img_size_str = ( + f"{img_size / 1024:.2f}KB" + if img_size < 1024 * 1024 + else f"{img_size / (1024 * 1024):.2f}MB" + ) + + if format_info.get("type") == "base64": + download_info = f"Base64 {img_size_str} ({parse_time:.1f}s)" + elif format_info.get("type") == "url": + speed = format_info.get("download_speed", 0) + download_info = f"URL {img_size_str} ({parse_time:.1f}s, {speed / (1024*1024):.1f}MB/s)" + else: + download_info = img_size_str + + print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → {download_info} ✓") + + total_time = time.time() - total_start + timing_info = { + "build_time": build_time, + "request_time": request_time, + "parse_time": parse_time, + "total_time": total_time, + "format_type": format_info.get("type", "unknown") + } + + return result_images, timing_info + + # ------------------------------------------------------------------ # + # 批量 & 同步接口（与 GeminiAPIClient 接口签名一致） # + # ------------------------------------------------------------------ # + + async def generate_batch_async( + self, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + batch_size: int, + images: Optional[List[Image.Image]] = None, + progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None, + debug: bool = False, + debug_request: bool = False, + enable_grounding: bool = False, + enable_image_search: bool = False + ) -> List[Image.Image]: + """批量全并发生成（单提示词 × batch_size 张）。""" + import asyncio + + all_images: List[Image.Image] = [] + completed = 0 + success_count = 0 + fail_count = 0 + first_error = None + + max_concurrent = 10 + num_batches = (batch_size + max_concurrent - 1) // max_concurrent + + print(f"OpenAIClient: 批量生成 {batch_size} 张，并发数: {max_concurrent}，分 {num_batches} 批") + + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + + async with aiohttp.ClientSession(connector=connector) as session: + for batch_idx in range(num_batches): + batch_start = batch_idx * max_concurrent + batch_end = min(batch_start + max_concurrent, batch_size) + batch_count = batch_end - batch_start + + if num_batches > 1: + print(f"OpenAIClient: 第 {batch_idx + 1}/{num_batches} 批 ({batch_start + 1}-{batch_end})") + + tasks = [ + asyncio.create_task( + self.generate_single_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=images, + session=session, + task_index=batch_start + i + 1, + total_tasks=batch_size, + debug=debug, + debug_request=debug_request + ), + name=f"task_{batch_start + i}" + ) + for i in range(batch_count) + ] + + batch_images: List[Image.Image] = [] + + for coro in asyncio.as_completed(tasks): + completed += 1 + try: + result_imgs, _ = await coro + for img in result_imgs: + batch_images.append(img) + all_images.append(img) + success_count += 1 + if progress_callback: + progress_callback(completed, batch_size, True, None) + print(f"OpenAIClient: 任务 {completed}/{batch_size} 成功 ✓") + except Exception as e: + fail_count += 1 + if first_error is None: + first_error = e + if progress_callback: + progress_callback(completed, batch_size, False, str(e)) + print(f"OpenAIClient: 任务 {completed}/{batch_size} 失败 ✗") + + if batch_images: + print(f"OpenAIClient: 第 {batch_idx + 1} 批完成，生成 {len(batch_images)} 张") + import gc + gc.collect() + await asyncio.sleep(0.1) + + batch_images = [] + + if not all_images: + if first_error: + raise first_error + raise RuntimeError(f"批量生成失败，{fail_count} 个请求全部失败") + + print(f"OpenAIClient: 批量完成，成功 {success_count}/{batch_size}，失败 {fail_count}") + return all_images + + def generate_sync( + self, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + batch_size: int, + images: Optional[List[Image.Image]] = None, + progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None, + debug: bool = False, + debug_request: bool = False, + enable_grounding: bool = False, + enable_image_search: bool = False + ) -> List[Image.Image]: + """同步生成接口（用于 ComfyUI 节点，接口与 GeminiAPIClient 完全一致）。""" + coro = self.generate_batch_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + batch_size=batch_size, + images=images, + progress_callback=progress_callback, + debug=debug, + debug_request=debug_request, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search + ) + return self.run_async_in_thread(coro) diff --git a/clients/sora_client.py b/clients/sora_client.py new file mode 100644 index 0000000..4d76bba --- /dev/null +++ b/clients/sora_client.py @@ -0,0 +1,530 @@ +""" +Sora 视频生成 API 客户端 +提供视频创建、状态轮询、视频下载功能 +""" + +import asyncio +import base64 +import json +import os +import time +from typing import Any, Callable, Dict, List, Optional + +import aiohttp + +from .base_client import BaseAPIClient +from ..utils.config import get_api_key_or_raise, get_api_base_url +from ..utils.image_utils import encode_image_to_base64 + + +def _translate_error_message(msg: str) -> str: + """将 API 返回的已知英文错误信息翻译为中文友好提示""" + if "people-in-user-uploads" in msg or ( + "moderation" in msg and "inputs" in msg + ): + return "上传的参考图片中包含了真实人物【官方风控】，请尝试使用其他办法绕开。" + return msg + + +class SoraClient(BaseAPIClient): + """ + Sora 视频生成客户端 + + 工作流程： + 1. create_video → POST /v1/videos (提交生成任务) + 2. poll_status → GET /v1/videos/{id} (轮询直到完成/失败) + 3. download_video→ GET /v1/videos/{id}/content (下载视频文件) + """ + + CREATE_ENDPOINT = "/v1/videos" + STATUS_ENDPOINT = "/v1/videos/{video_id}" + CONTENT_ENDPOINT = "/v1/videos/{video_id}/content" + + POLL_INITIAL_INTERVAL = 3 + POLL_MAX_INTERVAL = 15 + + def __init__(self): + api_key = get_api_key_or_raise() + base_url = get_api_base_url() + super().__init__(base_url=base_url, api_key=api_key) + + # ------------------------------------------------------------------ + # BaseAPIClient 抽象方法实现（本客户端主要使用自定义方法） + # ------------------------------------------------------------------ + + def get_endpoint(self, **kwargs) -> str: + return self.CREATE_ENDPOINT + + def build_request_body(self, **kwargs) -> Dict[str, Any]: + return {} + + def parse_response(self, response: Dict[str, Any]) -> Any: + return response + + # ------------------------------------------------------------------ + # 核心异步方法 + # ------------------------------------------------------------------ + + async def create_video_async( + self, + prompt: str, + model: str, + seconds: int = 4, + size: str = "720x1280", + input_reference_bytes: Optional[bytes] = None, + seed: Optional[int] = None, + session: Optional[aiohttp.ClientSession] = None, + ) -> Dict[str, Any]: + """ + 提交视频生成任务 + + 格式策略（根据抓包确认）： + - 无参考图片：application/json + - 有参考图片：multipart/form-data，input_reference 以 PNG 文件上传 + + 注意：seed 不被上游 API 接受，仅在 ComfyUI 节点侧用于缓存刷新 + + Returns: + API 响应 JSON，包含 video id 和初始状态 + """ + url = f"{self.base_url}{self.CREATE_ENDPOINT}" + headers = {"Authorization": f"Bearer {self.api_key}"} + + # ============================================================ + # ⚠️ 已验证可用的标准请求方案，请勿随意修改！（2026-02-28） + # ============================================================ + # 经多轮调试确认： + # - 有图片：必须使用 multipart/form-data，input_reference 以 PNG 文件上传 + # · filename="reference.png", content_type="image/png"（与抓包一致） + # · 不可改为 application/json + base64 → 400 "expected a file, got a string" + # · 不可改为 application/json + data URI → 500 upstream error + # · 不可改为 multipart + image/jpeg → 400 "Inpaint image must match..."（尺寸校验失败） + # - 无图片：使用 application/json，已验证成功 + # ============================================================ + if input_reference_bytes: + if len(input_reference_bytes) > self.max_request_size: + raise ValueError( + f"参考图片约 {len(input_reference_bytes) / 1024 / 1024:.1f}MB，" + f"超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制，请使用较小的图片" + ) + # ⚠️ 有图片：multipart/form-data + PNG 文件上传（唯一验证成功的方案） + form = aiohttp.FormData() + form.add_field("prompt", prompt) + form.add_field("model", model) + form.add_field("seconds", str(seconds)) + form.add_field("size", size) + form.add_field( + "input_reference", + input_reference_bytes, + filename="reference.png", # ⚠️ 不可改文件名/扩展名 + content_type="image/png", # ⚠️ 不可改为 image/jpeg + ) + send_kwargs: Dict[str, Any] = {"data": form, "headers": headers} + else: + # ⚠️ 无图片：application/json（已验证成功） + body: Dict[str, Any] = { + "model": model, + "prompt": prompt, + "seconds": str(seconds), + "size": size, + } + send_kwargs = {"json": body, "headers": headers} + + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + try: + async with session.post(url, **send_kwargs) as response: + if response.status != 200: + error_text = await response.text() + error_message = self._extract_error_message(error_text, response.status) + raise RuntimeError(error_message) + + resp_json = await response.json() + return resp_json + + finally: + if close_session: + await session.close() + + async def poll_video_status_async( + self, + video_id: str, + progress_callback: Optional[Callable[[int, float], None]] = None, + session: Optional[aiohttp.ClientSession] = None, + ) -> Dict[str, Any]: + """ + 轮询视频生成状态，直到完成或失败 + + Args: + video_id: 视频任务 ID + progress_callback: 进度回调 (progress_percent, elapsed_seconds) + session: aiohttp 会话 + + Returns: + 最终状态的 API 响应 + + Raises: + RuntimeError: 生成失败 + """ + url = f"{self.base_url}{self.STATUS_ENDPOINT.format(video_id=video_id)}" + headers = self.get_headers(use_bearer_token=True) + + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + interval = self.POLL_INITIAL_INTERVAL + + try: + while True: + async with session.get(url, headers=headers) as response: + if response.status != 200: + error_text = await response.text() + error_message = self._extract_error_message(error_text, response.status) + raise RuntimeError(error_message) + + data = await response.json() + + # status 兼容大小写：queued / in_progress / IN_PROGRESS / completed / COMPLETED + status = data.get("status", "").lower() + + # progress 兼容整数 (30) 和字符串 ("30%") 两种格式 + progress_raw = data.get("progress", 0) + if isinstance(progress_raw, str): + try: + progress = int(progress_raw.rstrip("%").strip()) + except ValueError: + progress = 0 + else: + progress = int(progress_raw) if progress_raw else 0 + + if progress_callback: + progress_callback(progress) + + if status == "completed": + return data + + if status == "failed": + error_info = data.get("error", {}) + error_msg = error_info.get("message", "未知错误") if isinstance(error_info, dict) else str(error_info) + error_msg = _translate_error_message(error_msg) + raise RuntimeError(f"视频生成失败: {error_msg}") + + await asyncio.sleep(interval) + interval = min(interval * 1.5, self.POLL_MAX_INTERVAL) + + finally: + if close_session: + await session.close() + + async def download_video_async( + self, + video_id: str, + save_path: str, + session: Optional[aiohttp.ClientSession] = None, + ) -> str: + """ + 下载生成的视频文件 + + 处理两种情况： + 1. 响应为重定向或 JSON 含下载 URL → 跟随下载 + 2. 响应为二进制视频流 → 直接保存 + + Returns: + 保存的文件路径 + """ + url = f"{self.base_url}{self.CONTENT_ENDPOINT.format(video_id=video_id)}" + headers = self.get_headers(use_bearer_token=True) + + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + try: + async with session.get(url, headers=headers, allow_redirects=True) as response: + if response.status != 200: + error_text = await response.text() + error_message = self._extract_error_message(error_text, response.status) + raise RuntimeError(f"视频下载失败: {error_message}") + + content_type = response.headers.get("Content-Type", "") + + if "application/json" in content_type: + data = await response.json() + download_url = data.get("url") or data.get("download_url") + if not download_url: + raise RuntimeError("视频下载失败: 响应中未找到下载链接") + await self._download_from_url(download_url, save_path, session) + else: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + with open(save_path, "wb") as f: + async for chunk in response.content.iter_chunked(8192): + f.write(chunk) + + return save_path + + finally: + if close_session: + await session.close() + + # ------------------------------------------------------------------ + # 同步包装 + # ------------------------------------------------------------------ + + def generate_video_sync( + self, + prompt: str, + model: str, + seconds: int, + size: str, + save_path: str, + input_reference_bytes: Optional[bytes] = None, + seed: Optional[int] = None, + progress_callback: Optional[Callable[[int, float], None]] = None, + on_stage: Optional[Callable[[str], None]] = None, + ) -> str: + """ + 同步执行完整的视频生成流程（创建 → 轮询 → 下载） + + Args: + on_stage: 阶段回调，用于打印状态切换信息 + + Returns: + 保存的视频文件路径 + """ + + async def _run(): + connector = aiohttp.TCPConnector(limit=0) + async with aiohttp.ClientSession(connector=connector) as session: + # 1. 提交任务 + if on_stage: + on_stage("submitting") + result = await self.create_video_async( + prompt=prompt, + model=model, + seconds=seconds, + size=size, + input_reference_bytes=input_reference_bytes, + seed=seed, + session=session, + ) + video_id = result.get("id") + if not video_id: + raise RuntimeError("API 未返回视频任务 ID") + + if on_stage: + on_stage(f"submitted:{video_id}") + + # 2. 轮询状态 + if on_stage: + on_stage("polling") + await self.poll_video_status_async( + video_id=video_id, + progress_callback=progress_callback, + session=session, + ) + + # 3. 下载视频 + if on_stage: + on_stage("downloading") + path = await self.download_video_async( + video_id=video_id, + save_path=save_path, + session=session, + ) + + if on_stage: + on_stage("done") + return path + + return self.run_async_in_thread(_run()) + + async def _generate_one_video_async( + self, + prompt: str, + model: str, + seconds: int, + size: str, + save_path: str, + input_reference_bytes: Optional[bytes] = None, + seed: Optional[int] = None, + session: Optional[aiohttp.ClientSession] = None, + ) -> str: + """ + 异步生成单个视频（创建 → 轮询 → 下载） + + Returns: + 保存的视频文件路径 + """ + result = await self.create_video_async( + prompt=prompt, + model=model, + seconds=seconds, + size=size, + input_reference_bytes=input_reference_bytes, + seed=seed, + session=session, + ) + video_id = result.get("id") + if not video_id: + raise RuntimeError("API 未返回视频任务 ID") + + await self.poll_video_status_async(video_id=video_id, session=session) + path = await self.download_video_async( + video_id=video_id, save_path=save_path, session=session + ) + return path + + async def generate_batch_videos_async( + self, + prompt: str, + model: str, + seconds: int, + size: str, + save_paths: List[str], + input_reference_bytes: Optional[bytes] = None, + seed: Optional[int] = None, + progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None, + ) -> List[str]: + """ + 并发生成多个视频 + + Args: + prompt: 提示词 + model: 模型名称 + seconds: 视频时长（秒） + size: 分辨率 + save_paths: 各视频的保存路径列表，长度决定并发数量 + input_reference_bytes: 参考图片字节（可选） + seed: 随机种子（仅节点侧使用） + progress_callback: 进度回调 (current, total, success, error_msg) + + Returns: + 成功生成的视频路径列表 + """ + batch_size = len(save_paths) + connector = aiohttp.TCPConnector(limit=0) + + async with aiohttp.ClientSession(connector=connector) as session: + tasks = [ + self._generate_one_video_async( + prompt=prompt, + model=model, + seconds=seconds, + size=size, + save_path=save_paths[i], + input_reference_bytes=input_reference_bytes, + seed=seed, + session=session, + ) + for i in range(batch_size) + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + + completed = 0 + paths: List[str] = [] + first_error = None + for i, result in enumerate(results): + if isinstance(result, Exception): + error_msg = str(result) + print(f"Sora: 第 {i + 1} 个视频生成失败") + print(f"原始错误详情:\n{error_msg}") + if first_error is None: + first_error = result + if progress_callback: + progress_callback(i + 1, batch_size, False, error_msg) + else: + completed += 1 + paths.append(result) + if progress_callback: + progress_callback(completed, batch_size, True, None) + + if not paths: + if first_error: + raise first_error + raise RuntimeError(f"批量视频生成失败，{batch_size} 个任务全部失败") + + return paths + + def generate_batch_videos_sync( + self, + prompt: str, + model: str, + seconds: int, + size: str, + save_paths: List[str], + input_reference_bytes: Optional[bytes] = None, + seed: Optional[int] = None, + progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None, + ) -> List[str]: + """ + 同步并发生成多个视频（用于 ComfyUI 节点） + + Args: + save_paths: 各视频的保存路径列表，长度决定并发数量 + + Returns: + 成功生成的视频路径列表 + """ + coro = self.generate_batch_videos_async( + prompt=prompt, + model=model, + seconds=seconds, + size=size, + save_paths=save_paths, + input_reference_bytes=input_reference_bytes, + seed=seed, + progress_callback=progress_callback, + ) + return self.run_async_in_thread(coro) + + # ------------------------------------------------------------------ + # 内部辅助方法 + # ------------------------------------------------------------------ + + async def _download_from_url( + self, + url: str, + save_path: str, + session: aiohttp.ClientSession, + ) -> None: + """从给定 URL 下载文件到本地路径""" + os.makedirs(os.path.dirname(save_path), exist_ok=True) + async with session.get(url) as response: + if response.status != 200: + raise RuntimeError(f"从下载链接获取视频失败 (状态码: {response.status})") + with open(save_path, "wb") as f: + async for chunk in response.content.iter_chunked(8192): + f.write(chunk) + + @staticmethod + def _extract_error_message(error_text: str, status_code: int) -> str: + """从错误响应中提取可读的错误信息""" + error_message = error_text + try: + error_json = json.loads(error_text) + if "error" in error_json: + if isinstance(error_json["error"], dict): + error_message = error_json["error"].get("message", error_text) + else: + error_message = str(error_json["error"]) + elif "message" in error_json: + error_message = error_json["message"] + except (json.JSONDecodeError, KeyError): + pass + + status_hints = { + 400: "请求参数错误 (400)", + 401: "认证失败 (401)，请检查 API 密钥", + 403: "权限不足 (403)，请检查账户权限或余额", + 429: "请求频率超限 (429)，请稍后重试", + 503: "服务暂时不可用 (503)，请稍后重试", + 504: "请求超时 (504)，请稍后重试", + } + hint = status_hints.get(status_code, f"API 请求失败 (状态码: {status_code})") + return f"{hint}\nAPI 返回: {error_message}" diff --git a/clients/veo_client.py b/clients/veo_client.py new file mode 100644 index 0000000..01d3e9d --- /dev/null +++ b/clients/veo_client.py @@ -0,0 +1,510 @@ +""" +Veo 视频生成 API 客户端 +提供视频创建、状态轮询、视频下载功能 +""" + +import asyncio +import base64 +import json +import os +import time +from typing import Any, Callable, Dict, List, Optional + +import aiohttp + +from .base_client import BaseAPIClient +from ..utils.config import get_api_key_or_raise, get_api_base_url +from ..utils.image_utils import encode_image_to_base64 + + +class VeoClient(BaseAPIClient): + """ + Veo 视频生成客户端 + + 工作流程： + 1. create_video → POST /v1/videos (提交生成任务) + 2. poll_status → GET /v1/videos/{id} (轮询直到完成/失败) + 3. download_video→ GET /v1/videos/{id}/content (下载视频文件) + """ + + CREATE_ENDPOINT = "/v1/videos" + STATUS_ENDPOINT = "/v1/videos/{video_id}" + CONTENT_ENDPOINT = "/v1/videos/{video_id}/content" + + POLL_INITIAL_INTERVAL = 3 + POLL_MAX_INTERVAL = 15 + + def __init__(self): + api_key = get_api_key_or_raise() + base_url = get_api_base_url() + super().__init__(base_url=base_url, api_key=api_key) + + # ------------------------------------------------------------------ + # BaseAPIClient 抽象方法实现 + # ------------------------------------------------------------------ + + def get_endpoint(self, **kwargs) -> str: + return self.CREATE_ENDPOINT + + def build_request_body(self, **kwargs) -> Dict[str, Any]: + return {} + + def parse_response(self, response: Dict[str, Any]) -> Any: + return response + + # ------------------------------------------------------------------ + # 核心异步方法 + # ------------------------------------------------------------------ + + async def create_video_async( + self, + prompt: str, + model: str, + seconds: int = 8, + size: str = "720x1280", + first_frame_bytes: Optional[bytes] = None, + last_frame_bytes: Optional[bytes] = None, + reference_bytes: Optional[bytes] = None, + seed: Optional[int] = None, + session: Optional[aiohttp.ClientSession] = None, + ) -> Dict[str, Any]: + """ + 提交视频生成任务 + + 格式策略： + - 无参考图片：application/json + - 有参考图片：multipart/form-data，图片以 PNG 文件上传 + + Args: + prompt: 提示词 + model: 模型名称 + seconds: 视频时长（秒） + size: 分辨率 + first_frame_bytes: 首帧图片字节 + last_frame_bytes: 尾帧图片字节 + reference_bytes: 参考图片字节 + seed: 随机种子 + session: aiohttp 会话 + + Returns: + API 响应 JSON，包含 video id 和初始状态 + """ + url = f"{self.base_url}{self.CREATE_ENDPOINT}" + headers = {"Authorization": f"Bearer {self.api_key}"} + + # 检查是否有图片 + has_images = any([first_frame_bytes, last_frame_bytes, reference_bytes]) + + if has_images: + # 有图片：multipart/form-data + PNG 文件上传 + if first_frame_bytes and len(first_frame_bytes) > self.max_request_size: + raise ValueError(f"首帧图片过大，超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制") + if last_frame_bytes and len(last_frame_bytes) > self.max_request_size: + raise ValueError(f"尾帧图片过大，超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制") + if reference_bytes and len(reference_bytes) > self.max_request_size: + raise ValueError(f"参考图片过大，超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制") + + form = aiohttp.FormData() + form.add_field("prompt", prompt) + form.add_field("model", model) + form.add_field("seconds", str(seconds)) + form.add_field("size", size) + # 注意：seed 不被上游 API 接受，仅在 ComfyUI 节点侧用于缓存刷新 + # if seed is not None: + # form.add_field("seed", str(seed)) + + # 使用 input_reference 字段（OpenAI兼容格式） + # 尝试支持多张图片：按顺序添加多个 input_reference 字段 + if first_frame_bytes: + form.add_field( + "input_reference", + first_frame_bytes, + filename="first_frame.png", + content_type="image/png", + ) + if last_frame_bytes: + form.add_field( + "input_reference", + last_frame_bytes, + filename="last_frame.png", + content_type="image/png", + ) + if reference_bytes: + form.add_field( + "input_reference", + reference_bytes, + filename="reference.png", + content_type="image/png", + ) + + send_kwargs: Dict[str, Any] = {"data": form, "headers": headers} + else: + # 无图片：application/json + body: Dict[str, Any] = { + "model": model, + "prompt": prompt, + "seconds": str(seconds), + "size": size, + } + # 注意：seed 不被上游 API 接受，仅在 ComfyUI 节点侧用于缓存刷新 + # if seed is not None: + # body["seed"] = str(seed) + send_kwargs = {"json": body, "headers": headers} + + # 打印请求调试信息 + import json + if has_images: + print(f"Veo: 使用 multipart/form-data 格式上传图片") + else: + print(f"Veo API 请求体: {json.dumps(body, ensure_ascii=False)}") + + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + try: + async with session.post(url, **send_kwargs) as response: + if response.status != 200: + error_text = await response.text() + error_message = self._extract_error_message(error_text, response.status) + raise RuntimeError(error_message) + + resp_json = await response.json() + return resp_json + + finally: + if close_session: + await session.close() + + async def poll_video_status_async( + self, + video_id: str, + progress_callback: Optional[Callable[[int, float], None]] = None, + session: Optional[aiohttp.ClientSession] = None, + ) -> Dict[str, Any]: + """ + 轮询视频生成状态，直到完成或失败 + + Args: + video_id: 视频任务 ID + progress_callback: 进度回调 (progress_percent, elapsed_seconds) + session: aiohttp 会话 + + Returns: + 最终状态的 API 响应 + + Raises: + RuntimeError: 生成失败 + """ + url = f"{self.base_url}{self.STATUS_ENDPOINT.format(video_id=video_id)}" + headers = self.get_headers(use_bearer_token=True) + + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + interval = self.POLL_INITIAL_INTERVAL + + try: + while True: + async with session.get(url, headers=headers) as response: + if response.status != 200: + error_text = await response.text() + error_message = self._extract_error_message(error_text, response.status) + raise RuntimeError(error_message) + + data = await response.json() + + # status 兼容大小写 + status = data.get("status", "").lower() + + # progress 兼容整数和字符串 + progress_raw = data.get("progress", 0) + if isinstance(progress_raw, str): + try: + progress = int(progress_raw.rstrip("%").strip()) + except ValueError: + progress = 0 + else: + progress = int(progress_raw) if progress_raw else 0 + + if progress_callback: + progress_callback(progress) + + if status == "completed": + return data + + if status == "failed": + error_info = data.get("error", {}) + error_msg = error_info.get("message", "未知错误") if isinstance(error_info, dict) else str(error_info) + raise RuntimeError(f"视频生成失败: {error_msg}") + + await asyncio.sleep(interval) + interval = min(interval * 1.5, self.POLL_MAX_INTERVAL) + + finally: + if close_session: + await session.close() + + async def download_video_async( + self, + video_id: str, + save_path: str, + session: Optional[aiohttp.ClientSession] = None, + ) -> str: + """ + 下载生成的视频文件 + + Returns: + 保存的文件路径 + """ + url = f"{self.base_url}{self.CONTENT_ENDPOINT.format(video_id=video_id)}" + headers = self.get_headers(use_bearer_token=True) + + close_session = False + if session is None: + session = aiohttp.ClientSession() + close_session = True + + try: + async with session.get(url, headers=headers, allow_redirects=True) as response: + if response.status != 200: + error_text = await response.text() + error_message = self._extract_error_message(error_text, response.status) + raise RuntimeError(f"视频下载失败: {error_message}") + + content_type = response.headers.get("Content-Type", "") + + if "application/json" in content_type: + data = await response.json() + download_url = data.get("url") or data.get("download_url") + if not download_url: + raise RuntimeError("视频下载失败: 响应中未找到下载链接") + await self._download_from_url(download_url, save_path, session) + else: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + with open(save_path, "wb") as f: + async for chunk in response.content.iter_chunked(8192): + f.write(chunk) + + return save_path + + finally: + if close_session: + await session.close() + + # ------------------------------------------------------------------ + # 同步包装 + # ------------------------------------------------------------------ + + def generate_video_sync( + self, + prompt: str, + model: str, + seconds: int, + size: str, + save_path: str, + first_frame_bytes: Optional[bytes] = None, + last_frame_bytes: Optional[bytes] = None, + reference_bytes: Optional[bytes] = None, + seed: Optional[int] = None, + progress_callback: Optional[Callable[[int], None]] = None, + on_stage: Optional[Callable[[str], None]] = None, + ) -> str: + """ + 同步执行完整的视频生成流程（创建 → 轮询 → 下载） + """ + + async def _run(): + connector = aiohttp.TCPConnector(limit=0) + async with aiohttp.ClientSession(connector=connector) as session: + # 1. 提交任务 + if on_stage: + on_stage("submitting") + result = await self.create_video_async( + prompt=prompt, + model=model, + seconds=seconds, + size=size, + first_frame_bytes=first_frame_bytes, + last_frame_bytes=last_frame_bytes, + reference_bytes=reference_bytes, + seed=seed, + session=session, + ) + video_id = result.get("id") + if not video_id: + raise RuntimeError("API 未返回视频任务 ID") + + if on_stage: + on_stage(f"submitted:{video_id}") + + # 2. 轮询状态 + if on_stage: + on_stage("polling") + await self.poll_video_status_async( + video_id=video_id, + progress_callback=progress_callback, + session=session, + ) + + # 3. 下载视频 + if on_stage: + on_stage("downloading") + path = await self.download_video_async( + video_id=video_id, + save_path=save_path, + session=session, + ) + + if on_stage: + on_stage("done") + return path + + return self.run_async_in_thread(_run()) + + def generate_batch_videos_sync( + self, + prompt: str, + model: str, + seconds: int, + size: str, + save_paths: List[str], + first_frame_bytes: Optional[bytes] = None, + last_frame_bytes: Optional[bytes] = None, + reference_bytes: Optional[bytes] = None, + seed: Optional[int] = None, + progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None, + ) -> List[str]: + """ + 同步并发生成多个视频 + """ + async def _run(): + batch_size = len(save_paths) + connector = aiohttp.TCPConnector(limit=0) + + async def generate_one(save_path: str): + return await self._generate_one_video_async( + prompt=prompt, + model=model, + seconds=seconds, + size=size, + save_path=save_path, + first_frame_bytes=first_frame_bytes, + last_frame_bytes=last_frame_bytes, + reference_bytes=reference_bytes, + seed=seed, + ) + + async with aiohttp.ClientSession(connector=connector) as session: + tasks = [generate_one(p) for p in save_paths] + results = await asyncio.gather(*tasks, return_exceptions=True) + + completed = 0 + paths: List[str] = [] + first_error = None + for i, result in enumerate(results): + if isinstance(result, Exception): + error_msg = str(result) + print(f"Veo: 第 {i + 1} 个视频生成失败") + if first_error is None: + first_error = result + if progress_callback: + progress_callback(i + 1, batch_size, False, error_msg) + else: + completed += 1 + paths.append(result) + if progress_callback: + progress_callback(completed, batch_size, True, None) + + if not paths: + if first_error: + raise first_error + raise RuntimeError(f"批量视频生成失败，{batch_size} 个任务全部失败") + + return paths + + return self.run_async_in_thread(_run()) + + async def _generate_one_video_async( + self, + prompt: str, + model: str, + seconds: int, + size: str, + save_path: str, + first_frame_bytes: Optional[bytes] = None, + last_frame_bytes: Optional[bytes] = None, + reference_bytes: Optional[bytes] = None, + seed: Optional[int] = None, + session: Optional[aiohttp.ClientSession] = None, + ) -> str: + """异步生成单个视频""" + result = await self.create_video_async( + prompt=prompt, + model=model, + seconds=seconds, + size=size, + first_frame_bytes=first_frame_bytes, + last_frame_bytes=last_frame_bytes, + reference_bytes=reference_bytes, + seed=seed, + session=session, + ) + video_id = result.get("id") + if not video_id: + raise RuntimeError("API 未返回视频任务 ID") + + await self.poll_video_status_async(video_id=video_id, session=session) + path = await self.download_video_async( + video_id=video_id, save_path=save_path, session=session + ) + return path + + # ------------------------------------------------------------------ + # 内部辅助方法 + # ------------------------------------------------------------------ + + async def _download_from_url( + self, + url: str, + save_path: str, + session: aiohttp.ClientSession, + ) -> None: + """从给定 URL 下载文件到本地路径""" + os.makedirs(os.path.dirname(save_path), exist_ok=True) + async with session.get(url) as response: + if response.status != 200: + raise RuntimeError(f"从下载链接获取视频失败 (状态码: {response.status})") + with open(save_path, "wb") as f: + async for chunk in response.content.iter_chunked(8192): + f.write(chunk) + + @staticmethod + def _extract_error_message(error_text: str, status_code: int) -> str: + """从错误响应中提取可读的错误信息""" + error_message = error_text + try: + error_json = json.loads(error_text) + if "error" in error_json: + if isinstance(error_json["error"], dict): + error_message = error_json["error"].get("message", error_text) + else: + error_message = str(error_json["error"]) + elif "message" in error_json: + error_message = error_json["message"] + except (json.JSONDecodeError, KeyError): + pass + + status_hints = { + 400: "请求参数错误 (400)", + 401: "认证失败 (401)，请检查 API 密钥", + 403: "权限不足 (403)，请检查账户权限或余额", + 429: "请求频率超限 (429)，请稍后重试", + 503: "服务暂时不可用 (503)，请稍后重试", + 504: "请求超时 (504)，请稍后重试", + } + hint = status_hints.get(status_code, f"API 请求失败 (状态码: {status_code})") + return f"{hint}\nAPI 返回: {error_message}" diff --git a/models_config.py b/models_config.py new file mode 100644 index 0000000..5f12f8a --- /dev/null +++ b/models_config.py @@ -0,0 +1,896 @@ +""" +模型配置中心 +用于集中管理所有支持的 Gemini 模型 + +使用方式: + 1. 添加新模型: 在对应的模型列表中添加新的模型字典 + 2. 临时关闭模型: 将模型的 enabled 字段设为 False + 3. 重新启用模型: 将模型的 enabled 字段改回 True + +模型类型: + - GEMINI_MODELS: Nano Banana 图像生成模型 + - GEMINI_FLASH_MODELS: Google Gemini Flash 文本生成模型 + +示例: + 添加新模型: + { + "id": "gemini-新模型名称", + "description": "模型说明和特点", + "enabled": True, + "endpoint_type": "standard", + "endpoint": "/v1beta/models/gemini-新模型名称:generateContent", + "thinking_config": { + "不思考": None, + "低": "low", + "中": None, + "高": "high" + } + } + + 临时关闭模型: + 将对应模型的 "enabled": True 改为 "enabled": False +""" + +from typing import List, Dict, Optional, Tuple + + +# ============================================================ +# 模型配置列表 +# ============================================================ + +# ============================================================ +# Nano Banana 图像生成模型 +# ============================================================ + +GEMINI_MODELS = [ + { + "id": "nano-banana-pro-限时特价", + "description": "Nano Banana Pro 限时特价,根据分辨率自动选择端点 (1K/2K/4K),高性能图像生成模型", + "enabled": True, + "endpoint_type": "dynamic", + "endpoint": None, # 动态端点，由代码根据分辨率选择 + "supported_aspect_ratios": [ + "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9" + ], + "supported_resolutions": ["1K", "2K", "4K"] + }, + { + "id": "nano-banana-pro-官方计费", + "description": "Nano Banana Pro 官方计费，按分辨率路由 (1K/2K/4K)，使用官方计费通道", + "enabled": True, + "endpoint_type": "dynamic", + "endpoint": None, # 动态端点，由代码根据分辨率选择 + "supported_aspect_ratios": [ + "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9" + ], + "supported_resolutions": ["1K", "2K", "4K"] + }, + { + "id": "nano-banana-2-限时特价", + "description": "Nano Banana 2 限时特价，固定端点，图像生成模型", + "enabled": True, + "endpoint_type": "standard", + "endpoint": "/v1beta/models/nano-banana-2:generateContent", + "supported_aspect_ratios": [ + "1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", + "8:1", "9:16", "16:9", "21:9" + ], + "supported_resolutions": ["512", "1K", "2K", "4K"] + }, + { + "id": "nano-banana-2-官方计费", + "description": "Nano Banana 2 官方计费，按分辨率路由 (512/1K/2K/4K)，使用官方计费通道", + "enabled": True, + "endpoint_type": "dynamic", + "endpoint": None, # 动态端点，由代码根据分辨率选择 + "supported_aspect_ratios": [ + "1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", + "8:1", "9:16", "16:9", "21:9" + ], + "supported_resolutions": ["512", "1K", "2K", "4K"] + }, + { + "id": "gemini-3-pro-image-preview", + "description": "标准模式,固定端点,适用于常规图像生成", + "enabled": False, + "endpoint_type": "standard", + "endpoint": "/v1beta/models/gemini-3-pro-image-preview:generateContent", + "supported_aspect_ratios": [ + "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9" + ], + "supported_resolutions": ["1K", "2K", "4K"] + }, + { + "id": "gemini-3.1-flash-image-preview", + "description": "Gemini 3.1 Flash 图像生成,固定端点,快速图像生成模型", + "enabled": False, + "endpoint_type": "standard", + "endpoint": "/v1beta/models/gemini-3.1-flash-image-preview:generateContent", + "supported_aspect_ratios": [ + "1:1", "1:4", "1:8", "2:3", "3:2", "3:4", "4:1", "4:3", "4:5", "5:4", + "8:1", "9:16", "16:9", "21:9" + ], + "supported_resolutions": ["512", "1K", "2K", "4K"] + } +] + + +# ============================================================ +# Google Gemini Flash 文本生成模型 +# ============================================================ + +GEMINI_FLASH_MODELS = [ + { + "id": "gemini-3-flash-preview", + "description": "Gemini 3 Flash,快速多模态文本生成,通过 thinkingConfig 控制思考等级", + "enabled": True, + "endpoint_type": "standard", + "endpoint": "/v1beta/models/gemini-3-flash-preview:generateContent", + "thinking_config": { + "低": "low", + "中": "medium", + "高": "high" + } + }, + + { + "id": "gemini-3.1-pro-preview", + "description": "Gemini 3.1 Pro,高性能多模态文本生成,通过 thinkingConfig 控制思考等级", + "enabled": True, + "endpoint_type": "standard", + "endpoint": "/v1beta/models/gemini-3.1-pro-preview:generateContent", + "thinking_config": { + "低": "low", + "中": "high" + } + }, + + { + "id": "gemini-3.1-flash-lite-preview", + "description": "Gemini 3.1 Flash Lite,轻量级多模态文本生成,通过 thinkingConfig 控制思考等级", + "enabled": True, + "endpoint_type": "standard", + "endpoint": "/v1beta/models/gemini-3.1-flash-lite-preview:generateContent", + "thinking_config": { + "低": "low", + "中": "medium", + "高": "high" + } + } +] + + +# ============================================================ +# 工具函数 +# ============================================================ + +def get_enabled_models() -> List[str]: + """ + 获取所有启用的模型 ID 列表 + + Returns: + 启用的模型 ID 列表 + + Example: + >>> get_enabled_models() + ['gemini-3-pro-image-preview-url', 'gemini-3-pro-image-preview', ...] + """ + return [model["id"] for model in GEMINI_MODELS if model.get("enabled", False)] + + +def get_all_models() -> List[str]: + """ + 获取所有模型 ID 列表（包括已禁用的） + + Returns: + 所有模型 ID 列表 + + Example: + >>> get_all_models() + ['gemini-3-pro-image-preview-url', 'gemini-3-pro-image-preview', ...] + """ + return [model["id"] for model in GEMINI_MODELS] + + +def get_model_config(model_id: str) -> Optional[Dict]: + """ + 根据模型 ID 获取完整的模型配置 + + Args: + model_id: 模型 ID + + Returns: + 模型配置字典，如果未找到则返回 None + + Example: + >>> config = get_model_config("gemini-3-pro-image-preview-url") + >>> print(config["description"]) + URL 模式,根据分辨率自动选择端点 (1K/2K/4K) + """ + for model in GEMINI_MODELS: + if model["id"] == model_id: + return model + return None + + +def is_model_enabled(model_id: str) -> bool: + """ + 检查指定模型是否启用 + + Args: + model_id: 模型 ID + + Returns: + True 如果模型启用，False 如果禁用或不存在 + + Example: + >>> is_model_enabled("gemini-3-pro-image-preview-url") + True + """ + config = get_model_config(model_id) + if config is None: + return False + return config.get("enabled", False) + + +def get_model_description(model_id: str) -> str: + """ + 获取模型的描述信息 + + Args: + model_id: 模型 ID + + Returns: + 模型描述，如果未找到则返回空字符串 + + Example: + >>> get_model_description("gemini-3-pro-image-preview") + '标准模式,固定端点,适用于常规图像生成' + """ + config = get_model_config(model_id) + if config is None: + return "" + return config.get("description", "") + + +def get_model_supported_aspect_ratios(model_id: str) -> List[str]: + """ + 获取模型支持的宽高比列表 + + Args: + model_id: 模型 ID + + Returns: + 支持的宽高比字符串列表，如果未配置则返回空列表 + + Example: + >>> get_model_supported_aspect_ratios("gemini-3-pro-image-preview") + ['1:1', '2:3', '3:2', ...] + """ + config = get_model_config(model_id) + if config is None: + return [] + return config.get("supported_aspect_ratios", []) + + +def get_all_supported_aspect_ratios() -> List[str]: + """ + 获取所有启用模型支持的宽高比（去重合并） + + Returns: + 所有启用模型支持的宽高比列表（保持顺序、去重） + + Example: + >>> get_all_supported_aspect_ratios() + ['1:1', '4:3', '3:4', '16:9', '9:16', '2:3', '3:2', '4:5', '5:4', '21:9', '1:4', '4:1', '1:8', '8:1'] + """ + seen = set() + result = [] + for model in GEMINI_MODELS: + if not model.get("enabled", False): + continue + for ratio in model.get("supported_aspect_ratios", []): + if ratio not in seen: + seen.add(ratio) + result.append(ratio) + return result + + +def get_model_supported_resolutions(model_id: str) -> List[str]: + """ + 获取模型支持的分辨率列表 + + Args: + model_id: 模型 ID + + Returns: + 支持的分辨率字符串列表，如果未配置则返回空列表 + + Example: + >>> get_model_supported_resolutions("gemini-3.1-flash-image-preview") + ['512', '1K', '2K', '4K'] + >>> get_model_supported_resolutions("gemini-3-pro-image-preview") + ['1K', '2K', '4K'] + """ + config = get_model_config(model_id) + if config is None: + return [] + return config.get("supported_resolutions", []) + + +def get_all_supported_resolutions() -> List[str]: + """ + 获取所有启用模型支持的分辨率（去重合并，按从小到大固定顺序排列） + + Returns: + 所有启用模型支持的分辨率列表（按 512 → 1K → 2K → 4K 顺序） + + Example: + >>> get_all_supported_resolutions() + ['512', '1K', '2K', '4K'] + """ + _ORDER = ["512", "1K", "2K", "4K"] + + seen = set() + for model in GEMINI_MODELS: + if not model.get("enabled", False): + continue + for res in model.get("supported_resolutions", []): + seen.add(res) + + return [res for res in _ORDER if res in seen] + + +def get_endpoint_type(model_id: str) -> Optional[str]: + """ + 获取模型的端点类型 + + Args: + model_id: 模型 ID + + Returns: + 端点类型 ("dynamic", "standard", "flatfee")，如果未找到则返回 None + + Example: + >>> get_endpoint_type("gemini-3-pro-image-preview-url") + 'dynamic' + """ + config = get_model_config(model_id) + if config is None: + return None + return config.get("endpoint_type") + + +def get_model_endpoint(model_id: str) -> Optional[str]: + """ + 获取模型的 API 端点 + + Args: + model_id: 模型 ID + + Returns: + API 端点路径，如果未找到或为动态端点则返回 None + + Example: + >>> get_model_endpoint("gemini-3-pro-image-preview") + '/v1beta/models/gemini-3-pro-image-preview:generateContent' + >>> get_model_endpoint("gemini-3-pro-image-preview-url") + None # 动态端点 + """ + config = get_model_config(model_id) + if config is None: + return None + return config.get("endpoint") + + +# ============================================================ +# Gemini Flash 模型工具函数 +# ============================================================ + +# ============================================================ +# Sora 视频生成模型 +# ============================================================ + +SORA_MODELS = [ + { + "id": "sora-2", + "description": "Sora 2 官方模型，支持标准时长和分辨率", + "enabled": True, + "supported_seconds": [4, 8, 10, 12, 15], + "supported_sizes": ["720x1280", "1280x720"], + "seconds_category": "官方", # 用于界面显示标签 + }, + { + "id": "sora-2-pro", + "description": "Sora 2 Pro 增强模型，支持扩展时长和竖屏/横屏高清分辨率", + "enabled": True, + "supported_seconds": [4, 8, 12, 15, 25], + "supported_sizes": ["720x1280", "1280x720", "1024x1792", "1792x1024"], + "seconds_category": "扩展", # Pro 模型支持全部时长 + }, +] + +# 秒数显示标签配置（用于界面下拉菜单） +# key: 实际秒数, value: 显示文本 +SECONDS_DISPLAY_MAP = { + 4: "4", + 8: "8", + 12: "12", + 10: "10", + 15: "15", + 25: "25(pro)", +} + +# 分辨率显示标签配置 +# key: 实际分辨率, value: (显示P数, 显示方向) +RESOLUTION_DISPLAY_MAP = { + "720x1280": ("720P", "竖屏"), + "1280x720": ("720P", "横屏"), + "1024x1792": ("1080P", "竖屏"), + "1792x1024": ("1080P", "横屏"), +} + + +# ============================================================ +# Sora 模型工具函数 +# ============================================================ + +def get_enabled_sora_models() -> List[str]: + """获取所有启用的 Sora 模型 ID 列表""" + return [model["id"] for model in SORA_MODELS if model.get("enabled", False)] + + +def get_sora_model_config(model_id: str) -> Optional[Dict]: + """根据模型 ID 获取 Sora 模型的完整配置""" + for model in SORA_MODELS: + if model["id"] == model_id: + return model + return None + + +def get_sora_supported_seconds(model_id: str) -> List[int]: + """获取 Sora 模型支持的视频时长列表（秒）""" + config = get_sora_model_config(model_id) + if config is None: + return [] + return config.get("supported_seconds", []) + + +def get_sora_supported_sizes(model_id: str) -> List[str]: + """获取 Sora 模型支持的分辨率列表""" + config = get_sora_model_config(model_id) + if config is None: + return [] + return config.get("supported_sizes", []) + + +def get_all_sora_seconds() -> List[int]: + """获取所有启用 Sora 模型支持的时长（去重、升序）""" + seen = set() + for model in SORA_MODELS: + if not model.get("enabled", False): + continue + for s in model.get("supported_seconds", []): + seen.add(s) + return sorted(seen) + + +def get_all_sora_sizes() -> List[str]: + """获取所有启用 Sora 模型支持的分辨率（去重、保持顺序）""" + seen = set() + result = [] + for model in SORA_MODELS: + if not model.get("enabled", False): + continue + for size in model.get("supported_sizes", []): + if size not in seen: + seen.add(size) + result.append(size) + return result + + +def get_sora_seconds_with_labels(model_id: str) -> List[Tuple[str, int]]: + """ + 获取指定模型支持的秒数列表（带标签显示） + + Returns: + 列表项为 (显示文本, 实际秒数)，如 [("4（官方）", 4), ("10（特殊）", 10)] + """ + config = get_sora_model_config(model_id) + if config is None: + return [] + + seconds_list = config.get("supported_seconds", []) + result = [] + for s in seconds_list: + category = SECONDS_CATEGORIES.get(s, "") + label = f"{s}（{category}）" if category else str(s) + result.append((label, s)) + return result + + +def get_sora_sizes_with_labels(model_id: str) -> List[Tuple[str, str]]: + """ + 获取指定模型支持的分辨率列表（带独占标识） + + Returns: + 列表项为 (显示文本, 实际分辨率)，如 [("720P 9:16 (720x1280)", "720x1280")] + """ + from math import gcd + + config = get_sora_model_config(model_id) + if config is None: + return [] + + sizes = config.get("supported_sizes", []) + result = [] + + # 检查哪些分辨率是独占的（仅该模型支持） + all_sizes_count = {} + for m in SORA_MODELS: + if not m.get("enabled", False): + continue + for size in m.get("supported_sizes", []): + all_sizes_count[size] = all_sizes_count.get(size, 0) + 1 + + for size in sizes: + # 解析分辨率 + parts = size.lower().split("x") + w, h = int(parts[0]), int(parts[1]) + short_side = min(w, h) + + # 分辨率等级 + if short_side >= 1792: + res = "2K+" + elif short_side >= 1080: + res = "1K+" + elif short_side >= 720: + res = "720P" + else: + res = f"{short_side}P" + + # 比例 + g = gcd(w, h) + ratio = f"{w // g}:{h // g}" + + # 检查是否独占 + exclusive = all_sizes_count.get(size, 0) == 1 + exclusive_tag = " [Pro独占]" if exclusive else "" + + # 方向 + orientation = "竖屏" if h > w else "横屏" if w > h else "方形" + + label = f"{res} {ratio} {orientation}{exclusive_tag} ({size})" + result.append((label, size)) + + return result + + +# ============================================================ +# Google Veo 视频生成模型 +# ============================================================ + +VEO_MODELS = [ + { + "id": "Veo3.1", + "description": "Google Veo 3.1 视频生成模型，支持文生视频和图生视频", + "enabled": True, + }, +] + +# Veo 分辨率映射表 +# key: "分辨率_宽高比", value: 实际分辨率字符串 +VEO_RESOLUTION_MAP = { + # 720p + "720p_9:16": "720x1280", + "720p_16:9": "1280x720", + # 1080p + "1080p_9:16": "1080x1920", + "1080p_16:9": "1920x1080", + # 4K + "4K_9:16": "2160x3840", + "4K_16:9": "3840x2160", +} + + +# ============================================================ +# Veo 模型工具函数 +# ============================================================ + +def get_enabled_veo_models() -> List[str]: + """获取所有启用的 Veo 模型 ID 列表""" + return [model["id"] for model in VEO_MODELS if model.get("enabled", False)] + + +def get_veo_model_config(model_id: str) -> Optional[Dict]: + """根据模型 ID 获取 Veo 模型的完整配置""" + for model in VEO_MODELS: + if model["id"] == model_id: + return model + return None + + +# ============================================================ +# Gemini Flash 模型工具函数 +# ============================================================ + +def get_enabled_flash_models() -> List[str]: + """ + 获取所有启用的 Flash 模型 ID 列表 + + Returns: + 启用的 Flash 模型 ID 列表 + + Example: + >>> get_enabled_flash_models() + ['gemini-3-flash-preview'] + """ + return [model["id"] for model in GEMINI_FLASH_MODELS if model.get("enabled", False)] + + +def get_all_flash_models() -> List[str]: + """ + 获取所有 Flash 模型 ID 列表（包括已禁用的） + + Returns: + 所有 Flash 模型 ID 列表 + """ + return [model["id"] for model in GEMINI_FLASH_MODELS] + + +def get_flash_model_config(model_id: str) -> Optional[Dict]: + """ + 根据模型 ID 获取 Flash 模型的完整配置 + + Args: + model_id: 模型 ID + + Returns: + 模型配置字典，如果未找到则返回 None + + Example: + >>> config = get_flash_model_config("gemini-3-flash-preview") + >>> print(config["description"]) + 'Gemini 3 Flash,快速多模态文本生成,支持图片和视频输入' + """ + for model in GEMINI_FLASH_MODELS: + if model["id"] == model_id: + return model + return None + + +def is_flash_model_enabled(model_id: str) -> bool: + """ + 检查指定 Flash 模型是否启用 + + Args: + model_id: 模型 ID + + Returns: + True 如果模型启用，False 如果禁用或不存在 + """ + config = get_flash_model_config(model_id) + if config is None: + return False + return config.get("enabled", False) + + +def get_flash_model_endpoint(model_id: str) -> Optional[str]: + """ + 获取 Flash 模型的 API 端点 + + Args: + model_id: 模型 ID + + Returns: + API 端点路径，如果未找到则返回 None + + Example: + >>> get_flash_model_endpoint("gemini-3-flash-preview") + '/v1beta/models/gemini-3-flash-preview:generateContent' + """ + config = get_flash_model_config(model_id) + if config is None: + return None + return config.get("endpoint") + + +def get_flash_model_description(model_id: str) -> str: + """ + 获取 Flash 模型的描述信息 + + Args: + model_id: 模型 ID + + Returns: + 模型描述，如果未找到则返回空字符串 + """ + config = get_flash_model_config(model_id) + if config is None: + return "" + return config.get("description", "") + + +def get_flash_model_thinking_level_value(model_id: str, thinking_level: str) -> Optional[str]: + """ + 获取指定模型在给定思考等级下应传入请求体的 thinkingLevel 值。 + + 仅对 endpoint_type="standard" 且配置了 thinking_config 的模型有效。 + 返回 None 表示该等级不受支持，请求体中不应包含 thinkingConfig。 + + Args: + model_id: 模型 ID + thinking_level: 思考等级中文名（不思考/低/中/高） + + Returns: + API thinkingLevel 值（如 "low"/"medium"/"high"），或 None（不传参） + + Example: + >>> get_flash_model_thinking_level_value("gemini-3-pro-preview", "低") + 'low' + >>> get_flash_model_thinking_level_value("gemini-3-pro-preview", "中") + None # 不受支持，省略 thinkingConfig + """ + config = get_flash_model_config(model_id) + if config is None: + return None + thinking_config = config.get("thinking_config") + if not thinking_config: + return None + return thinking_config.get(thinking_level) + + +# 已弃用：动态端点模式下不再需要这些函数 +# def get_flash_model_thinking_levels(model_id: str) -> List[str]: +# """ +# 获取 Flash 模型支持的思考等级列表 +# +# Args: +# model_id: 模型 ID +# +# Returns: +# 思考等级列表（中文），如果未找到则返回空列表 +# +# Example: +# >>> get_flash_model_thinking_levels("gemini-3-flash-preview") +# ['默认', '最低', '低', '中', '高'] +# """ +# config = get_flash_model_config(model_id) +# if config is None: +# return [] +# +# thinking_levels = config.get("thinking_levels", {}) +# return list(thinking_levels.keys()) + + +# def get_thinking_level_value(model_id: str, thinking_level: str) -> Optional[str]: +# """ +# 获取思考等级对应的 API 参数值 +# +# Args: +# model_id: 模型 ID +# thinking_level: 思考等级（中文） +# +# Returns: +# API 参数值（英文），如果未找到则返回 None +# +# Example: +# >>> get_thinking_level_value("gemini-3-flash-preview", "默认") +# 'high' +# >>> get_thinking_level_value("gemini-3-flash-preview", "最低") +# 'minimal' +# """ +# config = get_flash_model_config(model_id) +# if config is None: +# return None +# +# thinking_levels = config.get("thinking_levels", {}) +# return thinking_levels.get(thinking_level) + + +# ============================================================ +# 向后兼容性检查 +# ============================================================ + +def validate_models_config() -> None: + """ + 验证模型配置的完整性 + + 检查: + - 每个模型必须有 id, description, enabled, endpoint_type, endpoint 字段 + - 非动态端点模型必须配置有效的 endpoint + - 至少有一个模型是启用的 + + Raises: + ValueError: 如果配置不合法 + """ + if not GEMINI_MODELS: + raise ValueError("GEMINI_MODELS 列表不能为空") + + required_fields = ["id", "description", "enabled", "endpoint_type", "endpoint"] + valid_endpoint_types = ["dynamic", "standard", "flatfee"] + + for i, model in enumerate(GEMINI_MODELS): + # 检查必需字段 + for field in required_fields: + if field not in model: + raise ValueError(f"模型 #{i} 缺少必需字段: {field}") + + # 检查 endpoint_type 是否合法 + if model["endpoint_type"] not in valid_endpoint_types: + raise ValueError( + f"模型 {model['id']} 的 endpoint_type '{model['endpoint_type']}' 不合法。" + f"必须是: {', '.join(valid_endpoint_types)}" + ) + + # 检查非动态端点模型必须有有效的 endpoint + if model["endpoint_type"] != "dynamic" and not model.get("endpoint"): + raise ValueError( + f"模型 {model['id']} 的 endpoint_type 为 '{model['endpoint_type']}'，" + f"但未配置有效的 endpoint 字段" + ) + + # 检查端点格式（如果配置了） + endpoint = model.get("endpoint") + if endpoint and not endpoint.startswith("/v1beta/models/"): + raise ValueError( + f"模型 {model['id']} 的 endpoint '{endpoint}' 格式不正确。" + f"应以 '/v1beta/models/' 开头" + ) + + # 检查至少有一个启用的模型 + if not get_enabled_models(): + raise ValueError("至少需要启用一个模型") + + +def validate_flash_models_config() -> None: + """ + 验证 Flash 模型配置的完整性 + + 检查: + - 每个模型必须有 id, description, enabled 字段 + - 每个模型必须有 endpoint 字段且格式正确 + - 至少有一个模型是启用的 + + Raises: + ValueError: 如果配置不合法 + """ + if not GEMINI_FLASH_MODELS: + raise ValueError("GEMINI_FLASH_MODELS 列表不能为空") + + required_fields = ["id", "description", "enabled"] + + for i, model in enumerate(GEMINI_FLASH_MODELS): + # 检查必需字段 + for field in required_fields: + if field not in model: + raise ValueError(f"Flash 模型 #{i} 缺少必需字段: {field}") + + # 检查端点配置 + if "endpoint" not in model: + raise ValueError(f"Flash 模型 {model['id']} 缺少 'endpoint' 字段") + + endpoint = model.get("endpoint", "") + if not endpoint or not endpoint.startswith("/v1beta/models/"): + raise ValueError( + f"Flash 模型 {model['id']} 的 endpoint '{endpoint}' 格式不正确。" + f"应以 '/v1beta/models/' 开头" + ) + + # 检查至少有一个启用的模型 + if not get_enabled_flash_models(): + raise ValueError("至少需要启用一个 Flash 模型") + + +# 在模块加载时验证配置 +try: + validate_models_config() +except ValueError as e: + print(f"⚠️ 图像模型配置验证失败: {str(e)}") + print(f"⚠️ 请检查 models_config.py 文件") + +try: + validate_flash_models_config() +except ValueError as e: + print(f"⚠️ Flash 模型配置验证失败: {str(e)}") + print(f"⚠️ 请检查 models_config.py 文件") diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..4d6535b --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,24 @@ +""" +节点模块 +包含所有 ComfyUI 自定义节点的实现 +""" + +from .nano_banana_pro import NanoBananaPro +from .batch_nano_banana_pro import BatchNanoBananaPro +from .google_gemini import GoogleGemini +from .load_file import LoadFile +from .image_stitch_pro import ImageStitchPro +from .remove_metadata import SaveCleanImage, BatchCleanMetadata +from .video_preview import VideoPreview +from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset +from .veo_video import GoogleVeo +from .flux_edit import FluxImageEdit +from .universal_llm import UniversalLLMChat +from .quan_neng_sheng_tu import QuanNengShengTu +from .batch_quan_neng_sheng_tu import BatchQuanNengShengTu +from .multi_res_preview import MultiResPreview +from .batch_images_o1key import BatchImagesO1key +from .nano_banana_v2 import NanaBananaV2 +from .batch_nano_banana_v2 import BatchNanaBananaV2 + +__all__ = ['NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'QuanNengShengTu', 'BatchQuanNengShengTu', 'MultiResPreview', 'BatchImagesO1key', 'NanaBananaV2', 'BatchNanaBananaV2'] diff --git a/nodes/batch_images_o1key.py b/nodes/batch_images_o1key.py new file mode 100644 index 0000000..3970c1a --- /dev/null +++ b/nodes/batch_images_o1key.py @@ -0,0 +1,80 @@ +""" +批量图像（o1key）节点 +复刻 ComfyUI 原生「批量图像」节点的动态输入行为： + +- 默认显示 2 个图像输入端口（图1, 图2） +- 当最后一个端口连上图像后，自动追加新端口 +- 断开连线后，多余的端口自动消失，最少保留 2 个 + +与原生节点的区别： + 原生节点会把所有图像强制 resize 到第一张的分辨率再合并为单一 tensor。 + 本节点保留每张图的原始分辨率，以 list[Tensor] 形式输出（is_output_list）。 + 下游节点（如「多分辨率图像预览」）需开启 INPUT_IS_LIST 才能正确接收。 + +实现方式：使用 V3 API 的 io.Autogrow.TemplateNames， +框架原生支持动态 slot 增减，无需编写任何 JS 扩展。 +""" + +import torch +from comfy_api.latest import io + +# 预生成 50 个端口名：图1, 图2, ..., 图50 +_SLOT_NAMES = [f"图{i}" for i in range(1, 51)] + + +class BatchImagesO1key(io.ComfyNode): + """ + 批量图像（o1key） + + - 动态输入端口（默认 2 个，最多 50 个），端口名为 图1、图2、图3... + - 连接最后一个端口时自动增加新端口 + - 断开后自动减少，保持界面整洁 + - 保留每张图的原始分辨率，不做任何 resize / 裁剪 + - 输出为图像列表，可直接接入「多分辨率图像预览」节点 + """ + + @classmethod + def define_schema(cls): + autogrow_template = io.Autogrow.TemplateNames( + input=io.Image.Input("image"), + names=_SLOT_NAMES, + min=2, + ) + return io.Schema( + node_id="BatchImagesO1key", + display_name="加载图像（批量）", + category="image", + description=( + "将多个独立图像收集为图像列表输出，保留每张图的原始分辨率。\n" + "• 默认显示 2 个输入端口（图1、图2），连接最后一个后自动追加新端口\n" + "• 断开连线后端口自动减少，最少保留 2 个\n" + "• 不做任何 resize / 裁剪，原图尺寸原样输出\n" + "• 输出为图像列表，可直接接入「多分辨率图像预览」节点" + ), + search_aliases=["批量图像", "batch images", "合并图像", "图像合并", "stack images"], + inputs=[ + io.Autogrow.Input("images", template=autogrow_template) + ], + outputs=[ + io.Image.Output(display_name="图像", is_output_list=True), + ], + ) + + @classmethod + def execute(cls, images: io.Autogrow.Type) -> io.NodeOutput: + # images 是 dict，key 为 "图1", "图2", ... ；未连接的 slot 值为 None + tensors = [v for v in images.values() if v is not None] + + if not tensors: + raise ValueError("批量图像（o1key）：请至少连接一张图像") + + for i, t in enumerate(tensors): + h, w = t.shape[1], t.shape[2] + print(f"批量图像（o1key）：图{i + 1} → {w}×{h}，shape={list(t.shape)}") + + print(f"批量图像（o1key）：共收集 {len(tensors)} 张，原始分辨率原样输出") + + # 以 list[Tensor] 形式返回，每张图保持自身分辨率 + return io.NodeOutput(tensors) + + diff --git a/nodes/batch_nano_banana_pro.py b/nodes/batch_nano_banana_pro.py new file mode 100644 index 0000000..65bc01f --- /dev/null +++ b/nodes/batch_nano_banana_pro.py @@ -0,0 +1,1121 @@ +""" +批量 Nano Banana Pro 节点 +ComfyUI 自定义节点，用于批量处理图像生成任务 +支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存 +""" + +import time +import math +import random +import asyncio +import aiohttp +from concurrent.futures import ThreadPoolExecutor +from typing import Optional, Tuple, List +from PIL import Image + +import torch +import numpy as np + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts +from ..utils.file_utils import ( + ImageInfo, + load_images_from_folder, + pair_images_by_name, + pair_images_cartesian, + generate_timestamp_filename, + save_image, +) +from ..clients.gemini_client import GeminiAPIClient +from ..models_config import ( + get_enabled_models, + get_model_supported_aspect_ratios, get_all_supported_aspect_ratios, + get_model_supported_resolutions, get_all_supported_resolutions +) + +# 导入 ComfyUI 原生进度条 +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + PROGRESS_BAR_AVAILABLE = False + print("⚠️ BatchNanoBananaPro: comfy.utils.ProgressBar 不可用，将只使用终端进度显示") + +# 导入 ComfyUI 的文件夹路径管理 +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + print("⚠️ BatchNanoBananaPro: folder_paths 不可用，将无法使用默认保存路径") + +# 内存监控（可选） +try: + import psutil + MEMORY_MONITOR_AVAILABLE = True +except ImportError: + MEMORY_MONITOR_AVAILABLE = False + print("⚠️ BatchNanoBananaPro: psutil 不可用，内存监控功能禁用") + +# ============================================================================ +# 调试日志配置 +# ============================================================================ +# 是否启用调试日志（打印完整的 API 响应内容） +# 设置为 True 以启用调试日志，False 以禁用 +DEBUG_LOG_ENABLED = False +# 是否启用请求体日志（打印发送给 API 的请求体，base64 图片数据将自动截断） +# 设置为 True 以启用请求体日志，False 以禁用 +REQUEST_LOG_ENABLED = False +# ============================================================================ + +_NODE = "Nano Banana Pro" + + +def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor: + """ + 将 PIL Image 列表转换为 ComfyUI tensor，安全处理多张不同尺寸的情况。 + + ComfyUI 的 IMAGE tensor 格式为 [B, H, W, C]，要求 batch 内所有图尺寸相同。 + 当 API 返回多张不同分辨率的图时（主图 + 附图），直接 stack 会崩溃。 + + 策略： + - 所有图均已按原始分辨率保存到磁盘（调用此函数前已完成） + - 以第一张图的尺寸为基准，只将尺寸相同的图纳入 tensor 输出 + - 尺寸不同的图跳过（不 resize、不丢弃磁盘文件），并打印日志提示 + - 若没有任何图与第一张尺寸相同（极罕见），则只输出第一张 + """ + if not images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + return pil_to_tensor([placeholder]) + + base_size = images[0].size # PIL size = (W, H) + matched = [img for img in images if img.size == base_size] + skipped = [img for img in images if img.size != base_size] + + if skipped: + sizes_str = ", ".join(f"{img.size[0]}×{img.size[1]}" for img in skipped) + print( + f"{node_label}: API 额外返回了 {len(skipped)} 张不同尺寸的图 ({sizes_str})，" + f"已按原始分辨率保存到磁盘，tensor 输出仅包含与主图尺寸相同的 {len(matched)} 张 " + f"({base_size[0]}×{base_size[1]})" + ) + + return pil_to_tensor(matched if matched else [images[0]]) + + +class BatchNanoBananaPro: + """ + 批量 Nano Banana Pro 节点 + + 功能： + - 从多个文件夹加载图片 + - 支持三种配对模式： + * 1:1 - 索引配对（文件夹之间按位置配对） + * 1*N - 笛卡尔积配对（所有可能组合） + * 不配对 - 固定参考图模式（文件夹图片依次与所有参考图组合） + - 批量调用 API 生成图像 + - 智能命名保存（保留原始文件名） + - 并发控制（默认最大 100） + + 注意： + - 「不配对」模式只支持单个文件夹 + - 支持的模型列表从 models_config.py 动态加载 + - 要添加/禁用模型，请编辑 models_config.py 文件 + """ + + # 支持的模型列表（从配置文件动态加载） + MODELS = None # 将在 INPUT_TYPES 中动态获取 + + # 支持的宽高比列表（全量：所有启用模型的并集，动态加载） + # 实际渲染时通过 get_all_supported_aspect_ratios() 获取 + ASPECT_RATIOS = [ + "1:1", "4:3", "3:4", "16:9", "9:16", + "2:3", "3:2", "4:5", "5:4", "21:9", + "1:4", "4:1", "1:8", "8:1" + ] + + # 支持的分辨率列表（全量兜底，实际由 get_all_supported_resolutions() 动态生成） + RESOLUTIONS = ["512", "1K", "2K", "4K"] + + # 配对模式 + PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"] + + def __init__(self): + """初始化节点""" + self.client = None + + def resize_to_megapixels( + self, + image: Image.Image, + target_megapixels: float + ) -> Image.Image: + """ + 将图像缩放到指定的总像素数，保持纵横比 + + Args: + image: PIL Image 对象 + target_megapixels: 目标像素数（百万像素） + + Returns: + 缩放后的 PIL Image + + Example: + >>> resized = self.resize_to_megapixels(img, 2.0) # 缩放到2百万像素 + """ + # 计算当前像素数 + current_pixels = image.width * image.height + target_pixels = int(target_megapixels * 1_000_000) + + # 如果当前像素数已经接近目标，则不缩放 + if abs(current_pixels - target_pixels) / target_pixels < 0.05: + return image + + # 计算缩放比例 + scale = (target_pixels / current_pixels) ** 0.5 + + # 计算新尺寸 + new_width = int(image.width * scale) + new_height = int(image.height * scale) + + # 确保至少为1像素 + new_width = max(1, new_width) + new_height = max(1, new_height) + + # 使用 Lanczos 重采样 + resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) + + return resized_image + + @classmethod + def INPUT_TYPES(cls): + """ + 定义输入参数 + + ComfyUI 节点规范： + - required: 必选参数 + - optional: 可选参数 + """ + # 从配置文件动态获取启用的模型列表 + enabled_models = get_enabled_models() + + # 如果没有启用的模型，使用空列表（会导致节点不可用，提示用户配置） + if not enabled_models: + enabled_models = ["请在 models_config.py 中启用至少一个模型"] + + # 动态获取所有启用模型支持的宽高比（去重合并） + all_aspect_ratios = get_all_supported_aspect_ratios() + if not all_aspect_ratios: + all_aspect_ratios = cls.ASPECT_RATIOS + + # 动态获取所有启用模型支持的分辨率（去重合并） + all_resolutions = get_all_supported_resolutions() + if not all_resolutions: + all_resolutions = cls.RESOLUTIONS + + # 创建9个独立的图像输入 + optional_inputs = {} + for i in range(1, 10): # 1-9 + optional_inputs[f"参考图{i}"] = ("IMAGE",) + + # 图片配对模式移到可选参数 + optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, { + "default": "不配对" + }) + + return { + "required": { + "prompt": ("STRING", { + "default": "一个中国女子的OOTD", + "multiline": True + }), + "模型": (enabled_models, { + "default": enabled_models[0] + }), + "宽高比": (all_aspect_ratios, { + "default": "1:1" + }), + "分辨率": (all_resolutions, { + "default": "2K" + }), + "像素缩放": ("BOOLEAN", { + "default": False, + "label_on": "打开", + "label_off": "关闭" + }), + "分辨率像素": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 100.0, + "step": 0.1, + "display": "number" + }), + "谷歌搜索（联网）": (["关闭", "打开"], { + "default": "关闭" + }), + "图片搜索（联网）": (["关闭", "打开"], { + "default": "关闭" + }), + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 0xffffffffffffffff + }), + "文件夹1": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹2": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹3": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹4": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹5": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹6": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹7": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹8": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹9": ("STRING", { + "default": "", + "multiline": False + }), + "保存路径": ("STRING", { + "default": "", + "multiline": False + }) + }, + "optional": optional_inputs + } + + # 返回值类型 + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("输出图像",) + + # 执行函数名 + FUNCTION = "process_batch" + + # 节点分类 + CATEGORY = "image/batch" + + def _load_folders( + self, + folder1: str, + folder2: Optional[str], + folder3: Optional[str], + folder4: Optional[str], + enable_scaling: bool, + target_megapixels: float, + folder5: Optional[str] = None, + folder6: Optional[str] = None, + folder7: Optional[str] = None, + folder8: Optional[str] = None, + folder9: Optional[str] = None, + ) -> List[List[ImageInfo]]: + """ + 加载所有文件夹中的图片 + + Args: + folder1-9: 文件夹路径 + enable_scaling: 是否启用像素缩放 + target_megapixels: 目标像素数（百万像素） + + Returns: + 图片列表的列表 + """ + folders = [folder1, folder2, folder3, folder4, folder5, folder6, folder7, folder8, folder9] + all_images = [] + + for i, folder in enumerate(folders, 1): + if folder and folder.strip(): + try: + images = load_images_from_folder(folder) + if images: + # 应用像素缩放 + if enable_scaling: + scaled_images = [] + for img_info in images: + scaled_img = self.resize_to_megapixels( + img_info.image, + target_megapixels + ) + # 创建新的 ImageInfo，保留其他元数据 + scaled_info = ImageInfo( + image=scaled_img, + filename=img_info.filename, + extension=img_info.extension, + source_path=img_info.source_path + ) + scaled_images.append(scaled_info) + images = scaled_images + + all_images.append(images) + else: + # 空文件夹，静默跳过 + pass + except ValueError as e: + print(f"BatchNanoBananaPro: 文件夹{i} 加载失败 - {e}") + + return all_images + + def _create_pairs( + self, + image_lists: List[List[ImageInfo]], + pairing_mode: str, + manual_images: Optional[List[ImageInfo]] = None + ) -> List[Tuple[ImageInfo, ...]]: + """ + 根据配对模式创建图片组合 + + Args: + image_lists: 从文件夹加载的图片列表 + pairing_mode: 配对模式 (1:1, 1*N, 不配对) + manual_images: 手动输入的参考图 + + Returns: + 配对后的元组列表 + + Raises: + ValueError: 不配对模式下填入多个文件夹时 + """ + # === 不配对模式 === + if pairing_mode == "不配对": + # 验证：只支持单个文件夹 + if len(image_lists) > 1: + raise ValueError("「不配对」模式只支持单个文件夹，请清空其他文件夹路径") + + # 场景1：有文件夹 + 有参考图 → 每张文件夹图片 + 所有参考图 + if image_lists and manual_images: + folder_images = image_lists[0] + pairs = [] + for img in folder_images: + pair = (img,) + tuple(manual_images) + pairs.append(pair) + return pairs + + # 场景2：有文件夹 + 无参考图 → 每张图片单独成组 + elif image_lists: + return [(img,) for img in image_lists[0]] + + else: + return [] + + # === 1:1 和 1*N 模式 === + # 参考图不参与配对，仅在文件夹图片之间进行配对 + if not image_lists: + return [] + + # 文件夹图片配对 + if len(image_lists) == 1: + base_pairs = [(img,) for img in image_lists[0]] + elif pairing_mode == "按相同图片命名": + base_pairs = list(pair_images_by_name(*image_lists)) + else: # 1*N + base_pairs = list(pair_images_cartesian(*image_lists)) + + # 将所有参考图追加到每组末尾（不参与配对逻辑） + if manual_images: + manual_tuple = tuple(manual_images) + base_pairs = [pair + manual_tuple for pair in base_pairs] + + return base_pairs + + async def _generate_single_task( + self, + client: GeminiAPIClient, + session: aiohttp.ClientSession, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: List[ImageInfo], + output_folder: str, + task_index: int, + enable_grounding: bool = True, + enable_image_search: bool = False, + base_filename: str = None, + ) -> dict: + """ + 执行单个生成任务 + + Args: + client: API 客户端 + session: aiohttp 会话 + prompt: 提示词 + model: 模型名称 + resolution: 分辨率 + aspect_ratio: 宽高比 + images: 输入图片列表 + output_folder: 输出文件夹 + task_index: 任务索引 + + Returns: + 包含结果信息的字典 + """ + result = { + "task_index": task_index, + "prompt": prompt, + "success": False, + "generated_count": 0, + "saved_files": [], + "output_images": [], # 无保存路径时存储内存图片 + "error": None + } + + try: + # 准备输入图片 + input_pil_images = [info.image for info in images] + + # 调用 API 生成图片（固定生成1次） + generated_images = [] + try: + gen_result = await client.generate_single_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=input_pil_images, + session=session, + debug=DEBUG_LOG_ENABLED, + debug_request=REQUEST_LOG_ENABLED, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + if gen_result: + # 正确解包元组：第一个元素是图像列表，第二个是计时信息 + images_list, timing_info = gen_result + generated_images.extend(images_list) + except Exception as e: + import traceback + error_msg = str(e) + error_traceback = traceback.format_exc() + print(f"=" * 80) + print(f"🔍 【原始报错信息展示】") + print(f"=" * 80) + print(f"任务编号: {task_index + 1}") + print(f"失败时间: {time.strftime('%Y-%m-%d %H:%M:%S')}") + print(f"模型: {model}") + print(f"分辨率: {resolution}") + print(f"宽高比: {aspect_ratio}") + print(f"-" * 80) + print(f"错误信息: {error_msg}") + print(f"-" * 80) + print(f"完整堆栈追踪:") + print(error_traceback) + print(f"=" * 80) + result["error"] = error_msg + + # 保存生成的图片到磁盘（始终保存） + import os + for i, gen_img in enumerate(generated_images): + # 使用文件夹1图片的名称，如果重名则+1 + if base_filename: + base_name = base_filename + counter = 0 + while True: + if counter == 0: + filename = f"{base_name}.png" + else: + filename = f"{base_name}+{counter}.png" + output_path = os.path.join(output_folder, filename) + if not os.path.exists(output_path): + break + counter += 1 + else: + # 如果没有base_filename，使用时间戳 + output_path = generate_timestamp_filename( + output_folder=output_folder, + extension=".png" + ) + save_image(gen_img, output_path) + result["saved_files"].append(output_path) + # 立即释放内存 + gen_img = None + + # 只有生成了图片才标记为成功 + if len(generated_images) > 0: + result["success"] = True + result["generated_count"] = len(generated_images) + + except Exception as e: + result["error"] = str(e) + + return result + + async def _process_batch_async( + self, + pairs: List[Tuple[ImageInfo, ...]], + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + output_folder: str, + pbar=None, + prompts_per_task: Optional[List[str]] = None, + enable_grounding: bool = True, + enable_image_search: bool = False, + ) -> List[dict]: + """ + 异步批量处理所有任务 - 改进版：支持分批保存 + + Args: + pairs: 配对后的图片组合 + prompt: 提示词（单提示词模式时使用） + model: 模型名称 + resolution: 分辨率 + aspect_ratio: 宽高比 + output_folder: 输出文件夹 + pbar: ComfyUI 进度条 + prompts_per_task: 每个任务对应的提示词列表（批量提示词模式时使用） + + Returns: + 所有任务的结果列表 + """ + if self.client is None: + self.client = GeminiAPIClient() + + total_tasks = len(pairs) + + # 保持并发数为10不变（按用户要求） + max_concurrent = 10 + + # 分批保存的批次大小（与并发数一致） + save_batch_size = 10 + + print(f"BatchNanoBananaPro: 检测到 {total_tasks} 个任务") + + all_results = [] + completed = 0 + success_count = 0 + fail_count = 0 + + # 计算生成批次数量 + num_batches = math.ceil(total_tasks / max_concurrent) + + # 内存监控初始化 + if MEMORY_MONITOR_AVAILABLE and total_tasks > 50: + import psutil + process = psutil.Process() + initial_memory = process.memory_info().rss / 1024 / 1024 + print(f"BatchNanoBananaPro: 初始内存使用: {initial_memory:.1f} MB") + + # 进度打印配置：任务数 >= 50 时，额外显示百分比里程碑 + show_milestone = total_tasks >= 50 + milestones = [0.2, 0.4, 0.6, 0.8, 1.0] # 20%, 40%, 60%, 80%, 100% + milestone_index = 0 + + if num_batches > 1: + print(f"BatchNanoBananaPro: 任务数 {total_tasks} 超过并发上限 {max_concurrent}，将分 {num_batches} 批执行") + + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + + async with aiohttp.ClientSession(connector=connector) as session: + # 分批处理：每批最多10个任务 + for batch_idx in range(num_batches): + start_idx = batch_idx * max_concurrent + end_idx = min(start_idx + max_concurrent, total_tasks) + batch_pairs = pairs[start_idx:end_idx] + + if num_batches > 1: + print(f"BatchNanoBananaPro: 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...") + + # 创建当前批次的任务 + tasks = [] + for i, pair in enumerate(batch_pairs): + # 批量提示词模式时，每个任务使用对应的提示词；否则使用统一提示词 + task_prompt = prompts_per_task[start_idx + i] if prompts_per_task else prompt + + # 提取文件夹1图片的名称作为保存文件名 + base_filename = None + if pair and len(pair) > 0: + first_image = pair[0] + if hasattr(first_image, 'filename'): + base_filename = first_image.filename + + task = asyncio.create_task( + self._generate_single_task( + client=self.client, + session=session, + prompt=task_prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=list(pair), + output_folder=output_folder, + task_index=start_idx + i, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + base_filename=base_filename, + ) + ) + tasks.append(task) + + # 收集当前批次的结果 + batch_results = [] + for coro in asyncio.as_completed(tasks): + result_data = None + try: + result = await coro + if isinstance(result, Exception): + result_data = { + "success": False, + "error": str(result), + "generated_count": 0, + "saved_files": [] + } + batch_results.append(result_data) + else: + result_data = result + batch_results.append(result) + except Exception as e: + result_data = { + "success": False, + "error": str(e), + "generated_count": 0, + "saved_files": [] + } + batch_results.append(result_data) + + completed += 1 + + # 根据成功/失败状态打印不同信息 + if result_data and result_data.get("success", False): + success_count += 1 + print(f"BatchNanoBananaPro: 任务 {completed}/{total_tasks} 成功 ✓") + else: + fail_count += 1 + # 打印完整错误信息（用于排查问题） + error_msg = result_data.get("error", "未知错误") if result_data else "未知错误" + print(f"BatchNanoBananaPro: 任务 {completed}/{total_tasks} 失败 ✗") + print(f"=" * 80) + print(f"🔍 【原始报错信息展示】") + print(f"=" * 80) + print(f"任务编号: {completed}/{total_tasks}") + print(f"失败时间: {time.strftime('%Y-%m-%d %H:%M:%S')}") + print(f"-" * 80) + print(f"错误详情:") + print(error_msg) + print(f"=" * 80) + + # 更新 ComfyUI 原生进度条 + if pbar is not None: + pbar.update(1) + + # 大任务额外显示百分比里程碑 + if show_milestone and milestone_index < len(milestones): + progress = completed / total_tasks + if progress >= milestones[milestone_index]: + percentage = int(milestones[milestone_index] * 100) + print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<") + milestone_index += 1 + + # 当前批次完成后，立即保存结果并清理内存 + all_results.extend(batch_results) + + # 分批保存：每完成一批（10个任务），立即处理保存并清理内存 + print(f"BatchNanoBananaPro: 第 {batch_idx + 1} 批完成，开始分批保存...") + + # 统计当前批次的结果 + batch_success = sum(1 for r in batch_results if r.get("success", False)) + batch_fail = len(batch_results) - batch_success + batch_generated = sum(r.get("generated_count", 0) for r in batch_results) + + print(f"BatchNanoBananaPro: 本批结果 - 成功: {batch_success}/{len(batch_results)}，生成: {batch_generated} 张") + + # 强制垃圾回收，释放内存 + import gc + gc.collect() + + # 内存监控 + if MEMORY_MONITOR_AVAILABLE and total_tasks > 50: + current_memory = process.memory_info().rss / 1024 / 1024 + memory_increase = current_memory - initial_memory + print(f"BatchNanoBananaPro: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)") + + # 内存警告阈值（2GB） + if current_memory > 2000: + print(f"⚠️ BatchNanoBananaPro: 内存使用过高！但图片已分批保存，即使崩溃也不会丢失已完成的任务") + + # 短暂暂停，让系统有时间处理文件I/O + await asyncio.sleep(0.5) + + return all_results + + def process_batch( + self, + prompt: str, + 文件夹1: str, + 文件夹2: str, + 文件夹3: str, + 文件夹4: str, + 文件夹5: str, + 文件夹6: str, + 文件夹7: str, + 文件夹8: str, + 文件夹9: str, + 像素缩放: bool, + 分辨率像素: float, + seed: int, + 图片配对模式: str, + 模型: str, + 宽高比: str, + 分辨率: str, + 保存路径: str = "", + **kwargs + ) -> Tuple[torch.Tensor]: + """ + 批量处理图像生成任务 + + Args: + prompt: 提示词 + 文件夹1-9: 图片文件夹路径 + 像素缩放: 是否启用像素缩放 + 分辨率像素: 目标像素数（百万像素） + seed: 随机种子 + 保存路径: 输出保存路径 + 图片配对模式: 1:1 或 1*N + 模型: 模型名称 + 宽高比: 输出宽高比 + 分辨率: 输出分辨率 + **kwargs: 动态参考图输入 (参考图1-9) + + Returns: + 输出图像张量 + """ + start_time = time.time() + + # 从 kwargs 提取搜索参数（界面显示为「关闭/打开」，转为 bool 供调用） + enable_grounding: bool = (kwargs.pop("谷歌搜索（联网）", "关闭") == "打开") + enable_image_search: bool = (kwargs.pop("图片搜索（联网）", "关闭") == "打开") + + + try: + # 设置随机种子（用于本地随机操作） + random.seed(seed) + np.random.seed(seed % (2**32)) + + # 验证：至少需要填写一个文件夹路径 + has_any_folder = any( + f and f.strip() + for f in [文件夹1, 文件夹2, 文件夹3, 文件夹4, 文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9] + ) + if not has_any_folder: + raise ValueError("请至少填写一个文件夹路径，该节点专为批量文件夹处理设计") + + # 校验分辨率与模型的兼容性 + supported_resolutions = get_model_supported_resolutions(模型) + if supported_resolutions and 分辨率 not in supported_resolutions: + raise ValueError( + f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的分辨率：{', '.join(supported_resolutions)}" + ) + + # 校验宽高比与模型的兼容性 + supported_ratios = get_model_supported_aspect_ratios(模型) + if supported_ratios and 宽高比 not in supported_ratios: + raise ValueError( + f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的宽高比：{', '.join(supported_ratios)}" + ) + + # 校验图片搜索（联网）与模型的兼容性 + # 仅 nano-banana-2-限时特价 和 gemini-3.1-flash-image-preview 支持图片搜索 + IMAGE_SEARCH_UNSUPPORTED_MODELS = ["nano-banana-pro-限时特价", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview"] + if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS: + raise ValueError( + f"模型 \"{模型}\" 不支持【图片搜索（联网）】功能！" + f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用" + ) + + # 加载文件夹图片 + print("BatchNanoBananaPro: 开始加载图片...") + image_lists = self._load_folders( + 文件夹1, 文件夹2, 文件夹3, 文件夹4, + 像素缩放, 分辨率像素, + 文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9 + ) + + # 验证文件夹是否有可用图片 + total_folder_images = sum(len(lst) for lst in image_lists) + if total_folder_images == 0: + raise ValueError("文件夹中未找到任何图片，请检查文件夹路径是否正确") + + # 处理独立的参考图输入 + manual_images = [] + for i in range(1, 10): # 1-9 + key = f"参考图{i}" + if key in kwargs and kwargs[key] is not None: + pil_images = tensor_to_pil(kwargs[key]) + for j, img in enumerate(pil_images): + # 如果启用像素缩放，也对参考图进行缩放 + if 像素缩放: + img = self.resize_to_megapixels(img, 分辨率像素) + + manual_images.append( + ImageInfo( + image=img, + filename=f"manual_{i}_{j}", + extension=".png", + source_path="" + ) + ) + + # 创建配对 + pairs = self._create_pairs(image_lists, 图片配对模式, manual_images if manual_images else None) + + if not pairs: + raise ValueError("配对结果为空，请检查输入") + + # 解析批量提示词（使用 --- 分隔多个提示词） + batch_prompts = parse_batch_prompts(prompt) + prompts_per_task = None + if batch_prompts: + # 展开 pairs × prompts：每个图片组合 × 每个提示词 = 一个任务 + expanded_pairs = [] + expanded_prompts = [] + for pair in pairs: + for bp in batch_prompts: + expanded_pairs.append(pair) + expanded_prompts.append(bp) + pairs = expanded_pairs + prompts_per_task = expanded_prompts + + total_tasks = len(pairs) + + # 打印首行概览 + # 图片搜索（联网）开启时隐含谷歌搜索接地，与客户端请求逻辑保持一致 + grounding_str = "" + if enable_image_search: + grounding_str = " | 谷歌图片搜索接地" + elif enable_grounding: + grounding_str = " | 谷歌搜索接地" + + if batch_prompts: + print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}") + else: + print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}") + + # 创建 ComfyUI 原生进度条 + pbar = None + if PROGRESS_BAR_AVAILABLE: + pbar = ProgressBar(total_tasks) + + # 检查保存路径（重要！） + has_save_path = bool(保存路径 and 保存路径.strip()) + if not has_save_path: + # 使用 ComfyUI 默认 output 目录作为保存路径 + if FOLDER_PATHS_AVAILABLE: + 保存路径 = folder_paths.get_output_directory() + has_save_path = True + print(f"BatchNanoBananaPro: 未设置保存路径，将使用 ComfyUI 默认 output 目录: {保存路径}") + else: + print("BatchNanoBananaPro: 未设置保存路径，图片将输出到节点") + + if has_save_path: + # 验证保存路径 + import os + try: + os.makedirs(保存路径, exist_ok=True) + # 测试写入权限 + test_file = os.path.join(保存路径, ".write_test") + with open(test_file, 'w') as f: + f.write("test") + os.remove(test_file) + print(f"BatchNanoBananaPro: 保存路径验证通过: {保存路径}") + except Exception as e: + raise ValueError(f"保存路径无效或无写入权限: {保存路径} - {str(e)}") + + # 初始化 API 客户端 + if self.client is None: + try: + self.client = GeminiAPIClient() + except ValueError as e: + raise ValueError(f"初始化 API 客户端失败: {str(e)}") + + # 判断是否使用默认 output 目录 + original_save_path = kwargs.get('保存路径', '') + user_set_save_path = bool(original_save_path and original_save_path.strip()) + + # 执行批量生成 + # 在新线程中运行异步代码，避免事件循环冲突 + def run_async_in_thread(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete( + self._process_batch_async( + pairs=pairs, + prompt=prompt, + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + output_folder=保存路径, + pbar=pbar, + prompts_per_task=prompts_per_task, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + ) + except Exception as e: + # 即使崩溃，也记录错误 + print(f"BatchNanoBananaPro: 异步任务执行异常: {str(e)}") + raise + finally: + loop.close() + + # 使用线程池在新线程中运行事件循环 + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_async_in_thread) + try: + results = future.result(timeout=3600) # 1小时超时 + except TimeoutError: + print("BatchNanoBananaPro: 任务执行超时（1小时）") + raise RuntimeError("任务执行超时，请减少任务数量或检查网络连接") + except Exception as e: + print(f"BatchNanoBananaPro: 任务执行失败: {str(e)}") + # 即使失败，也尝试返回部分结果 + if 'all_saved_files' in locals(): + print(f"BatchNanoBananaPro: 部分保存的图片: {len(all_saved_files)} 张") + raise + + # 统计结果 + success_count = sum(1 for r in results if r.get("success", False)) + fail_count = len(results) - success_count + total_generated = sum(r.get("generated_count", 0) for r in results) + all_saved_files = [] + for r in results: + all_saved_files.extend(r.get("saved_files", [])) + + elapsed = time.time() - start_time + + # 格式化时间 + if elapsed < 1: + time_str = f"{elapsed:.3f}s" + else: + time_str = f"{elapsed:.2f}s" + + # 计算平均耗时 + avg_time = elapsed / success_count if success_count > 0 else 0 + avg_time_str = f"{avg_time:.1f}s/张" if success_count > 0 else "N/A" + + # 精简统计信息 + has_save_path = bool(保存路径 and 保存路径.strip()) + is_default_path = not bool(kwargs.get('保存路径', '').strip() if '保存路径' in locals() else False) + print("=" * 60) + print(f"完成！总耗时 {time_str} | 成功: {success_count}/{total_tasks} | 生成 {total_generated} 张 | 平均 {avg_time_str}") + if has_save_path: + if is_default_path: + print(f"保存路径: {保存路径} (ComfyUI 默认 output 目录)") + else: + print(f"保存路径: {保存路径}") + else: + print("保存路径: 未设置（仅输出到节点）") + + # 失败详情（如果有） + failed_results = [r for r in results if not r.get("success", False)] + if failed_results: + print(f"-" * 60) + print(f"❌ 失败任务汇总: {len(failed_results)} 个") + print(f"-" * 60) + + # 显示前3个失败任务的详细信息 + for idx, failed in enumerate(failed_results[:3], 1): + task_num = failed.get('task_index', '?') + 1 + error_msg = failed.get('error', '未知错误') + print(f"\n【失败任务 #{task_num}】") + print(f"错误信息: {error_msg}") + + if len(failed_results) > 3: + remaining = [str(r.get('task_index', '?') + 1) for r in failed_results[3:]] + print(f"\n其他失败任务编号: {', '.join(remaining)}") + + print(f"-" * 60) + + # 收集最后几张图片用于 ComfyUI 节点输出 + output_images = [] + max_output_images = 10 + + if all_saved_files: + # 从磁盘加载最近的图片 + recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):] + for file_path in recent_files: + try: + img = Image.open(file_path) + output_images.append(img) + except Exception as e: + print(f"BatchNanoBananaPro: 无法加载图片 {file_path} - {e}") + + # 策略3：如果还是没有图片，创建一个占位图 + if not output_images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + output_images = [placeholder] + + # 转换为张量 + output_tensor = _images_to_tensor_safe(output_images, _NODE) + + # 最终内存清理 + import gc + gc.collect() + + # 打印最终统计信息 + total_saved = len(all_saved_files) + print(f"BatchNanoBananaPro: 任务完成！共保存 {total_saved} 张图片到磁盘") + if total_saved > 0: + print(f"BatchNanoBananaPro: 最新保存的文件: {all_saved_files[-1]}") + + return (output_tensor,) + + except ValueError as e: + # 检测是否为授权错误 + if str(e) == "未授权！": + print("请联系作者授权后方可使用！") + raise ValueError("未授权！") from None + else: + # 用户输入错误 - 打印完整错误信息 + error_msg = str(e) + print(f"BatchNanoBananaPro: ❌ {error_msg}") + raise ValueError(error_msg) from None + + except RuntimeError as e: + # 打印完整错误信息 + error_full = str(e) + print(f"BatchNanoBananaPro: ❌ {error_full}") + raise RuntimeError(error_full) from None + + except Exception as e: + # 其他未知错误 - 打印完整错误信息 + error_msg = str(e) + print(f"BatchNanoBananaPro: ❌ {error_msg}") + + raise type(e)(error_msg) from None + + finally: + # 查询余额 + if self.client is not None: + try: + balance_data = self.client.query_balance_sync() + balance_info = self.client.format_balance_info(balance_data) + print(f"BatchNanaBananaPro: {balance_info}") + print("=" * 60) + except Exception: + pass + + # 最终内存清理 + import gc + gc.collect() + print(f"BatchNanoBananaPro: 最终内存清理完成") \ No newline at end of file diff --git a/nodes/batch_nano_banana_v2.py b/nodes/batch_nano_banana_v2.py new file mode 100644 index 0000000..3168829 --- /dev/null +++ b/nodes/batch_nano_banana_v2.py @@ -0,0 +1,781 @@ +""" +批量 Nano Banana v2 节点 +BatchNanoBananaPro 的完全复刻，唯一改动： + + 将原来 9 个独立「参考图1~9」输入端 + 改为 1 个「参考图」输入端（可选），配合「加载图像（批量）」节点使用。 + + 「加载图像（批量）」输出 is_output_list=True（list[Tensor]）， + 本节点声明 INPUT_IS_LIST = True 来整体接收该列表， + 然后在 process_batch() 开头对所有参数统一解包，其余业务逻辑与原节点完全一致。 +""" + +import os +import gc +import time +import math +import random +import asyncio +import aiohttp +from concurrent.futures import ThreadPoolExecutor +from typing import Optional, Tuple, List +from PIL import Image + +import torch +import numpy as np + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts +from ..utils.file_utils import ( + ImageInfo, + load_images_from_folder, + pair_images_by_name, + pair_images_cartesian, + generate_timestamp_filename, + save_image, +) +from ..clients.gemini_client import GeminiAPIClient +from ..models_config import ( + get_enabled_models, + get_model_supported_aspect_ratios, get_all_supported_aspect_ratios, + get_model_supported_resolutions, get_all_supported_resolutions +) + +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + PROGRESS_BAR_AVAILABLE = False + +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +try: + import psutil + MEMORY_MONITOR_AVAILABLE = True +except ImportError: + MEMORY_MONITOR_AVAILABLE = False + +DEBUG_LOG_ENABLED = False +REQUEST_LOG_ENABLED = False + +_NODE = "BatchNanoBananaV2" + + +def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor: + """ + 将 PIL Image 列表转换为 ComfyUI tensor，安全处理多张不同尺寸的情况。 + + ComfyUI 的 IMAGE tensor 格式为 [B, H, W, C]，要求 batch 内所有图尺寸相同。 + 当 API 返回多张不同分辨率的图时（主图 + 附图），直接 stack 会崩溃。 + + 策略： + - 所有图均已按原始分辨率保存到磁盘（调用此函数前已完成） + - 以第一张图的尺寸为基准，只将尺寸相同的图纳入 tensor 输出 + - 尺寸不同的图跳过（不 resize、不丢弃磁盘文件），并打印日志提示 + - 若没有任何图与第一张尺寸相同（极罕见），则只输出第一张 + """ + if not images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + return pil_to_tensor([placeholder]) + + base_size = images[0].size # PIL size = (W, H) + matched = [img for img in images if img.size == base_size] + skipped = [img for img in images if img.size != base_size] + + if skipped: + sizes_str = ", ".join(f"{img.size[0]}×{img.size[1]}" for img in skipped) + print( + f"{node_label}: API 额外返回了 {len(skipped)} 张不同尺寸的图 ({sizes_str})，" + f"已按原始分辨率保存到磁盘，tensor 输出仅包含与主图尺寸相同的 {len(matched)} 张 " + f"({base_size[0]}×{base_size[1]})" + ) + + return pil_to_tensor(matched if matched else [images[0]]) + + +class BatchNanaBananaV2: + """ + 批量 Nano Banana v2 + + 与 BatchNanoBananaPro 完全一致，参考图输入方式不同： + - 原版：9 个独立可选端口（参考图1~9） + - v2：1 个可选端口「参考图」，配合「加载图像（批量）」可传入任意数量图片 + """ + + ASPECT_RATIOS = [ + "1:1", "4:3", "3:4", "16:9", "9:16", + "2:3", "3:2", "4:5", "5:4", "21:9", + "1:4", "4:1", "1:8", "8:1" + ] + RESOLUTIONS = ["512", "1K", "2K", "4K"] + PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"] + + def __init__(self): + self.client = None + + def resize_to_megapixels(self, image: Image.Image, target_megapixels: float) -> Image.Image: + current_pixels = image.width * image.height + target_pixels = int(target_megapixels * 1_000_000) + if abs(current_pixels - target_pixels) / target_pixels < 0.05: + return image + scale = (target_pixels / current_pixels) ** 0.5 + new_width = max(1, int(image.width * scale)) + new_height = max(1, int(image.height * scale)) + return image.resize((new_width, new_height), Image.Resampling.LANCZOS) + + @classmethod + def INPUT_TYPES(cls): + enabled_models = get_enabled_models() + if not enabled_models: + enabled_models = ["请在 models_config.py 中启用至少一个模型"] + + all_aspect_ratios = get_all_supported_aspect_ratios() or cls.ASPECT_RATIOS + all_resolutions = get_all_supported_resolutions() or cls.RESOLUTIONS + + return { + "required": { + "prompt": ("STRING", {"default": "一个中国女子的OOTD", "multiline": True}), + "模型": (enabled_models, {"default": enabled_models[0]}), + "宽高比": (all_aspect_ratios, {"default": "1:1"}), + "分辨率": (all_resolutions, {"default": "2K"}), + "像素缩放": ("BOOLEAN", {"default": False, "label_on": "打开", "label_off": "关闭"}), + "分辨率像素": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.1, "display": "number"}), + "谷歌搜索（联网）": (["关闭", "打开"], {"default": "关闭"}), + "图片搜索（联网）": (["关闭", "打开"], {"default": "关闭"}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "文件夹1": ("STRING", {"default": "", "multiline": False}), + "文件夹2": ("STRING", {"default": "", "multiline": False}), + "文件夹3": ("STRING", {"default": "", "multiline": False}), + "文件夹4": ("STRING", {"default": "", "multiline": False}), + "文件夹5": ("STRING", {"default": "", "multiline": False}), + "文件夹6": ("STRING", {"default": "", "multiline": False}), + "文件夹7": ("STRING", {"default": "", "multiline": False}), + "文件夹8": ("STRING", {"default": "", "multiline": False}), + "文件夹9": ("STRING", {"default": "", "multiline": False}), + "保存路径": ("STRING", {"default": "", "multiline": False}), + }, + "optional": { + # 单个参考图端口，接受普通 IMAGE 或「加载图像（批量）」输出的列表 + "参考图": ("IMAGE",), + "图片配对模式": (cls.PAIRING_MODES, {"default": "不配对"}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("输出图像",) + FUNCTION = "process_batch" + CATEGORY = "image/batch" + + # 声明 INPUT_IS_LIST，使 ComfyUI 将「加载图像（批量）」的 list[Tensor] + # 整体传入而非逐张迭代执行，同时其余所有参数也会被包进 list，需解包。 + INPUT_IS_LIST = True + + # ------------------------------------------------------------------ # + # 以下方法与 BatchNanoBananaPro 完全相同 + # ------------------------------------------------------------------ # + + def _load_folders( + self, + folder1, folder2, folder3, folder4, + enable_scaling, target_megapixels, + folder5=None, folder6=None, folder7=None, folder8=None, folder9=None, + ) -> List[List[ImageInfo]]: + folders = [folder1, folder2, folder3, folder4, + folder5, folder6, folder7, folder8, folder9] + all_images = [] + for i, folder in enumerate(folders, 1): + if folder and folder.strip(): + try: + images = load_images_from_folder(folder) + if images: + if enable_scaling: + scaled = [] + for info in images: + scaled_img = self.resize_to_megapixels(info.image, target_megapixels) + scaled.append(ImageInfo( + image=scaled_img, + filename=info.filename, + extension=info.extension, + source_path=info.source_path + )) + images = scaled + all_images.append(images) + except ValueError as e: + print(f"{_NODE}: 文件夹{i} 加载失败 - {e}") + return all_images + + def _create_pairs( + self, + image_lists: List[List[ImageInfo]], + pairing_mode: str, + manual_images: Optional[List[ImageInfo]] = None + ) -> List[Tuple[ImageInfo, ...]]: + if pairing_mode == "不配对": + if len(image_lists) > 1: + raise ValueError("「不配对」模式只支持单个文件夹，请清空其他文件夹路径") + if image_lists and manual_images: + return [(img,) + tuple(manual_images) for img in image_lists[0]] + elif image_lists: + return [(img,) for img in image_lists[0]] + else: + return [] + + if not image_lists: + return [] + + if len(image_lists) == 1: + base_pairs = [(img,) for img in image_lists[0]] + elif pairing_mode == "按相同图片命名": + base_pairs = list(pair_images_by_name(*image_lists)) + else: + base_pairs = list(pair_images_cartesian(*image_lists)) + + if manual_images: + manual_tuple = tuple(manual_images) + base_pairs = [pair + manual_tuple for pair in base_pairs] + + return base_pairs + + async def _generate_single_task( + self, + client: GeminiAPIClient, + session: aiohttp.ClientSession, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: List[ImageInfo], + output_folder: str, + task_index: int, + enable_grounding: bool = True, + enable_image_search: bool = False, + base_filename: str = None, + ) -> dict: + result = { + "task_index": task_index, + "prompt": prompt, + "success": False, + "generated_count": 0, + "saved_files": [], + "output_images": [], + "error": None + } + try: + input_pil_images = [info.image for info in images] + generated_images = [] + try: + gen_result = await client.generate_single_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=input_pil_images, + session=session, + debug=DEBUG_LOG_ENABLED, + debug_request=REQUEST_LOG_ENABLED, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + if gen_result: + images_list, timing_info = gen_result + generated_images.extend(images_list) + except Exception as e: + import traceback + error_msg = str(e) + error_traceback = traceback.format_exc() + print(f"=" * 80) + print(f"🔍 【原始报错信息展示】") + print(f"=" * 80) + print(f"任务编号: {task_index + 1}") + print(f"失败时间: {time.strftime('%Y-%m-%d %H:%M:%S')}") + print(f"模型: {model}") + print(f"分辨率: {resolution}") + print(f"宽高比: {aspect_ratio}") + print(f"-" * 80) + print(f"错误信息: {error_msg}") + print(f"-" * 80) + print(f"完整堆栈追踪:") + print(error_traceback) + print(f"=" * 80) + result["error"] = error_msg + + for i, gen_img in enumerate(generated_images): + if base_filename: + base_name = base_filename + counter = 0 + while True: + filename = f"{base_name}.png" if counter == 0 else f"{base_name}+{counter}.png" + output_path = os.path.join(output_folder, filename) + if not os.path.exists(output_path): + break + counter += 1 + else: + output_path = generate_timestamp_filename( + output_folder=output_folder, extension=".png" + ) + save_image(gen_img, output_path) + result["saved_files"].append(output_path) + gen_img = None + + if len(generated_images) > 0: + result["success"] = True + result["generated_count"] = len(generated_images) + + except Exception as e: + result["error"] = str(e) + + return result + + async def _process_batch_async( + self, + pairs: List[Tuple[ImageInfo, ...]], + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + output_folder: str, + pbar=None, + prompts_per_task: Optional[List[str]] = None, + enable_grounding: bool = True, + enable_image_search: bool = False, + ) -> List[dict]: + if self.client is None: + self.client = GeminiAPIClient() + + total_tasks = len(pairs) + max_concurrent = 10 + + print(f"{_NODE}: 检测到 {total_tasks} 个任务") + + all_results = [] + completed = 0 + success_count = 0 + fail_count = 0 + num_batches = math.ceil(total_tasks / max_concurrent) + + if MEMORY_MONITOR_AVAILABLE and total_tasks > 50: + process = psutil.Process() + initial_memory = process.memory_info().rss / 1024 / 1024 + print(f"{_NODE}: 初始内存使用: {initial_memory:.1f} MB") + + show_milestone = total_tasks >= 50 + milestones = [0.2, 0.4, 0.6, 0.8, 1.0] + milestone_index = 0 + + if num_batches > 1: + print(f"{_NODE}: 任务数 {total_tasks} 超过并发上限 {max_concurrent}，将分 {num_batches} 批执行") + + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + async with aiohttp.ClientSession(connector=connector) as session: + for batch_idx in range(num_batches): + start_idx = batch_idx * max_concurrent + end_idx = min(start_idx + max_concurrent, total_tasks) + batch_pairs = pairs[start_idx:end_idx] + + if num_batches > 1: + print(f"{_NODE}: 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...") + + tasks = [] + for i, pair in enumerate(batch_pairs): + task_prompt = prompts_per_task[start_idx + i] if prompts_per_task else prompt + base_filename = None + if pair and len(pair) > 0: + first_image = pair[0] + if hasattr(first_image, 'filename'): + base_filename = first_image.filename + + task = asyncio.create_task( + self._generate_single_task( + client=self.client, + session=session, + prompt=task_prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=list(pair), + output_folder=output_folder, + task_index=start_idx + i, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + base_filename=base_filename, + ) + ) + tasks.append(task) + + batch_results = [] + for coro in asyncio.as_completed(tasks): + result_data = None + try: + result = await coro + if isinstance(result, Exception): + result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": []} + batch_results.append(result_data) + else: + result_data = result + batch_results.append(result) + except Exception as e: + result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": []} + batch_results.append(result_data) + + completed += 1 + if result_data and result_data.get("success", False): + success_count += 1 + print(f"{_NODE}: 任务 {completed}/{total_tasks} 成功 ✓") + else: + fail_count += 1 + error_msg = result_data.get("error", "未知错误") if result_data else "未知错误" + print(f"{_NODE}: 任务 {completed}/{total_tasks} 失败 ✗") + print(f"=" * 80) + print(f"🔍 【原始报错信息展示】") + print(f"=" * 80) + print(f"任务编号: {completed}/{total_tasks}") + print(f"失败时间: {time.strftime('%Y-%m-%d %H:%M:%S')}") + print(f"-" * 80) + print(f"错误详情:") + print(error_msg) + print(f"=" * 80) + + if pbar is not None: + pbar.update(1) + + if show_milestone and milestone_index < len(milestones): + progress = completed / total_tasks + if progress >= milestones[milestone_index]: + percentage = int(milestones[milestone_index] * 100) + print(f"{_NODE}: >>> 进度 {percentage}% <<<") + milestone_index += 1 + + all_results.extend(batch_results) + print(f"{_NODE}: 第 {batch_idx + 1} 批完成，开始分批保存...") + + batch_success = sum(1 for r in batch_results if r.get("success", False)) + batch_fail = len(batch_results) - batch_success + batch_generated = sum(r.get("generated_count", 0) for r in batch_results) + print(f"{_NODE}: 本批结果 - 成功: {batch_success}/{len(batch_results)}，生成: {batch_generated} 张") + + gc.collect() + + if MEMORY_MONITOR_AVAILABLE and total_tasks > 50: + current_memory = process.memory_info().rss / 1024 / 1024 + memory_increase = current_memory - initial_memory + print(f"{_NODE}: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)") + if current_memory > 2000: + print(f"⚠️ {_NODE}: 内存使用过高！但图片已分批保存，即使崩溃也不会丢失已完成的任务") + + await asyncio.sleep(0.5) + + return all_results + + def process_batch( + self, + prompt, + 文件夹1, 文件夹2, 文件夹3, 文件夹4, + 文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9, + 像素缩放, + 分辨率像素, + seed, + 模型, + 宽高比, + 分辨率, + 保存路径, + **kwargs + ) -> Tuple[torch.Tensor]: + + # ---------------------------------------------------------------- + # INPUT_IS_LIST=True 时，所有参数均为 list，先统一解包为标量 + # ---------------------------------------------------------------- + def _unpack(v): + return v[0] if isinstance(v, list) else v + + prompt = _unpack(prompt) + 文件夹1 = _unpack(文件夹1) + 文件夹2 = _unpack(文件夹2) + 文件夹3 = _unpack(文件夹3) + 文件夹4 = _unpack(文件夹4) + 文件夹5 = _unpack(文件夹5) + 文件夹6 = _unpack(文件夹6) + 文件夹7 = _unpack(文件夹7) + 文件夹8 = _unpack(文件夹8) + 文件夹9 = _unpack(文件夹9) + 像素缩放 = _unpack(像素缩放) + 分辨率像素 = _unpack(分辨率像素) + seed = _unpack(seed) + 模型 = _unpack(模型) + 宽高比 = _unpack(宽高比) + 分辨率 = _unpack(分辨率) + 保存路径 = _unpack(保存路径) + + # 含全角括号的参数名无法作为形参，从 kwargs 中提取 + enable_grounding: bool = (_unpack(kwargs.pop("谷歌搜索（联网）", "关闭"))) == "打开" + enable_image_search: bool = (_unpack(kwargs.pop("图片搜索（联网）", "关闭"))) == "打开" + + # 图片配对模式（可选参数） + 图片配对模式 = _unpack(kwargs.pop("图片配对模式", "不配对")) + + # ---------------------------------------------------------------- + # 收集参考图：兼容两种来源 + # 1. 「加载图像（批量）」→ is_output_list=True → list[Tensor] + # INPUT_IS_LIST 下传入的是 list[list[Tensor]] 或 list[Tensor]，需展平 + # 2. 普通 IMAGE 端口（单 tensor 或 batch tensor）→ list 中只有 1 个元素 + # ---------------------------------------------------------------- + ref_raw = kwargs.pop("参考图", None) + manual_images: List[ImageInfo] = [] + + if ref_raw is not None: + items = ref_raw if isinstance(ref_raw, list) else [ref_raw] + idx = 0 + for item in items: + if item is None: + continue + if isinstance(item, list): + sub_tensors = item + elif isinstance(item, torch.Tensor): + sub_tensors = [item] + else: + continue + for tensor in sub_tensors: + if tensor is None or not isinstance(tensor, torch.Tensor): + continue + pil_images = tensor_to_pil(tensor) + for j, img in enumerate(pil_images): + if 像素缩放: + img = self.resize_to_megapixels(img, 分辨率像素) + manual_images.append(ImageInfo( + image=img, + filename=f"manual_{idx}_{j}", + extension=".png", + source_path="" + )) + idx += 1 + + # ---------------------------------------------------------------- + # 以下逻辑与 BatchNanoBananaPro.process_batch() 完全一致 + # ---------------------------------------------------------------- + start_time = time.time() + + try: + random.seed(seed) + np.random.seed(seed % (2 ** 32)) + + has_any_folder = any( + f and f.strip() + for f in [文件夹1, 文件夹2, 文件夹3, 文件夹4, + 文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9] + ) + if not has_any_folder: + raise ValueError("请至少填写一个文件夹路径，该节点专为批量文件夹处理设计") + + supported_resolutions = get_model_supported_resolutions(模型) + if supported_resolutions and 分辨率 not in supported_resolutions: + raise ValueError( + f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的分辨率：{', '.join(supported_resolutions)}" + ) + + supported_ratios = get_model_supported_aspect_ratios(模型) + if supported_ratios and 宽高比 not in supported_ratios: + raise ValueError( + f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的宽高比：{', '.join(supported_ratios)}" + ) + + IMAGE_SEARCH_UNSUPPORTED_MODELS = [ + "nano-banana-pro-限时特价", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview" + ] + if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS: + raise ValueError( + f"模型 \"{模型}\" 不支持【图片搜索（联网）】功能！" + f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用" + ) + + print(f"{_NODE}: 开始加载图片...") + image_lists = self._load_folders( + 文件夹1, 文件夹2, 文件夹3, 文件夹4, + 像素缩放, 分辨率像素, + 文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9 + ) + + total_folder_images = sum(len(lst) for lst in image_lists) + if total_folder_images == 0: + raise ValueError("文件夹中未找到任何图片，请检查文件夹路径是否正确") + + pairs = self._create_pairs(image_lists, 图片配对模式, manual_images if manual_images else None) + + if not pairs: + raise ValueError("配对结果为空，请检查输入") + + batch_prompts = parse_batch_prompts(prompt) + prompts_per_task = None + if batch_prompts: + expanded_pairs = [] + expanded_prompts = [] + for pair in pairs: + for bp in batch_prompts: + expanded_pairs.append(pair) + expanded_prompts.append(bp) + pairs = expanded_pairs + prompts_per_task = expanded_prompts + + total_tasks = len(pairs) + + grounding_str = "" + if enable_image_search: + grounding_str = " | 谷歌图片搜索接地" + elif enable_grounding: + grounding_str = " | 谷歌搜索接地" + + if batch_prompts: + print(f"{_NODE}: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}") + else: + print(f"{_NODE}: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}") + + pbar = None + if PROGRESS_BAR_AVAILABLE: + pbar = ProgressBar(total_tasks) + + has_save_path = bool(保存路径 and 保存路径.strip()) + if not has_save_path: + if FOLDER_PATHS_AVAILABLE: + 保存路径 = folder_paths.get_output_directory() + has_save_path = True + print(f"{_NODE}: 未设置保存路径，将使用 ComfyUI 默认 output 目录: {保存路径}") + else: + print(f"{_NODE}: 未设置保存路径，图片将输出到节点") + + if has_save_path: + try: + os.makedirs(保存路径, exist_ok=True) + test_file = os.path.join(保存路径, ".write_test") + with open(test_file, 'w') as f: + f.write("test") + os.remove(test_file) + print(f"{_NODE}: 保存路径验证通过: {保存路径}") + except Exception as e: + raise ValueError(f"保存路径无效或无写入权限: {保存路径} - {str(e)}") + + if self.client is None: + try: + self.client = GeminiAPIClient() + except ValueError as e: + raise ValueError(f"初始化 API 客户端失败: {str(e)}") + + def run_async_in_thread(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete( + self._process_batch_async( + pairs=pairs, + prompt=prompt, + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + output_folder=保存路径, + pbar=pbar, + prompts_per_task=prompts_per_task, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + ) + except Exception as e: + print(f"{_NODE}: 异步任务执行异常: {str(e)}") + raise + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_async_in_thread) + try: + results = future.result(timeout=3600) + except TimeoutError: + print(f"{_NODE}: 任务执行超时（1小时）") + raise RuntimeError("任务执行超时，请减少任务数量或检查网络连接") + except Exception as e: + print(f"{_NODE}: 任务执行失败: {str(e)}") + raise + + success_count = sum(1 for r in results if r.get("success", False)) + fail_count = len(results) - success_count + total_generated = sum(r.get("generated_count", 0) for r in results) + all_saved_files = [f for r in results for f in r.get("saved_files", [])] + + elapsed = time.time() - start_time + time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" + avg_time = elapsed / success_count if success_count > 0 else 0 + avg_time_str = f"{avg_time:.1f}s/张" if success_count > 0 else "N/A" + + print("=" * 60) + print(f"完成！总耗时 {time_str} | 成功: {success_count}/{total_tasks} | 生成 {total_generated} 张 | 平均 {avg_time_str}") + if has_save_path: + print(f"保存路径: {保存路径}") + else: + print("保存路径: 未设置（仅输出到节点）") + + failed_results = [r for r in results if not r.get("success", False)] + if failed_results: + print(f"-" * 60) + print(f"❌ 失败任务汇总: {len(failed_results)} 个") + print(f"-" * 60) + for idx, failed in enumerate(failed_results[:3], 1): + task_num = failed.get('task_index', '?') + 1 + error_msg = failed.get('error', '未知错误') + print(f"\n【失败任务 #{task_num}】") + print(f"错误信息: {error_msg}") + if len(failed_results) > 3: + remaining = [str(r.get('task_index', '?') + 1) for r in failed_results[3:]] + print(f"\n其他失败任务编号: {', '.join(remaining)}") + print(f"-" * 60) + + output_images = [] + if all_saved_files: + for fp in all_saved_files[-min(10, len(all_saved_files)):]: + try: + output_images.append(Image.open(fp)) + except Exception as e: + print(f"{_NODE}: 无法加载图片 {fp} - {e}") + + if not output_images: + output_images = [Image.new('RGB', (512, 512), color=(128, 128, 128))] + + output_tensor = _images_to_tensor_safe(output_images, _NODE) + gc.collect() + + total_saved = len(all_saved_files) + print(f"{_NODE}: 任务完成！共保存 {total_saved} 张图片到磁盘") + if total_saved > 0: + print(f"{_NODE}: 最新保存的文件: {all_saved_files[-1]}") + + + return (output_tensor,) + + except ValueError as e: + if str(e) == "未授权！": + print("请联系作者授权后方可使用！") + raise ValueError("未授权！") from None + error_msg = str(e) + print(f"{_NODE}: ❌ {error_msg}") + raise ValueError(error_msg) from None + + except RuntimeError as e: + error_full = str(e) + print(f"{_NODE}: ❌ {error_full}") + raise RuntimeError(error_full) from None + + except Exception as e: + error_msg = str(e) + print(f"{_NODE}: ❌ {error_msg}") + raise type(e)(error_msg) from None + + finally: + if self.client is not None: + try: + balance_data = self.client.query_balance_sync() + balance_info = self.client.format_balance_info(balance_data) + print(f"{_NODE}: {balance_info}") + print("=" * 60) + except Exception: + pass + gc.collect() diff --git a/nodes/batch_quan_neng_sheng_tu.py b/nodes/batch_quan_neng_sheng_tu.py new file mode 100644 index 0000000..b61fdd9 --- /dev/null +++ b/nodes/batch_quan_neng_sheng_tu.py @@ -0,0 +1,642 @@ +""" +全能生图（批量）节点 +ComfyUI 自定义节点，用于批量处理图像生成任务 +支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存 +""" + +import time +import math +import random +import asyncio +import aiohttp +from concurrent.futures import ThreadPoolExecutor +from typing import Optional, Tuple, List +from PIL import Image + +import torch +import numpy as np + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts +from ..utils.file_utils import ( + ImageInfo, + load_images_from_folder, + pair_images_by_name, + pair_images_cartesian, + generate_timestamp_filename, + save_image, +) +from ..clients.openai_client import OpenAIAPIClient +from ..models_config import ( + get_enabled_models, + get_model_supported_aspect_ratios, get_all_supported_aspect_ratios, + get_model_supported_resolutions, get_all_supported_resolutions +) + +# 导入 ComfyUI 原生进度条 +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + PROGRESS_BAR_AVAILABLE = False + print("⚠️ 全能生图（批量）: comfy.utils.ProgressBar 不可用，将只使用终端进度显示") + +# 导入 ComfyUI 的文件夹路径管理 +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + print("⚠️ 全能生图（批量）: folder_paths 不可用，将无法使用默认保存路径") + +# 内存监控（可选） +try: + import psutil + MEMORY_MONITOR_AVAILABLE = True +except ImportError: + MEMORY_MONITOR_AVAILABLE = False + print("⚠️ 全能生图（批量）: psutil 不可用，内存监控功能禁用") + +# ============================================================================ +# 调试日志配置 +# ============================================================================ +DEBUG_LOG_ENABLED = False +REQUEST_LOG_ENABLED = False +# ============================================================================ + + +class BatchQuanNengShengTu: + """ + 全能生图（批量）节点 + + 功能： + - 从多个文件夹加载图片 + - 支持三种配对模式： + * 按相同图片命名 - 索引配对（文件夹之间按位置配对） + * 1*N - 笛卡尔积配对（所有可能组合） + * 不配对 - 固定参考图模式（文件夹图片依次与所有参考图组合） + - 批量调用 API 生成图像 + - 智能命名保存（保留原始文件名） + - 并发控制（默认最大 10） + + 注意： + - 「不配对」模式只支持单个文件夹 + - 支持的模型列表从 models_config.py 动态加载 + """ + + MODELS = None + ASPECT_RATIOS = [ + "1:1", "4:3", "3:4", "16:9", "9:16", + "2:3", "3:2", "4:5", "5:4", "21:9", + "1:4", "4:1", "1:8", "8:1" + ] + RESOLUTIONS = ["512", "1K", "2K", "4K"] + PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"] + + def __init__(self): + """初始化节点""" + self.client = None + + def resize_to_megapixels( + self, + image: Image.Image, + target_megapixels: float + ) -> Image.Image: + """将图像缩放到指定的总像素数，保持纵横比""" + current_pixels = image.width * image.height + target_pixels = int(target_megapixels * 1_000_000) + + if abs(current_pixels - target_pixels) / target_pixels < 0.05: + return image + + scale = (target_pixels / current_pixels) ** 0.5 + new_width = max(1, int(image.width * scale)) + new_height = max(1, int(image.height * scale)) + + return image.resize((new_width, new_height), Image.Resampling.LANCZOS) + + @classmethod + def INPUT_TYPES(cls): + """定义输入参数""" + enabled_models = get_enabled_models() + enabled_models = [m for m in enabled_models if "限时特价" not in m] + + if not enabled_models: + enabled_models = ["请在 models_config.py 中启用至少一个模型"] + + all_aspect_ratios = get_all_supported_aspect_ratios() + if not all_aspect_ratios: + all_aspect_ratios = cls.ASPECT_RATIOS + + all_resolutions = get_all_supported_resolutions() + if not all_resolutions: + all_resolutions = cls.RESOLUTIONS + + optional_inputs = {} + for i in range(1, 10): + optional_inputs[f"参考图{i}"] = ("IMAGE",) + + optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, { + "default": "不配对" + }) + + return { + "required": { + "提示词": ("STRING", { + "default": "一个中国女子的OOTD", + "multiline": True + }), + "模型": (enabled_models, { + "default": enabled_models[0] + }), + "宽高比": (all_aspect_ratios, { + "default": "1:1" + }), + "分辨率": (all_resolutions, { + "default": "2K" + }), + "像素缩放": ("BOOLEAN", { + "default": False, + "label_on": "打开", + "label_off": "关闭" + }), + "分辨率像素": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 100.0, + "step": 0.1, + "display": "number" + }), + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 0xffffffffffffffff + }), + "文件夹1": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹2": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹3": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹4": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹5": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹6": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹7": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹8": ("STRING", { + "default": "", + "multiline": False + }), + "文件夹9": ("STRING", { + "default": "", + "multiline": False + }), + "保存路径": ("STRING", { + "default": "", + "multiline": False + }) + }, + "optional": optional_inputs + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("输出图像",) + FUNCTION = "process_batch" + CATEGORY = "image/batch" + + def _load_folders( + self, + folder1: str, + folder2: Optional[str], + folder3: Optional[str], + folder4: Optional[str], + enable_scaling: bool, + target_megapixels: float, + folder5: Optional[str] = None, + folder6: Optional[str] = None, + folder7: Optional[str] = None, + folder8: Optional[str] = None, + folder9: Optional[str] = None, + ) -> List[List[ImageInfo]]: + """加载所有文件夹中的图片""" + folders = [folder1, folder2, folder3, folder4, folder5, folder6, folder7, folder8, folder9] + all_images = [] + + for i, folder in enumerate(folders, 1): + if folder and folder.strip(): + try: + images = load_images_from_folder(folder) + if images: + if enable_scaling: + scaled_images = [] + for img_info in images: + scaled_img = self.resize_to_megapixels( + img_info.image, + target_megapixels + ) + scaled_info = ImageInfo( + image=scaled_img, + filename=img_info.filename, + extension=img_info.extension, + source_path=img_info.source_path + ) + scaled_images.append(scaled_info) + images = scaled_images + all_images.append(images) + except ValueError as e: + print(f"全能生图（批量）: 文件夹{i} 加载失败 - {e}") + + return all_images + + def _create_pairs( + self, + image_lists: List[List[ImageInfo]], + pairing_mode: str, + manual_images: Optional[List[ImageInfo]] = None + ) -> List[Tuple[ImageInfo, ...]]: + """根据配对模式创建图片组合""" + if pairing_mode == "不配对": + if len(image_lists) > 1: + raise ValueError("「不配对」模式只支持单个文件夹，请清空其他文件夹路径") + + if image_lists and manual_images: + folder_images = image_lists[0] + pairs = [] + for img in folder_images: + pair = (img,) + tuple(manual_images) + pairs.append(pair) + return pairs + elif image_lists: + return [(img,) for img in image_lists[0]] + else: + return [] + + if not image_lists: + return [] + + if len(image_lists) == 1: + base_pairs = [(img,) for img in image_lists[0]] + elif pairing_mode == "按相同图片命名": + base_pairs = list(pair_images_by_name(*image_lists)) + else: + base_pairs = list(pair_images_cartesian(*image_lists)) + + if manual_images: + manual_tuple = tuple(manual_images) + base_pairs = [pair + manual_tuple for pair in base_pairs] + + return base_pairs + + async def _generate_single_task( + self, + session: aiohttp.ClientSession, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: List[ImageInfo], + output_folder: str, + task_index: int, + base_filename: str = None, + ) -> dict: + """执行单个生成任务""" + result = { + "task_index": task_index, + "prompt": prompt, + "success": False, + "generated_count": 0, + "saved_files": [], + "error": None + } + + try: + input_pil_images = [info.image for info in images] + + gen_result = await self.client.generate_single_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=input_pil_images, + session=session, + debug=DEBUG_LOG_ENABLED, + debug_request=REQUEST_LOG_ENABLED, + enable_grounding=False, + enable_image_search=False + ) + + if gen_result: + images_list, _ = gen_result + + import os + for gen_img in images_list: + if base_filename: + base_name = base_filename + counter = 0 + while True: + filename = f"{base_name}.png" if counter == 0 else f"{base_name}+{counter}.png" + output_path = os.path.join(output_folder, filename) + if not os.path.exists(output_path): + break + counter += 1 + else: + output_path = generate_timestamp_filename( + output_folder=output_folder, + extension=".png" + ) + save_image(gen_img, output_path) + result["saved_files"].append(output_path) + gen_img = None + + result["success"] = True + result["generated_count"] = len(images_list) + + except Exception as e: + result["error"] = str(e) + + return result + + async def _process_batch_async( + self, + pairs: List[Tuple[ImageInfo, ...]], + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + output_folder: str, + pbar=None, + prompts_per_task: Optional[List[str]] = None, + ) -> List[dict]: + """异步批量处理所有任务""" + if self.client is None: + self.client = OpenAIAPIClient() + + total_tasks = len(pairs) + max_concurrent = 10 + + print(f"全能生图（批量）: 检测到 {total_tasks} 个任务") + + all_results = [] + completed = 0 + success_count = 0 + fail_count = 0 + + num_batches = math.ceil(total_tasks / max_concurrent) + + if num_batches > 1: + print(f"全能生图（批量）: 任务数 {total_tasks} 超过并发上限 {max_concurrent}，将分 {num_batches} 批执行") + + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + + async with aiohttp.ClientSession(connector=connector) as session: + for batch_idx in range(num_batches): + start_idx = batch_idx * max_concurrent + end_idx = min(start_idx + max_concurrent, total_tasks) + batch_pairs = pairs[start_idx:end_idx] + + if num_batches > 1: + print(f"全能生图（批量）: 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...") + + tasks = [] + for i, pair in enumerate(batch_pairs): + task_prompt = prompts_per_task[start_idx + i] if prompts_per_task else prompt + + base_filename = None + if pair and len(pair) > 0: + first_image = pair[0] + if hasattr(first_image, 'filename'): + base_filename = first_image.filename + + task = asyncio.create_task( + self._generate_single_task( + session=session, + prompt=task_prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=list(pair), + output_folder=output_folder, + task_index=start_idx + i, + base_filename=base_filename, + ) + ) + tasks.append(task) + + batch_results = [] + for coro in asyncio.as_completed(tasks): + result_data = None + try: + result = await coro + if isinstance(result, Exception): + result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": []} + else: + result_data = result + batch_results.append(result_data) + except Exception as e: + result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": []} + batch_results.append(result_data) + + completed += 1 + + if result_data and result_data.get("success", False): + success_count += 1 + print(f"全能生图（批量）: 任务 {completed}/{total_tasks} 成功 ✓") + else: + fail_count += 1 + error_msg = result_data.get("error", "未知错误") if result_data else "未知错误" + print(f"全能生图（批量）: 任务 {completed}/{total_tasks} 失败 ✗ - {error_msg}") + + if pbar is not None: + pbar.update(1) + + all_results.extend(batch_results) + + import gc + gc.collect() + + await asyncio.sleep(0.1) + + return all_results + + def process_batch( + self, + 提示词: str, + 模型: str, + 宽高比: str, + 分辨率: str, + 像素缩放: bool, + 分辨率像素: float, + seed: int, + 文件夹1: str, + 文件夹2: str, + 文件夹3: str, + 文件夹4: str, + 文件夹5: str, + 文件夹6: str, + 文件夹7: str, + 文件夹8: str, + 文件夹9: str, + 保存路径: str, + **kwargs + ) -> Tuple[torch.Tensor]: + """批量处理图像生成""" + start_time = time.time() + + try: + random.seed(seed) + np.random.seed(seed % (2**32)) + + if self.client is None: + self.client = OpenAIAPIClient() + + supported_resolutions = get_model_supported_resolutions(模型) + if supported_resolutions and 分辨率 not in supported_resolutions: + raise ValueError( + f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的分辨率：{', '.join(supported_resolutions)}" + ) + + supported_ratios = get_model_supported_aspect_ratios(模型) + if supported_ratios and 宽高比 not in supported_ratios: + raise ValueError( + f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的宽高比：{', '.join(supported_ratios)}" + ) + + manual_images = [] + for i in range(1, 10): + key = f"参考图{i}" + if key in kwargs and kwargs[key] is not None: + pil_imgs = tensor_to_pil(kwargs[key]) + for pil_img in pil_imgs: + manual_images.append(ImageInfo( + image=pil_img, + filename=f"manual_{i}", + extension=".png", + source_path="" + )) + + if 像素缩放 and manual_images: + scaled_manual = [] + for img_info in manual_images: + scaled_img = self.resize_to_megapixels(img_info.image, 分辨率像素) + scaled_manual.append(ImageInfo( + image=scaled_img, + filename=img_info.filename, + extension=img_info.extension, + source_path=img_info.source_path + )) + manual_images = scaled_manual + + folder_images = self._load_folders( + 文件夹1, 文件夹2, 文件夹3, 文件夹4, + 像素缩放, 分辨率像素, + 文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9 + ) + + pairing_mode = kwargs.get("图片配对模式", "不配对") + pairs = self._create_pairs(folder_images, pairing_mode, manual_images if manual_images else None) + + if not pairs: + raise ValueError("没有可处理的图片组合，请检查文件夹路径和参考图输入") + + batch_prompts = parse_batch_prompts(提示词) + prompts_per_task = None + + if batch_prompts: + if len(batch_prompts) != len(pairs): + raise ValueError( + f"批量提示词数量 ({len(batch_prompts)}) 与任务数量 ({len(pairs)}) 不匹配！\n" + f"请确保提示词数量与图片组合数量一致" + ) + prompts_per_task = batch_prompts + print(f"全能生图（批量）: 批量提示词模式 - {len(batch_prompts)} 个提示词") + + output_folder = 保存路径.strip() if 保存路径 else "" + if not output_folder and FOLDER_PATHS_AVAILABLE: + output_folder = folder_paths.get_output_directory() + + if not output_folder: + raise ValueError("无法确定保存路径，请指定保存路径或确保 folder_paths 可用") + + import os + os.makedirs(output_folder, exist_ok=True) + print(f"全能生图（批量）: 保存路径 → {output_folder}") + + pbar = None + if PROGRESS_BAR_AVAILABLE: + pbar = ProgressBar(len(pairs)) + + def run_async(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete( + self._process_batch_async( + pairs=pairs, + prompt=提示词, + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + output_folder=output_folder, + pbar=pbar, + prompts_per_task=prompts_per_task, + ) + ) + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_async) + results = future.result(timeout=3600) + + success_count = sum(1 for r in results if r.get("success", False)) + fail_count = len(results) - success_count + all_saved_files = [] + for r in results: + all_saved_files.extend(r.get("saved_files", [])) + + elapsed = time.time() - start_time + print(f"全能生图（批量）: 完成！总耗时 {elapsed:.2f}s | 成功: {success_count}/{len(pairs)} | 失败: {fail_count}") + + output_images = [] + max_output = 10 + recent_files = all_saved_files[-min(max_output, len(all_saved_files)):] + for file_path in recent_files: + try: + img = Image.open(file_path) + output_images.append(img) + except Exception as e: + print(f"全能生图（批量）: 无法加载 {file_path} - {e}") + + if not output_images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + output_images = [placeholder] + + output_tensor = pil_to_tensor(output_images) + print(f"全能生图（批量）: 共保存 {len(all_saved_files)} 张图片，节点输出最后 {len(output_images)} 张") + + import gc + gc.collect() + return (output_tensor,) + + except Exception as e: + print(f"全能生图（批量）: ❌ {str(e)}") + raise + diff --git a/nodes/flux_edit.py b/nodes/flux_edit.py new file mode 100644 index 0000000..5df7206 --- /dev/null +++ b/nodes/flux_edit.py @@ -0,0 +1,170 @@ +""" +Flux2 图像编辑节点 +通过 vip.o1key.com 调用 Flux2 + SeedVR2 远程服务进行图像编辑和超分辨率 + +功能： +- 接收主图和参考图 +- 上传到远程服务器执行图像编辑 +- 轮询等待 SeedVR2 超分辨率结果 +- 返回最终放大后的图像 +""" + +import time +from io import BytesIO +from typing import Tuple + +import torch +from PIL import Image + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor +from ..clients.flux_edit_client import FluxEditClient + + +class FluxImageEdit: + """ + Flux2 图像编辑节点 + + 通过远程 API 将主图与参考图结合，按照提示词进行图像编辑， + 并经 SeedVR2 超分辨率放大后返回最终结果。 + """ + + SIZES = ["2K", "4K"] + + def __init__(self): + self.client = None + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "主图": ("IMAGE",), + "参考图": ("IMAGE",), + "提示词": ("STRING", { + "default": "Replace the woman's underwear in Figure 1 with the strapless bra in Figure 2", + "multiline": True, + }), + "分辨率": (cls.SIZES, { + "default": "4K", + }), + "轮询间隔": ("INT", { + "default": 15, + "min": 5, + "max": 60, + "step": 5, + }), + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 0xffffffffffffffff, + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("输出图像",) + FUNCTION = "generate" + CATEGORY = "image/edit" + + def _image_to_jpeg_bytes(self, image: Image.Image, quality: int = 92) -> bytes: + """将 PIL Image 转为 JPEG 二进制""" + if image.mode in ("RGBA", "P", "LA"): + image = image.convert("RGB") + buf = BytesIO() + image.save(buf, format="JPEG", quality=quality) + return buf.getvalue() + + def generate( + self, + 主图: torch.Tensor, + 参考图: torch.Tensor, + 提示词: str, + 分辨率: str, + 轮询间隔: int, + seed: int, + ) -> Tuple[torch.Tensor]: + """ + 执行图像编辑 + + Args: + 主图: 要编辑的原始图像 (ComfyUI tensor, [B, H, W, C]) + 参考图: 参考/风格图像 (ComfyUI tensor, [B, H, W, C]) + 提示词: 编辑指令 + 分辨率: 超分辨率目标 ("2K" 或 "4K"，会自动映射为 2048/4096) + 轮询间隔: 轮询秒数 + seed: 随机种子 + + Returns: + 输出图像 tensor (IMAGE,) + """ + start_time = time.time() + + try: + # 初始化客户端 + if self.client is None: + self.client = FluxEditClient() + + # Tensor → PIL（取第一张） + main_pils = tensor_to_pil(主图) + ref_pils = tensor_to_pil(参考图) + + if not main_pils: + raise ValueError("主图不能为空") + if not ref_pils: + raise ValueError("参考图不能为空") + + main_img = main_pils[0] + ref_img = ref_pils[0] + + # PIL → JPEG bytes + main_bytes = self._image_to_jpeg_bytes(main_img) + ref_bytes = self._image_to_jpeg_bytes(ref_img) + + print(f"Flux Edit: 开始处理 | 主图 {main_img.size} | 参考图 {ref_img.size} | 分辨率 {分辨率} | seed {seed}") + + # 进度回调 + def progress_callback(status_str: str): + print(f"Flux Edit: {status_str}") + + # 提交任务并等待结果 + result_bytes = self.client.submit_and_wait( + image_bytes=main_bytes, + mask_bytes=ref_bytes, + prompt=提示词, + size=分辨率, + poll_interval=轮询间隔, + progress_callback=progress_callback, + ) + + # 解码结果 + result_img = Image.open(BytesIO(result_bytes)) + if result_img.mode != "RGB": + result_img = result_img.convert("RGB") + + print(f"Flux Edit: 结果图像尺寸 {result_img.size}") + + # 转为 tensor + output_tensor = pil_to_tensor([result_img]) + + # 打印耗时 + elapsed = time.time() - start_time + if elapsed < 60: + time_str = f"{elapsed:.1f}s" + else: + minutes = int(elapsed // 60) + seconds = elapsed % 60 + time_str = f"{minutes}m {seconds:.0f}s" + print(f"Flux Edit: 完成！总耗时 {time_str}") + + return (output_tensor,) + + except ValueError as e: + if str(e) == "未授权！": + print("请联系作者授权后方可使用！") + raise ValueError("未授权！") from None + print(f"Flux Edit: ❌ {e}") + raise + + except Exception as e: + error_msg = str(e) + print(f"Flux Edit: ❌ {error_msg}") + raise RuntimeError(error_msg) from None diff --git a/nodes/google_gemini.py b/nodes/google_gemini.py new file mode 100644 index 0000000..6c4f1cd --- /dev/null +++ b/nodes/google_gemini.py @@ -0,0 +1,755 @@ +""" +Google Gemini 节点 +ComfyUI 自定义节点,用于调用 Gemini Flash 模型进行多模态文本生成 +""" + +import base64 +import os +import time +import tempfile +from typing import Dict, List, Optional, Tuple +from io import BytesIO + +import torch +from PIL import Image + +from ..utils.image_utils import tensor_to_pil, encode_image_to_base64 +from ..utils.file_types import FileData +from ..clients.gemini_flash_client import GeminiFlashClient +from ..models_config import get_enabled_flash_models + +# 文件大小限制（20MB） +MAX_FILE_SIZE = 20 * 1024 * 1024 + +# 图片缩放后最大尺寸（1K分辨率 = 1024像素） +MAX_IMAGE_DIMENSION = 1024 + +# 视频压缩目标大小（1-10MB） +TARGET_VIDEO_SIZE_MIN = 1 * 1024 * 1024 +TARGET_VIDEO_SIZE_MAX = 10 * 1024 * 1024 + + +# 支持的视频 MIME 类型映射 +VIDEO_MIME_TYPES = { + ".mp4": "video/mp4", + ".mpeg": "video/mpeg", + ".mpg": "video/mpg", + ".mov": "video/quicktime", + ".avi": "video/x-msvideo", + ".flv": "video/x-flv", + ".webm": "video/webm", + ".wmv": "video/x-ms-wmv", + ".3gp": "video/3gpp", + ".3gpp": "video/3gpp" +} + +# 尝试导入视频处理库 +try: + import cv2 + CV2_AVAILABLE = True +except ImportError: + CV2_AVAILABLE = False + print("⚠️ Google Gemini: OpenCV (cv2) 不可用，视频压缩功能将受限") + +try: + import subprocess + FFMPEG_AVAILABLE = True +except ImportError: + FFMPEG_AVAILABLE = False + + +class GoogleGemini: + """ + Google Gemini 节点 + + 功能： + - 支持多个 Gemini Flash 模型 + - 支持图片、视频和文件输入 + - 支持不同思考等级（不思考/低/中/高）- 通过 thinkingConfig.thinkingLevel 控制 + - 输出生成的文本内容（主要内容 + 思考内容） + """ + + # 支持的思考等级选项 + THINKING_LEVELS = ["不思考", "低", "中", "高"] + + def __init__(self): + """初始化节点""" + self.client = None + + @classmethod + def INPUT_TYPES(cls): + """ + 定义输入参数 + """ + # 从配置获取启用的模型列表 + enabled_models = get_enabled_flash_models() + default_model = enabled_models[0] if enabled_models else "gemini-3-flash-preview" + + return { + "required": { + "模型": (enabled_models, { + "default": default_model + }), + "提示词": ("STRING", { + "default": "", + "multiline": True + }), + "思考等级": (cls.THINKING_LEVELS, { + "default": "不思考" + }) + }, + "optional": { + "图片": ("IMAGE",), + "视频": ("VIDEO",), + "文件": ("FILE",) + } + } + + # 返回值类型 + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("主要内容",) + + # 执行函数名 + FUNCTION = "generate" + + # 节点分类 + CATEGORY = "text/generation" + + # 允许输出到 UI + OUTPUT_NODE = True + + def _resize_image_if_needed(self, img: Image.Image) -> Image.Image: + """ + 如果图片过大，缩放到1K分辨率 + + Args: + img: PIL Image 对象 + + Returns: + 缩放后的 PIL Image + """ + width, height = img.size + max_dim = max(width, height) + + if max_dim > MAX_IMAGE_DIMENSION: + # 计算缩放比例 + scale = MAX_IMAGE_DIMENSION / max_dim + new_width = int(width * scale) + new_height = int(height * scale) + + print(f"Google Gemini: 图片尺寸 {width}x{height} 超过限制，缩放至 {new_width}x{new_height}") + img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) + + return img + + def _check_and_compress_image(self, img: Image.Image) -> str: + """ + 检查图片大小，如果超过20MB则进行压缩 + + Args: + img: PIL Image 对象 + + Returns: + base64 编码的字符串 + """ + # 先进行尺寸缩放（如果需要） + img = self._resize_image_if_needed(img) + + # 尝试不同的压缩质量 + qualities = [95, 85, 75, 65, 55, 45] + + for quality in qualities: + buffer = BytesIO() + # 转换为RGB模式（去除alpha通道）以减小体积 + if img.mode in ('RGBA', 'P'): + img_rgb = img.convert('RGB') + else: + img_rgb = img + + img_rgb.save(buffer, format='JPEG', quality=quality, optimize=True) + buffer.seek(0) + data = buffer.getvalue() + + if len(data) <= MAX_FILE_SIZE: + print(f"Google Gemini: 图片压缩后大小 {len(data) / 1024 / 1024:.2f}MB (质量{quality})") + return base64.b64encode(data).decode('utf-8') + + # 如果所有质量都无法满足，使用最低质量 + print(f"Google Gemini: 警告 - 即使最低质量仍超过20MB，将使用最低质量发送") + return base64.b64encode(data).decode('utf-8') + + def _prepare_image_data( + self, + images: Optional[torch.Tensor] + ) -> Optional[List[Dict[str, str]]]: + """ + 准备图片数据 + + 如果图片超过20MB，会自动进行缩放和压缩 + + Args: + images: ComfyUI 图片张量 [B, H, W, C] + + Returns: + 图片数据列表，每个元素包含 mime_type 和 data + """ + if images is None: + return None + + pil_images = tensor_to_pil(images) + if not pil_images: + return None + + # 将所有图片转为 RGB PIL Image 并首次编码 + processed = [] # [(pil_img_rgb, b64_data, mime_type)] + for img in pil_images: + buffer = BytesIO() + img.save(buffer, format='PNG') + original_size = buffer.tell() + buffer.close() + + if original_size > MAX_FILE_SIZE: + print(f"Google Gemini: 检测到图片过大 ({original_size / 1024 / 1024:.2f}MB)，正在进行压缩...") + img_rgb = img.convert('RGB') if img.mode != 'RGB' else img.copy() + b64_str = self._check_and_compress_image(img_rgb) + processed.append((img_rgb, b64_str, "image/jpeg")) + else: + b64_str = encode_image_to_base64(img) + processed.append((None, b64_str, "image/png")) + + # 多图总体积控制 + def calc_total_bytes(): + return sum(len(base64.b64decode(item[1])) for item in processed) + + total = calc_total_bytes() + if total > MAX_FILE_SIZE and len(processed) > 1: + print(f"Google Gemini: 图片总体积 {total / 1024 / 1024:.2f}MB 超过 {MAX_FILE_SIZE // 1024 // 1024}MB 限制，正在压缩...") + + # 降质量 + for quality in range(70, 19, -10): + new_processed = [] + for pil_img, _, _ in processed: + if pil_img is None: + # PNG 原图需要转 RGB + continue + buf = BytesIO() + pil_img.save(buf, format='JPEG', quality=quality, optimize=True) + data = buf.getvalue() + new_processed.append((pil_img, base64.b64encode(data).decode('utf-8'), "image/jpeg")) + if not new_processed: + break + processed = new_processed + total = calc_total_bytes() + if total <= MAX_FILE_SIZE: + print(f"Google Gemini: 图片压缩完成，总体积 {total / 1024 / 1024:.2f}MB ({len(processed)}张图片，质量{quality})") + break + + # 降分辨率 + if total > MAX_FILE_SIZE: + for scale in [0.75, 0.5, 0.35]: + new_processed = [] + for pil_img, _, _ in processed: + if pil_img is None: + continue + w, h = pil_img.size + resized = pil_img.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS) + buf = BytesIO() + resized.save(buf, format='JPEG', quality=20, optimize=True) + data = buf.getvalue() + new_processed.append((resized, base64.b64encode(data).decode('utf-8'), "image/jpeg")) + if not new_processed: + break + processed = new_processed + total = calc_total_bytes() + if total <= MAX_FILE_SIZE: + print(f"Google Gemini: 图片压缩完成，总体积 {total / 1024 / 1024:.2f}MB ({len(processed)}张图片，缩放{int(scale*100)}%)") + break + + if total > MAX_FILE_SIZE: + print(f"Google Gemini: 无法将 {len(processed)} 张图片压缩到 {MAX_FILE_SIZE // 1024 // 1024}MB 以内，请减少图片数量或降低分辨率") + raise ValueError(f"图片总体积 {total / 1024 / 1024:.2f}MB 超过限制，无法压缩到 {MAX_FILE_SIZE // 1024 // 1024}MB 以内") + + image_data = [{"mime_type": mt, "data": b64} for _, b64, mt in processed] + return image_data + + def _compress_video_with_ffmpeg(self, input_path: str, output_path: str, target_size: int) -> bool: + """ + 使用 FFmpeg 压缩视频到目标大小 + + Args: + input_path: 输入视频路径 + output_path: 输出视频路径 + target_size: 目标文件大小（字节） + + Returns: + 是否压缩成功 + """ + try: + # 获取视频时长（秒） + probe_cmd = ['ffprobe', '-v', 'error', '-show_entries', 'format=duration', + '-of', 'default=noprint_wrappers=1:nokey=1', input_path] + duration = float(subprocess.check_output(probe_cmd).decode().strip()) + + # 计算目标比特率（bit/s），预留一些余量 + target_bitrate = int((target_size * 8) / duration * 0.9) + + # 使用 FFmpeg 压缩视频 + # -c:v libx264: 使用 H.264 编码器 + # -b:v: 视频比特率 + # -maxrate 和 -bufsize: 控制码率波动 + # -c:a aac: 音频使用 AAC 编码 + # -b:a 128k: 音频比特率 128k + # -movflags +faststart: 优化网络播放 + cmd = [ + 'ffmpeg', '-y', '-i', input_path, + '-c:v', 'libx264', + '-b:v', f'{target_bitrate}', + '-maxrate', f'{int(target_bitrate * 1.5)}', + '-bufsize', f'{target_bitrate * 2}', + '-c:a', 'aac', + '-b:a', '128k', + '-movflags', '+faststart', + '-preset', 'fast', + output_path + ] + + print(f"Google Gemini: 正在压缩视频到 {target_size / 1024 / 1024:.1f}MB...") + result = subprocess.run(cmd, capture_output=True, text=True) + + if result.returncode == 0 and os.path.exists(output_path): + final_size = os.path.getsize(output_path) + print(f"Google Gemini: 视频压缩完成，最终大小 {final_size / 1024 / 1024:.2f}MB") + return True + else: + print(f"Google Gemini: FFmpeg 压缩失败: {result.stderr}") + return False + + except Exception as e: + print(f"Google Gemini: 视频压缩异常: {str(e)}") + return False + + def _compress_video_with_opencv(self, input_path: str, output_path: str, scale: float = 0.5) -> bool: + """ + 使用 OpenCV 压缩视频（备用方案） + + Args: + input_path: 输入视频路径 + output_path: 输出视频路径 + scale: 尺寸缩放比例 + + Returns: + 是否压缩成功 + """ + if not CV2_AVAILABLE: + return False + + try: + cap = cv2.VideoCapture(input_path) + if not cap.isOpened(): + return False + + # 获取原视频参数 + fps = cap.get(cv2.CAP_PROP_FPS) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + # 计算新尺寸 + new_width = int(width * scale) + new_height = int(height * scale) + + # 创建视频写入器 + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter(output_path, fourcc, fps, (new_width, new_height)) + + print(f"Google Gemini: 使用 OpenCV 压缩视频，分辨率 {width}x{height} -> {new_width}x{new_height}") + + while True: + ret, frame = cap.read() + if not ret: + break + + # 缩放帧 + resized = cv2.resize(frame, (new_width, new_height)) + out.write(resized) + + cap.release() + out.release() + + if os.path.exists(output_path): + final_size = os.path.getsize(output_path) + print(f"Google Gemini: 视频压缩完成，最终大小 {final_size / 1024 / 1024:.2f}MB") + return True + return False + + except Exception as e: + print(f"Google Gemini: OpenCV 压缩失败: {str(e)}") + return False + + def _compress_video(self, video_path: str) -> str: + """ + 压缩视频到 1-10MB 之间 + + Args: + video_path: 原视频路径 + + Returns: + 压缩后的视频路径（临时文件） + """ + original_size = os.path.getsize(video_path) + print(f"Google Gemini: 视频文件过大 ({original_size / 1024 / 1024:.2f}MB)，正在压缩...") + + # 创建临时文件 + temp_dir = tempfile.gettempdir() + _, ext = os.path.splitext(video_path) + output_path = os.path.join(temp_dir, f"compressed_{int(time.time())}{ext}") + + # 确定目标大小（优先尝试 10MB，如果不行再降低） + target_sizes = [ + TARGET_VIDEO_SIZE_MAX, # 10MB + int(TARGET_VIDEO_SIZE_MAX * 0.8), # 8MB + int(TARGET_VIDEO_SIZE_MAX * 0.6), # 6MB + int(TARGET_VIDEO_SIZE_MAX * 0.5), # 5MB + TARGET_VIDEO_SIZE_MIN * 5, # 5MB + TARGET_VIDEO_SIZE_MIN * 3, # 3MB + TARGET_VIDEO_SIZE_MIN * 2, # 2MB + ] + + # 优先尝试 FFmpeg + if FFMPEG_AVAILABLE: + for target_size in target_sizes: + if self._compress_video_with_ffmpeg(video_path, output_path, target_size): + # 检查最终大小 + final_size = os.path.getsize(output_path) + if TARGET_VIDEO_SIZE_MIN <= final_size <= MAX_FILE_SIZE: + return output_path + # 如果仍然太大，继续降低目标 + os.remove(output_path) + + # FFmpeg 失败或不可用，尝试 OpenCV + if CV2_AVAILABLE: + scales = [0.7, 0.5, 0.4, 0.3, 0.25] + for scale in scales: + if self._compress_video_with_opencv(video_path, output_path, scale): + final_size = os.path.getsize(output_path) + if final_size <= MAX_FILE_SIZE: + return output_path + # 如果仍然太大，继续降低分辨率 + os.remove(output_path) + + # 所有压缩方法都失败 + raise ValueError( + f"视频文件过大 ({original_size / 1024 / 1024:.2f}MB) 且无法压缩到 20MB 以下。" + f"请安装 FFmpeg 以获得更好的压缩效果，或手动压缩视频。" + ) + + def _prepare_video_data( + self, + video + ) -> Optional[Dict[str, str]]: + """ + 准备视频数据 + + ComfyUI VIDEO 类型包含视频文件路径信息。 + 读取视频文件并转换为 base64。 + 如果视频超过 20MB，会自动进行压缩。 + + Args: + video: ComfyUI VIDEO 类型数据 + + Returns: + 视频数据字典，包含 mime_type 和 data + """ + if video is None: + return None + + # VIDEO 类型处理：支持多种格式 + video_path = None + temp_compressed_path = None + + if isinstance(video, dict): + # 字典格式：尝试常见的键名 + video_path = video.get("video") or video.get("path") or video.get("file") or video.get("filename") + # 如果还是找不到，遍历所有键找到有效路径 + if not video_path: + for key, val in video.items(): + if isinstance(val, str) and os.path.exists(val): + video_path = val + break + elif isinstance(video, str): + # 字符串格式：直接作为路径 + video_path = video + else: + # 对象格式：尝试常见属性 + # 1. 尝试 __file 属性（VideoFromFile 对象） + if hasattr(video, "__file"): + video_path = video.__file + # 2. 尝试其他常见属性 + elif hasattr(video, "video"): + video_path = video.video + elif hasattr(video, "path"): + video_path = video.path + elif hasattr(video, "filename"): + video_path = video.filename + # 3. 尝试从 __dict__ 中查找路径（支持私有属性如 _VideoFromFile__file） + elif hasattr(video, "__dict__"): + for attr_name, attr_value in video.__dict__.items(): + # 查找字符串类型的属性，且包含 file 或 path 关键字 + if isinstance(attr_value, str): + if "file" in attr_name.lower() or "path" in attr_name.lower(): + # 验证路径是否有效 + if os.path.exists(attr_value): + video_path = attr_value + break + # 如果属性值本身看起来像文件路径，也尝试使用 + elif os.path.exists(attr_value) and os.path.isfile(attr_value): + video_path = attr_value + break + + if not video_path or not os.path.exists(video_path): + print(f"Google Gemini: 视频文件不存在或路径无效: {video_path}") + return None + + # 获取文件扩展名和 MIME 类型 + _, ext = os.path.splitext(video_path) + ext = ext.lower() + + mime_type = VIDEO_MIME_TYPES.get(ext, "video/mp4") + + try: + # 检查文件大小 + file_size = os.path.getsize(video_path) + + # 如果超过 20MB，进行压缩 + if file_size > MAX_FILE_SIZE: + video_path = self._compress_video(video_path) + temp_compressed_path = video_path + # 压缩后统一使用 mp4 格式 + mime_type = "video/mp4" + + # 读取并编码视频 + with open(video_path, "rb") as f: + video_bytes = f.read() + + b64_str = base64.b64encode(video_bytes).decode("utf-8") + + # 清理临时文件 + if temp_compressed_path and os.path.exists(temp_compressed_path): + try: + os.remove(temp_compressed_path) + print(f"Google Gemini: 临时压缩文件已清理") + except: + pass + + return { + "mime_type": mime_type, + "data": b64_str + } + + except Exception as e: + # 清理临时文件 + if temp_compressed_path and os.path.exists(temp_compressed_path): + try: + os.remove(temp_compressed_path) + except: + pass + + print(f"Google Gemini: 处理视频文件失败 - {str(e)}") + return None + + def _prepare_file_data( + self, + file: Optional[FileData] + ) -> Optional[Dict[str, str]]: + """ + 准备文件数据 + + 从 FILE 类型提取文件数据 + + Args: + file: FileData 对象（来自 LoadFile 节点） + + Returns: + 文件数据字典，包含 mime_type 和 data + """ + if file is None: + return None + + return { + "mime_type": file.mime_type, + "data": file.data + } + + def _parse_dual_output(self, raw_response: Dict) -> Tuple[str, str]: + """ + 解析包含思考内容和主要内容的响应 + + Args: + raw_response: API 原始响应字典 + + Returns: + (主要内容, 思考内容) + """ + candidates = raw_response.get("candidates", []) + if not candidates: + return ("", "") + + parts = candidates[0].get("content", {}).get("parts", []) + + thought_text = "" + main_text = "" + + for part in parts: + if part.get("thought") is True: + # 思考部分 + thought_text = part.get("text", "") + elif "thoughtSignature" in part or "text" in part: + # 主要内容 + main_text = part.get("text", "") + + return main_text + + def generate( + self, + 模型: str, + 提示词: str, + 思考等级: str, + 图片: Optional[torch.Tensor] = None, + 视频=None, + 文件: Optional[FileData] = None + ) -> Tuple[str]: + """ + 生成文本 + + Args: + 模型: 使用的模型名称 + 提示词: 用户提示词 + 思考等级: 思考等级选项 + 图片: 输入图片 + 视频: 输入视频 + 文件: 输入文件（PDF/TXT） + + Returns: + (主要内容, 思考内容) + """ + start_time = time.time() + + try: + # 初始化 API 客户端 + if self.client is None: + try: + self.client = GeminiFlashClient() + except ValueError as e: + raise ValueError(f"初始化失败: {str(e)}") + + # 准备图片数据 + image_data = self._prepare_image_data(图片) + if image_data: + print(f"Google Gemini: 输入 {len(image_data)} 张图片") + + # 准备视频数据 + video_data = self._prepare_video_data(视频) + if video_data: + print(f"Google Gemini: 输入视频 ({video_data['mime_type']})") + + # 准备文件数据 + document_data = self._prepare_file_data(文件) + if document_data: + file_type = "PDF" if document_data['mime_type'] == "application/pdf" else "TXT" + print(f"Google Gemini: 输入文件 ({file_type})") + + # 构建输入描述 + input_desc = [] + if 提示词: + input_desc.append("文本") + if image_data: + input_desc.append(f"{len(image_data)}张图片") + if video_data: + input_desc.append("视频") + if document_data: + input_desc.append("文件") + + print(f"Google Gemini: 模型 = {模型}") + print(f"Google Gemini: 多模态输入 ({', '.join(input_desc)})") + print(f"Google Gemini: 思考等级 = {思考等级}") + + # 获取端点和构建请求体 + endpoint = self.client.get_endpoint(model=模型) + request_body = self.client.build_request_body( + prompt=提示词, + model=模型, + thinking_level=思考等级, + image_data=image_data, + video_data=video_data, + document_data=document_data + ) + + print(f"Google Gemini: 发送请求...") + + # 调用底层 API 获取原始响应 + async def get_raw_response(): + return await self.client.request_async( + endpoint, + request_body, + session=None + ) + + # 在独立线程中执行异步请求 + raw_response = self.client.run_async_in_thread(get_raw_response()) + + # 计算耗时 + elapsed = time.time() - start_time + + # 解析响应，分离主要内容和思考内容 + main_text = self._parse_dual_output(raw_response) + + # 打印响应 token 用量 + usage = raw_response.get("usageMetadata", {}) + prompt_tokens = usage.get("promptTokenCount", 0) + candidates_tokens = usage.get("candidatesTokenCount", 0) + thoughts_tokens = usage.get("thoughtsTokenCount", 0) + total_tokens = usage.get("totalTokenCount", 0) + finish_reason = "" + candidates = raw_response.get("candidates", []) + if candidates: + finish_reason = candidates[0].get("finishReason", "") + + print(f"Google Gemini: 生成完成 (耗时: {elapsed:.2f}s)") + print(f"Google Gemini: finishReason = {finish_reason}") + print(f"Google Gemini: Token 用量 — 输入: {prompt_tokens}, 输出: {candidates_tokens}, 思考: {thoughts_tokens}, 合计: {total_tokens}") + print(f"Google Gemini: 主要内容长度: {len(main_text)} 字符") + + # 输出预览 + if main_text: + preview = main_text[:100] + "..." if len(main_text) > 100 else main_text + print(f"Google Gemini: 主要内容预览: {preview}") + + return (main_text,) + if str(e) == "未授权！": + print("请联系作者授权后方可使用！") + raise ValueError("未授权！") from None + else: + # 用户输入错误 - 只显示简洁信息 + error_msg = str(e).split('\n')[0] # 只取第一行 + print(f"Google Gemini: ❌ {error_msg}") + raise ValueError(error_msg) from None + + except RuntimeError as e: + # 日志只打第一行；报错框展示完整多行 + error_full = str(e) + print(f"Google Gemini: ❌ {error_full.split('\n')[0]}") + raise RuntimeError(error_full) from None + + except Exception as e: + # 其他未知错误 - 只显示简洁信息 + error_msg = str(e).split('\n')[0] + print(f"Google Gemini: ❌ {error_msg}") + raise type(e)(error_msg) from None + + finally: + if self.client is not None: + try: + balance_data = self.client.query_balance_sync() + balance_info = self.client.format_balance_info(balance_data) + print(f"Google Gemini: {balance_info}") + except Exception: + pass diff --git a/nodes/image_stitch_pro.py b/nodes/image_stitch_pro.py new file mode 100644 index 0000000..018c568 --- /dev/null +++ b/nodes/image_stitch_pro.py @@ -0,0 +1,241 @@ +""" +高级图像拼接节点 +支持最多 10 张图像按指定方向（上、下、左、右）依次拼接， +支持调整图像大小匹配和添加间隔。 +""" + +from typing import Optional, Tuple, List + +import torch +from PIL import Image + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor +from ..utils.file_utils import load_images_from_folder + + +# 间隔颜色映射 +SPACING_COLOR_MAP = { + "white": (255, 255, 255), + "black": (0, 0, 0), + "red": (255, 0, 0), + "green": (0, 255, 0), + "blue": (0, 0, 255), +} + + +def _resize_to_match(img: Image.Image, ref: Image.Image, direction: str) -> Image.Image: + """ + 按拼接方向将 img 缩放，使其与 ref 在垂直于拼接轴的尺寸上一致。 + + - 水平拼接 (right/left)：统一高度 + - 垂直拼接 (down/up)：统一宽度 + """ + ref_w, ref_h = ref.size + img_w, img_h = img.size + + if direction in ("right", "left"): + if img_h != ref_h: + scale = ref_h / img_h + new_w = max(1, int(img_w * scale)) + img = img.resize((new_w, ref_h), Image.LANCZOS) + else: + if img_w != ref_w: + scale = ref_w / img_w + new_h = max(1, int(img_h * scale)) + img = img.resize((ref_w, new_h), Image.LANCZOS) + + return img + + +def _make_spacer(ref: Image.Image, spacing_width: int, + direction: str, color: Tuple[int, int, int]) -> Image.Image: + """创建间隔色块""" + if direction in ("right", "left"): + return Image.new("RGB", (spacing_width, ref.size[1]), color) + else: + return Image.new("RGB", (ref.size[0], spacing_width), color) + + +def _stitch_two(img_a: Image.Image, img_b: Image.Image, + direction: str, match_size: bool, + spacing_width: int, spacing_color: Tuple[int, int, int]) -> Image.Image: + """ + 将两张 PIL 图像按指定方向拼接。 + img_a 为基准图像，img_b 拼接在 img_a 的指定方向侧。 + direction="right" → img_b 在 img_a 右侧 + direction="left" → img_b 在 img_a 左侧 + direction="down" → img_b 在 img_a 下方 + direction="up" → img_b 在 img_a 上方 + """ + if img_a.mode != "RGB": + img_a = img_a.convert("RGB") + if img_b.mode != "RGB": + img_b = img_b.convert("RGB") + + if match_size: + img_b = _resize_to_match(img_b, img_a, direction) + + if direction == "right": + pieces = [img_a, img_b] + elif direction == "left": + pieces = [img_b, img_a] + elif direction == "down": + pieces = [img_a, img_b] + else: # up + pieces = [img_b, img_a] + + if spacing_width > 0: + interleaved: List[Image.Image] = [] + for idx, piece in enumerate(pieces): + interleaved.append(piece) + if idx < len(pieces) - 1: + interleaved.append(_make_spacer(piece, spacing_width, direction, spacing_color)) + pieces = interleaved + + if direction in ("right", "left"): + total_w = sum(p.size[0] for p in pieces) + max_h = max(p.size[1] for p in pieces) + canvas = Image.new("RGB", (total_w, max_h), spacing_color) + x = 0 + for piece in pieces: + canvas.paste(piece, (x, 0)) + x += piece.size[0] + else: + max_w = max(p.size[0] for p in pieces) + total_h = sum(p.size[1] for p in pieces) + canvas = Image.new("RGB", (max_w, total_h), spacing_color) + y = 0 + for piece in pieces: + canvas.paste(piece, (0, y)) + y += piece.size[1] + + return canvas + + +def _natural_sort_key(filename: str): + """按数字优先的文件名排序，使 1, 2, 3, 10 而非 1, 10, 2, 3""" + try: + return (0, int(filename)) + except ValueError: + return (1, filename.lower()) + + +class ImageStitchPro: + """ + 高级图像拼接节点 + + 在 ComfyUI 原生拼接节点基础上扩展，支持同时输入最多 10 张图像， + 按指定方向依次拼接，并可在图像间添加任意颜色的间隔。 + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "方向": (["right", "down", "left", "up"], {"default": "down"}), + "匹配图像尺寸": ("BOOLEAN", {"default": True}), + "间距宽度": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 2}), + "间距颜色": (["white", "black", "red", "green", "blue"], {"default": "white"}), + }, + "optional": { + "图1": ("IMAGE",), + "图2": ("IMAGE",), + "图3": ("IMAGE",), + "图4": ("IMAGE",), + "图5": ("IMAGE",), + "图6": ("IMAGE",), + "图7": ("IMAGE",), + "图8": ("IMAGE",), + "图9": ("IMAGE",), + "图10": ("IMAGE",), + "图11": ("IMAGE",), + "图12": ("IMAGE",), + "图片路径（可选）": ("STRING", {"default": "", "multiline": False}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("拼接图像",) + FUNCTION = "stitch" + CATEGORY = "image" + + DESCRIPTION = ( + "高级图像拼接节点，支持最多 12 张图像按指定方向（右/下/左/上）依次拼接。\n" + "可选择是否将后续图像缩放以匹配第一张图像的尺寸，并可在图像间添加彩色间隔。\n" + "可选填「图片路径」：仅处理该文件夹内图片，按文件名顺序依次拼接；与输入端图片不可同时使用。" + ) + + def stitch( + self, + 方向: str = "down", + 匹配图像尺寸: bool = True, + 间距宽度: int = 0, + 间距颜色: str = "white", + 图1: Optional[torch.Tensor] = None, + 图2: Optional[torch.Tensor] = None, + 图3: Optional[torch.Tensor] = None, + 图4: Optional[torch.Tensor] = None, + 图5: Optional[torch.Tensor] = None, + 图6: Optional[torch.Tensor] = None, + 图7: Optional[torch.Tensor] = None, + 图8: Optional[torch.Tensor] = None, + 图9: Optional[torch.Tensor] = None, + 图10: Optional[torch.Tensor] = None, + 图11: Optional[torch.Tensor] = None, + 图12: Optional[torch.Tensor] = None, + **kwargs: object, + ) -> Tuple[torch.Tensor]: + + color = SPACING_COLOR_MAP.get(间距颜色, (255, 255, 255)) + raw_tensors = [图1, 图2, 图3, 图4, 图5, 图6, 图7, 图8, 图9, 图10, 图11, 图12] + tensors = [t for t in raw_tensors if t is not None] + has_input_images = len(tensors) > 0 + image_folder = (kwargs.get("图片路径（可选）") or "").strip() + + if image_folder and has_input_images: + raise ValueError("不可同时使用「图片路径（可选）」与输入端图片，请二选一。") + + if image_folder: + infos = load_images_from_folder(image_folder) + if not infos: + raise ValueError(f"文件夹中未找到可用的图片，或路径无效: {image_folder}") + infos.sort(key=lambda x: _natural_sort_key(x.filename)) + pil_list = [info.image for info in infos] + if len(pil_list) == 1: + return (pil_to_tensor(pil_list),) + base = pil_list[0] + for next_img in pil_list[1:]: + base = _stitch_two( + base, next_img, + direction=方向, + match_size=匹配图像尺寸, + spacing_width=间距宽度, + spacing_color=color, + ) + return (pil_to_tensor([base]),) + else: + if not has_input_images: + raise ValueError("请至少接入一张图片，或填写「图片路径（可选）」中的文件夹路径。") + + if len(tensors) == 1: + return (tensors[0],) + + pil_batches: List[List[Image.Image]] = [tensor_to_pil(t) for t in tensors] + + batch_size = min(len(b) for b in pil_batches) + result_images: List[Image.Image] = [] + + for i in range(batch_size): + frames = [batch[i] for batch in pil_batches] + base = frames[0] + for next_img in frames[1:]: + base = _stitch_two( + base, next_img, + direction=方向, + match_size=匹配图像尺寸, + spacing_width=间距宽度, + spacing_color=color, + ) + result_images.append(base) + + return (pil_to_tensor(result_images),) diff --git a/nodes/kling_video.py b/nodes/kling_video.py new file mode 100644 index 0000000..9e1d6f4 --- /dev/null +++ b/nodes/kling_video.py @@ -0,0 +1,736 @@ +""" +Kling 3.0 Video Nodes +""" + +import os +import re + +from ..clients.kling_client import KlingClient +from ..clients.gemini_client import GeminiAPIClient +from ..utils.image_utils import tensor_to_pil, encode_image_to_base64 + +from comfy_api.latest import InputImpl + +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + + +def _get_video_output_dir() -> str: + if FOLDER_PATHS_AVAILABLE: + base = folder_paths.get_output_directory() + else: + plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + base = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output") + video_dir = os.path.join(base, "video") + os.makedirs(video_dir, exist_ok=True) + return video_dir + + +def _get_next_counter(directory: str, prefix: str) -> int: + if not os.path.exists(directory): + return 1 + pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)") + max_counter = 0 + for f in os.listdir(directory): + m = pattern.match(f) + if m: + max_counter = max(max_counter, int(m.group(1))) + return max_counter + 1 + + +def _tensor_to_base64(tensor) -> str: + """ComfyUI IMAGE tensor → base64 PNG 字符串""" + pil_images = tensor_to_pil(tensor) + return encode_image_to_base64(pil_images[0], format="PNG") + + +def _validate_prompt(prompt: str, *, required: bool = True) -> None: + """校验单条提示词。 + + Args: + prompt: 提示词字符串。 + required: 为 True 时不允许为空（多镜头关闭或 shot_type 为 intelligence 时适用）。 + """ + if required and not prompt.strip(): + raise ValueError("提示词不能为空（非多镜头模式下必填）。") + if len(prompt) > 2500: + raise ValueError( + f"提示词长度 ({len(prompt)}) 超过上限 2500 个字符，请缩短后重试。" + ) + + +def _validate_multi_prompt(multi_prompt_list: list, total_duration: int) -> None: + """校验多镜头分镜列表。 + + 规则： + - 分镜数量：1 ~ 6； + - 每个分镜提示词不超过 512 个字符； + - 每个分镜时长 ≥ 1 且 ≤ total_duration； + - 所有分镜时长之和必须等于 total_duration。 + """ + count = len(multi_prompt_list) + if count < 1 or count > 6: + raise ValueError( + f"多镜头分镜数量须在 1~6 之间，当前为 {count}。" + ) + + duration_sum = 0 + for entry in multi_prompt_list: + idx = entry["index"] + p = entry.get("prompt", "") + dur = entry.get("duration", 0) + + if len(p) > 512: + raise ValueError( + f"镜头 {idx} 提示词长度 ({len(p)}) 超过上限 512 个字符。" + ) + if dur < 1: + raise ValueError( + f"镜头 {idx} 时长 ({dur}s) 不能小于 1 秒。" + ) + if dur > total_duration: + raise ValueError( + f"镜头 {idx} 时长 ({dur}s) 超过任务总时长 ({total_duration}s)。" + ) + duration_sum += dur + + if duration_sum != total_duration: + raise ValueError( + f"所有分镜时长之和 ({duration_sum}s) 必须等于任务总时长 ({total_duration}s)。" + ) + + +def _validate_image(tensor, label: str = "图片") -> None: + """校验图片张量。 + + 规则： + - 文件大小（PNG）不超过 10MB； + - 宽、高均不小于 300px； + - 宽高比介于 1:2.5 ~ 2.5:1 之间（即 ratio ∈ [0.4, 2.5]）。 + """ + import io + + pil_images = tensor_to_pil(tensor) + img = pil_images[0] + w, h = img.size + + # ── 最小尺寸 ────────────────────────────────────────────────────── + if w < 300 or h < 300: + raise ValueError( + f"{label} 宽高不得小于 300px，当前为 {w}×{h}px。" + ) + + # ── 宽高比 ──────────────────────────────────────────────────────── + ratio = w / h + if ratio < 1 / 2.5 or ratio > 2.5: + raise ValueError( + f"{label} 宽高比须在 1:2.5 ~ 2.5:1 之间，" + f"当前为 {w}:{h}（比值 {ratio:.2f}）。" + ) + + # ── 文件大小 ────────────────────────────────────────────────────── + buf = io.BytesIO() + img.save(buf, format="PNG") + size_mb = buf.tell() / (1024 * 1024) + if size_mb > 10: + raise ValueError( + f"{label} PNG 大小 ({size_mb:.1f}MB) 超过上限 10MB。" + ) + + +class KlingVideo: + """Kling 3.0 视频生成节点（支持多镜头）""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "提示词": ("STRING", {"multiline": True, "default": ""}), + "反向提示词": ("STRING", {"multiline": True, "default": ""}), + "时长": ([5, 10, 15],), + "分辨率": (["1080p", "720p"],), + "宽高比": (["智能", "16:9", "9:16", "1:1"], {"default": "智能"}), + "生成音频": (["打开", "关闭"], {"default": "打开"}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}), + }, + "optional": { + "起始帧": ("IMAGE",), + "镜头1_提示词": ("STRING", {"multiline": True, "default": ""}), + "镜头1_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}), + "镜头2_提示词": ("STRING", {"multiline": True, "default": ""}), + "镜头2_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}), + "镜头3_提示词": ("STRING", {"multiline": True, "default": ""}), + "镜头3_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}), + "镜头4_提示词": ("STRING", {"multiline": True, "default": ""}), + "镜头4_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}), + "镜头5_提示词": ("STRING", {"multiline": True, "default": ""}), + "镜头5_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}), + "镜头6_提示词": ("STRING", {"multiline": True, "default": ""}), + "镜头6_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}), + } + } + + RETURN_TYPES = ("VIDEO",) + RETURN_NAMES = ("视频",) + FUNCTION = "generate" + CATEGORY = "comfyui_o1key/Kling" + + async def generate(self, **kwargs): + """生成视频（支持多镜头）""" + prompt = kwargs["提示词"] + negative_prompt = kwargs["反向提示词"] + duration = kwargs["时长"] + resolution = kwargs["分辨率"] + aspect_ratio = kwargs["宽高比"] + generate_audio = kwargs["生成音频"] + start_frame = kwargs.get("起始帧", None) + seed = kwargs.get("seed", 0) # noqa: F841 — 触发 ComfyUI 缓存刷新 + + mode = "pro" if resolution == "1080p" else "std" + voice = "voice" if generate_audio == "打开" else "novoice" + + # ── 多镜头检测 ──────────────────────────────────────────────── + multi_prompt_list = [] + for i in range(1, 7): + sb_prompt = kwargs.get(f"镜头{i}_提示词", "").strip() + if sb_prompt: + sb_duration = kwargs.get(f"镜头{i}_时长", 5) + multi_prompt_list.append({ + "index": i, + "prompt": sb_prompt, + "duration": sb_duration, + }) + + multi_shot_enabled = len(multi_prompt_list) > 0 + + if multi_shot_enabled: + total_duration = sum(e["duration"] for e in multi_prompt_list) + if total_duration < 3 or total_duration > 15: + raise ValueError( + f"多镜头总时长 ({total_duration}s) 必须在 3~15 秒之间。" + ) + _validate_multi_prompt(multi_prompt_list, total_duration) + duration = total_duration + else: + _validate_prompt(prompt, required=True) + + # ── 构建模型名 & 请求体 ─────────────────────────────────────── + import json, base64, copy + model_name = f"kling-v3-{mode}-{duration}s-{voice}" + + body = { + "model": model_name, + "mode": mode, + "duration": duration, + } + + sound = "on" if generate_audio == "打开" else "off" + + if multi_shot_enabled or sound == "on": + ms_payload = {} + ms_payload["prompt"] = prompt + + if sound == "on": + ms_payload["sound"] = "on" + + if multi_shot_enabled: + ms_payload["multi_shot"] = True + ms_payload["shot_type"] = "customize" + ms_payload["multi_prompt"] = multi_prompt_list + + encoded = base64.b64encode( + json.dumps(ms_payload, ensure_ascii=False).encode("utf-8") + ).decode("utf-8") + body["prompt"] = f"__MS__:{encoded}" + else: + body["prompt"] = prompt + + if negative_prompt.strip(): + body["negative_prompt"] = negative_prompt + + if start_frame is not None: + _validate_image(start_frame, "起始帧") + body["image"] = _tensor_to_base64(start_frame) + endpoint_type = "image2video" + else: + body["metadata"] = {"aspect_ratio": aspect_ratio} + endpoint_type = "text2video" + + # ── 保存路径 ────────────────────────────────────────────────── + video_dir = _get_video_output_dir() + counter = _get_next_counter(video_dir, "kling") + save_path = os.path.join(video_dir, f"kling_{counter:05d}.mp4") + + client = KlingClient() + + # ── 进度条 ──────────────────────────────────────────────────── + try: + from comfy.utils import ProgressBar + pbar = ProgressBar(100) + except Exception: + pbar = None + + def on_stage(stage: str): + if stage == "submitting": + print("[视频生成] 提交中...") + if pbar: pbar.update_absolute(0, 100) + elif stage.startswith("submitted:"): + print(f"[视频生成] 任务已提交 → {stage.split(':',1)[1]}") + if pbar: pbar.update_absolute(5, 100) + elif stage == "downloading": + print("[视频生成] 下载视频...") + if pbar: pbar.update_absolute(99, 100) + elif stage == "done": + print("[视频生成] 完成") + if pbar: pbar.update_absolute(100, 100) + + def on_progress(pct: int): + mapped = 5 + int(pct * 0.94) + if pbar: pbar.update_absolute(mapped, 100) + + try: + result_path = await client.generate_async( + endpoint_type=endpoint_type, + body=body, + save_path=save_path, + on_stage=on_stage, + on_progress=on_progress, + ) + return (InputImpl.VideoFromFile(result_path),) + finally: + # 查询余额 + try: + _balance_client = GeminiAPIClient() + balance_data = _balance_client.query_balance_sync() + balance_info = _balance_client.format_balance_info(balance_data) + print(f"自研视频模型: {balance_info}") + except Exception: + pass + + +class KlingFirstLastFrame: + """Kling 3.0 首尾帧到视频节点""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "首帧": ("IMAGE",), + "尾帧": ("IMAGE",), + "提示词": ("STRING", {"multiline": True, "default": ""}), + "时长": ([5, 10, 15],), + "生成音频": (["打开", "关闭"], {"default": "打开"}), + "模型": (["v3"],), + "分辨率": (["1080p", "720p"],), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}), + } + } + + RETURN_TYPES = ("VIDEO",) + RETURN_NAMES = ("视频",) + FUNCTION = "generate" + CATEGORY = "comfyui_o1key/Kling" + + async def generate(self, **kwargs): + first_frame = kwargs["首帧"] + end_frame = kwargs["尾帧"] + prompt = kwargs["提示词"] + duration = kwargs["时长"] + generate_audio = kwargs["生成音频"] + model_base = kwargs["模型"] + model_base = "kling-" + model_base # v3 → kling-v3（后端值还原） + resolution = kwargs["分辨率"] + seed = kwargs.get("seed", 0) # noqa: F841 — 触发 ComfyUI 缓存刷新 + + _validate_prompt(prompt, required=True) + + # 时长校验 + if duration not in (5, 10, 15): + raise ValueError(f"时长仅支持 5、10、15 秒，当前值为 {duration}，请重新选择。") + + # 拼接模型名：kling-v3-{mode}-{dur}s-{voice} + mode = "pro" if resolution == "1080p" else "std" + voice = "voice" if generate_audio == "打开" else "novoice" + model_name = f"{model_base}-{mode}-{duration}s-{voice}" + + # 图片校验 & 转 base64 + _validate_image(first_frame, "首帧") + _validate_image(end_frame, "尾帧") + image_b64 = _tensor_to_base64(first_frame) + image_tail_b64 = _tensor_to_base64(end_frame) + + # ── 按规范编码 prompt 和 sound ────────────────────────── + import json, base64 + sound = "on" if generate_audio == "打开" else "off" + + body = { + "model": model_name, + "image": image_b64, + "mode": mode, + "duration": duration, + "metadata": { + "image_tail": image_tail_b64, + }, + } + + if sound == "on": + ms_payload = { + "prompt": prompt, + "sound": "on", + } + encoded = base64.b64encode( + json.dumps(ms_payload, ensure_ascii=False).encode("utf-8") + ).decode("utf-8") + body["prompt"] = f"__MS__:{encoded}" + else: + body["prompt"] = prompt + + # 保存路径 + video_dir = _get_video_output_dir() + counter = _get_next_counter(video_dir, "kling") + save_path = os.path.join(video_dir, f"kling_{counter:05d}.mp4") + + client = KlingClient() + + # 进度条：0~100 步 + try: + from comfy.utils import ProgressBar + pbar = ProgressBar(100) + except Exception: + pbar = None + + def on_stage(stage: str): + if stage == "submitting": + print("[视频生成] 提交中...") + if pbar: + pbar.update_absolute(0, 100) + elif stage.startswith("submitted:"): + print(f"[视频生成] 任务已提交 → {stage.split(':',1)[1]}") + if pbar: + pbar.update_absolute(5, 100) + elif stage == "downloading": + print("[视频生成] 下载视频...") + if pbar: + pbar.update_absolute(99, 100) + elif stage == "done": + print("[视频生成] 完成") + if pbar: + pbar.update_absolute(100, 100) + + def on_progress(pct: int): + # pct 来自 API progress 字段，如 50 表示 50% + # 生成阶段占 5~99 区间 + mapped = 5 + int(pct * 0.94) + if pbar: + pbar.update_absolute(mapped, 100) + + try: + result_path = await client.generate_async( + endpoint_type="image2video", + body=body, + save_path=save_path, + on_stage=on_stage, + on_progress=on_progress, + ) + return (InputImpl.VideoFromFile(result_path),) + finally: + # 查询余额 + try: + _balance_client = GeminiAPIClient() + balance_data = _balance_client.query_balance_sync() + balance_info = _balance_client.format_balance_info(balance_data) + print(f"自研视频模型: {balance_info}") + except Exception: + pass + + +class KlingMotionControlTest: + """Kling 动作控制（测试）节点 —— reference_video 接受 VIDEO 类型输入""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "提示词": ("STRING", {"multiline": True, "default": ""}), + "参考图片": ("IMAGE",), + "参考视频": ("VIDEO",), + }, + "optional": { + "保留原声": ("BOOLEAN", {"default": True}), + "人物朝向": (["video", "image"],), + "画质模式": (["专家", "标准"],), + "模型版本": (["v3"],), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}), + }, + } + + RETURN_TYPES = ("VIDEO",) + RETURN_NAMES = ("视频",) + FUNCTION = "generate" + CATEGORY = "comfyui_o1key/Kling" + + async def generate(self, **kwargs): + """动作控制（测试）：VIDEO 类型参考视频 + 图片人物动作迁移""" + import base64 + + prompt = kwargs["提示词"] + reference_image = kwargs["参考图片"] + reference_video = kwargs["参考视频"] + keep_original_sound = kwargs.get("保留原声", True) + character_orientation = kwargs.get("人物朝向", "video") + mode = kwargs.get("画质模式", "专家") + mode = "pro" if mode == "专家" else "std" # 映射为 API 参数值 + model = kwargs.get("模型版本", "v3") + model = "kling-" + model # v3 → kling-v3（后端值还原） + seed = kwargs.get("seed", 0) # noqa: F841 — 触发 ComfyUI 缓存刷新 + + # ── 校验提示词 ──────────────────────────────────────────────── + _validate_prompt(prompt, required=True) + + # ── 校验参考图片 ────────────────────────────────────────────── + _validate_image(reference_image, "参考图片") + image_b64 = _tensor_to_base64(reference_image) + + # ── 从 VIDEO 对象获取本地文件路径并读取 ─────────────────────── + # ComfyUI VIDEO 对象有 .source_path 或通过 VideoFromFile 构造 + video_path = None + if hasattr(reference_video, "source_path"): + video_path = reference_video.source_path + elif hasattr(reference_video, "path"): + video_path = reference_video.path + elif isinstance(reference_video, str): + video_path = reference_video.strip() + + if not video_path or not os.path.isfile(video_path): + raise ValueError( + f"无法获取参考视频文件路径，请确保连接的是本地视频文件。" + f"（当前路径：{video_path}）" + ) + + # ── 校验视频时长约束 ────────────────────────────────────────── + # 人物朝向="video" → 3~30 秒；人物朝向="image" → 3~10 秒 + try: + import subprocess, json as _json + ffprobe_cmd = [ + "ffprobe", "-v", "quiet", + "-print_format", "json", + "-show_format", + video_path, + ] + result_proc = subprocess.run(ffprobe_cmd, capture_output=True, text=True, timeout=30) + if result_proc.returncode == 0: + info = _json.loads(result_proc.stdout) + duration_sec = float(info.get("format", {}).get("duration", 0)) + if character_orientation == "video": + if not (3 <= duration_sec <= 30): + raise ValueError( + f"当人物朝向为 'video' 时，" + f"参考视频时长须在 3~30 秒之间，当前为 {duration_sec:.1f}s。" + ) + else: # "image" + if not (3 <= duration_sec <= 10): + raise ValueError( + f"当人物朝向为 'image' 时，" + f"参考视频时长须在 3~10 秒之间，当前为 {duration_sec:.1f}s。" + ) + except FileNotFoundError: + # ffprobe 不可用时跳过时长校验，但打印提示 + print("[动作控制] 警告：ffprobe 未找到，跳过视频时长校验。") + except ValueError: + raise + except Exception as e: + print(f"[动作控制] 时长校验异常（已跳过）：{e}") + + # ── 视频转 base64 ───────────────────────────────────────────── + with open(video_path, "rb") as f: + video_b64 = base64.b64encode(f.read()).decode("utf-8") + + # ── 构建请求体 ──────────────────────────────────────────────── + body = { + "prompt": prompt, + "character_orientation": character_orientation, + "mode": mode, + "model": model, + "keep_original_sound": "yes" if keep_original_sound else "no", + "image": image_b64, + "video": video_b64, + } + + # ── 保存路径 ────────────────────────────────────────────────── + video_dir = _get_video_output_dir() + counter = _get_next_counter(video_dir, "kling_motion_test") + save_path = os.path.join(video_dir, f"kling_motion_test_{counter:05d}.mp4") + + client = KlingClient() + + # ── 进度条 ──────────────────────────────────────────────────── + try: + from comfy.utils import ProgressBar + pbar = ProgressBar(100) + except Exception: + pbar = None + + def on_stage(stage: str): + if stage == "submitting": + print("[动作控制] 提交中...") + if pbar: pbar.update_absolute(0, 100) + elif stage.startswith("submitted:"): + print(f"[动作控制] 任务已提交 → {stage.split(':',1)[1]}") + if pbar: pbar.update_absolute(5, 100) + elif stage == "downloading": + print("[动作控制] 下载视频...") + if pbar: pbar.update_absolute(99, 100) + elif stage == "done": + print("[动作控制] 完成") + if pbar: pbar.update_absolute(100, 100) + + def on_progress(pct: int): + mapped = 5 + int(pct * 0.94) + if pbar: pbar.update_absolute(mapped, 100) + + try: + result_path = await client.generate_async( + endpoint_type="motion_control", + body=body, + save_path=save_path, + on_stage=on_stage, + on_progress=on_progress, + ) + return (InputImpl.VideoFromFile(result_path),) + finally: + # 查询余额 + try: + _balance_client = GeminiAPIClient() + balance_data = _balance_client.query_balance_sync() + balance_info = _balance_client.format_balance_info(balance_data) + print(f"自研视频模型: {balance_info}") + except Exception: + pass + + +class AspectRatioPreset: + """图片宽高比预设节点""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "图像": ("IMAGE",), + "宽高比": (["智能", "16:9", "9:16", "4:3", "3:4", "1:1"], {"default": "智能"}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("图像",) + FUNCTION = "resize" + CATEGORY = "comfyui_o1key/Utils" + + def resize(self, 图像, 宽高比): + import torch + from PIL import Image + import numpy as np + + pil_images = tensor_to_pil(图像) + img = pil_images[0] + w, h = img.size + img_ratio = w / h + + # 确定原图所属的宽高比家族 + ratios = {"16:9": 16/9, "9:16": 9/16, "4:3": 4/3, "3:4": 3/4, "1:1": 1.0} + closest_ratio = min(ratios.keys(), key=lambda k: abs(ratios[k] - img_ratio)) + + # 智能模式：使用最接近的比例 + if 宽高比 == "智能": + 宽高比 = closest_ratio + + # 解析目标比例 + target_w, target_h = map(int, 宽高比.split(":")) + target_ratio = target_w / target_h + + # 确定分辨率级别（1K/2K） + max_dim = max(w, h) + if max_dim <= 1080: + base = 1080 + elif max_dim <= 2160: + base = 2160 + else: + base = 2160 + + # 计算目标尺寸 + if target_ratio >= 1: + target_width = base + target_height = int(base / target_ratio) + else: + target_height = base + target_width = int(base * target_ratio) + + # 判断是否同家族（横向家族：16:9, 4:3；纵向家族：9:16, 3:4；正方形：1:1） + horizontal_family = ["16:9", "4:3"] + vertical_family = ["9:16", "3:4"] + + same_family = False + if closest_ratio in horizontal_family and 宽高比 in horizontal_family: + same_family = True + elif closest_ratio in vertical_family and 宽高比 in vertical_family: + same_family = True + elif closest_ratio == "1:1" and 宽高比 == "1:1": + same_family = True + + # 同家族：直接缩放或裁剪（无白底） + if same_family: + if img_ratio > target_ratio: + # 图像更宽，以高度为准缩放后裁剪 + scale = target_height / h + scaled_w = int(w * scale) + scaled_h = target_height + scaled = img.resize((scaled_w, scaled_h), Image.LANCZOS) + left = (scaled_w - target_width) // 2 + result = scaled.crop((left, 0, left + target_width, target_height)) + else: + # 图像更高，以宽度为准缩放后裁剪 + scale = target_width / w + scaled_w = target_width + scaled_h = int(h * scale) + scaled = img.resize((scaled_w, scaled_h), Image.LANCZOS) + top = (scaled_h - target_height) // 2 + result = scaled.crop((0, top, target_width, top + target_height)) + + # 不同家族：保持宽高比 + 白底填充 + else: + if img_ratio > target_ratio: + scaled_w = target_width + scaled_h = int(target_width / img_ratio) + else: + scaled_h = target_height + scaled_w = int(target_height * img_ratio) + + scaled = img.resize((scaled_w, scaled_h), Image.LANCZOS) + canvas = Image.new("RGB", (target_width, target_height), (255, 255, 255)) + paste_x = (target_width - scaled_w) // 2 + paste_y = (target_height - scaled_h) // 2 + canvas.paste(scaled, (paste_x, paste_y)) + result = canvas + + # 转回 tensor + arr = np.array(result).astype(np.float32) / 255.0 + tensor = torch.from_numpy(arr).unsqueeze(0) + + return (tensor,) + + +NODE_CLASS_MAPPINGS = { + "KlingVideo": KlingVideo, + "KlingFirstLastFrame": KlingFirstLastFrame, + "KlingMotionControlTest": KlingMotionControlTest, + "AspectRatioPreset": AspectRatioPreset, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "KlingVideo": "自研模型 3.0 视频", + "KlingFirstLastFrame": "自研模型 3.0 首尾帧到视频", + "KlingMotionControlTest": "自研模型 动作控制（测试）", + "AspectRatioPreset": "图片宽高比预设", +} diff --git a/nodes/load_file.py b/nodes/load_file.py new file mode 100644 index 0000000..7d432b3 --- /dev/null +++ b/nodes/load_file.py @@ -0,0 +1,146 @@ +""" +LoadFile 节点 +ComfyUI 自定义节点,用于加载文件并转换为 FILE 类型数据 +""" + +import base64 +import os +from pathlib import Path +from typing import Tuple + +from ..utils.file_types import FileData, DOCUMENT_MIME_TYPES, FILE_SIZE_LIMITS + + +class LoadFile: + """ + LoadFile 节点 + + 功能： + - 从文件系统加载文件 + - 支持 PDF 和 TXT 文件 + - 转换为 FILE 类型数据（包含 base64 编码内容） + - 验证文件大小和格式 + """ + + @classmethod + def INPUT_TYPES(cls): + """ + 定义输入参数 + """ + return { + "required": { + "文件路径": ("STRING", { + "default": "", + "multiline": False + }) + } + } + + # 返回值类型 + RETURN_TYPES = ("FILE", "STRING") + RETURN_NAMES = ("文件", "文件信息") + + # 执行函数名 + FUNCTION = "load_file" + + # 节点分类 + CATEGORY = "file/input" + + def load_file(self, 文件路径: str) -> Tuple[FileData, str]: + """ + 加载文件并转换为 FILE 类型 + + Args: + 文件路径: 文件的完整路径（支持绝对路径和相对路径） + + Returns: + (FileData, 文件信息预览) + + Raises: + ValueError: 文件不存在、不支持的文件类型或文件过大 + """ + try: + # 清理路径（去除空格和引号） + file_path = 文件路径.strip().strip('"').strip("'") + + if not file_path: + raise ValueError("文件路径不能为空") + + # 转换为 Path 对象 + path = Path(file_path) + + # 如果是相对路径，转换为绝对路径 + if not path.is_absolute(): + # 相对于当前工作目录 + path = Path.cwd() / path + + # 验证文件是否存在 + if not path.exists(): + raise ValueError(f"文件不存在: {file_path}") + + if not path.is_file(): + raise ValueError(f"路径不是文件: {file_path}") + + # 获取文件信息 + extension = path.suffix.lower() + filename = path.stem + file_size = path.stat().st_size + + # 验证文件类型 + if extension not in DOCUMENT_MIME_TYPES: + supported_types = ", ".join(DOCUMENT_MIME_TYPES.keys()) + raise ValueError( + f"不支持的文件类型: {extension}\n" + f"支持的类型: {supported_types}" + ) + + # 获取 MIME 类型 + mime_type = DOCUMENT_MIME_TYPES[extension] + + # 验证文件大小 + size_limit = FILE_SIZE_LIMITS.get(extension, 20 * 1024 * 1024) + if file_size > size_limit: + raise ValueError( + f"文件过大 ({file_size / 1024 / 1024:.2f}MB)，" + f"最大支持 {size_limit / 1024 / 1024:.0f}MB" + ) + + # 读取文件并转换为 base64 + print(f"LoadFile: 正在加载文件 {filename}{extension}") + print(f"LoadFile: 文件大小 = {file_size / 1024:.2f}KB") + + with open(path, "rb") as f: + file_bytes = f.read() + + # Base64 编码 + b64_str = base64.b64encode(file_bytes).decode("utf-8") + + # 创建 FileData 对象 + file_data = FileData( + path=str(path), + filename=filename, + extension=extension, + mime_type=mime_type, + data=b64_str, + size=file_size + ) + + # 生成文件信息预览 + file_info = ( + f"文件名: {filename}{extension}\n" + f"类型: {mime_type}\n" + f"大小: {file_size / 1024:.2f}KB\n" + f"路径: {path}" + ) + + print(f"LoadFile: 加载成功") + + return (file_data, file_info) + + except ValueError as e: + print(f"LoadFile: 输入错误 - {str(e)}") + raise + + except Exception as e: + print(f"LoadFile: 未知错误 - {str(e)}") + raise diff --git a/nodes/multi_res_preview.py b/nodes/multi_res_preview.py new file mode 100644 index 0000000..17fb2f1 --- /dev/null +++ b/nodes/multi_res_preview.py @@ -0,0 +1,165 @@ +""" +多分辨率图像预览节点 +ComfyUI 自定义节点，支持同时预览多张不同分辨率的图像 + +背景： + ComfyUI 原生「预览图像」节点要求 batch 内所有图片分辨率相同（因为它们被 + stack 成一个 [B, H, W, C] tensor）。当 API 返回多张不同尺寸的图片时 + （例如 nano-banana-2 同时返回 1K + 2K），原生节点会报错。 + +解决方案： + 声明 INPUT_IS_LIST = True，ComfyUI 会将连入的所有图像作为 + Python list[Tensor] 传入，而不是强行 stack 成单个 tensor。 + 节点逐张单独保存为临时 PNG，再通过 ui.images 列表返回给前端并列展示， + 完全不受分辨率一致性的限制。 +""" + +import os +import uuid +import json +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo + +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + + +def _get_temp_dir() -> str: + """获取 ComfyUI temp 目录，不可用时回退到系统临时目录""" + if FOLDER_PATHS_AVAILABLE: + return folder_paths.get_temp_directory() + import tempfile + return tempfile.gettempdir() + + +def _tensor_to_pil(tensor) -> list: + """ + 将单个 IMAGE tensor 转换为 PIL Image 列表。 + + ComfyUI IMAGE tensor 格式：[B, H, W, C]，float32，值域 [0, 1] + 支持： + - 单张图 tensor: shape [H, W, C] 或 [1, H, W, C] + - batch tensor: shape [B, H, W, C]（B 张相同尺寸图） + """ + import torch + if not isinstance(tensor, torch.Tensor): + return [] + + if tensor.ndim == 3: + tensor = tensor.unsqueeze(0) + + results = [] + for i in range(tensor.shape[0]): + img_np = tensor[i].cpu().numpy() + img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8) + results.append(Image.fromarray(img_np)) + return results + + +class MultiResPreview: + """ + 多分辨率图像预览节点 + + 功能： + - 单个「图像」输入端口，支持接入批次图像 + - INPUT_IS_LIST = True：ComfyUI 将每张图作为独立 tensor 传入， + 不强制要求尺寸相同，彻底解决不同分辨率无法共存的问题 + - 每张图像独立保存为临时 PNG，在节点上并列展示所有图像 + + 用法： + 将 Nano Banana 节点的输出直接连入「图像」端口即可， + 无论返回几张、分辨率是否相同，都能正确展示。 + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "图像": ("IMAGE",), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + }, + } + + # 关键：告知 ComfyUI 以 list[Tensor] 而非 stacked Tensor 传入图像 + # 这样不同分辨率的图片可以共存于同一个输入中 + INPUT_IS_LIST = True + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "preview" + CATEGORY = "image" + + DESCRIPTION = ( + "多分辨率图像预览节点。\n" + "单个图像输入端口，支持任意数量、任意分辨率的批次图像。\n" + "解决了原生「预览图像」节点要求 batch 内图片尺寸相同的限制。\n" + "常用场景：nano-banana-2 同时返回 1K + 2K 图时，直接连入本节点即可。" + ) + + def preview(self, 图像, prompt=None, extra_pnginfo=None) -> dict: + """ + 逐张将图像保存到 temp 目录，返回 ui.images 供前端展示。 + + Args: + 图像: list[Tensor]，每个元素是一张或一批图（INPUT_IS_LIST） + prompt: ComfyUI 注入的 prompt 元数据（可选） + extra_pnginfo: ComfyUI 注入的额外 PNG 信息（可选） + + Returns: + {"ui": {"images": [...]}} 格式，每项对应一张图 + """ + temp_dir = _get_temp_dir() + os.makedirs(temp_dir, exist_ok=True) + + # 构建 PNG 元数据（与原生预览节点行为一致） + metadata = PngInfo() + # INPUT_IS_LIST 时 hidden 值也会被包装成 list，取第一个元素 + _prompt = prompt[0] if isinstance(prompt, list) else prompt + _extra = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo + if _prompt is not None: + try: + metadata.add_text("prompt", json.dumps(_prompt)) + except Exception: + pass + if _extra is not None: + try: + for k, v in _extra.items(): + metadata.add_text(k, json.dumps(v)) + except Exception: + pass + + saved = [] + total_input = 0 + total_saved = 0 + + # 图像 是 list[Tensor]，逐个处理（每个 Tensor 可能自身是个 batch） + for tensor in 图像: + pil_images = _tensor_to_pil(tensor) + total_input += len(pil_images) + + for pil_img in pil_images: + try: + filename = f"multi_res_preview_{uuid.uuid4().hex[:12]}.png" + filepath = os.path.join(temp_dir, filename) + pil_img.save(filepath, pnginfo=metadata, compress_level=1) + + saved.append({ + "filename": filename, + "subfolder": "", + "type": "temp", + }) + total_saved += 1 + except Exception as e: + print(f"多分辨率预览: ⚠️ 保存图像失败 - {e}") + + if total_input == 0: + print("多分辨率预览: ⚠️ 没有接收到任何图像") + + return {"ui": {"images": saved}} diff --git a/nodes/nano_banana_pro.py b/nodes/nano_banana_pro.py new file mode 100644 index 0000000..996172b --- /dev/null +++ b/nodes/nano_banana_pro.py @@ -0,0 +1,903 @@ +""" +Nano Banana Pro 节点 +ComfyUI 自定义节点，用于调用 Gemini 模型生成图像 +""" + +import time +import math +import random +import asyncio +import aiohttp +from concurrent.futures import ThreadPoolExecutor +from typing import Optional, Tuple, List + +import torch +import numpy as np +from PIL import Image + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts +from ..utils.file_utils import ImageInfo, generate_timestamp_filename, save_image +from ..clients.gemini_client import GeminiAPIClient +from ..models_config import ( + get_enabled_models, get_model_description, + get_model_supported_aspect_ratios, get_all_supported_aspect_ratios, + get_model_supported_resolutions, get_all_supported_resolutions +) + +# 检查 folder_paths 是否可用 +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +# 导入 ComfyUI 原生进度条 +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + PROGRESS_BAR_AVAILABLE = False + print("⚠️ NanoBananaPro: comfy.utils.ProgressBar 不可用，将只使用终端进度显示") + +# 内存监控（可选） +try: + import psutil + MEMORY_MONITOR_AVAILABLE = True +except ImportError: + MEMORY_MONITOR_AVAILABLE = False + print("⚠️ NanoBananaPro: psutil 不可用，内存监控功能禁用") + +# ============================================================================ +# 调试日志配置 +# ============================================================================ +# 是否启用调试日志（打印完整的 API 响应内容） +# 设置为 True 以启用调试日志，False 以禁用 +DEBUG_LOG_ENABLED = False +# 是否启用请求体日志（打印发送给 API 的请求体，base64 图片数据将自动截断） +# 设置为 True 以启用请求体日志，False 以禁用 +REQUEST_LOG_ENABLED = False +# ============================================================================ + +_NODE = "Nano Banana Pro" + + +def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor: + """ + 将 PIL Image 列表转换为 ComfyUI tensor，安全处理多张不同尺寸的情况。 + + ComfyUI 的 IMAGE tensor 格式为 [B, H, W, C]，要求 batch 内所有图尺寸相同。 + 当 API 返回多张不同分辨率的图时（主图 + 附图），直接 stack 会崩溃。 + + 策略： + - 所有图均已按原始分辨率保存到磁盘（调用此函数前已完成） + - 以第一张图的尺寸为基准，只将尺寸相同的图纳入 tensor 输出 + - 尺寸不同的图跳过（不 resize、不丢弃磁盘文件），并打印日志提示 + - 若没有任何图与第一张尺寸相同（极罕见），则只输出第一张 + """ + if not images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + return pil_to_tensor([placeholder]) + + base_size = images[0].size # PIL size = (W, H) + matched = [img for img in images if img.size == base_size] + skipped = [img for img in images if img.size != base_size] + + if skipped: + sizes_str = ", ".join(f"{img.size[0]}×{img.size[1]}" for img in skipped) + print( + f"{node_label}: API 额外返回了 {len(skipped)} 张不同尺寸的图 ({sizes_str})，" + f"已按原始分辨率保存到磁盘，tensor 输出仅包含与主图尺寸相同的 {len(matched)} 张 " + f"({base_size[0]}×{base_size[1]})" + ) + + return pil_to_tensor(matched if matched else [images[0]]) + + +class NanoBananaPro: + """ + Nano Banana Pro 节点 + + 功能： + - 文生图：基于提示词生成图像 + - 图生图：基于输入图像和提示词生成新图像 + - 批量生成：支持并发生成多张图像 + + 注意： + - 支持的模型列表从 models_config.py 动态加载 + - 要添加/禁用模型，请编辑 models_config.py 文件 + """ + + # 支持的模型列表（从配置文件动态加载） + MODELS = None # 将在 INPUT_TYPES 中动态获取 + + # 支持的宽高比列表（全量：所有启用模型的并集，动态加载） + # 实际渲染时通过 get_all_supported_aspect_ratios() 获取 + ASPECT_RATIOS = [ + "1:1", "4:3", "3:4", "16:9", "9:16", + "2:3", "3:2", "4:5", "5:4", "21:9", + "1:4", "4:1", "1:8", "8:1" + ] + + # 支持的分辨率列表（全量兜底，实际由 get_all_supported_resolutions() 动态生成） + RESOLUTIONS = ["512", "1K", "2K", "4K"] + + def __init__(self): + """初始化节点""" + self.client = None + + @classmethod + def INPUT_TYPES(cls): + """ + 定义输入参数 + + ComfyUI 节点规范： + - required: 必选参数 + - optional: 可选参数 + """ + # 从配置文件动态获取启用的模型列表 + enabled_models = get_enabled_models() + + # 如果没有启用的模型，使用空列表（会导致节点不可用，提示用户配置） + if not enabled_models: + enabled_models = ["请在 models_config.py 中启用至少一个模型"] + + # 动态获取所有启用模型支持的宽高比（去重合并） + all_aspect_ratios = get_all_supported_aspect_ratios() + if not all_aspect_ratios: + all_aspect_ratios = cls.ASPECT_RATIOS + + # 动态获取所有启用模型支持的分辨率（去重合并） + all_resolutions = get_all_supported_resolutions() + if not all_resolutions: + all_resolutions = cls.RESOLUTIONS + + # 创建9个独立的图像输入 + optional_inputs = {} + for i in range(1, 10): # 1-9 + optional_inputs[f"参考图{i}"] = ("IMAGE",) + + return { + "required": { + "prompt": ("STRING", { + "default": "一个中国女子的OOTD", + "multiline": True + }), + "模型": (enabled_models, { + "default": enabled_models[0] + }), + "宽高比": (all_aspect_ratios, { + "default": "1:1" + }), + "分辨率": (all_resolutions, { + "default": "2K" + }), + "生图数量": ("INT", { + "default": 1, + "min": 1, + "max": 1000, + "step": 1 + }), + "像素缩放": ("BOOLEAN", { + "default": True, + "label_on": "打开", + "label_off": "关闭" + }), + "分辨率像素": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 100.0, + "step": 0.1, + "display": "number" + }), + "谷歌搜索（联网）": (["关闭", "打开"], { + "default": "关闭" + }), + "图片搜索（联网）": (["关闭", "打开"], { + "default": "关闭" + }), + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 0xffffffffffffffff + }) + }, + "optional": optional_inputs + } + + # 返回值类型 + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("输出图像",) + + # 导入 ComfyUI 的文件夹路径管理 + try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True + except ImportError: + FOLDER_PATHS_AVAILABLE = False + + # 执行函数名 + FUNCTION = "generate" + + # 节点分类 + CATEGORY = "image/generation" + + def resize_to_megapixels( + self, + image: Image.Image, + target_megapixels: float + ) -> Image.Image: + """ + 将图像缩放到指定的总像素数，保持纵横比 + + Args: + image: PIL Image 对象 + target_megapixels: 目标像素数（百万像素） + + Returns: + 缩放后的 PIL Image + + Example: + >>> resized = self.resize_to_megapixels(img, 2.0) # 缩放到2百万像素 + """ + # 计算当前像素数 + current_pixels = image.width * image.height + target_pixels = int(target_megapixels * 1_000_000) + + # 如果当前像素数已经接近目标，则不缩放 + if abs(current_pixels - target_pixels) / target_pixels < 0.05: + return image + + # 计算缩放比例 + scale = (target_pixels / current_pixels) ** 0.5 + + # 计算新尺寸 + new_width = int(image.width * scale) + new_height = int(image.height * scale) + + # 确保至少为1像素 + new_width = max(1, new_width) + new_height = max(1, new_height) + + # 使用 Lanczos 重采样 + resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) + + return resized_image + + def validate_inputs( + self, + images: Optional[torch.Tensor], + batch_size: int + ) -> None: + """ + 验证输入参数 + + Args: + images: 输入图像张量（可选） + batch_size: 批次大小 + + Raises: + ValueError: 如果输入参数不合法 + """ + # 检查图像数量 + if images is not None: + num_images = images.shape[0] + if num_images > 14: + raise ValueError( + f"输入图像数量 {num_images} 超过限制 14 张，请减少输入图像数量" + ) + + # 检查批次大小 + if batch_size < 1 or batch_size > 1000: + raise ValueError( + f"批次大小 {batch_size} 超出范围 [1, 1000]" + ) + + async def _generate_single_task( + self, + session: aiohttp.ClientSession, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: List[Image.Image], + output_folder: str, + global_task_index: int, + enable_grounding: bool = False, + enable_image_search: bool = False, + ) -> dict: + """执行单个生成任务，生成后立即保存到磁盘""" + result = { + "global_task_index": global_task_index, + "prompt": prompt, + "success": False, + "generated_count": 0, + "saved_files": [], + "error": None + } + + try: + gen_result = await self.client.generate_single_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=images if images else None, + session=session, + debug=DEBUG_LOG_ENABLED, + debug_request=REQUEST_LOG_ENABLED, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + if gen_result: + images_list, _ = gen_result + for gen_img in images_list: + output_path = generate_timestamp_filename( + output_folder=output_folder, + extension=".png" + ) + save_image(gen_img, output_path) + result["saved_files"].append(output_path) + gen_img = None # 释放内存 + + result["success"] = True + result["generated_count"] = len(images_list) + except Exception as e: + result["error"] = str(e) + + return result + + async def _process_batch_async( + self, + prompts: List[str], + model: str, + resolution: str, + aspect_ratio: str, + images_per_prompt: int, + input_images: List[Image.Image], + output_folder: str, + pbar=None, + enable_grounding: bool = False, + enable_image_search: bool = False, + ) -> List[dict]: + """异步批量处理：每个提示词独立调用 API，生成后立即写磁盘""" + # 构建任务列表：(prompt, sub_index) 用于 images_per_prompt > 1 的情况 + tasks_def = [] + for p_idx, prompt in enumerate(prompts): + for sub_idx in range(images_per_prompt): + tasks_def.append((p_idx, sub_idx, prompt)) + + total_tasks = len(tasks_def) + num_prompts = len(prompts) + print(f"Nano Banana Pro: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务") + + max_concurrent = 10 + num_batches = math.ceil(total_tasks / max_concurrent) + + all_results = [] + completed = 0 + success_count = 0 + fail_count = 0 + + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + + async with aiohttp.ClientSession(connector=connector) as session: + for batch_idx in range(num_batches): + start_idx = batch_idx * max_concurrent + end_idx = min(start_idx + max_concurrent, total_tasks) + + tasks = [] + for i in range(start_idx, end_idx): + _, _, prompt = tasks_def[i] + task = asyncio.create_task( + self._generate_single_task( + session=session, + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=input_images, + output_folder=output_folder, + global_task_index=i, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + ) + tasks.append(task) + + batch_results = [] + for coro in asyncio.as_completed(tasks): + result_data = None + try: + result = await coro + if isinstance(result, Exception): + result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": [], "prompt": ""} + else: + result_data = result + batch_results.append(result_data) + except Exception as e: + result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": [], "prompt": ""} + batch_results.append(result_data) + + completed += 1 + prompt_snippet = (result_data.get("prompt", "") or "")[:30] + + if result_data and result_data.get("success", False): + success_count += 1 + count = result_data.get("generated_count", 1) + print(f"Nano Banana Pro: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✓成功({count}张)") + else: + fail_count += 1 + error_msg = result_data.get("error", "未知错误") if result_data else "未知错误" + print(f"Nano Banana Pro: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}") + + if pbar is not None: + pbar.update(1) + + all_results.extend(batch_results) + + import gc + gc.collect() + + await asyncio.sleep(0.1) + + return all_results + + def generate( + self, + prompt: str, + 模型: str, + 宽高比: str, + 分辨率: str, + 生图数量: int, + 像素缩放: bool, + 分辨率像素: float, + seed: int, + **kwargs + ) -> Tuple[torch.Tensor]: + """ + 生成图像 + + Args: + prompt: 提示词 + 模型: 模型名称 + 宽高比: 宽高比 + 分辨率: 分辨率 + 生图数量: 批次大小 + 像素缩放: 是否启用像素缩放 + 分辨率像素: 目标像素数（百万像素） + seed: 随机种子 + **kwargs: 搜索开关（谷歌搜索（联网）/ 图片搜索（联网））及动态参考图输入 (参考图1-9) + 注：两个搜索参数名含全角括号，不能作为 Python 形参，从 kwargs 中提取 + + 注意： + 调试日志功能已移至文件顶部配置，通过修改 DEBUG_LOG_ENABLED 常量控制 + + Returns: + 生成的图像张量 (IMAGE,) + """ + start_time = time.time() + + # 从 kwargs 提取搜索参数（界面显示为「关闭/打开」，转为 bool 供调用） + enable_grounding: bool = (kwargs.pop("谷歌搜索（联网）", "关闭") == "打开") + enable_image_search: bool = (kwargs.pop("图片搜索（联网）", "关闭") == "打开") + + # 创建 ComfyUI 原生进度条 + pbar = None + if PROGRESS_BAR_AVAILABLE: + pbar = ProgressBar(生图数量) + + try: + # 设置随机种子（用于本地随机操作） + random.seed(seed) + np.random.seed(seed % (2**32)) + + # 内存监控初始化 + if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50: + import psutil + process = psutil.Process() + initial_memory = process.memory_info().rss / 1024 / 1024 + print(f"Nano Banana Pro: 初始内存使用: {initial_memory:.1f} MB") + + # 初始化 API 客户端 + if self.client is None: + try: + self.client = GeminiAPIClient() + except ValueError as e: + raise ValueError(f"初始化失败: {str(e)}") + + # 校验分辨率与模型的兼容性 + supported_resolutions = get_model_supported_resolutions(模型) + if supported_resolutions and 分辨率 not in supported_resolutions: + raise ValueError( + f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的分辨率：{', '.join(supported_resolutions)}" + ) + + # 校验宽高比与模型的兼容性 + supported_ratios = get_model_supported_aspect_ratios(模型) + if supported_ratios and 宽高比 not in supported_ratios: + raise ValueError( + f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的宽高比：{', '.join(supported_ratios)}" + ) + + # 校验图片搜索（联网）与模型的兼容性 + # 仅 nano-banana-2-限时特价 和 gemini-3.1-flash-image-preview 支持图片搜索 + IMAGE_SEARCH_UNSUPPORTED_MODELS = ["nano-banana-pro-限时特价", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview"] + if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS: + raise ValueError( + f"模型 \"{模型}\" 不支持【图片搜索（联网）】功能！" + f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用" + ) + + # 收集独立输入的参考图 + input_images = [] + for i in range(1, 10): # 1-9 + key = f"参考图{i}" + if key in kwargs and kwargs[key] is not None: + pil_imgs = tensor_to_pil(kwargs[key]) + input_images.extend(pil_imgs) + + # 验证输入图像数量 + if input_images: + if len(input_images) > 14: + raise ValueError( + f"输入图像数量 {len(input_images)} 超过限制 14 张，请减少输入图像数量" + ) + + # 应用像素缩放（如果启用） + if input_images and 像素缩放: + scaled_images = [] + for img in input_images: + scaled = self.resize_to_megapixels(img, 分辨率像素) + scaled_images.append(scaled) + input_images = scaled_images + + # 解析批量提示词 + batch_prompts = parse_batch_prompts(prompt) + + # 打印首行概览 + # 图片搜索（联网）开启时隐含谷歌搜索接地，与客户端请求逻辑保持一致 + grounding_str = "" + if enable_image_search: + grounding_str = " | 谷歌图片搜索接地" + elif enable_grounding: + grounding_str = " | 谷歌搜索接地" + + if batch_prompts: + # 批量提示词模式 + num_prompts = len(batch_prompts) + total_images = num_prompts * 生图数量 + mode_str = f"批量提示词模式 ({num_prompts}个提示词)" + if input_images: + mode_str += f" (输入{len(input_images)}张)" + print(f"Nano Banana Pro: {mode_str} | {分辨率} {宽高比} | 共{total_images}张{grounding_str}") + + # 大批量警告 + if total_images > 100: + print(f"⚠️ Nano Banana Pro: 警告！批量生成 {total_images} 张图片，内存占用可能较高") + print(f"⚠️ 建议：分批执行或减少生图数量") + else: + # 单提示词模式 + mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式" + print(f"Nano Banana Pro: {mode_str} | {分辨率} {宽高比} | {生图数量}张{grounding_str}") + + # 大批量警告 + if 生图数量 > 100: + print(f"⚠️ Nano Banana Pro: 警告！批量生成 {生图数量} 张图片，内存占用可能较高") + print(f"⚠️ 建议：分批执行或减少生图数量") + + # 统计变量 + success_count = 0 + fail_count = 0 + + # 进度回调 - 打印错误信息并更新进度条，添加内存监控 + def progress_callback(current, total, success, error_msg=None): + nonlocal success_count, fail_count + if success: + success_count += 1 + print(f"Nano Banana Pro: 任务 {current}/{total} 成功 ✓") + else: + fail_count += 1 + # 打印完整的错误信息（用于排查问题） + if error_msg: + print(f"Nano Banana Pro: 任务 {current}/{total} 失败 ✗") + print(f"原始错误详情:\n{error_msg}") + else: + print(f"Nano Banana Pro: 任务 {current}/{total} 失败 ✗") + + # 更新 ComfyUI 原生进度条 + if pbar is not None: + pbar.update(1) + + # 内存监控（每完成10个任务检查一次） + if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0: + import gc + gc.collect() # 强制垃圾回收 + current_memory = process.memory_info().rss / 1024 / 1024 + memory_increase = current_memory - initial_memory + print(f"Nano Banana Pro: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)") + + # 内存警告阈值（2GB） + if current_memory > 2000: + print(f"⚠️ Nano Banana Pro: 内存使用过高！建议减少生图数量或分批执行") + + # 根据是否有批量提示词选择生成模式 + if batch_prompts: + num_prompts = len(batch_prompts) + total_images = num_prompts * 生图数量 + + # ===== 批量提示词模式：异步并发+磁盘保存 ===== + if pbar is not None: + pbar = ProgressBar(total_images) + + # 确定保存路径 + output_folder = "" + if FOLDER_PATHS_AVAILABLE: + output_folder = folder_paths.get_output_directory() + print(f"Nano Banana Pro: 磁盘保存模式 → {output_folder}") + else: + raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用") + + import os + os.makedirs(output_folder, exist_ok=True) + + def run_async_in_thread(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete( + self._process_batch_async( + prompts=batch_prompts, + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + images_per_prompt=生图数量, + input_images=input_images, + output_folder=output_folder, + pbar=pbar, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + ) + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_async_in_thread) + try: + results = future.result(timeout=3600) + except TimeoutError: + raise RuntimeError("任务执行超时（1小时），请减少提示词数量或检查网络连接") + + # 统计结果 + success_count = sum(1 for r in results if r.get("success", False)) + fail_count = len(results) - success_count + total_generated = sum(r.get("generated_count", 0) for r in results) + all_saved_files = [] + for r in results: + all_saved_files.extend(r.get("saved_files", [])) + + elapsed = time.time() - start_time + time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" + + print(f"完成！总耗时 {time_str} | 成功: {success_count}/{total_images} | 失败: {fail_count}") + + # 失败详情 + failed_results = [r for r in results if not r.get("success", False)] + if failed_results: + for fr in failed_results: + idx = fr.get("global_task_index", -1) + 1 + prompt_snippet = (fr.get("prompt", "") or "")[:30] + error_msg = fr.get("error", "未知错误") + print(f" 失败 #{idx}: {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → {error_msg}") + + # 从磁盘加载最后 10 张图片 + output_images = [] + max_output_images = 10 + recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):] + for file_path in recent_files: + try: + img = Image.open(file_path) + output_images.append(img) + except Exception as e: + print(f"Nano Banana Pro: 无法加载 {file_path} - {e}") + + if not output_images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + output_images = [placeholder] + + output_tensor = _images_to_tensor_safe(output_images, _NODE) + print(f"Nano Banana Pro: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张") + + import gc + gc.collect() + return (output_tensor,) + else: + # 单提示词模式 + if 生图数量 == 1: + # 单张：同步生成 + 保存到磁盘 + 输出 tensor + generated_images = self.client.generate_sync( + prompt=prompt, + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + batch_size=1, + images=input_images, + progress_callback=progress_callback, + debug=DEBUG_LOG_ENABLED, + debug_request=REQUEST_LOG_ENABLED, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + # 单张：保存到磁盘 + import os + output_folder = "" + if FOLDER_PATHS_AVAILABLE: + output_folder = folder_paths.get_output_directory() + print(f"Nano Banana Pro: 磁盘保存模式 → {output_folder}") + else: + raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用") + os.makedirs(output_folder, exist_ok=True) + for gen_img in generated_images: + output_path = generate_timestamp_filename(output_folder=output_folder) + save_image(gen_img, output_path) + else: + # 多张：异步并发 + 磁盘保存（与批量提示词逻辑一致） + print(f"Nano Banana Pro: 单提示词×{生图数量}张 → 异步并发模式") + + if pbar is not None: + pbar = ProgressBar(生图数量) + + output_folder = "" + if FOLDER_PATHS_AVAILABLE: + output_folder = folder_paths.get_output_directory() + print(f"Nano Banana Pro: 磁盘保存模式 → {output_folder}") + else: + raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用") + + import os + os.makedirs(output_folder, exist_ok=True) + + def run_async_in_thread(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete( + self._process_batch_async( + prompts=[prompt], + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + images_per_prompt=生图数量, + input_images=input_images, + output_folder=output_folder, + pbar=pbar, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + ) + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_async_in_thread) + try: + results = future.result(timeout=3600) + except TimeoutError: + raise RuntimeError("任务执行超时（1小时），请减少生图数量或检查网络连接") + + success_count = sum(1 for r in results if r.get("success", False)) + fail_count = len(results) - success_count + total_generated = sum(r.get("generated_count", 0) for r in results) + all_saved_files = [] + for r in results: + all_saved_files.extend(r.get("saved_files", [])) + + elapsed = time.time() - start_time + time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" + print(f"完成！总耗时 {time_str} | 成功: {success_count}/{生图数量} | 失败: {fail_count}") + + # 失败详情 + failed_results = [r for r in results if not r.get("success", False)] + if failed_results: + for fr in failed_results: + idx = fr.get("global_task_index", -1) + 1 + error_msg = fr.get("error", "未知错误") + print(f" 失败 #{idx}: {prompt[:30]}{'...' if len(prompt) >= 30 else ''} → {error_msg}") + + # 从磁盘加载最后 10 张图片 + output_images = [] + max_output_images = 10 + recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):] + for file_path in recent_files: + try: + img = Image.open(file_path) + output_images.append(img) + except Exception as e: + print(f"Nano Banana Pro: 无法加载 {file_path} - {e}") + + if not output_images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + output_images = [placeholder] + + output_tensor = _images_to_tensor_safe(output_images, _NODE) + print(f"Nano Banana Pro: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张") + # 不生成 prompts_map.txt（单提示词无需映射） + + import gc + gc.collect() + return (output_tensor,) + + + # 优化：限制输出图片数量，避免内存爆炸 + max_output_images = 20 # 最多输出20张图片到ComfyUI + + if len(generated_images) > max_output_images: + print(f"Nano Banana Pro: 生成 {len(generated_images)} 张图片，限制输出前 {max_output_images} 张到ComfyUI") + output_images = generated_images[:max_output_images] + else: + output_images = generated_images + + # 转换输出图像 + output_tensor = _images_to_tensor_safe(output_images, _NODE) + + # 计算耗时并打印最终统计 + elapsed = time.time() - start_time + if elapsed < 1: + time_str = f"{elapsed:.3f}s" + else: + time_str = f"{elapsed:.2f}s" + + # 打印最终汇总 + if fail_count > 0: + print(f"[4/4] 完成！总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张") + else: + print(f"[4/4] 完成！总耗时 {time_str} | 成功 {len(generated_images)}张") + + # 最终内存清理 + import gc + gc.collect() + if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50: + final_memory = process.memory_info().rss / 1024 / 1024 + print(f"Nano Banana Pro: 最终内存使用: {final_memory:.1f} MB") + + return (output_tensor,) + + except ValueError as e: + # 检测是否为授权错误 + if str(e) == "未授权！": + print("请联系作者授权后方可使用！") + raise ValueError("未授权！") from None + else: + # 用户输入错误 - 打印完整错误信息 + error_msg = str(e) + print(f"Nano Banana Pro: ❌ {error_msg}") + raise ValueError(error_msg) from None + + except RuntimeError as e: + # 打印完整错误信息 + error_full = str(e) + print(f"Nano Banana Pro: ❌ {error_full}") + raise RuntimeError(error_full) from None + + except Exception as e: + # 其他未知错误 - 打印完整错误信息 + error_msg = str(e) + print(f"Nano Banana Pro: ❌ {error_msg}") + raise type(e)(error_msg) from None + + finally: + # 查询余额 + if self.client is not None: + try: + balance_data = self.client.query_balance_sync() + balance_info = self.client.format_balance_info(balance_data) + print(f"Nano Banana Pro: {balance_info}") + except Exception: + pass + + # 最终内存清理 + import gc + gc.collect() + print(f"Nano Banana Pro: 最终内存清理完成") \ No newline at end of file diff --git a/nodes/nano_banana_v2.py b/nodes/nano_banana_v2.py new file mode 100644 index 0000000..0b653db --- /dev/null +++ b/nodes/nano_banana_v2.py @@ -0,0 +1,656 @@ +""" +Nano Banana v2 节点 +NanoBananaPro 的完全复刻，唯一改动： + + 将原来 9 个独立「参考图1~9」输入端 + 改为 1 个「参考图」输入端（可选），配合「加载图像（批量）」节点使用。 + + 「加载图像（批量）」输出 is_output_list=True（list[Tensor]）， + 本节点声明 INPUT_IS_LIST = True 来整体接收该列表， + 然后在 generate() 开头对所有参数统一解包，其余业务逻辑与原节点完全一致。 +""" + +import os +import gc +import time +import math +import random +import asyncio +import aiohttp +from concurrent.futures import ThreadPoolExecutor +from typing import Optional, Tuple, List + +import torch +import numpy as np +from PIL import Image + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts +from ..utils.file_utils import ImageInfo, generate_timestamp_filename, save_image +from ..clients.gemini_client import GeminiAPIClient +from ..models_config import ( + get_enabled_models, get_model_description, + get_model_supported_aspect_ratios, get_all_supported_aspect_ratios, + get_model_supported_resolutions, get_all_supported_resolutions +) + +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + PROGRESS_BAR_AVAILABLE = False + +try: + import psutil + MEMORY_MONITOR_AVAILABLE = True +except ImportError: + MEMORY_MONITOR_AVAILABLE = False + +DEBUG_LOG_ENABLED = False +REQUEST_LOG_ENABLED = False + +_NODE = "Nano Banana v2" + + +def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor: + """ + 将 PIL Image 列表转换为 ComfyUI tensor，安全处理多张不同尺寸的情况。 + + ComfyUI 的 IMAGE tensor 格式为 [B, H, W, C]，要求 batch 内所有图尺寸相同。 + 当 API 返回多张不同分辨率的图时（主图 + 附图），直接 stack 会崩溃。 + + 策略： + - 所有图均已按原始分辨率保存到磁盘（调用此函数前已完成） + - 以第一张图的尺寸为基准，只将尺寸相同的图纳入 tensor 输出 + - 尺寸不同的图跳过（不 resize、不丢弃磁盘文件），并打印日志提示 + - 若没有任何图与第一张尺寸相同（极罕见），则只输出第一张 + """ + if not images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + return pil_to_tensor([placeholder]) + + base_size = images[0].size # PIL size = (W, H) + matched = [img for img in images if img.size == base_size] + skipped = [img for img in images if img.size != base_size] + + if skipped: + sizes_str = ", ".join(f"{img.size[0]}×{img.size[1]}" for img in skipped) + print( + f"{node_label}: API 额外返回了 {len(skipped)} 张不同尺寸的图 ({sizes_str})，" + f"已按原始分辨率保存到磁盘，tensor 输出仅包含与主图尺寸相同的 {len(matched)} 张 " + f"({base_size[0]}×{base_size[1]})" + ) + + return pil_to_tensor(matched if matched else [images[0]]) + + +class NanaBananaV2: + """ + Nano Banana v2 + + 与 NanoBananaPro 完全一致，参考图输入方式不同： + - 原版：9 个独立可选端口（参考图1~9） + - v2：1 个可选端口「参考图」，配合「加载图像（批量）」可传入任意数量图片 + """ + + ASPECT_RATIOS = [ + "1:1", "4:3", "3:4", "16:9", "9:16", + "2:3", "3:2", "4:5", "5:4", "21:9", + "1:4", "4:1", "1:8", "8:1" + ] + RESOLUTIONS = ["512", "1K", "2K", "4K"] + + def __init__(self): + self.client = None + + @classmethod + def INPUT_TYPES(cls): + enabled_models = get_enabled_models() + if not enabled_models: + enabled_models = ["请在 models_config.py 中启用至少一个模型"] + + all_aspect_ratios = get_all_supported_aspect_ratios() or cls.ASPECT_RATIOS + all_resolutions = get_all_supported_resolutions() or cls.RESOLUTIONS + + return { + "required": { + "prompt": ("STRING", { + "default": "一个中国女子的OOTD", + "multiline": True + }), + "模型": (enabled_models, {"default": enabled_models[0]}), + "宽高比": (all_aspect_ratios, {"default": "1:1"}), + "分辨率": (all_resolutions, {"default": "2K"}), + "生图数量": ("INT", {"default": 1, "min": 1, "max": 1000, "step": 1}), + "像素缩放": ("BOOLEAN", {"default": True, "label_on": "打开", "label_off": "关闭"}), + "分辨率像素": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.1, "display": "number"}), + "谷歌搜索（联网）": (["关闭", "打开"], {"default": "关闭"}), + "图片搜索（联网）": (["关闭", "打开"], {"default": "关闭"}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }, + "optional": { + # 单个参考图端口，接受普通 IMAGE 或「加载图像（批量）」输出的列表 + "参考图": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("输出图像",) + FUNCTION = "generate" + CATEGORY = "image/generation" + + # 声明 INPUT_IS_LIST，使 ComfyUI 将「加载图像（批量）」的 list[Tensor] + # 整体传入而非逐张迭代执行，同时其余所有参数也会被包进 list，需解包。 + INPUT_IS_LIST = True + + # ------------------------------------------------------------------ # + # 以下方法与 NanoBananaPro 完全相同，仅 generate() 开头增加了解包逻辑 + # ------------------------------------------------------------------ # + + def resize_to_megapixels(self, image: Image.Image, target_megapixels: float) -> Image.Image: + current_pixels = image.width * image.height + target_pixels = int(target_megapixels * 1_000_000) + if abs(current_pixels - target_pixels) / target_pixels < 0.05: + return image + scale = (target_pixels / current_pixels) ** 0.5 + new_width = max(1, int(image.width * scale)) + new_height = max(1, int(image.height * scale)) + return image.resize((new_width, new_height), Image.Resampling.LANCZOS) + + async def _generate_single_task( + self, + session: aiohttp.ClientSession, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: List[Image.Image], + output_folder: str, + global_task_index: int, + enable_grounding: bool = False, + enable_image_search: bool = False, + ) -> dict: + result = { + "global_task_index": global_task_index, + "prompt": prompt, + "success": False, + "generated_count": 0, + "saved_files": [], + "error": None + } + try: + gen_result = await self.client.generate_single_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=images if images else None, + session=session, + debug=DEBUG_LOG_ENABLED, + debug_request=REQUEST_LOG_ENABLED, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + if gen_result: + images_list, _ = gen_result + for gen_img in images_list: + output_path = generate_timestamp_filename( + output_folder=output_folder, + extension=".png" + ) + save_image(gen_img, output_path) + result["saved_files"].append(output_path) + gen_img = None + result["success"] = True + result["generated_count"] = len(images_list) + except Exception as e: + result["error"] = str(e) + return result + + async def _process_batch_async( + self, + prompts: List[str], + model: str, + resolution: str, + aspect_ratio: str, + images_per_prompt: int, + input_images: List[Image.Image], + output_folder: str, + pbar=None, + enable_grounding: bool = False, + enable_image_search: bool = False, + ) -> List[dict]: + tasks_def = [] + for p_idx, prompt in enumerate(prompts): + for sub_idx in range(images_per_prompt): + tasks_def.append((p_idx, sub_idx, prompt)) + + total_tasks = len(tasks_def) + num_prompts = len(prompts) + print(f"{_NODE}: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务") + + max_concurrent = 10 + num_batches = math.ceil(total_tasks / max_concurrent) + all_results = [] + completed = 0 + success_count = 0 + fail_count = 0 + + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + async with aiohttp.ClientSession(connector=connector) as session: + for batch_idx in range(num_batches): + start_idx = batch_idx * max_concurrent + end_idx = min(start_idx + max_concurrent, total_tasks) + tasks = [] + for i in range(start_idx, end_idx): + _, _, prompt = tasks_def[i] + task = asyncio.create_task( + self._generate_single_task( + session=session, + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=input_images, + output_folder=output_folder, + global_task_index=i, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + ) + tasks.append(task) + + batch_results = [] + for coro in asyncio.as_completed(tasks): + result_data = None + try: + result = await coro + if isinstance(result, Exception): + result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": [], "prompt": ""} + else: + result_data = result + batch_results.append(result_data) + except Exception as e: + result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": [], "prompt": ""} + batch_results.append(result_data) + + completed += 1 + prompt_snippet = (result_data.get("prompt", "") or "")[:30] + if result_data and result_data.get("success", False): + success_count += 1 + count = result_data.get("generated_count", 1) + print(f"{_NODE}: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✓成功({count}张)") + else: + fail_count += 1 + error_msg = result_data.get("error", "未知错误") if result_data else "未知错误" + print(f"{_NODE}: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}") + + if pbar is not None: + pbar.update(1) + + all_results.extend(batch_results) + gc.collect() + await asyncio.sleep(0.1) + + return all_results + + def generate( + self, + prompt, + 模型, + 宽高比, + 分辨率, + 生图数量, + 像素缩放, + 分辨率像素, + **kwargs + ) -> Tuple[torch.Tensor]: + # ---------------------------------------------------------------- + # INPUT_IS_LIST=True 时，所有参数均为 list，先统一解包为标量 + # ---------------------------------------------------------------- + prompt = prompt[0] if isinstance(prompt, list) else prompt + 模型 = 模型[0] if isinstance(模型, list) else 模型 + 宽高比 = 宽高比[0] if isinstance(宽高比, list) else 宽高比 + 分辨率 = 分辨率[0] if isinstance(分辨率, list) else 分辨率 + 生图数量 = 生图数量[0] if isinstance(生图数量, list) else 生图数量 + 像素缩放 = 像素缩放[0] if isinstance(像素缩放, list) else 像素缩放 + 分辨率像素 = 分辨率像素[0] if isinstance(分辨率像素, list) else 分辨率像素 + + # seed 也在 kwargs 里（含全角括号的参数名无法作为形参） + seed_raw = kwargs.pop("seed", [0]) + seed: int = seed_raw[0] if isinstance(seed_raw, list) else seed_raw + + # 搜索开关同理 + grounding_raw = kwargs.pop("谷歌搜索（联网）", ["关闭"]) + image_search_raw = kwargs.pop("图片搜索（联网）", ["关闭"]) + enable_grounding: bool = (grounding_raw[0] if isinstance(grounding_raw, list) else grounding_raw) == "打开" + enable_image_search: bool = (image_search_raw[0] if isinstance(image_search_raw, list) else image_search_raw) == "打开" + + # ---------------------------------------------------------------- + # 收集参考图：兼容两种来源 + # 1. 「加载图像（批量）」→ is_output_list=True → list[Tensor] + # INPUT_IS_LIST 下传入的是 list[list[Tensor]] 或 list[Tensor]，需展平 + # 2. 普通 IMAGE 端口（单 tensor 或 batch tensor）→ list 中只有 1 个元素 + # ---------------------------------------------------------------- + ref_raw = kwargs.pop("参考图", None) + input_images: List[Image.Image] = [] + + if ref_raw is not None: + # INPUT_IS_LIST 下，可选端口若连接则为 list；元素可能是 Tensor 或 list[Tensor] + items = ref_raw if isinstance(ref_raw, list) else [ref_raw] + for item in items: + if item is None: + continue + if isinstance(item, list): + # 来自 is_output_list 的嵌套 list，继续展平 + for sub in item: + if sub is not None and isinstance(sub, torch.Tensor): + input_images.extend(tensor_to_pil(sub)) + elif isinstance(item, torch.Tensor): + input_images.extend(tensor_to_pil(item)) + + # ---------------------------------------------------------------- + # 以下逻辑与 NanoBananaPro.generate() 完全一致 + # ---------------------------------------------------------------- + start_time = time.time() + + pbar = None + if PROGRESS_BAR_AVAILABLE: + pbar = ProgressBar(生图数量) + + try: + random.seed(seed) + np.random.seed(seed % (2 ** 32)) + + if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50: + process = psutil.Process() + initial_memory = process.memory_info().rss / 1024 / 1024 + print(f"{_NODE}: 初始内存使用: {initial_memory:.1f} MB") + + if self.client is None: + try: + self.client = GeminiAPIClient() + except ValueError as e: + raise ValueError(f"初始化失败: {str(e)}") + + # 校验分辨率 + supported_resolutions = get_model_supported_resolutions(模型) + if supported_resolutions and 分辨率 not in supported_resolutions: + raise ValueError( + f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的分辨率：{', '.join(supported_resolutions)}" + ) + + # 校验宽高比 + supported_ratios = get_model_supported_aspect_ratios(模型) + if supported_ratios and 宽高比 not in supported_ratios: + raise ValueError( + f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的宽高比：{', '.join(supported_ratios)}" + ) + + # 校验图片搜索与模型兼容性 + IMAGE_SEARCH_UNSUPPORTED_MODELS = [ + "nano-banana-pro-限时特价", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview" + ] + if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS: + raise ValueError( + f"模型 \"{模型}\" 不支持【图片搜索（联网）】功能！" + f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用" + ) + + # 验证输入图像数量上限 + if len(input_images) > 14: + raise ValueError( + f"输入图像数量 {len(input_images)} 超过限制 14 张，请减少输入图像数量" + ) + + # 像素缩放 + if input_images and 像素缩放: + input_images = [self.resize_to_megapixels(img, 分辨率像素) for img in input_images] + + # 解析批量提示词 + batch_prompts = parse_batch_prompts(prompt) + + # 打印概览 + grounding_str = "" + if enable_image_search: + grounding_str = " | 谷歌图片搜索接地" + elif enable_grounding: + grounding_str = " | 谷歌搜索接地" + + if batch_prompts: + num_prompts = len(batch_prompts) + total_images = num_prompts * 生图数量 + mode_str = f"批量提示词模式 ({num_prompts}个提示词)" + if input_images: + mode_str += f" (输入{len(input_images)}张)" + print(f"{_NODE}: {mode_str} | {分辨率} {宽高比} | 共{total_images}张{grounding_str}") + if total_images > 100: + print(f"⚠️ {_NODE}: 警告！批量生成 {total_images} 张图片，内存占用可能较高") + print(f"⚠️ 建议：分批执行或减少生图数量") + else: + mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式" + print(f"{_NODE}: {mode_str} | {分辨率} {宽高比} | {生图数量}张{grounding_str}") + if 生图数量 > 100: + print(f"⚠️ {_NODE}: 警告！批量生成 {生图数量} 张图片，内存占用可能较高") + print(f"⚠️ 建议：分批执行或减少生图数量") + + success_count = 0 + fail_count = 0 + + def progress_callback(current, total, success, error_msg=None): + nonlocal success_count, fail_count + if success: + success_count += 1 + print(f"{_NODE}: 任务 {current}/{total} 成功 ✓") + else: + fail_count += 1 + if error_msg: + print(f"{_NODE}: 任务 {current}/{total} 失败 ✗") + print(f"原始错误详情:\n{error_msg}") + else: + print(f"{_NODE}: 任务 {current}/{total} 失败 ✗") + if pbar is not None: + pbar.update(1) + if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0: + gc.collect() + current_memory = process.memory_info().rss / 1024 / 1024 + memory_increase = current_memory - initial_memory + print(f"{_NODE}: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)") + if current_memory > 2000: + print(f"⚠️ {_NODE}: 内存使用过高！建议减少生图数量或分批执行") + + def _get_output_folder(): + if FOLDER_PATHS_AVAILABLE: + folder = folder_paths.get_output_directory() + return folder + raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用") + + def run_async_in_thread(coro_fn): + def _run(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete(coro_fn()) + finally: + loop.close() + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(_run) + try: + return future.result(timeout=3600) + except TimeoutError: + raise RuntimeError("任务执行超时（1小时），请减少数量或检查网络连接") + + # ── 批量提示词模式 ────────────────────────────────────── + if batch_prompts: + num_prompts = len(batch_prompts) + total_images = num_prompts * 生图数量 + if pbar is not None: + pbar = ProgressBar(total_images) + + output_folder = _get_output_folder() + os.makedirs(output_folder, exist_ok=True) + + results = run_async_in_thread(lambda: self._process_batch_async( + prompts=batch_prompts, + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + images_per_prompt=生图数量, + input_images=input_images, + output_folder=output_folder, + pbar=pbar, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + )) + + success_count = sum(1 for r in results if r.get("success", False)) + fail_count = len(results) - success_count + all_saved_files = [f for r in results for f in r.get("saved_files", [])] + + elapsed = time.time() - start_time + time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" + print(f"完成！总耗时 {time_str} | 成功: {success_count}/{total_images} | 失败: {fail_count}") + + failed_results = [r for r in results if not r.get("success", False)] + for fr in failed_results: + idx = fr.get("global_task_index", -1) + 1 + snippet = (fr.get("prompt", "") or "")[:30] + print(f" 失败 #{idx}: {snippet}{'...' if len(snippet) >= 30 else ''} → {fr.get('error', '未知错误')}") + + output_images = [] + for fp in all_saved_files[-min(10, len(all_saved_files)):]: + try: + output_images.append(Image.open(fp)) + except Exception as e: + print(f"{_NODE}: 无法加载 {fp} - {e}") + + if not output_images: + output_images = [Image.new('RGB', (512, 512), color=(128, 128, 128))] + + output_tensor = _images_to_tensor_safe(output_images, _NODE) + print(f"{_NODE}: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张") + gc.collect() + return (output_tensor,) + + # ── 单提示词模式 ──────────────────────────────────────── + if 生图数量 == 1: + generated_images = self.client.generate_sync( + prompt=prompt, + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + batch_size=1, + images=input_images, + progress_callback=progress_callback, + debug=DEBUG_LOG_ENABLED, + debug_request=REQUEST_LOG_ENABLED, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + ) + output_folder = _get_output_folder() + os.makedirs(output_folder, exist_ok=True) + for gen_img in generated_images: + output_path = generate_timestamp_filename(output_folder=output_folder) + save_image(gen_img, output_path) + else: + print(f"{_NODE}: 单提示词×{生图数量}张 → 异步并发模式") + if pbar is not None: + pbar = ProgressBar(生图数量) + + output_folder = _get_output_folder() + os.makedirs(output_folder, exist_ok=True) + + results = run_async_in_thread(lambda: self._process_batch_async( + prompts=[prompt], + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + images_per_prompt=生图数量, + input_images=input_images, + output_folder=output_folder, + pbar=pbar, + enable_grounding=enable_grounding, + enable_image_search=enable_image_search, + )) + + success_count = sum(1 for r in results if r.get("success", False)) + fail_count = len(results) - success_count + all_saved_files = [f for r in results for f in r.get("saved_files", [])] + + elapsed = time.time() - start_time + time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" + print(f"完成！总耗时 {time_str} | 成功: {success_count}/{生图数量} | 失败: {fail_count}") + + failed_results = [r for r in results if not r.get("success", False)] + for fr in failed_results: + idx = fr.get("global_task_index", -1) + 1 + print(f" 失败 #{idx}: {prompt[:30]}{'...' if len(prompt) >= 30 else ''} → {fr.get('error', '未知错误')}") + + output_images = [] + for fp in all_saved_files[-min(10, len(all_saved_files)):]: + try: + output_images.append(Image.open(fp)) + except Exception as e: + print(f"{_NODE}: 无法加载 {fp} - {e}") + + if not output_images: + output_images = [Image.new('RGB', (512, 512), color=(128, 128, 128))] + + output_tensor = _images_to_tensor_safe(output_images, _NODE) + print(f"{_NODE}: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张") + gc.collect() + return (output_tensor,) + + # 单张同步模式的输出路径（生图数量==1 走到这里） + max_output_images = 20 + if len(generated_images) > max_output_images: + print(f"{_NODE}: 生成 {len(generated_images)} 张图片，限制输出前 {max_output_images} 张到ComfyUI") + output_images = generated_images[:max_output_images] + else: + output_images = generated_images + + output_tensor = _images_to_tensor_safe(output_images, _NODE) + elapsed = time.time() - start_time + time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" + if fail_count > 0: + print(f"[4/4] 完成！总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张") + else: + print(f"[4/4] 完成！总耗时 {time_str} | 成功 {len(generated_images)}张") + + gc.collect() + return (output_tensor,) + + except ValueError as e: + if str(e) == "未授权！": + print("请联系作者授权后方可使用！") + raise ValueError("未授权！") from None + error_msg = str(e) + print(f"{_NODE}: ❌ {error_msg}") + raise ValueError(error_msg) from None + + except RuntimeError as e: + error_full = str(e) + print(f"{_NODE}: ❌ {error_full}") + raise RuntimeError(error_full) from None + + except Exception as e: + error_msg = str(e) + print(f"{_NODE}: ❌ {error_msg}") + raise type(e)(error_msg) from None + + finally: + if self.client is not None: + try: + balance_data = self.client.query_balance_sync() + balance_info = self.client.format_balance_info(balance_data) + print(f"{_NODE}: {balance_info}") + except Exception: + pass + gc.collect() diff --git a/nodes/quan_neng_sheng_tu.py b/nodes/quan_neng_sheng_tu.py new file mode 100644 index 0000000..dc1e190 --- /dev/null +++ b/nodes/quan_neng_sheng_tu.py @@ -0,0 +1,826 @@ +""" +全能生图 节点 +ComfyUI 自定义节点，用于调用 Gemini 模型生成图像 +""" + +import time +import math +import random +import asyncio +import aiohttp +from concurrent.futures import ThreadPoolExecutor +from typing import Optional, Tuple, List + +import torch +import numpy as np +from PIL import Image + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts +from ..utils.file_utils import ImageInfo, generate_timestamp_filename, save_image +from ..clients.openai_client import OpenAIAPIClient +from ..models_config import ( + get_enabled_models, get_model_description, + get_model_supported_aspect_ratios, get_all_supported_aspect_ratios, + get_model_supported_resolutions, get_all_supported_resolutions +) + +# 检查 folder_paths 是否可用 +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +# 导入 ComfyUI 原生进度条 +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + PROGRESS_BAR_AVAILABLE = False + print("⚠️ 全能生图: comfy.utils.ProgressBar 不可用，将只使用终端进度显示") + +# 内存监控（可选） +try: + import psutil + MEMORY_MONITOR_AVAILABLE = True +except ImportError: + MEMORY_MONITOR_AVAILABLE = False + print("⚠️ 全能生图: psutil 不可用，内存监控功能禁用") + +# ============================================================================ +# 调试日志配置 +# ============================================================================ +# 是否启用调试日志（打印完整的 API 响应内容） +# 设置为 True 以启用调试日志，False 以禁用 +DEBUG_LOG_ENABLED = False +# 是否启用请求体日志（打印发送给 API 的请求体，base64 图片数据将自动截断） +# 设置为 True 以启用请求体日志，False 以禁用 +REQUEST_LOG_ENABLED = False +# ============================================================================ + + +class QuanNengShengTu: + """ + 全能生图 节点 + + 功能： + - 文生图：基于提示词生成图像 + - 图生图：基于输入图像和提示词生成新图像 + - 批量生成：支持并发生成多张图像 + + 注意： + - 支持的模型列表从 models_config.py 动态加载 + - 要添加/禁用模型，请编辑 models_config.py 文件 + """ + + # 支持的模型列表（从配置文件动态加载） + MODELS = None # 将在 INPUT_TYPES 中动态获取 + + # 支持的宽高比列表（全量：所有启用模型的并集，动态加载） + ASPECT_RATIOS = [ + "1:1", "4:3", "3:4", "16:9", "9:16", + "2:3", "3:2", "4:5", "5:4", "21:9", + "1:4", "4:1", "1:8", "8:1" + ] + + # 支持的分辨率列表（全量兜底，实际由 get_all_supported_resolutions() 动态生成） + RESOLUTIONS = ["512", "1K", "2K", "4K"] + + def __init__(self): + """初始化节点""" + self.client = None + + @classmethod + def INPUT_TYPES(cls): + """ + 定义输入参数 + + ComfyUI 节点规范： + - required: 必选参数 + - optional: 可选参数 + """ + # 从配置文件动态获取启用的模型列表 + enabled_models = get_enabled_models() + + # 过滤掉包含"限时特价"的模型 + enabled_models = [m for m in enabled_models if "限时特价" not in m] + + # 如果没有启用的模型，使用空列表（会导致节点不可用，提示用户配置） + if not enabled_models: + enabled_models = ["请在 models_config.py 中启用至少一个模型"] + + # 动态获取所有启用模型支持的宽高比（去重合并） + all_aspect_ratios = get_all_supported_aspect_ratios() + if not all_aspect_ratios: + all_aspect_ratios = cls.ASPECT_RATIOS + + # 动态获取所有启用模型支持的分辨率（去重合并） + all_resolutions = get_all_supported_resolutions() + if not all_resolutions: + all_resolutions = cls.RESOLUTIONS + + # 创建9个独立的图像输入 + optional_inputs = {} + for i in range(1, 10): # 1-9 + optional_inputs[f"参考图{i}"] = ("IMAGE",) + + return { + "required": { + "提示词": ("STRING", { + "default": "一个中国女子的OOTD", + "multiline": True + }), + "模型": (enabled_models, { + "default": enabled_models[0] + }), + "宽高比": (all_aspect_ratios, { + "default": "1:1" + }), + "分辨率": (all_resolutions, { + "default": "2K" + }), + "生图数量": ("INT", { + "default": 1, + "min": 1, + "max": 1000, + "step": 1 + }), + "像素缩放": ("BOOLEAN", { + "default": True, + "label_on": "打开", + "label_off": "关闭" + }), + "分辨率像素": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 100.0, + "step": 0.1, + "display": "number" + }), + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 0xffffffffffffffff + }) + }, + "optional": optional_inputs + } + + # 返回值类型 + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("输出图像",) + + # 导入 ComfyUI 的文件夹路径管理 + try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True + except ImportError: + FOLDER_PATHS_AVAILABLE = False + + # 执行函数名 + FUNCTION = "generate" + + # 节点分类 + CATEGORY = "image/generation" + + def resize_to_megapixels( + self, + image: Image.Image, + target_megapixels: float + ) -> Image.Image: + """ + 将图像缩放到指定的总像素数，保持纵横比 + + Args: + image: PIL Image 对象 + target_megapixels: 目标像素数（百万像素） + + Returns: + 缩放后的 PIL Image + """ + # 计算当前像素数 + current_pixels = image.width * image.height + target_pixels = int(target_megapixels * 1_000_000) + + # 如果当前像素数已经接近目标，则不缩放 + if abs(current_pixels - target_pixels) / target_pixels < 0.05: + return image + + # 计算缩放比例 + scale = (target_pixels / current_pixels) ** 0.5 + + # 计算新尺寸 + new_width = int(image.width * scale) + new_height = int(image.height * scale) + + # 确保至少为1像素 + new_width = max(1, new_width) + new_height = max(1, new_height) + + # 使用 Lanczos 重采样 + resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) + + return resized_image + + def validate_inputs( + self, + images: Optional[torch.Tensor], + batch_size: int + ) -> None: + """ + 验证输入参数 + + Args: + images: 输入图像张量（可选） + batch_size: 批次大小 + + Raises: + ValueError: 如果输入参数不合法 + """ + # 检查图像数量 + if images is not None: + num_images = images.shape[0] + if num_images > 14: + raise ValueError( + f"输入图像数量 {num_images} 超过限制 14 张，请减少输入图像数量" + ) + + # 检查批次大小 + if batch_size < 1 or batch_size > 1000: + raise ValueError( + f"批次大小 {batch_size} 超出范围 [1, 1000]" + ) + + async def _generate_single_task( + self, + session: aiohttp.ClientSession, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: List[Image.Image], + output_folder: str, + global_task_index: int, + ) -> dict: + """执行单个生成任务，生成后立即保存到磁盘""" + result = { + "global_task_index": global_task_index, + "prompt": prompt, + "success": False, + "generated_count": 0, + "saved_files": [], + "error": None + } + + try: + gen_result = await self.client.generate_single_async( + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=images if images else None, + session=session, + debug=DEBUG_LOG_ENABLED, + debug_request=REQUEST_LOG_ENABLED, + enable_grounding=False, + enable_image_search=False + ) + if gen_result: + images_list, _ = gen_result + for gen_img in images_list: + output_path = generate_timestamp_filename( + output_folder=output_folder, + extension=".png" + ) + save_image(gen_img, output_path) + result["saved_files"].append(output_path) + gen_img = None # 释放内存 + + result["success"] = True + result["generated_count"] = len(images_list) + except Exception as e: + result["error"] = str(e) + + return result + + async def _process_batch_async( + self, + prompts: List[str], + model: str, + resolution: str, + aspect_ratio: str, + images_per_prompt: int, + input_images: List[Image.Image], + output_folder: str, + pbar=None, + ) -> List[dict]: + """异步批量处理：每个提示词独立调用 API，生成后立即写磁盘""" + # 构建任务列表：(prompt, sub_index) 用于 images_per_prompt > 1 的情况 + tasks_def = [] + for p_idx, prompt in enumerate(prompts): + for sub_idx in range(images_per_prompt): + tasks_def.append((p_idx, sub_idx, prompt)) + + total_tasks = len(tasks_def) + num_prompts = len(prompts) + print(f"全能生图: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务") + + max_concurrent = 10 + num_batches = math.ceil(total_tasks / max_concurrent) + + all_results = [] + completed = 0 + success_count = 0 + fail_count = 0 + + connector = aiohttp.TCPConnector(limit=0, limit_per_host=0) + + async with aiohttp.ClientSession(connector=connector) as session: + for batch_idx in range(num_batches): + start_idx = batch_idx * max_concurrent + end_idx = min(start_idx + max_concurrent, total_tasks) + + tasks = [] + for i in range(start_idx, end_idx): + _, _, prompt = tasks_def[i] + task = asyncio.create_task( + self._generate_single_task( + session=session, + prompt=prompt, + model=model, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=input_images, + output_folder=output_folder, + global_task_index=i, + ) + ) + tasks.append(task) + + batch_results = [] + for coro in asyncio.as_completed(tasks): + result_data = None + try: + result = await coro + if isinstance(result, Exception): + result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": [], "prompt": ""} + else: + result_data = result + batch_results.append(result_data) + except Exception as e: + result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": [], "prompt": ""} + batch_results.append(result_data) + + completed += 1 + prompt_snippet = (result_data.get("prompt", "") or "")[:30] + + if result_data and result_data.get("success", False): + success_count += 1 + count = result_data.get("generated_count", 1) + print(f"全能生图: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✓成功({count}张)") + else: + fail_count += 1 + error_msg = result_data.get("error", "未知错误") if result_data else "未知错误" + print(f"全能生图: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}") + + if pbar is not None: + pbar.update(1) + + all_results.extend(batch_results) + + import gc + gc.collect() + + await asyncio.sleep(0.1) + + return all_results + + def generate( + self, + 提示词: str, + 模型: str, + 宽高比: str, + 分辨率: str, + 生图数量: int, + 像素缩放: bool, + 分辨率像素: float, + seed: int, + **kwargs + ) -> Tuple[torch.Tensor]: + """ + 生成图像 + + Args: + prompt: 提示词 + 模型: 模型名称 + 宽高比: 宽高比 + 分辨率: 分辨率 + 生图数量: 批次大小 + 像素缩放: 是否启用像素缩放 + 分辨率像素: 目标像素数（百万像素） + seed: 随机种子 + **kwargs: 动态参考图输入 (参考图1-9) + + 注意： + 调试日志功能已移至文件顶部配置，通过修改 DEBUG_LOG_ENABLED 常量控制 + + Returns: + 生成的图像张量 (IMAGE,) + """ + start_time = time.time() + + # 创建 ComfyUI 原生进度条 + pbar = None + if PROGRESS_BAR_AVAILABLE: + pbar = ProgressBar(生图数量) + + try: + # 设置随机种子（用于本地随机操作） + random.seed(seed) + np.random.seed(seed % (2**32)) + + # 内存监控初始化 + if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50: + import psutil + process = psutil.Process() + initial_memory = process.memory_info().rss / 1024 / 1024 + print(f"全能生图: 初始内存使用: {initial_memory:.1f} MB") + + # 初始化 API 客户端 + if self.client is None: + try: + self.client = OpenAIAPIClient() + except ValueError as e: + raise ValueError(f"初始化失败: {str(e)}") + + # 校验分辨率与模型的兼容性 + supported_resolutions = get_model_supported_resolutions(模型) + if supported_resolutions and 分辨率 not in supported_resolutions: + raise ValueError( + f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的分辨率：{', '.join(supported_resolutions)}" + ) + + # 校验宽高比与模型的兼容性 + supported_ratios = get_model_supported_aspect_ratios(模型) + if supported_ratios and 宽高比 not in supported_ratios: + raise ValueError( + f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的宽高比：{', '.join(supported_ratios)}" + ) + + # 收集独立输入的参考图 + input_images = [] + for i in range(1, 10): # 1-9 + key = f"参考图{i}" + if key in kwargs and kwargs[key] is not None: + pil_imgs = tensor_to_pil(kwargs[key]) + input_images.extend(pil_imgs) + + # 验证输入图像数量 + if input_images: + if len(input_images) > 14: + raise ValueError( + f"输入图像数量 {len(input_images)} 超过限制 14 张，请减少输入图像数量" + ) + + # 应用像素缩放（如果启用） + if input_images and 像素缩放: + scaled_images = [] + for img in input_images: + scaled = self.resize_to_megapixels(img, 分辨率像素) + scaled_images.append(scaled) + input_images = scaled_images + + # 解析批量提示词 + batch_prompts = parse_batch_prompts(提示词) + + # 打印首行概览 + if batch_prompts: + # 批量提示词模式 + num_prompts = len(batch_prompts) + total_images = num_prompts * 生图数量 + mode_str = f"批量提示词模式 ({num_prompts}个提示词)" + if input_images: + mode_str += f" (输入{len(input_images)}张)" + print(f"全能生图: {mode_str} | {分辨率} {宽高比} | 共{total_images}张") + + # 大批量警告 + if total_images > 100: + print(f"⚠️ 全能生图: 警告！批量生成 {total_images} 张图片，内存占用可能较高") + print(f"⚠️ 建议：分批执行或减少生图数量") + else: + # 单提示词模式 + mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式" + print(f"全能生图: {mode_str} | {分辨率} {宽高比} | {生图数量}张") + + # 大批量警告 + if 生图数量 > 100: + print(f"⚠️ 全能生图: 警告！批量生成 {生图数量} 张图片，内存占用可能较高") + print(f"⚠️ 建议：分批执行或减少生图数量") + + # 统计变量 + success_count = 0 + fail_count = 0 + + # 进度回调 - 打印错误信息并更新进度条，添加内存监控 + def progress_callback(current, total, success, error_msg=None): + nonlocal success_count, fail_count + if success: + success_count += 1 + print(f"全能生图: 任务 {current}/{total} 成功 ✓") + else: + fail_count += 1 + if error_msg: + print(f"全能生图: 任务 {current}/{total} 失败 ✗") + print(f"原始错误详情:\n{error_msg}") + else: + print(f"全能生图: 任务 {current}/{total} 失败 ✗") + + # 更新 ComfyUI 原生进度条 + if pbar is not None: + pbar.update(1) + + # 内存监控（每完成10个任务检查一次） + if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0: + import gc + gc.collect() # 强制垃圾回收 + current_memory = process.memory_info().rss / 1024 / 1024 + memory_increase = current_memory - initial_memory + print(f"全能生图: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)") + + # 内存警告阈值（2GB） + if current_memory > 2000: + print(f"⚠️ 全能生图: 内存使用过高！建议减少生图数量或分批执行") + + # 根据是否有批量提示词选择生成模式 + if batch_prompts: + num_prompts = len(batch_prompts) + total_images = num_prompts * 生图数量 + + # ===== 批量提示词模式：异步并发+磁盘保存 ===== + if pbar is not None: + pbar = ProgressBar(total_images) + + # 确定保存路径 + output_folder = "" + if FOLDER_PATHS_AVAILABLE: + output_folder = folder_paths.get_output_directory() + print(f"全能生图: 磁盘保存模式 → {output_folder}") + else: + raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用") + + import os + os.makedirs(output_folder, exist_ok=True) + + def run_async_in_thread(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete( + self._process_batch_async( + prompts=batch_prompts, + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + images_per_prompt=生图数量, + input_images=input_images, + output_folder=output_folder, + pbar=pbar, + ) + ) + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_async_in_thread) + try: + results = future.result(timeout=3600) + except TimeoutError: + raise RuntimeError("任务执行超时（1小时），请减少提示词数量或检查网络连接") + + # 统计结果 + success_count = sum(1 for r in results if r.get("success", False)) + fail_count = len(results) - success_count + total_generated = sum(r.get("generated_count", 0) for r in results) + all_saved_files = [] + for r in results: + all_saved_files.extend(r.get("saved_files", [])) + + elapsed = time.time() - start_time + time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" + + print(f"完成！总耗时 {time_str} | 成功: {success_count}/{total_images} | 失败: {fail_count}") + + # 失败详情 + failed_results = [r for r in results if not r.get("success", False)] + if failed_results: + for fr in failed_results: + idx = fr.get("global_task_index", -1) + 1 + prompt_snippet = (fr.get("prompt", "") or "")[:30] + error_msg = fr.get("error", "未知错误") + print(f" 失败 #{idx}: {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → {error_msg}") + + # 从磁盘加载最后 10 张图片 + output_images = [] + max_output_images = 10 + recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):] + for file_path in recent_files: + try: + img = Image.open(file_path) + output_images.append(img) + except Exception as e: + print(f"全能生图: 无法加载 {file_path} - {e}") + + if not output_images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + output_images = [placeholder] + + output_tensor = pil_to_tensor(output_images) + print(f"全能生图: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张") + + import gc + gc.collect() + return (output_tensor,) + else: + # 单提示词模式 + if 生图数量 == 1: + # 单张：同步生成 + 保存到磁盘 + 输出 tensor + generated_images = self.client.generate_sync( + prompt=提示词, + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + batch_size=1, + images=input_images, + progress_callback=progress_callback, + debug=DEBUG_LOG_ENABLED, + debug_request=REQUEST_LOG_ENABLED, + enable_grounding=False, + enable_image_search=False + ) + # 单张：保存到磁盘 + import os + output_folder = "" + if FOLDER_PATHS_AVAILABLE: + output_folder = folder_paths.get_output_directory() + print(f"全能生图: 磁盘保存模式 → {output_folder}") + else: + raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用") + os.makedirs(output_folder, exist_ok=True) + for gen_img in generated_images: + output_path = generate_timestamp_filename(output_folder=output_folder) + save_image(gen_img, output_path) + else: + # 多张：异步并发 + 磁盘保存（与批量提示词逻辑一致） + print(f"全能生图: 单提示词×{生图数量}张 → 异步并发模式") + + if pbar is not None: + pbar = ProgressBar(生图数量) + + output_folder = "" + if FOLDER_PATHS_AVAILABLE: + output_folder = folder_paths.get_output_directory() + print(f"全能生图: 磁盘保存模式 → {output_folder}") + else: + raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用") + + import os + os.makedirs(output_folder, exist_ok=True) + + def run_async_in_thread(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete( + self._process_batch_async( + prompts=[提示词], + model=模型, + resolution=分辨率, + aspect_ratio=宽高比, + images_per_prompt=生图数量, + input_images=input_images, + output_folder=output_folder, + pbar=pbar, + ) + ) + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(run_async_in_thread) + try: + results = future.result(timeout=3600) + except TimeoutError: + raise RuntimeError("任务执行超时（1小时），请减少生图数量或检查网络连接") + + success_count = sum(1 for r in results if r.get("success", False)) + fail_count = len(results) - success_count + total_generated = sum(r.get("generated_count", 0) for r in results) + all_saved_files = [] + for r in results: + all_saved_files.extend(r.get("saved_files", [])) + + elapsed = time.time() - start_time + time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" + print(f"完成！总耗时 {time_str} | 成功: {success_count}/{生图数量} | 失败: {fail_count}") + + # 失败详情 + failed_results = [r for r in results if not r.get("success", False)] + if failed_results: + for fr in failed_results: + idx = fr.get("global_task_index", -1) + 1 + error_msg = fr.get("error", "未知错误") + print(f" 失败 #{idx}: {提示词[:30]}{'...' if len(提示词) >= 30 else ''} → {error_msg}") + + # 从磁盘加载最后 10 张图片 + output_images = [] + max_output_images = 10 + recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):] + for file_path in recent_files: + try: + img = Image.open(file_path) + output_images.append(img) + except Exception as e: + print(f"全能生图: 无法加载 {file_path} - {e}") + + if not output_images: + placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) + output_images = [placeholder] + + output_tensor = pil_to_tensor(output_images) + print(f"全能生图: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张") + + import gc + gc.collect() + return (output_tensor,) + + + # 优化：限制输出图片数量，避免内存爆炸 + max_output_images = 20 # 最多输出20张图片到ComfyUI + + if len(generated_images) > max_output_images: + print(f"全能生图: 生成 {len(generated_images)} 张图片，限制输出前 {max_output_images} 张到ComfyUI") + output_images = generated_images[:max_output_images] + else: + output_images = generated_images + + # 转换输出图像 + output_tensor = pil_to_tensor(output_images) + + # 计算耗时并打印最终统计 + elapsed = time.time() - start_time + if elapsed < 1: + time_str = f"{elapsed:.3f}s" + else: + time_str = f"{elapsed:.2f}s" + + # 打印最终汇总 + if fail_count > 0: + print(f"[4/4] 完成！总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张") + else: + print(f"[4/4] 完成！总耗时 {time_str} | 成功 {len(generated_images)}张") + + # 最终内存清理 + import gc + gc.collect() + if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50: + final_memory = process.memory_info().rss / 1024 / 1024 + print(f"全能生图: 最终内存使用: {final_memory:.1f} MB") + + return (output_tensor,) + + except ValueError as e: + # 检测是否为授权错误 + if str(e) == "未授权！": + print("请联系作者授权后方可使用！") + raise ValueError("未授权！") from None + else: + error_msg = str(e) + print(f"全能生图: ❌ {error_msg}") + raise ValueError(error_msg) from None + + except RuntimeError as e: + error_full = str(e) + print(f"全能生图: ❌ {error_full}") + raise RuntimeError(error_full) from None + + except Exception as e: + error_msg = str(e) + print(f"全能生图: ❌ {error_msg}") + raise type(e)(error_msg) from None + + finally: + # 查询余额 + if self.client is not None: + try: + balance_data = self.client.query_balance_sync() + balance_info = self.client.format_balance_info(balance_data) + print(f"全能生图: {balance_info}") + except Exception: + pass + + # 最终内存清理 + import gc + gc.collect() + print(f"全能生图: 最终内存清理完成") diff --git a/nodes/remove_metadata.py b/nodes/remove_metadata.py new file mode 100644 index 0000000..568daf4 --- /dev/null +++ b/nodes/remove_metadata.py @@ -0,0 +1,371 @@ +""" +图像元数据去除节点 +替代 ComfyUI 原生"保存图像"节点，保存时不写入提示词、工作流等 AI 元数据 + +提供两种节点： +1. SaveCleanImage - 接收 IMAGE 张量，去除元数据后直接保存到 output 目录 +2. BatchCleanMetadata - 指定文件夹路径，批量去除已有图片中的元数据 +""" + +import os +from datetime import datetime +import random +from typing import List + +import numpy as np +import torch +from PIL import Image +from PIL.PngImagePlugin import PngInfo + +from ..utils.image_utils import tensor_to_pil +from ..utils.file_utils import _get_port_suffix + +# 尝试导入 ComfyUI 的 folder_paths +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +# 支持的图片格式 +SUPPORTED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.webp', '.bmp', '.tiff', '.tif'} + + +def _get_output_dir() -> str: + """ + 获取 ComfyUI output 目录 + + Returns: + output 目录的绝对路径 + """ + if FOLDER_PATHS_AVAILABLE: + return folder_paths.get_output_directory() + # fallback: 相对于插件目录推断 + plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + return os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output") + + +def _get_next_counter(directory: str, prefix: str) -> int: + """ + 扫描目录，获取下一个可用的文件计数器 + + Args: + directory: 目标目录 + prefix: 文件名前缀 + + Returns: + 下一个计数器值 + """ + if not os.path.exists(directory): + return 1 + + if prefix: + pattern = re.compile(rf'^{re.escape(prefix)}_(\d+)') + else: + pattern = re.compile(rf'^(\d+)\.') + max_counter = 0 + + for f in os.listdir(directory): + m = pattern.match(f) + if m: + counter = int(m.group(1)) + max_counter = max(max_counter, counter) + + return max_counter + 1 + + +def _save_image_clean(image: Image.Image, path: str, fmt: str = None, quality: int = 95) -> None: + """ + 保存图像，不包含任何元数据 + + 通过提取纯像素数据并重建全新的 Image 对象，确保没有任何元数据残留。 + + Args: + image: PIL Image 对象 + path: 保存路径 + fmt: 图像格式（PNG/JPEG/WEBP），为 None 时根据扩展名推断 + quality: JPEG/WEBP 质量（1-100） + """ + # 确保 RGB 模式 + if image.mode != 'RGB': + image = image.convert('RGB') + + # 提取纯像素数据，重建全新的 Image 对象 + # 使用 tobytes() + frombytes() 确保只保留像素数据，彻底断开与原图像的关联 + pixel_data = image.tobytes() + clean = Image.frombytes('RGB', image.size, pixel_data) + + # 显式清空 info 字典，确保不会有任何残留元数据 + clean.info = {} + + # 推断格式 + if fmt is None: + ext = os.path.splitext(path)[1].lower() + format_map = { + '.png': 'PNG', + '.jpg': 'JPEG', + '.jpeg': 'JPEG', + '.webp': 'WEBP', + '.bmp': 'BMP', + '.tiff': 'TIFF', + '.tif': 'TIFF', + } + fmt = format_map.get(ext, 'PNG') + + # 构建保存参数（确保不写入任何元数据） + save_kwargs = {} + if fmt == 'PNG': + save_kwargs['pnginfo'] = PngInfo() # 空的 PngInfo，不包含任何文本块 + elif fmt == 'JPEG': + save_kwargs['quality'] = quality + # 不传 exif 参数，自然不会写入 EXIF 数据 + elif fmt == 'WEBP': + save_kwargs['quality'] = quality + save_kwargs['exif'] = b"" # 显式清空 EXIF + + clean.save(path, format=fmt, **save_kwargs) + + +# ============================================================================ +# 节点 1：保存干净图像 +# ============================================================================ + +class SaveCleanImage: + """ + 保存干净图像节点（不含元数据） + + 功能： + - 接收 IMAGE 张量（支持单图和批次） + - 去除所有元数据后保存到 ComfyUI/output 目录 + - 文件名自动添加 nometa 标识，方便辨认 + - 支持 PNG/JPEG/WEBP 格式 + - 作为终端节点，替代 ComfyUI 原生"保存图像"节点 + + 使用场景： + - 生图完成后，直接保存不含 AI 元数据的干净图像 + - 分享图像时不暴露提示词和工作流 + """ + + SAVE_FORMATS = ["PNG", "JPEG", "WEBP"] + + @classmethod + def INPUT_TYPES(cls): + """ + 定义输入参数 + + Returns: + 输入参数配置字典 + """ + return { + "required": { + "图像": ("IMAGE",), + "文件名前缀": ("STRING", {"default": "ComfyUI_nometa"}), + "保存格式": (cls.SAVE_FORMATS, {"default": "PNG"}), + }, + "optional": { + "JPEG/WEBP质量": ("INT", { + "default": 95, + "min": 1, + "max": 100, + "step": 1 + }), + } + } + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "save_clean" + CATEGORY = "image" + + DESCRIPTION = ( + "保存干净图像（不含元数据）。\n" + "替代 ComfyUI 原生'保存图像'节点，保存时不写入提示词、工作流等 AI 元数据。\n" + "文件保存到 ComfyUI/output 目录。" + ) + + def save_clean( + self, + 图像: torch.Tensor, + 文件名前缀: str = "ComfyUI_nometa", + 保存格式: str = "PNG", + **kwargs + ) -> dict: + """ + 去除元数据并保存图像 + + Args: + 图像: ComfyUI 图像张量 [B, H, W, C] + 文件名前缀: 保存文件名前缀 + 保存格式: 图像格式（PNG/JPEG/WEBP） + **kwargs: 可选参数（JPEG/WEBP质量） + + Returns: + UI 结果字典，包含保存的图像信息用于前端预览 + """ + quality = kwargs.get("JPEG/WEBP质量", 95) + + output_dir = _get_output_dir() + port_suffix = _get_port_suffix() + os.makedirs(output_dir, exist_ok=True) + + # 格式与扩展名映射 + ext_map = {"PNG": ".png", "JPEG": ".jpg", "WEBP": ".webp"} + ext = ext_map.get(保存格式, ".png") + + # 转换为 PIL 图像 + pil_images = tensor_to_pil(图像) + + results = [] + saved_paths = [] + for img in pil_images: + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + ms = random.randint(0, 999) + + while True: + if 文件名前缀: + filename = f"{文件名前缀}_{ts}_{ms:03d}{port_suffix}{ext}" + else: + filename = f"{ts}_{ms:03d}{port_suffix}{ext}" + filepath = os.path.join(output_dir, filename) + if not os.path.exists(filepath): + break + ms = (ms + 1) % 1000 + + _save_image_clean(img, filepath, fmt=保存格式, quality=quality) + + results.append({ + "filename": filename, + "subfolder": "", + "type": "output" + }) + saved_paths.append(filepath) + + # 打印详细日志，方便用户定位保存的文件 + print(f"保存干净图像: 已保存 {len(pil_images)} 张无元数据图像 (格式: {保存格式})") + for p in saved_paths: + print(f" → {p}") + + return {"ui": {"images": results}} + + +# ============================================================================ +# 节点 2：批量去除元数据 +# ============================================================================ + +class BatchCleanMetadata: + """ + 批量去除文件夹中图片元数据的节点 + + 功能： + - 指定文件夹路径，批量处理其中所有图片 + - 去除 EXIF、PNG tEXt 块、ComfyUI 工作流等所有元数据 + - 支持保存到原目录（添加 _nometa 后缀）或覆盖原文件 + - 支持 PNG/JPG/JPEG/WEBP/BMP/TIFF 格式 + + 使用场景： + - 已经保存了一批含有 AI 元数据的图片，需要批量清理 + - 批量处理指定文件夹中的所有图片 + """ + + @classmethod + def INPUT_TYPES(cls): + """ + 定义输入参数 + + Returns: + 输入参数配置字典 + """ + return { + "required": { + "文件夹路径": ("STRING", {"default": ""}), + "覆盖原文件": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("处理结果",) + OUTPUT_NODE = True + FUNCTION = "batch_clean" + CATEGORY = "image" + + DESCRIPTION = ( + "批量去除文件夹中图片的元数据。\n" + "支持 PNG/JPG/JPEG/WEBP/BMP/TIFF 格式。\n" + "默认在原文件名后添加 _nometa 后缀保存，也可选择覆盖原文件。" + ) + + def batch_clean( + self, + 文件夹路径: str, + 覆盖原文件: bool = False, + ) -> tuple: + """ + 批量去除文件夹中图片的元数据 + + Args: + 文件夹路径: 待处理图片所在的文件夹路径 + 覆盖原文件: 是否覆盖原文件（False 则添加 _nometa 后缀） + + Returns: + 处理结果字符串 + + Raises: + ValueError: 文件夹路径无效 + """ + if not 文件夹路径 or not 文件夹路径.strip(): + raise ValueError("请输入文件夹路径") + + folder = 文件夹路径.strip() + + if not os.path.isdir(folder): + raise ValueError(f"文件夹路径无效或不存在: {folder}") + + # 扫描支持的图片文件 + files = [] + for f in sorted(os.listdir(folder)): + ext = os.path.splitext(f)[1].lower() + if ext in SUPPORTED_EXTENSIONS: + files.append(f) + + if not files: + msg = f"文件夹中未找到支持的图片文件 ({', '.join(SUPPORTED_EXTENSIONS)})" + print(f"批量去除元数据: {msg}") + return (msg,) + + print(f"批量去除元数据: 找到 {len(files)} 张图片，开始处理...") + + success_count = 0 + fail_count = 0 + + for f in files: + try: + src_path = os.path.join(folder, f) + img = Image.open(src_path) + + if 覆盖原文件: + dst_path = src_path + else: + name, ext = os.path.splitext(f) + dst_path = os.path.join(folder, f"{name}_nometa{ext}") + + _save_image_clean(img, dst_path) + success_count += 1 + + except Exception as e: + print(f"批量去除元数据: 处理 {f} 失败 - {str(e)}") + fail_count += 1 + + # 构建结果消息 + if fail_count > 0: + msg = f"处理完成: 成功 {success_count} 张, 失败 {fail_count} 张" + else: + msg = f"处理完成: 全部 {success_count} 张成功" + + if not 覆盖原文件: + msg += " (已添加 _nometa 后缀)" + else: + msg += " (已覆盖原文件)" + + print(f"批量去除元数据: {msg}") + + return (msg,) diff --git a/nodes/sora_video.py b/nodes/sora_video.py new file mode 100644 index 0000000..8508a6c --- /dev/null +++ b/nodes/sora_video.py @@ -0,0 +1,526 @@ +""" +Sora 视频生成节点 +ComfyUI 自定义节点，调用 Sora API 生成视频 +""" + +import os +import re +import time +from math import gcd +from typing import Optional, Tuple + +import torch + +from ..utils.image_utils import tensor_to_pil +from ..clients.sora_client import SoraClient +from ..models_config import ( + get_enabled_sora_models, + get_all_sora_seconds, + get_all_sora_sizes, + get_sora_supported_seconds, + get_sora_supported_sizes, + get_sora_seconds_with_labels, + get_sora_sizes_with_labels, + SORA_MODELS, +) + +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + PROGRESS_BAR_AVAILABLE = False + print("⚠️ SoraVideo: comfy.utils.ProgressBar 不可用，将只使用终端进度显示") + + +def _size_to_display(size: str) -> str: + """ + 将 'WxH' 格式的分辨率转换为友好显示名。 + + 例如： + "720x1280" → "720P 9:16" + "1280x720" → "720P 16:9" + "1024x1792" → "1K 4:7" + "1792x1024" → "1K 7:4" + + Args: + size: 分辨率字符串，格式 "WxH" + + Returns: + 友好显示名字符串 + """ + parts = size.lower().split("x") + w, h = int(parts[0]), int(parts[1]) + short_side = min(w, h) + if short_side >= 3840: + res = "4K" + elif short_side >= 1920: + res = "2K" + elif short_side >= 1080: + res = "1K" + elif short_side >= 720: + res = "720P" + elif short_side >= 480: + res = "480P" + else: + res = f"{short_side}P" + g = gcd(w, h) + ratio = f"{w // g}:{h // g}" + return f"{res} {ratio} ({size})" + + +def _build_size_display_map(sizes: list) -> dict: + """ + 构建 显示名 → 实际值 映射字典。 + + Args: + sizes: 实际分辨率列表，如 ["720x1280", "1280x720"] + + Returns: + 字典，key 为显示名，value 为实际分辨率字符串 + """ + mapping = {} + for size in sizes: + display = _size_to_display(size) + if display in mapping: + # 极少数情况下防止重名 + display = f"{display} ({size})" + mapping[display] = size + return mapping + + +def _get_video_output_dir() -> str: + """获取视频输出目录: ComfyUI/output/video""" + if FOLDER_PATHS_AVAILABLE: + base = folder_paths.get_output_directory() + else: + plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + base = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output") + video_dir = os.path.join(base, "video") + os.makedirs(video_dir, exist_ok=True) + return video_dir + + +def _get_next_counter(directory: str, prefix: str) -> int: + """扫描目录，获取下一个可用的文件计数器""" + if not os.path.exists(directory): + return 1 + pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)") + max_counter = 0 + for f in os.listdir(directory): + m = pattern.match(f) + if m: + max_counter = max(max_counter, int(m.group(1))) + return max_counter + 1 + + +def _fit_image_to_target(image, target_size: str): + """ + 将参考图片按 "等比缩放覆盖 + 居中裁剪" 策略适配到目标分辨率。 + + 策略 (Cover Crop)： + 1. 比较图片宽高比和目标宽高比 + 2. 等比缩放，使图片最短边刚好覆盖目标对应边（图片完全覆盖目标区域） + 3. 居中裁剪多余部分，得到精确目标尺寸 + + Args: + image: PIL Image 对象 + target_size: 目标分辨率字符串，格式 "WxH"（如 "720x1280"） + + Returns: + 适配后的 PIL Image 对象 + """ + from PIL import Image as PILImage + + # 解析目标尺寸 + parts = target_size.lower().split("x") + target_w, target_h = int(parts[0]), int(parts[1]) + + src_w, src_h = image.size + src_ratio = src_w / src_h + target_ratio = target_w / target_h + + # 宽高比一致且尺寸不超过目标，无需处理 + if abs(src_ratio - target_ratio) < 0.01 and src_w <= target_w and src_h <= target_h: + return image + + print(f"Sora: 参考图片 {src_w}x{src_h} (比例 {src_ratio:.2f}) → 目标 {target_w}x{target_h} (比例 {target_ratio:.2f})") + + # 获取高质量重采样滤波器 + resample = PILImage.Resampling.LANCZOS if hasattr(PILImage, "Resampling") else PILImage.LANCZOS + + # Cover Crop: 缩放使图片完全覆盖目标区域，然后居中裁剪 + if src_ratio > target_ratio: + # 图片更宽：以高度为基准缩放，裁左右 + scale = target_h / src_h + new_w = round(src_w * scale) + new_h = target_h + image = image.resize((new_w, new_h), resample=resample) + # 居中裁剪宽度 + left = (new_w - target_w) // 2 + image = image.crop((left, 0, left + target_w, target_h)) + else: + # 图片更高（或一样）：以宽度为基准缩放，裁上下 + scale = target_w / src_w + new_w = target_w + new_h = round(src_h * scale) + image = image.resize((new_w, new_h), resample=resample) + # 居中裁剪高度 + top = (new_h - target_h) // 2 + image = image.crop((0, top, target_w, top + target_h)) + + print(f"Sora: 参考图片已适配为 {image.size[0]}x{image.size[1]}") + return image + + +def _compress_image_for_upload( + image, + target_size: Optional[str] = None, +) -> bytes: + """ + 将 PIL Image 适配目标分辨率并编码为 PNG 字节，用于上传。 + + ============================================================ + ⚠️ 已验证可用的标准做法，请勿随意修改以下编码逻辑！ + ============================================================ + 经过多轮调试（2026-02-28），以下参数组合为唯一验证成功的方案： + + 1. 图片格式：PNG（format="PNG"） + - 不可改为 JPEG —— API 会校验 Content-Type，抓包确认服务端使用 image/png + - 不可使用 base64 字符串 —— 会报 "expected a file, got a string" + - 不可使用 data URI —— 服务端不识别，返回 500 + + 2. 图片尺寸：必须与视频分辨率完全一致（target_size） + - 不可缩放降采样 —— 会报 "Inpaint image must match the requested width and height" + - 尺寸由 _fit_image_to_target() 保证（等比缩放 + 居中裁剪） + + 3. 上传方式：由调用方（sora_client.py）以 multipart/form-data 文件字段上传 + - filename="reference.png", content_type="image/png" + - 不可改回 application/json —— 服务端校验 input_reference 必须为 file 类型 + ============================================================ + + Args: + image: PIL Image 对象 + target_size: 目标分辨率字符串 "WxH"（如 "720x1280"） + + Returns: + PNG 格式的二进制字节 + """ + from io import BytesIO + + # 统一转换为 RGB（去除透明通道及其他模式） + if image.mode != "RGB": + image = image.convert("RGB") + + # 适配到目标分辨率（等比缩放 + 居中裁剪） + # ⚠️ 必须保持此尺寸不变，API 强制要求参考图片与视频分辨率完全一致 + if target_size: + image = _fit_image_to_target(image, target_size) + + # ⚠️ 必须使用 PNG 格式，不可改为 JPEG 或其他格式 + buffered = BytesIO() + image.save(buffered, format="PNG") + size_kb = buffered.tell() / 1024 + print(f"Sora: 参考图片编码为 PNG，{size_kb:.0f} KB ({image.size[0]}x{image.size[1]})") + return buffered.getvalue() + + +class SoraVideo: + """ + Sora 视频生成节点 + + 功能： + - 文生视频：基于提示词生成视频 + - 图生视频：基于参考图片和提示词生成视频 + - 异步轮询：自动等待生成完成并下载 + """ + + def __init__(self): + self.client = None + + @classmethod + def INPUT_TYPES(cls): + from ..models_config import SECONDS_DISPLAY_MAP, RESOLUTION_DISPLAY_MAP + + enabled_models = get_enabled_sora_models() + if not enabled_models: + enabled_models = ["请在 models_config.py 中启用至少一个 Sora 模型"] + + # 构建秒数选项列表（按数字顺序排序） + # 格式: ["4", "8", "10", "12", "15", "25(pro)"] + all_seconds_display = [] + seen_seconds = set() + for model_id in enabled_models: + supported = get_sora_supported_seconds(model_id) + for s in supported: + if s not in seen_seconds: + seen_seconds.add(s) + display = SECONDS_DISPLAY_MAP.get(s, str(s)) + all_seconds_display.append((s, display)) + # 按秒数数值排序 + all_seconds_display = sorted(all_seconds_display, key=lambda x: x[0]) + seconds_options = [d for _, d in all_seconds_display] if all_seconds_display else ["4", "8", "12"] + + # 构建分辨率选项列表（去重） + # 格式: ["720P", "1080P"] + seen_resolutions = set() + for model_id in enabled_models: + supported = get_sora_supported_sizes(model_id) + for size in supported: + if size in RESOLUTION_DISPLAY_MAP: + res_name, _ = RESOLUTION_DISPLAY_MAP[size] + seen_resolutions.add(res_name) + resolution_options = sorted(list(seen_resolutions)) if seen_resolutions else ["720P"] + + return { + "required": { + "prompt": ("STRING", { + "default": "A calico cat playing a piano on stage", + "multiline": True, + }), + "模型": (enabled_models, { + "default": enabled_models[0], + }), + "分辨率": (resolution_options, { + "default": resolution_options[0] if resolution_options else "720P", + }), + "宽高比": (["竖屏", "横屏"], { + "default": "竖屏", + }), + "视频时长": (seconds_options, { + "default": seconds_options[0] if seconds_options else "4", + }), + "生成数量": ("INT", { + "default": 1, + "min": 1, + "max": 10, + "step": 1, + }), + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 0xffffffffffffffff + }), + }, + "optional": { + "参考图片": ("IMAGE",), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("预览视频",) + FUNCTION = "generate_video" + CATEGORY = "video/generation" + + DESCRIPTION = ( + "Sora 视频生成节点。\n" + "支持文生视频和图生视频，自动轮询任务状态并下载视频。\n" + "视频保存到 ComfyUI/output/video/ 目录。\n\n" + "【模型说明】\n" + "• sora-2：官方模型，支持 4/8/12秒、720P 分辨率\n" + "• sora-2-pro：增强模型，支持全时长（含25秒）、1080P 分辨率\n\n" + "【时长说明】\n" + "• 25(pro)：仅 sora-2-pro 支持的25秒时长\n\n" + "【分辨率说明】\n" + "• 720P：sora-2 和 sora-2-pro 均支持\n" + "• 1080P：仅 sora-2-pro 支持的高清分辨率" + ) + + def generate_video( + self, + prompt: str, + 模型: str, + **kwargs, + ) -> Tuple[str]: + from ..models_config import SECONDS_DISPLAY_MAP, RESOLUTION_DISPLAY_MAP + + 视频时长_display = kwargs.pop("视频时长", "4") + 分辨率_display = kwargs.pop("分辨率", "720P") + 宽高比 = kwargs.pop("宽高比", "竖屏") + 生成数量 = kwargs.pop("生成数量", 1) + seed = kwargs.pop("seed", 0) + start_time = time.time() + + # 解析秒数显示值（如 "25(pro)" → 25） + seconds = 4 # 默认 + for actual, display in SECONDS_DISPLAY_MAP.items(): + if display == 视频时长_display: + seconds = actual + break + # 如果找不到映射，尝试直接解析数字 + if seconds == 4 and 视频时长_display != "4": + try: + seconds = int(视频时长_display.replace("(pro)", "")) + except ValueError: + seconds = 4 + + # 根据分辨率和宽高比确定实际分辨率值 + 分辨率 = "720x1280" # 默认 + for actual, (res_name, orientation) in RESOLUTION_DISPLAY_MAP.items(): + if res_name == 分辨率_display and orientation == 宽高比: + 分辨率 = actual + break + + # 检查参考图片 + ref_image = kwargs.get("参考图片") + ref_image_bytes = None + if ref_image is not None: + pil_images = tensor_to_pil(ref_image) + if pil_images: + ref_image_bytes = _compress_image_for_upload(pil_images[0], target_size=分辨率) + + mode_str = "图生视频 (含参考图)" if ref_image_bytes else "文生视频" + # 获取用户友好的显示值用于日志 + seconds_display = SECONDS_DISPLAY_MAP.get(seconds, str(seconds)) + res_display = f"{分辨率_display} {宽高比}" + if 生成数量 > 1: + print(f"Sora: {mode_str} | 并发{生成数量}个 | {模型} | {seconds_display} | {res_display}") + else: + print(f"Sora: {mode_str} | {模型} | {seconds_display} | {res_display}") + + # 校验参数兼容性 + supported_seconds = get_sora_supported_seconds(模型) + if supported_seconds and seconds not in supported_seconds: + # 构建带标签的支持时长列表 + supported_labels = [] + for s in supported_seconds: + display = SECONDS_DISPLAY_MAP.get(s, str(s)) + supported_labels.append(display) + raise ValueError( + f"时长 {SECONDS_DISPLAY_MAP.get(seconds, str(seconds))} 与模型 \"{模型}\" 不兼容！\n" + f"该模型支持的时长: {', '.join(supported_labels)}" + ) + + supported_sizes = get_sora_supported_sizes(模型) + if supported_sizes and 分辨率 not in supported_sizes: + # 检查该分辨率是否为Pro独占 + pro_only_sizes = ["1024x1792", "1792x1024"] + _, orientation = RESOLUTION_DISPLAY_MAP.get(分辨率, (分辨率, "")) + extra_hint = f"\n提示：1080P {orientation} 为 sora-2-pro 独占，请切换模型或选择720P。" if 分辨率 in pro_only_sizes else "" + raise ValueError( + f"分辨率 \"{分辨率_display} {宽高比}\" 与模型 \"{模型}\" 不兼容！" + f"支持的分辨率: {', '.join(supported_sizes)}" + extra_hint + ) + + # 准备保存路径 + video_dir = _get_video_output_dir() + counter = _get_next_counter(video_dir, "sora") + + # ProgressBar + pbar = None + if PROGRESS_BAR_AVAILABLE: + pbar = ProgressBar(生成数量 if 生成数量 > 1 else 100) + + try: + if self.client is None: + self.client = SoraClient() + + if 生成数量 == 1: + # ── 单个视频：保留详细进度（提交→轮询→下载） + save_path = os.path.join(video_dir, f"sora_{counter:05d}.mp4") + last_progress = [0] + + def progress_callback(progress_pct: int): + print( + f"\rSora: 生成中... 进度: {progress_pct}%", + end="", flush=True + ) + if pbar is not None and progress_pct > last_progress[0]: + pbar.update(progress_pct - last_progress[0]) + last_progress[0] = progress_pct + + def on_stage(stage: str): + if stage == "submitting": + print("Sora: 正在提交视频生成任务...") + elif stage.startswith("submitted:"): + vid = stage.split(":", 1)[1] + print(f"Sora: 视频任务已提交，ID: {vid}") + elif stage == "polling": + print("Sora: 等待视频生成...") + elif stage == "downloading": + print("") # 换行（结束 \r 行） + print("Sora: 视频生成完成，正在下载...") + + result_path = self.client.generate_video_sync( + prompt=prompt, + model=模型, + seconds=seconds, + size=分辨率, + save_path=save_path, + input_reference_bytes=ref_image_bytes, + seed=seed, + progress_callback=progress_callback, + on_stage=on_stage, + ) + result_paths = [result_path] + + else: + # ── 批量并发：同时提交多个任务 + save_paths = [ + os.path.join(video_dir, f"sora_{counter + i:05d}.mp4") + for i in range(生成数量) + ] + success_count = [0] + fail_count = [0] + + def batch_progress_callback(current: int, total: int, success: bool, error_msg): + if success: + success_count[0] += 1 + print(f"Sora: 第 {current}/{total} 个视频完成 ✓") + else: + fail_count[0] += 1 + print(f"Sora: 第 {current}/{total} 个视频失败 ✗") + if error_msg: + print(f"原始错误详情:\n{error_msg}") + if pbar is not None: + pbar.update(1) + + print(f"Sora: 正在并发提交 {生成数量} 个视频任务，请耐心等待...") + result_paths = self.client.generate_batch_videos_sync( + prompt=prompt, + model=模型, + seconds=seconds, + size=分辨率, + save_paths=save_paths, + input_reference_bytes=ref_image_bytes, + seed=seed, + progress_callback=batch_progress_callback, + ) + + elapsed = time.time() - start_time + time_str = f"{elapsed:.2f}s" if elapsed >= 1 else f"{elapsed:.3f}s" + print(f"Sora: 完成！总耗时 {time_str} | 已生成 {len(result_paths)} 个视频") + for p in result_paths: + print(f" → {p}") + + output_path = "\n".join(result_paths) + return (output_path,) + + except ValueError as e: + error_msg = str(e) + print(f"\nSora: ❌ {error_msg}") + raise ValueError(error_msg) from None + + except RuntimeError as e: + error_msg = str(e) + print(f"\nSora: ❌ {error_msg}") + raise RuntimeError(error_msg) from None + + except Exception as e: + error_msg = str(e) + print(f"\nSora: ❌ {error_msg}") + raise type(e)(error_msg) from None + + finally: + if self.client is not None: + try: + balance_data = self.client.query_balance_sync() + balance_info = self.client.format_balance_info(balance_data) + print(f"Sora: {balance_info}") + except Exception: + pass diff --git a/nodes/universal_llm.py b/nodes/universal_llm.py new file mode 100644 index 0000000..0f445cb --- /dev/null +++ b/nodes/universal_llm.py @@ -0,0 +1,304 @@ +""" +全能LLM对话助手节点 +ComfyUI 自定义节点，通过 OpenAI 兼容协议调用市面上主流的 AI 对话大模型 +支持多模态（图片输入），单轮对话，非流式输出 + +API 密钥和地址通过插件统一配置（环境变量或 .config 文件），与 Google Gemini 节点一致 +""" + +import time +import base64 +import json +from io import BytesIO +from typing import Optional, Tuple + +import torch +from PIL import Image + +from ..utils.image_utils import tensor_to_pil +from ..utils.config import get_api_key_or_raise, get_api_base_url + +# ============================================================================ +# 模型配置 +# ============================================================================ + +SUPPORTED_MODELS = [ + "gpt-5.4", + "gemini-3-flash-preview", + "gemini-3.1-flash-lite-preview", + "gemini-3.1-pro-preview", + "deepseek-v3.2", + "kimi-k2.5", + "doubao-seed-2-0-pro-260215", + "qwen3.5-plus-2026-02-15", + "qwen3.5-plus", +] + +# 图片缩放最大尺寸 +MAX_IMAGE_DIMENSION = 1568 + +# 图片最大文件大小（20MB） +MAX_IMAGE_SIZE = 20 * 1024 * 1024 + + +class UniversalLLMChat: + """ + 全能LLM对话助手 + + 功能： + - 通过 OpenAI 兼容协议调用主流大模型 + - 支持多模态（图片输入） + - 单轮对话，非流式输出 + - API 密钥和地址继承插件统一配置 + """ + + def __init__(self): + self._api_key = None + self._base_url = None + + def _ensure_config(self): + """延迟加载配置，首次调用时初始化""" + if self._api_key is None: + self._api_key = get_api_key_or_raise("O1KEY_API_KEY") + self._base_url = get_api_base_url() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "模型": (SUPPORTED_MODELS, { + "default": SUPPORTED_MODELS[0] + }), + "提示词": ("STRING", { + "default": "", + "multiline": True, + }), + }, + "optional": { + "图片": ("IMAGE",), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("回复",) + FUNCTION = "generate" + CATEGORY = "text/generation" + OUTPUT_NODE = True + + def _resize_image(self, img: Image.Image) -> Image.Image: + """如果图片过长边超过限制，等比缩放""" + w, h = img.size + max_dim = max(w, h) + if max_dim > MAX_IMAGE_DIMENSION: + scale = MAX_IMAGE_DIMENSION / max_dim + new_w, new_h = int(w * scale), int(h * scale) + print(f"全能LLM: 图片缩放 {w}x{h} -> {new_w}x{new_h}") + return img.resize((new_w, new_h), Image.Resampling.LANCZOS) + return img + + def _image_to_data_url(self, img: Image.Image) -> str: + """将 PIL Image 转为 data URL（JPEG base64）""" + img = self._resize_image(img) + if img.mode in ('RGBA', 'P'): + img = img.convert('RGB') + + for quality in [92, 82, 72, 60, 45]: + buf = BytesIO() + img.save(buf, format='JPEG', quality=quality, optimize=True) + data = buf.getvalue() + if len(data) <= MAX_IMAGE_SIZE: + b64 = base64.b64encode(data).decode('utf-8') + return f"data:image/jpeg;base64,{b64}" + + b64 = base64.b64encode(data).decode('utf-8') + return f"data:image/jpeg;base64,{b64}" + + def _build_messages( + self, + prompt: str, + images: Optional[torch.Tensor] = None, + ) -> list: + """构建 OpenAI 格式的 messages 数组""" + image_data_urls = [] + pil_images_cache = [] # 保留 PIL Image 用于总体积重新编码 + + if images is not None: + pil_images = tensor_to_pil(images) + for img in pil_images: + img_resized = self._resize_image(img) + if img_resized.mode in ('RGBA', 'P'): + img_resized = img_resized.convert('RGB') + pil_images_cache.append(img_resized) + image_data_urls.append(self._image_to_data_url(img_resized)) + + # 多图总体积控制 + if pil_images_cache and len(pil_images_cache) > 1: + total_bytes = sum( + len(base64.b64decode(url.split(',', 1)[1])) for url in image_data_urls + ) + if total_bytes > MAX_IMAGE_SIZE: + print(f"全能LLM: 图片总体积 {total_bytes / 1024 / 1024:.2f}MB 超过 {MAX_IMAGE_SIZE // 1024 // 1024}MB 限制，正在压缩...") + + # 降质量 + compressed = False + for quality in [80, 70, 60, 50, 40, 30, 20]: + new_urls = [] + for img in pil_images_cache: + buf = BytesIO() + img.save(buf, format='JPEG', quality=quality, optimize=True) + b64 = base64.b64encode(buf.getvalue()).decode('utf-8') + new_urls.append(f"data:image/jpeg;base64,{b64}") + total_bytes = sum(len(base64.b64decode(u.split(',', 1)[1])) for u in new_urls) + if total_bytes <= MAX_IMAGE_SIZE: + image_data_urls = new_urls + print(f"全能LLM: 图片压缩完成，总体积 {total_bytes / 1024 / 1024:.2f}MB ({len(pil_images_cache)}张图片，质量{quality})") + compressed = True + break + + # 降分辨率 + if not compressed: + for scale in [0.75, 0.5, 0.35]: + new_urls = [] + for img in pil_images_cache: + w, h = img.size + resized = img.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS) + buf = BytesIO() + resized.save(buf, format='JPEG', quality=20, optimize=True) + b64 = base64.b64encode(buf.getvalue()).decode('utf-8') + new_urls.append(f"data:image/jpeg;base64,{b64}") + total_bytes = sum(len(base64.b64decode(u.split(',', 1)[1])) for u in new_urls) + if total_bytes <= MAX_IMAGE_SIZE: + image_data_urls = new_urls + print(f"全能LLM: 图片压缩完成，总体积 {total_bytes / 1024 / 1024:.2f}MB ({len(pil_images_cache)}张图片，缩放{int(scale*100)}%)") + compressed = True + break + + if not compressed: + print(f"全能LLM: 无法将 {len(pil_images_cache)} 张图片压缩到 {MAX_IMAGE_SIZE // 1024 // 1024}MB 以内，请减少图片数量或降低分辨率") + raise ValueError(f"图片总体积 {total_bytes / 1024 / 1024:.2f}MB 超过限制，无法压缩到 {MAX_IMAGE_SIZE // 1024 // 1024}MB 以内") + + if not image_data_urls: + return [{"role": "user", "content": prompt}] + + content_parts = [] + for url in image_data_urls: + content_parts.append({ + "type": "image_url", + "image_url": {"url": url} + }) + content_parts.append({ + "type": "text", + "text": prompt + }) + + return [{"role": "user", "content": content_parts}] + + def generate( + self, + 模型: str, + 提示词: str, + 图片: Optional[torch.Tensor] = None, + ) -> Tuple[str]: + start_time = time.time() + + try: + self._ensure_config() + + # 构建 messages + messages = self._build_messages(提示词, 图片) + + img_count = len(tensor_to_pil(图片)) if 图片 is not None else 0 + input_desc = "文本" + (f" + {img_count}张图片" if img_count > 0 else "") + + print(f"全能LLM: 模型 = {模型}") + print(f"全能LLM: 输入 = {input_desc}") + + # 构建请求体 + request_body = { + "model": 模型, + "messages": messages, + "stream": False, + } + + # 发送请求（在独立线程中运行异步请求，避免与 ComfyUI 事件循环冲突） + import aiohttp + import asyncio + from concurrent.futures import ThreadPoolExecutor + + async def _do_request(): + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {self._api_key}", + } + url = f"{self._base_url}/v1/chat/completions" + timeout = aiohttp.ClientTimeout(total=120) + + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.post(url, headers=headers, json=request_body) as resp: + status = resp.status + body = await resp.text() + + if status != 200: + try: + err_data = json.loads(body) + err_msg = err_data.get("error", {}).get("message", body[:200]) + except Exception: + err_msg = body[:200] + + if status == 401: + raise ValueError(f"认证失败：API Key 无效或已过期") + elif status == 403: + raise ValueError(f"无权访问模型 {模型}") + elif status == 429: + raise ValueError(f"请求频率超限，请稍后重试") + elif status == 404: + raise ValueError(f"模型 {模型} 不存在或 API 地址错误") + else: + raise RuntimeError(f"API 错误 ({status}): {err_msg}") + + return json.loads(body) + + def _run_in_thread(): + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete(_do_request()) + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as pool: + response_data = pool.submit(_run_in_thread).result() + + # 解析响应 + choices = response_data.get("choices", []) + if not choices: + raise RuntimeError("API 返回了空响应（无 choices）") + + reply = choices[0].get("message", {}).get("content", "") + + # Token 用量 + usage = response_data.get("usage", {}) + prompt_tokens = usage.get("prompt_tokens", 0) + completion_tokens = usage.get("completion_tokens", 0) + total_tokens = usage.get("total_tokens", 0) + + elapsed = time.time() - start_time + print(f"全能LLM: 生成完成 (耗时: {elapsed:.2f}s)") + print(f"全能LLM: Token 用量 — 输入: {prompt_tokens}, 输出: {completion_tokens}, 合计: {total_tokens}") + if reply: + preview = reply[:100] + "..." if len(reply) > 100 else reply + print(f"全能LLM: 回复预览: {preview}") + + return (reply,) + + except ValueError as e: + if str(e) == "未授权！": + print("全能LLM: 请联系作者授权后方可使用！") + raise ValueError("未授权！") from None + error_msg = str(e).split('\n')[0] + print(f"全能LLM: ❌ {error_msg}") + raise + + except Exception as e: + error_msg = str(e).split('\n')[0] + print(f"全能LLM: ❌ {error_msg}") + raise RuntimeError(error_msg) from None diff --git a/nodes/veo_video.py b/nodes/veo_video.py new file mode 100644 index 0000000..922774c --- /dev/null +++ b/nodes/veo_video.py @@ -0,0 +1,422 @@ +""" +Google Veo 视频生成节点 +ComfyUI 自定义节点，调用 Veo API 生成视频 +""" + +import os +import re +import time +from typing import Optional, Tuple + +import torch + +from ..utils.image_utils import tensor_to_pil +from ..clients.veo_client import VeoClient +from ..models_config import ( + get_enabled_veo_models, + VEO_MODELS, + VEO_RESOLUTION_MAP, +) + +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + PROGRESS_BAR_AVAILABLE = False + print("⚠️ GoogleVeo: comfy.utils.ProgressBar 不可用，将只使用终端进度显示") + + +def _get_video_output_dir() -> str: + """获取视频输出目录: ComfyUI/output/video""" + if FOLDER_PATHS_AVAILABLE: + base = folder_paths.get_output_directory() + else: + plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + base = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output") + video_dir = os.path.join(base, "video") + os.makedirs(video_dir, exist_ok=True) + return video_dir + + +def _get_next_counter(directory: str, prefix: str) -> int: + """扫描目录，获取下一个可用的文件计数器""" + if not os.path.exists(directory): + return 1 + pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)") + max_counter = 0 + for f in os.listdir(directory): + m = pattern.match(f) + if m: + max_counter = max(max_counter, int(m.group(1))) + return max_counter + 1 + + +def _fit_image_to_target(image, target_size: str): + """ + 将参考图片按 "等比缩放覆盖 + 居中裁剪" 策略适配到目标分辨率。 + """ + from PIL import Image as PILImage + + parts = target_size.lower().split("x") + target_w, target_h = int(parts[0]), int(parts[1]) + + src_w, src_h = image.size + src_ratio = src_w / src_h + target_ratio = target_w / target_h + + if abs(src_ratio - target_ratio) < 0.01 and src_w <= target_w and src_h <= target_h: + return image + + print(f"Veo: 参考图片 {src_w}x{src_h} (比例 {src_ratio:.2f}) → 目标 {target_w}x{target_h} (比例 {target_ratio:.2f})") + + resample = PILImage.Resampling.LANCZOS if hasattr(PILImage, "Resampling") else PILImage.LANCZOS + + if src_ratio > target_ratio: + scale = target_h / src_h + new_w = round(src_w * scale) + new_h = target_h + image = image.resize((new_w, new_h), resample=resample) + left = (new_w - target_w) // 2 + image = image.crop((left, 0, left + target_w, target_h)) + else: + scale = target_w / src_w + new_w = target_w + new_h = round(src_h * scale) + image = image.resize((new_w, new_h), resample=resample) + top = (new_h - target_h) // 2 + image = image.crop((0, top, target_w, top + target_h)) + + print(f"Veo: 参考图片已适配为 {image.size[0]}x{image.size[1]}") + return image + + +def _compress_image_to_bytes(image, target_size: Optional[str] = None) -> bytes: + """ + 将 PIL Image 适配目标分辨率并编码为 PNG 字节 + """ + from io import BytesIO + + if image.mode != "RGB": + image = image.convert("RGB") + + if target_size: + image = _fit_image_to_target(image, target_size) + + buffered = BytesIO() + image.save(buffered, format="PNG") + size_kb = buffered.tell() / 1024 + print(f"Veo: 参考图片编码为 PNG，{size_kb:.0f} KB ({image.size[0]}x{image.size[1]})") + return buffered.getvalue() + + +class GoogleVeo: + """ + Google Veo 视频生成节点 + + 功能： + - 文生视频：基于提示词生成视频 + - 图生视频：基于首帧/尾帧/参考图生成视频 + - 异步轮询：自动等待生成完成并下载 + """ + + def __init__(self): + self.client = None + + @classmethod + def INPUT_TYPES(cls): + enabled_models = get_enabled_veo_models() + if not enabled_models: + enabled_models = ["请在 models_config.py 中启用 Veo 模型"] + + # 分辨率选项 + resolution_options = ["720p", "1080p", "4K"] + + # 宽高比选项 + aspect_ratio_options = ["16:9", "9:16"] + + # 视频秒数选项 + seconds_options = ["4", "6", "8"] + + return { + "required": { + "prompt": ("STRING", { + "default": "A calico cat playing a piano on stage", + "multiline": True, + }), + "模型": (enabled_models, { + "default": enabled_models[0] if enabled_models else "Veo3.1", + }), + "分辨率": (resolution_options, { + "default": "720p", + }), + "宽高比": (aspect_ratio_options, { + "default": "9:16", + }), + "视频时长": (seconds_options, { + "default": "8", + }), + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 0xffffffffffffffff, + }), + "生成数量": ("INT", { + "default": 1, + "min": 1, + "max": 10, + "step": 1, + }), + }, + "optional": { + "首帧": ("IMAGE",), + "尾帧": ("IMAGE",), + "参考图": ("IMAGE",), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("预览视频",) + FUNCTION = "generate_video" + CATEGORY = "video/generation" + + DESCRIPTION = ( + "Google Veo 视频生成节点。\n" + "支持文生视频和图生视频（图生视频支持首帧、尾帧、参考图）。\n" + "视频保存到 ComfyUI/output/video/ 目录。\n\n" + "【模型说明】\n" + "• Veo3.1：Google 最新视频生成模型\n\n" + "【分辨率说明】\n" + "• 720p：标清\n" + "• 1080p：高清\n" + "• 4K：超高清\n\n" + "【时长说明】\n" + "• 4秒：短视频\n" + "• 6秒：标准\n" + "• 8秒：长视频（默认）\n\n" + "【图生视频说明】\n" + "• 首帧：视频开始的第一帧图像\n" + "• 尾帧：视频结束时的最后一帧图像\n" + "• 参考图：参考图像（与首帧/尾帧配合使用）\n" + "• 至少需要提供首帧或参考图之一" + ) + + def generate_video( + self, + prompt: str, + 模型: str, + **kwargs, + ) -> Tuple[str]: + 分辨率 = kwargs.pop("分辨率", "720p") + 宽高比 = kwargs.pop("宽高比", "9:16") + 视频时长 = kwargs.pop("视频时长", "8") + seed = kwargs.pop("seed", 0) + 生成数量 = kwargs.pop("生成数量", 1) + start_time = time.time() + + # 解析视频时长 + seconds = int(视频时长) + + # 解析分辨率和宽高比，映射到模型名称 + size_key = f"{分辨率}_{宽高比}" + actual_size = VEO_RESOLUTION_MAP.get(size_key) + if not actual_size: + # 默认值 + actual_size = "720x1280" # 720p 9:16 + + # 检查是否有参考图输入 + 首帧 = kwargs.get("首帧") + 尾帧 = kwargs.get("尾帧") + 参考图 = kwargs.get("参考图") + + has_image = 首帧 is not None or 尾帧 is not None or 参考图 is not None + + # 根据是否有图片选择模型前缀 + if has_image: + model_prefix = "veo3.1" + else: + model_prefix = "veo3.1" + + # 构建完整模型名称 + # 格式: veo3.1-portrait / veo3.1-landscape / veo3.1-portrait-fl / veo3.1-landscape-fl 等 + if 分辨率 == "720p": + res_suffix = "" + if 宽高比 == "9:16": + orientation = "portrait" + else: + orientation = "landscape" + elif 分辨率 == "1080p": + res_suffix = "-hd" + if 宽高比 == "9:16": + orientation = "portrait" + else: + orientation = "landscape" + else: # 4K + res_suffix = "-4k" + if 宽高比 == "9:16": + orientation = "portrait" + else: + orientation = "landscape" + + # 图生视频添加 -fl 后缀 + if has_image: + model_suffix = f"-{orientation}-fl{res_suffix}" + else: + model_suffix = f"-{orientation}{res_suffix}" + + model = f"{model_prefix}{model_suffix}" + + # 准备图片字节 + first_frame_bytes = None + last_frame_bytes = None + reference_bytes = None + + if 首帧 is not None: + pil_images = tensor_to_pil(首帧) + if pil_images: + first_frame_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size) + + if 尾帧 is not None: + pil_images = tensor_to_pil(尾帧) + if pil_images: + last_frame_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size) + + if 参考图 is not None: + pil_images = tensor_to_pil(参考图) + if pil_images: + reference_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size) + + mode_str = "图生视频" if has_image else "文生视频" + print(f"Veo: {mode_str} | 并发{生成数量}个 | 模型: {model} | {seconds}秒 | {分辨率} {宽高比}") + + # 准备保存路径 + video_dir = _get_video_output_dir() + counter = _get_next_counter(video_dir, "veo") + + # ProgressBar + pbar = None + if PROGRESS_BAR_AVAILABLE: + pbar = ProgressBar(生成数量 if 生成数量 > 1 else 100) + + try: + if self.client is None: + self.client = VeoClient() + + if 生成数量 == 1: + save_path = os.path.join(video_dir, f"veo_{counter:05d}.mp4") + last_progress = [0] + + def progress_callback(progress_pct: int): + print( + f"\rVeo: 生成中... 进度: {progress_pct}%", + end="", flush=True + ) + if pbar is not None and progress_pct > last_progress[0]: + pbar.update(progress_pct - last_progress[0]) + last_progress[0] = progress_pct + + def on_stage(stage: str): + if stage == "submitting": + print("Veo: 正在提交视频生成任务...") + elif stage.startswith("submitted:"): + vid = stage.split(":", 1)[1] + print(f"Veo: 视频任务已提交，ID: {vid}") + elif stage == "polling": + print("Veo: 等待视频生成...") + elif stage == "downloading": + print("") + print("Veo: 视频生成完成，正在下载...") + + result_path = self.client.generate_video_sync( + prompt=prompt, + model=model, + seconds=seconds, + size=actual_size, + save_path=save_path, + first_frame_bytes=first_frame_bytes, + last_frame_bytes=last_frame_bytes, + reference_bytes=reference_bytes, + seed=seed, + progress_callback=progress_callback, + on_stage=on_stage, + ) + result_paths = [result_path] + + else: + save_paths = [ + os.path.join(video_dir, f"veo_{counter + i:05d}.mp4") + for i in range(生成数量) + ] + success_count = [0] + + def batch_progress_callback(current: int, total: int, success: bool, error_msg): + if success: + success_count[0] += 1 + print(f"Veo: 第 {current}/{total} 个视频完成 ✓") + else: + print(f"Veo: 第 {current}/{total} 个视频失败 ✗") + if error_msg: + print(f"原始错误详情:\n{error_msg}") + if pbar is not None: + pbar.update(1) + + print(f"Veo: 正在并发提交 {生成数量} 个视频任务，请耐心等待...") + result_paths = self.client.generate_batch_videos_sync( + prompt=prompt, + model=model, + seconds=seconds, + size=actual_size, + save_paths=save_paths, + first_frame_bytes=first_frame_bytes, + last_frame_bytes=last_frame_bytes, + reference_bytes=reference_bytes, + seed=seed, + progress_callback=batch_progress_callback, + ) + + elapsed = time.time() - start_time + time_str = f"{elapsed:.2f}s" if elapsed >= 1 else f"{elapsed:.3f}s" + print(f"Veo: 完成！总耗时 {time_str} | 已生成 {len(result_paths)} 个视频") + for p in result_paths: + print(f" → {p}") + + output_path = "\n".join(result_paths) + return (output_path,) + + except ValueError as e: + error_msg = str(e) + print(f"\nVeo: ❌ {error_msg}") + raise ValueError(error_msg) from None + + except RuntimeError as e: + error_msg = str(e) + print(f"\nVeo: ❌ {error_msg}") + raise RuntimeError(error_msg) from None + + except Exception as e: + error_msg = str(e) + print(f"\nVeo: ❌ {error_msg}") + raise type(e)(error_msg) from None + + finally: + if self.client is not None: + try: + balance_data = self.client.query_balance_sync() + balance_info = self.client.format_balance_info(balance_data) + print(f"Veo: {balance_info}") + except Exception: + pass + + +NODE_CLASS_MAPPINGS = { + "GoogleVeo": GoogleVeo, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "GoogleVeo": "Google Veo - ab", +} diff --git a/nodes/video_preview.py b/nodes/video_preview.py new file mode 100644 index 0000000..d660598 --- /dev/null +++ b/nodes/video_preview.py @@ -0,0 +1,107 @@ +""" +通用视频预览节点 +ComfyUI 自定义节点，接收视频文件路径并在前端展示预览 +""" + +import os + +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +SUPPORTED_VIDEO_EXTENSIONS = {".mp4", ".webm", ".mov", ".avi", ".mkv", ".flv", ".wmv", ".3gp"} + + +def _get_output_dir() -> str: + if FOLDER_PATHS_AVAILABLE: + return folder_paths.get_output_directory() + plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + return os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output") + + +class VideoPreview: + """ + 通用视频预览节点 + + 功能： + - 接收视频文件路径（STRING） + - 在 ComfyUI 前端节点上内嵌