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 <[email protected]>
This commit is contained in:
o1key
+756
@@ -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 格式：`<type>: <description>`
|
||||
|
||||
### 常用类型
|
||||
- `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`
|
||||
|
||||
更新脚本会自动：
|
||||
- 检查远程更新
|
||||
- 备份配置文件
|
||||
- 拉取最新代码
|
||||
- 更新依赖包
|
||||
- 显示更新日志
|
||||
|
||||
---
|
||||
@@ -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
|
||||
@@ -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
|
||||
+27
@@ -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
|
||||
+601
@@ -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 密钥
|
||||
- 完善的错误处理和日志输出
|
||||
@@ -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.
|
||||
@@ -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**
|
||||
+108
@@ -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']
|
||||
@@ -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']
|
||||
@@ -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}"
|
||||
@@ -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}
|
||||
@@ -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 格式: 
|
||||
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"<base64 data, {len(obj)} chars>"
|
||||
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"<base64 data, {len(obj)} chars>"
|
||||
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)
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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/<model-name>:generateContent
|
||||
# 提取 <model-name> 部分
|
||||
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,<data>
|
||||
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 图片链接：
|
||||
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},<base64 {len(data)} chars>"
|
||||
if len(obj) > 200 and all(
|
||||
c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
|
||||
for c in obj[:64]
|
||||
):
|
||||
return f"<base64 {len(obj)} chars>"
|
||||
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},<base64 {len(data)} chars>"
|
||||
if len(obj) > 200 and all(
|
||||
c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
|
||||
for c in obj[:64]
|
||||
):
|
||||
return f"<base64 {len(obj)} chars>"
|
||||
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)
|
||||
@@ -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}"
|
||||
@@ -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}"
|
||||
@@ -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 文件")
|
||||
@@ -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']
|
||||
@@ -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)
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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()
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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),)
|
||||
@@ -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": "图片宽高比预设",
|
||||
}
|
||||
@@ -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
|
||||
@@ -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}}
|
||||
@@ -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: 最终内存清理完成")
|
||||
@@ -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()
|
||||
@@ -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"全能生图: 最终内存清理完成")
|
||||
@@ -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,)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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 前端节点上内嵌 <video> 播放器进行预览
|
||||
- 支持 mp4, webm, mov, avi, mkv 等主流格式
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"预览视频": ("STRING", {"forceInput": True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "preview"
|
||||
CATEGORY = "video"
|
||||
|
||||
DESCRIPTION = (
|
||||
"通用视频预览节点。\n"
|
||||
"接收视频文件路径，在节点上显示视频播放器。\n"
|
||||
"支持 mp4, webm, mov, avi, mkv 等主流视频格式。"
|
||||
)
|
||||
|
||||
def preview(self, **kwargs) -> dict:
|
||||
video_path = kwargs.get("预览视频", "")
|
||||
if not video_path or not video_path.strip():
|
||||
raise ValueError("视频路径为空")
|
||||
|
||||
video_path = video_path.strip()
|
||||
|
||||
if not os.path.isfile(video_path):
|
||||
raise ValueError(f"视频文件不存在: {video_path}")
|
||||
|
||||
ext = os.path.splitext(video_path)[1].lower()
|
||||
if ext not in SUPPORTED_VIDEO_EXTENSIONS:
|
||||
raise ValueError(
|
||||
f"不支持的视频格式 '{ext}'，"
|
||||
f"支持: {', '.join(sorted(SUPPORTED_VIDEO_EXTENSIONS))}"
|
||||
)
|
||||
|
||||
output_dir = _get_output_dir()
|
||||
abs_video = os.path.abspath(video_path)
|
||||
abs_output = os.path.abspath(output_dir)
|
||||
|
||||
if abs_video.startswith(abs_output):
|
||||
rel_path = os.path.relpath(abs_video, abs_output)
|
||||
subfolder = os.path.dirname(rel_path).replace("\\", "/")
|
||||
filename = os.path.basename(rel_path)
|
||||
file_type = "output"
|
||||
else:
|
||||
filename = os.path.basename(abs_video)
|
||||
subfolder = ""
|
||||
file_type = "output"
|
||||
|
||||
# 如果文件不在 output 目录下，复制一份到 output/video/
|
||||
target_dir = os.path.join(output_dir, "video")
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
target_path = os.path.join(target_dir, filename)
|
||||
|
||||
if not os.path.exists(target_path) or abs_video != os.path.abspath(target_path):
|
||||
import shutil
|
||||
shutil.copy2(abs_video, target_path)
|
||||
|
||||
subfolder = "video"
|
||||
|
||||
file_size = os.path.getsize(abs_video)
|
||||
size_mb = file_size / (1024 * 1024)
|
||||
print(f"视频预览: {filename} ({size_mb:.1f}MB)")
|
||||
|
||||
return {
|
||||
"ui": {
|
||||
"videos": [{
|
||||
"filename": filename,
|
||||
"subfolder": subfolder,
|
||||
"type": file_type,
|
||||
}],
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
aiohttp>=3.9.0
|
||||
Pillow>=10.0.0
|
||||
requests>=2.31.0
|
||||
@@ -0,0 +1,37 @@
|
||||
"""
|
||||
工具模块
|
||||
包含图像处理、配置管理、文件处理等通用工具函数
|
||||
"""
|
||||
|
||||
from .image_utils import (
|
||||
tensor_to_pil,
|
||||
pil_to_tensor,
|
||||
encode_image_to_base64,
|
||||
decode_base64_to_pil
|
||||
)
|
||||
from .config import load_config, get_api_key
|
||||
from .file_utils import (
|
||||
ImageInfo,
|
||||
load_images_from_folder,
|
||||
pair_images_indexed,
|
||||
pair_images_cartesian,
|
||||
generate_timestamp_filename,
|
||||
save_image,
|
||||
get_folder_image_count
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
'tensor_to_pil',
|
||||
'pil_to_tensor',
|
||||
'encode_image_to_base64',
|
||||
'decode_base64_to_pil',
|
||||
'load_config',
|
||||
'get_api_key',
|
||||
'ImageInfo',
|
||||
'load_images_from_folder',
|
||||
'pair_images_indexed',
|
||||
'pair_images_cartesian',
|
||||
'generate_timestamp_filename',
|
||||
'save_image',
|
||||
'get_folder_image_count'
|
||||
]
|
||||
+117
@@ -0,0 +1,117 @@
|
||||
"""
|
||||
配置管理模块
|
||||
处理环境变量和 API 密钥管理
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Dict, Optional
|
||||
|
||||
|
||||
# 获取插件根目录
|
||||
PLUGIN_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
CONFIG_FILE = os.path.join(PLUGIN_ROOT, ".config")
|
||||
|
||||
|
||||
# ============ API 基础配置 ============
|
||||
# 所有 API 客户端的统一基础 URL
|
||||
# 可通过环境变量 O1KEY_API_BASE_URL 覆盖
|
||||
DEFAULT_API_BASE_URL = "https://vip.o1key.com"
|
||||
|
||||
|
||||
def load_config(config_path: Optional[str] = None) -> Dict[str, str]:
|
||||
"""
|
||||
从配置文件加载所有配置项
|
||||
|
||||
Args:
|
||||
config_path: 配置文件路径，默认为插件目录下的 .config
|
||||
|
||||
Returns:
|
||||
配置字典 {key: value}
|
||||
|
||||
Example:
|
||||
>>> config = load_config()
|
||||
>>> api_key = config.get('O1KEY_API_KEY')
|
||||
"""
|
||||
if config_path is None:
|
||||
config_path = CONFIG_FILE
|
||||
|
||||
config = {}
|
||||
|
||||
if not os.path.exists(config_path):
|
||||
return config
|
||||
|
||||
try:
|
||||
with open(config_path, 'r', encoding='utf-8') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
|
||||
# 跳过空行和注释
|
||||
if not line or line.startswith('#'):
|
||||
continue
|
||||
|
||||
# 解析 KEY=VALUE 格式
|
||||
if '=' in line:
|
||||
key, value = line.split('=', 1)
|
||||
key = key.strip()
|
||||
value = value.strip().strip('"').strip("'")
|
||||
|
||||
if key and value:
|
||||
config[key] = value
|
||||
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取配置文件失败: {e}")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def get_api_key(key_name: str = "O1KEY_API_KEY") -> Optional[str]:
|
||||
"""
|
||||
获取 API 密钥
|
||||
从 .config 文件读取
|
||||
|
||||
Args:
|
||||
key_name: 密钥名称，默认为 O1KEY_API_KEY
|
||||
|
||||
Returns:
|
||||
API 密钥字符串，如果未找到则返回 None
|
||||
"""
|
||||
config = load_config()
|
||||
return config.get(key_name)
|
||||
|
||||
|
||||
def get_api_key_or_raise(key_name: str = "O1KEY_API_KEY") -> str:
|
||||
"""
|
||||
获取 API 密钥，如果未找到则抛出异常
|
||||
|
||||
Args:
|
||||
key_name: 密钥名称
|
||||
|
||||
Returns:
|
||||
API 密钥字符串
|
||||
|
||||
Raises:
|
||||
ValueError: 如果未找到 API 密钥
|
||||
"""
|
||||
api_key = get_api_key(key_name)
|
||||
|
||||
if not api_key:
|
||||
raise ValueError("未授权！")
|
||||
|
||||
return api_key
|
||||
|
||||
|
||||
def get_api_base_url() -> str:
|
||||
"""
|
||||
获取 API 基础 URL
|
||||
从 .config 文件读取，如果未配置则使用默认值
|
||||
|
||||
Returns:
|
||||
API 基础 URL 字符串
|
||||
"""
|
||||
config = load_config()
|
||||
base_url = config.get("O1KEY_API_BASE_URL")
|
||||
|
||||
if base_url:
|
||||
return base_url.rstrip('/')
|
||||
|
||||
return DEFAULT_API_BASE_URL
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
文件数据类型定义
|
||||
用于在 ComfyUI 节点间传递文件数据
|
||||
"""
|
||||
|
||||
from typing import NamedTuple
|
||||
|
||||
|
||||
class FileData(NamedTuple):
|
||||
"""
|
||||
文件数据类型，用于在节点间传递
|
||||
|
||||
Attributes:
|
||||
path: 文件完整路径
|
||||
filename: 文件名（不含扩展名）
|
||||
extension: 文件扩展名（如 .pdf）
|
||||
mime_type: MIME 类型
|
||||
data: Base64 编码的文件内容
|
||||
size: 文件大小（字节）
|
||||
"""
|
||||
path: str
|
||||
filename: str
|
||||
extension: str
|
||||
mime_type: str
|
||||
data: str
|
||||
size: int
|
||||
|
||||
|
||||
# 支持的文档 MIME 类型映射
|
||||
DOCUMENT_MIME_TYPES = {
|
||||
".pdf": "application/pdf",
|
||||
".txt": "text/plain"
|
||||
}
|
||||
|
||||
|
||||
# 文件大小限制（字节）
|
||||
FILE_SIZE_LIMITS = {
|
||||
".pdf": 50 * 1024 * 1024, # 50MB (Gemini API 官方限制)
|
||||
".txt": 20 * 1024 * 1024 # 20MB (保守限制)
|
||||
}
|
||||
@@ -0,0 +1,375 @@
|
||||
"""
|
||||
文件处理工具模块
|
||||
提供文件夹图片加载、智能命名、图片配对等功能
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
import time
|
||||
import random
|
||||
from datetime import datetime
|
||||
from itertools import product
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Optional, NamedTuple
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def _get_server_port() -> Optional[int]:
|
||||
"""获取当前 ComfyUI 实例的端口号，失败返回 None"""
|
||||
try:
|
||||
import comfy.cli_args
|
||||
port = getattr(comfy.cli_args.args, 'port', None) or getattr(comfy.cli_args, 'server_port', None) or getattr(comfy.cli_args, 'port', None)
|
||||
if port is not None:
|
||||
return int(port)
|
||||
except Exception:
|
||||
pass
|
||||
# 备用：从 listen 环境变量或命令行参数尝试
|
||||
try:
|
||||
import sys
|
||||
for arg in sys.argv:
|
||||
if '--port' in arg or '--listen-port' in arg:
|
||||
parts = arg.split('=')
|
||||
if len(parts) == 2:
|
||||
return int(parts[1].strip())
|
||||
elif arg in ('--port', '--listen-port'):
|
||||
idx = sys.argv.index(arg)
|
||||
if idx + 1 < len(sys.argv):
|
||||
return int(sys.argv[idx + 1])
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _get_port_suffix() -> str:
|
||||
"""
|
||||
返回非默认端口的后缀字符串（如 "_8189"），默认端口 8188 或获取失败时返回空字符串。
|
||||
"""
|
||||
try:
|
||||
port = _get_server_port()
|
||||
if port is not None and port != 8188:
|
||||
return f"_{port}"
|
||||
except Exception:
|
||||
pass
|
||||
return ""
|
||||
|
||||
|
||||
# 支持的图片格式
|
||||
SUPPORTED_IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
|
||||
|
||||
|
||||
class ImageInfo(NamedTuple):
|
||||
"""图片信息结构"""
|
||||
image: Image.Image
|
||||
filename: str # 不含扩展名的文件名
|
||||
extension: str # 扩展名（如 .png）
|
||||
source_path: str # 原始文件路径
|
||||
|
||||
|
||||
def load_images_from_folder(
|
||||
folder_path: str,
|
||||
recursive: bool = False
|
||||
) -> List[ImageInfo]:
|
||||
"""
|
||||
从文件夹加载所有图片
|
||||
|
||||
Args:
|
||||
folder_path: 文件夹路径
|
||||
recursive: 是否递归加载子文件夹
|
||||
|
||||
Returns:
|
||||
ImageInfo 列表，包含图片和元数据
|
||||
|
||||
Raises:
|
||||
ValueError: 文件夹不存在或为空
|
||||
|
||||
Example:
|
||||
>>> images = load_images_from_folder("D:/images")
|
||||
>>> for info in images:
|
||||
... print(f"{info.filename}: {info.image.size}")
|
||||
"""
|
||||
folder_path = folder_path.strip()
|
||||
|
||||
if not folder_path:
|
||||
return []
|
||||
|
||||
path = Path(folder_path)
|
||||
|
||||
if not path.exists():
|
||||
raise ValueError(f"文件夹不存在: {folder_path}")
|
||||
|
||||
if not path.is_dir():
|
||||
raise ValueError(f"路径不是文件夹: {folder_path}")
|
||||
|
||||
images = []
|
||||
|
||||
# 获取文件列表
|
||||
if recursive:
|
||||
files = list(path.rglob("*"))
|
||||
else:
|
||||
files = list(path.iterdir())
|
||||
|
||||
# 按文件名排序，确保顺序一致
|
||||
files = sorted(files, key=lambda x: x.name.lower())
|
||||
|
||||
for file_path in files:
|
||||
if not file_path.is_file():
|
||||
continue
|
||||
|
||||
ext = file_path.suffix.lower()
|
||||
if ext not in SUPPORTED_IMAGE_EXTENSIONS:
|
||||
continue
|
||||
|
||||
try:
|
||||
img = Image.open(file_path)
|
||||
img.load() # 确保图片完全加载
|
||||
|
||||
# 转换为 RGB 模式
|
||||
if img.mode != 'RGB':
|
||||
img = img.convert('RGB')
|
||||
|
||||
images.append(ImageInfo(
|
||||
image=img,
|
||||
filename=file_path.stem,
|
||||
extension=ext,
|
||||
source_path=str(file_path)
|
||||
))
|
||||
except Exception as e:
|
||||
print(f"警告: 无法加载图片 {file_path}: {e}")
|
||||
continue
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def pair_images_indexed(
|
||||
*image_lists: List[ImageInfo]
|
||||
) -> List[Tuple[ImageInfo, ...]]:
|
||||
"""
|
||||
1:1 索引配对
|
||||
|
||||
按索引位置配对多个图片列表，以最短列表长度为准。
|
||||
|
||||
Args:
|
||||
*image_lists: 多个 ImageInfo 列表
|
||||
|
||||
Returns:
|
||||
配对后的元组列表
|
||||
|
||||
Example:
|
||||
>>> list_a = [a1, a2, a3]
|
||||
>>> list_b = [b1, b2, b3]
|
||||
>>> pairs = pair_images_indexed(list_a, list_b)
|
||||
>>> # [(a1, b1), (a2, b2), (a3, b3)]
|
||||
"""
|
||||
if not image_lists:
|
||||
return []
|
||||
|
||||
# 过滤空列表
|
||||
non_empty_lists = [lst for lst in image_lists if lst]
|
||||
|
||||
if not non_empty_lists:
|
||||
return []
|
||||
|
||||
# 使用 zip 进行索引配对（以最短列表为准）
|
||||
return list(zip(*non_empty_lists))
|
||||
|
||||
|
||||
def pair_images_by_name(
|
||||
*image_lists: List[ImageInfo]
|
||||
) -> List[Tuple[ImageInfo, ...]]:
|
||||
"""
|
||||
按文件名配对（同名匹配）
|
||||
|
||||
取所有文件夹中文件名（不含扩展名）的交集，按文件名字母升序排列后配对。
|
||||
只有在所有文件夹中都存在同名文件，该文件名才会被纳入配对。
|
||||
扩展名不同的文件（如 1.jpg 与 1.png）视为同名。
|
||||
|
||||
Args:
|
||||
*image_lists: 多个 ImageInfo 列表
|
||||
|
||||
Returns:
|
||||
配对后的元组列表，按文件名字母升序排列
|
||||
|
||||
Raises:
|
||||
ValueError: 所有文件夹之间没有任何相同文件名时抛出
|
||||
|
||||
Example:
|
||||
>>> list_a = [ImageInfo(filename="1", ...), ImageInfo(filename="2", ...)]
|
||||
>>> list_b = [ImageInfo(filename="1", ...), ImageInfo(filename="3", ...)]
|
||||
>>> pairs = pair_images_by_name(list_a, list_b)
|
||||
>>> # [(list_a[0], list_b[0])] # 只有 "1" 匹配
|
||||
"""
|
||||
if not image_lists:
|
||||
return []
|
||||
|
||||
non_empty_lists = [lst for lst in image_lists if lst]
|
||||
if not non_empty_lists:
|
||||
return []
|
||||
|
||||
# 单文件夹直接返回（无需配对）
|
||||
if len(non_empty_lists) == 1:
|
||||
return [(img,) for img in non_empty_lists[0]]
|
||||
|
||||
# 为每个文件夹建立 filename（stem）-> ImageInfo 的映射
|
||||
name_maps = [
|
||||
{img.filename: img for img in lst}
|
||||
for lst in non_empty_lists
|
||||
]
|
||||
|
||||
# 取所有文件夹文件名的交集
|
||||
common_names = set(name_maps[0].keys())
|
||||
for nm in name_maps[1:]:
|
||||
common_names &= set(nm.keys())
|
||||
|
||||
if not common_names:
|
||||
# 收集各文件夹的文件名示例，帮助用户排查问题
|
||||
folder_samples = []
|
||||
for i, nm in enumerate(name_maps):
|
||||
sample = sorted(nm.keys())[:3]
|
||||
sample_str = "、".join(f'"{n}"' for n in sample)
|
||||
folder_samples.append(f"文件夹{i + 1}：{sample_str}")
|
||||
samples_info = "\n".join(folder_samples)
|
||||
raise ValueError(
|
||||
f"所有文件夹中没有找到任何同名图片，无法进行配对！\n"
|
||||
f"请确保各文件夹内存在文件名相同的图片后重试。\n"
|
||||
f"（文件名比较不含扩展名，例如「1.jpg」与「1.png」视为同名）\n\n"
|
||||
f"各文件夹当前文件名示例：\n{samples_info}"
|
||||
)
|
||||
|
||||
# 按文件名字母升序排列，保证顺序稳定
|
||||
sorted_names = sorted(common_names, key=lambda x: x.lower())
|
||||
|
||||
return [
|
||||
tuple(nm[name] for nm in name_maps)
|
||||
for name in sorted_names
|
||||
]
|
||||
|
||||
|
||||
def pair_images_cartesian(
|
||||
*image_lists: List[ImageInfo]
|
||||
) -> List[Tuple[ImageInfo, ...]]:
|
||||
"""
|
||||
笛卡尔积配对
|
||||
|
||||
生成多个图片列表的所有组合。
|
||||
|
||||
Args:
|
||||
*image_lists: 多个 ImageInfo 列表
|
||||
|
||||
Returns:
|
||||
配对后的元组列表
|
||||
|
||||
Example:
|
||||
>>> list_a = [a1, a2]
|
||||
>>> list_b = [b1, b2]
|
||||
>>> pairs = pair_images_cartesian(list_a, list_b)
|
||||
>>> # [(a1, b1), (a1, b2), (a2, b1), (a2, b2)]
|
||||
"""
|
||||
if not image_lists:
|
||||
return []
|
||||
|
||||
# 过滤空列表
|
||||
non_empty_lists = [lst for lst in image_lists if lst]
|
||||
|
||||
if not non_empty_lists:
|
||||
return []
|
||||
|
||||
# 使用 itertools.product 生成笛卡尔积
|
||||
return list(product(*non_empty_lists))
|
||||
|
||||
|
||||
def generate_timestamp_filename(output_folder: str, prefix: str = "", extension: str = ".png", port_suffix: str = "") -> str:
|
||||
"""
|
||||
生成基于时间戳的文件名，确保按文件名排序 = 按生成时间排序。
|
||||
|
||||
格式：{prefix}{HHMMSS_YYYYMMDD_mmm}{port_suffix}{extension}
|
||||
例如：161700_20260322_001.png 或 去除ai_161700_20260322_001.png
|
||||
|
||||
Args:
|
||||
output_folder: 输出目录
|
||||
prefix: 文件名前缀（如 "去除ai_"）
|
||||
extension: 文件扩展名（如 ".png"）
|
||||
port_suffix: 端口后缀（如 "_8189"），为空时自动获取
|
||||
|
||||
Returns:
|
||||
完整文件路径
|
||||
"""
|
||||
Path(output_folder).mkdir(parents=True, exist_ok=True)
|
||||
if not port_suffix:
|
||||
port_suffix = _get_port_suffix()
|
||||
|
||||
date_part = datetime.now().strftime("%Y%m%d")
|
||||
time_part = datetime.now().strftime("%H%M%S")
|
||||
ms = random.randint(0, 999)
|
||||
|
||||
while True:
|
||||
filename = f"{prefix}{time_part}_{date_part}_{ms:03d}{port_suffix}{extension}"
|
||||
full_path = Path(output_folder) / filename
|
||||
if not full_path.exists():
|
||||
return str(full_path)
|
||||
ms = (ms + 1) % 1000
|
||||
|
||||
|
||||
def save_image(
|
||||
image: Image.Image,
|
||||
output_path: str,
|
||||
quality: int = 95
|
||||
) -> str:
|
||||
"""
|
||||
保存图片到指定路径
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
output_path: 输出文件路径
|
||||
quality: JPEG 质量（仅对 JPEG 格式有效）
|
||||
|
||||
Returns:
|
||||
实际保存的文件路径
|
||||
"""
|
||||
# 确保目录存在
|
||||
output_dir = Path(output_path).parent
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 根据扩展名选择保存参数
|
||||
ext = Path(output_path).suffix.lower()
|
||||
|
||||
if ext in {'.jpg', '.jpeg'}:
|
||||
# 转换为 RGB（JPEG 不支持 alpha 通道）
|
||||
if image.mode != 'RGB':
|
||||
image = image.convert('RGB')
|
||||
image.save(output_path, quality=quality)
|
||||
elif ext == '.webp':
|
||||
image.save(output_path, quality=quality)
|
||||
else:
|
||||
image.save(output_path)
|
||||
|
||||
return output_path
|
||||
|
||||
|
||||
def get_folder_image_count(folder_path: str) -> int:
|
||||
"""
|
||||
获取文件夹中的图片数量（不加载图片）
|
||||
|
||||
Args:
|
||||
folder_path: 文件夹路径
|
||||
|
||||
Returns:
|
||||
图片数量
|
||||
"""
|
||||
folder_path = folder_path.strip()
|
||||
|
||||
if not folder_path:
|
||||
return 0
|
||||
|
||||
path = Path(folder_path)
|
||||
|
||||
if not path.exists() or not path.is_dir():
|
||||
return 0
|
||||
|
||||
count = 0
|
||||
for file_path in path.iterdir():
|
||||
if file_path.is_file() and file_path.suffix.lower() in SUPPORTED_IMAGE_EXTENSIONS:
|
||||
count += 1
|
||||
|
||||
return count
|
||||
@@ -0,0 +1,193 @@
|
||||
"""
|
||||
图像处理工具模块
|
||||
提供 ComfyUI Tensor 与 PIL Image 之间的转换功能
|
||||
"""
|
||||
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def tensor_to_pil(tensor: torch.Tensor) -> List[Image.Image]:
|
||||
"""
|
||||
将 ComfyUI 的 Tensor 转换为 PIL Image 列表
|
||||
|
||||
Args:
|
||||
tensor: 形状为 [B, H, W, C] 的张量，值范围 [0, 1]
|
||||
|
||||
Returns:
|
||||
PIL Image 列表
|
||||
|
||||
Example:
|
||||
>>> images = tensor_to_pil(input_tensor)
|
||||
>>> for img in images:
|
||||
... img.save(f"output_{i}.png")
|
||||
"""
|
||||
images = []
|
||||
|
||||
# 转换为 numpy 数组
|
||||
np_images = tensor.cpu().numpy()
|
||||
|
||||
# 处理每张图像
|
||||
for i in range(np_images.shape[0]):
|
||||
img_array = np_images[i]
|
||||
|
||||
# 转换值范围从 [0, 1] 到 [0, 255]
|
||||
img_array = (img_array * 255).astype(np.uint8)
|
||||
|
||||
# 创建 PIL Image
|
||||
img = Image.fromarray(img_array)
|
||||
images.append(img)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def pil_to_tensor(images: List[Image.Image]) -> torch.Tensor:
|
||||
"""
|
||||
将 PIL Image 列表转换为 ComfyUI 的 Tensor
|
||||
|
||||
Args:
|
||||
images: PIL Image 列表
|
||||
|
||||
Returns:
|
||||
形状为 [B, H, W, C] 的张量，值范围 [0, 1]
|
||||
|
||||
Example:
|
||||
>>> pil_images = [Image.open("test.png")]
|
||||
>>> tensor = pil_to_tensor(pil_images)
|
||||
>>> print(tensor.shape) # [1, H, W, 3]
|
||||
"""
|
||||
tensors = []
|
||||
|
||||
for img in images:
|
||||
# 确保是 RGB 模式
|
||||
if img.mode != 'RGB':
|
||||
img = img.convert('RGB')
|
||||
|
||||
# 转换为 numpy 数组
|
||||
img_array = np.array(img).astype(np.float32)
|
||||
|
||||
# 转换值范围从 [0, 255] 到 [0, 1]
|
||||
img_array = img_array / 255.0
|
||||
|
||||
tensors.append(img_array)
|
||||
|
||||
# 堆叠为批次
|
||||
batch_tensor = np.stack(tensors, axis=0)
|
||||
|
||||
# 转换为 torch tensor
|
||||
return torch.from_numpy(batch_tensor)
|
||||
|
||||
|
||||
def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
|
||||
"""
|
||||
将 PIL Image 编码为 base64 字符串
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
format: 图像格式，默认 PNG
|
||||
|
||||
Returns:
|
||||
base64 编码的字符串
|
||||
|
||||
Example:
|
||||
>>> img = Image.open("test.png")
|
||||
>>> b64_str = encode_image_to_base64(img)
|
||||
"""
|
||||
buffered = BytesIO()
|
||||
|
||||
# 转换为 RGB 模式（如果是 RGBA）
|
||||
if image.mode == 'RGBA':
|
||||
image = image.convert('RGB')
|
||||
|
||||
image.save(buffered, format=format)
|
||||
img_bytes = buffered.getvalue()
|
||||
|
||||
return base64.b64encode(img_bytes).decode('utf-8')
|
||||
|
||||
|
||||
def decode_base64_to_pil(base64_string: str) -> Image.Image:
|
||||
"""
|
||||
将 base64 字符串解码为 PIL Image
|
||||
|
||||
Args:
|
||||
base64_string: base64 编码的图像字符串
|
||||
|
||||
Returns:
|
||||
PIL Image 对象
|
||||
|
||||
Example:
|
||||
>>> img = decode_base64_to_pil(b64_str)
|
||||
>>> img.save("decoded.png")
|
||||
"""
|
||||
img_bytes = base64.b64decode(base64_string)
|
||||
img = Image.open(BytesIO(img_bytes))
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def parse_batch_prompts(prompt: str) -> List[str]:
|
||||
"""
|
||||
解析批量提示词
|
||||
|
||||
检测单独行的 --- 分隔符，分割提示词。
|
||||
如果 --- 不是单独占据一行，则返回空列表（表示单提示词模式）。
|
||||
|
||||
Args:
|
||||
prompt: 用户输入的提示词文本
|
||||
|
||||
Returns:
|
||||
提示词列表。如果未检测到单独行的 ---，返回空列表（表示单提示词模式）
|
||||
|
||||
Raises:
|
||||
ValueError: 如果所有提示词都为空
|
||||
|
||||
Example:
|
||||
>>> prompts = parse_batch_prompts("a woman\\n---\\na man")
|
||||
>>> print(prompts) # ['a woman', 'a man']
|
||||
|
||||
>>> prompts = parse_batch_prompts("a woman --- a man")
|
||||
>>> print(prompts) # [] (单提示词模式)
|
||||
"""
|
||||
lines = prompt.split('\n')
|
||||
|
||||
# 检查是否存在单独行的 ---
|
||||
has_separator = False
|
||||
for line in lines:
|
||||
if line.strip() == '---':
|
||||
has_separator = True
|
||||
break
|
||||
|
||||
# 如果没有单独行的 ---，返回空列表（单提示词模式）
|
||||
if not has_separator:
|
||||
return []
|
||||
|
||||
# 按单独行的 --- 分割
|
||||
# 先将所有单独行的 --- 替换为特殊标记
|
||||
processed_lines = []
|
||||
for line in lines:
|
||||
if line.strip() == '---':
|
||||
processed_lines.append('<<<SEPARATOR>>>')
|
||||
else:
|
||||
processed_lines.append(line)
|
||||
|
||||
# 重新组合并分割
|
||||
processed_text = '\n'.join(processed_lines)
|
||||
raw_prompts = processed_text.split('<<<SEPARATOR>>>')
|
||||
|
||||
# 过滤空提示词
|
||||
filtered_prompts = []
|
||||
for p in raw_prompts:
|
||||
stripped = p.strip()
|
||||
if stripped:
|
||||
filtered_prompts.append(stripped)
|
||||
|
||||
# 如果所有提示词都为空，抛出错误
|
||||
if not filtered_prompts:
|
||||
raise ValueError("批量提示词模式下，所有提示词都为空，请至少提供一个有效的提示词")
|
||||
|
||||
return filtered_prompts
|
||||
@@ -0,0 +1,86 @@
|
||||
"""
|
||||
更新检查工具
|
||||
在插件加载时检查是否有新版本
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def get_current_version() -> Optional[str]:
|
||||
"""
|
||||
获取当前版本号
|
||||
|
||||
Returns:
|
||||
版本号字符串，如果读取失败返回 None
|
||||
"""
|
||||
version_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "version.txt")
|
||||
try:
|
||||
with open(version_file, 'r', encoding='utf-8') as f:
|
||||
return f.read().strip()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def check_for_updates() -> bool:
|
||||
"""
|
||||
检查是否有更新
|
||||
|
||||
Returns:
|
||||
True 如果有更新，False 如果已是最新或检查失败
|
||||
"""
|
||||
try:
|
||||
# 获取当前目录
|
||||
plugin_dir = os.path.dirname(os.path.dirname(__file__))
|
||||
|
||||
# 检查是否是 Git 仓库
|
||||
git_dir = os.path.join(plugin_dir, '.git')
|
||||
if not os.path.exists(git_dir):
|
||||
return False
|
||||
|
||||
# 执行 git fetch（禁止弹出认证弹框，失败时静默处理）
|
||||
env = os.environ.copy()
|
||||
env['GIT_TERMINAL_PROMPT'] = '0'
|
||||
subprocess.run(
|
||||
['git', 'fetch', 'origin'],
|
||||
cwd=plugin_dir,
|
||||
capture_output=True,
|
||||
timeout=10,
|
||||
env=env
|
||||
)
|
||||
|
||||
# 检查本地和远程版本
|
||||
local = subprocess.run(
|
||||
['git', 'rev-parse', '@'],
|
||||
cwd=plugin_dir,
|
||||
capture_output=True,
|
||||
text=True
|
||||
).stdout.strip()
|
||||
|
||||
remote = subprocess.run(
|
||||
['git', 'rev-parse', '@{u}'],
|
||||
cwd=plugin_dir,
|
||||
capture_output=True,
|
||||
text=True
|
||||
).stdout.strip()
|
||||
|
||||
return local != remote
|
||||
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def notify_update_available():
|
||||
"""通知用户有更新可用"""
|
||||
current_version = get_current_version()
|
||||
version_str = f" (当前版本: {current_version})" if current_version else ""
|
||||
|
||||
print("\n" + "="*60)
|
||||
print(f"🎉 Comfyui_o1key 有新版本可用{version_str}")
|
||||
print("="*60)
|
||||
print("更新方法：")
|
||||
print(" Windows: 双击运行 update.bat")
|
||||
print(" Linux/Mac: 运行 ./update.sh")
|
||||
print("或手动执行: git pull origin main")
|
||||
print("="*60 + "\n")
|
||||
@@ -0,0 +1 @@
|
||||
v1.10.2
|
||||
@@ -0,0 +1,149 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { api } from "../../../scripts/api.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "comfyui_o1key.videoPreview",
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, _app) {
|
||||
if (nodeData.name !== "VideoPreview") return;
|
||||
|
||||
const origOnExecuted = nodeType.prototype.onExecuted;
|
||||
|
||||
nodeType.prototype.onExecuted = function (message) {
|
||||
if (origOnExecuted) {
|
||||
origOnExecuted.apply(this, arguments);
|
||||
}
|
||||
|
||||
const videos = message?.videos;
|
||||
if (!videos || videos.length === 0) return;
|
||||
|
||||
const videoInfo = videos[0];
|
||||
const params = new URLSearchParams();
|
||||
params.set("filename", videoInfo.filename);
|
||||
if (videoInfo.subfolder) params.set("subfolder", videoInfo.subfolder);
|
||||
params.set("type", videoInfo.type || "output");
|
||||
|
||||
const videoUrl = api.apiURL(`/view?${params.toString()}`);
|
||||
|
||||
// ── 首次创建 DOM 结构 ──────────────────────────────
|
||||
if (!this._videoContainer) {
|
||||
this._videoContainer = document.createElement("div");
|
||||
this._videoContainer.style.cssText =
|
||||
"width:100%;display:flex;flex-direction:column;align-items:center;" +
|
||||
"padding:4px;box-sizing:border-box;";
|
||||
|
||||
this._videoEl = document.createElement("video");
|
||||
this._videoEl.controls = true;
|
||||
this._videoEl.loop = true;
|
||||
this._videoEl.autoplay = true;
|
||||
this._videoEl.muted = true;
|
||||
this._videoEl.playsInline = true;
|
||||
// 宽度铺满容器，高度由 object-fit 自适应，不限制 max-height
|
||||
this._videoEl.style.cssText =
|
||||
"width:100%;display:block;border-radius:4px;" +
|
||||
"background:#000;object-fit:contain;";
|
||||
|
||||
this._videoLabel = document.createElement("div");
|
||||
this._videoLabel.style.cssText =
|
||||
"font-size:10px;color:#aaa;margin-top:2px;" +
|
||||
"text-align:center;word-break:break-all;";
|
||||
|
||||
this._videoResLabel = document.createElement("div");
|
||||
this._videoResLabel.style.cssText =
|
||||
"font-size:10px;color:#888;margin-top:1px;" +
|
||||
"text-align:center;";
|
||||
|
||||
this._videoContainer.appendChild(this._videoEl);
|
||||
this._videoContainer.appendChild(this._videoLabel);
|
||||
this._videoContainer.appendChild(this._videoResLabel);
|
||||
|
||||
// ── 视频元数据加载后，根据真实宽高比重新调整节点大小 ──
|
||||
this._videoEl.addEventListener("loadedmetadata", () => {
|
||||
const vw = this._videoEl.videoWidth;
|
||||
const vh = this._videoEl.videoHeight;
|
||||
if (!vw || !vh) return;
|
||||
|
||||
// 存储宽高比（高/宽），供 computeSize 使用
|
||||
this._videoAspectRatio = vh / vw;
|
||||
|
||||
// 显示分辨率
|
||||
if (this._videoResLabel) {
|
||||
this._videoResLabel.textContent = `${vw} × ${vh}`;
|
||||
}
|
||||
|
||||
// 用真实比例重新计算节点高度
|
||||
this._resizeToVideo();
|
||||
});
|
||||
}
|
||||
|
||||
this._videoEl.src = videoUrl;
|
||||
this._videoLabel.textContent = videoInfo.filename;
|
||||
|
||||
// ── 注册 DOM Widget（仅第一次）──────────────────────
|
||||
if (!this.widgets?.find((w) => w.name === "video_preview_widget")) {
|
||||
const self = this;
|
||||
const widget = this.addDOMWidget(
|
||||
"video_preview_widget",
|
||||
"div",
|
||||
this._videoContainer,
|
||||
{ serialize: false, hideOnZoom: false }
|
||||
);
|
||||
|
||||
// computeSize 在 LiteGraph 布局时被调用，返回 [宽, 高]
|
||||
widget.computeSize = function (width) {
|
||||
const w = width ?? self.size?.[0] ?? 300;
|
||||
if (self._videoAspectRatio) {
|
||||
const innerW = Math.max(w - 16, 10); // 减去左右 padding
|
||||
const videoH = Math.round(innerW * self._videoAspectRatio);
|
||||
return [w, videoH + 40]; // +40 = 文件名 + 分辨率标签高度
|
||||
}
|
||||
// 元数据未就绪时给一个合理默认值
|
||||
return [w, 260];
|
||||
};
|
||||
}
|
||||
|
||||
// 初次渲染（元数据尚未加载）给出合理初始尺寸
|
||||
if (!this._videoAspectRatio) {
|
||||
const w = Math.max(this.size[0], 320);
|
||||
const h = Math.max(this.size[1], 300);
|
||||
this.setSize([w, h]);
|
||||
}
|
||||
|
||||
this.setDirtyCanvas(true, true);
|
||||
};
|
||||
|
||||
// ── 辅助方法：按视频真实比例自适应节点大小 ──────────────
|
||||
nodeType.prototype._resizeToVideo = function () {
|
||||
if (!this._videoAspectRatio) return;
|
||||
|
||||
const nodeWidth = Math.max(this.size[0], 320);
|
||||
const innerW = nodeWidth - 16;
|
||||
const videoH = Math.round(innerW * this._videoAspectRatio);
|
||||
const labelH = 40; // 文件名 + 分辨率两行
|
||||
|
||||
// 节点头部 + 其他 widget 的高度
|
||||
// LiteGraph 节点头部约 30px，每个普通 widget 约 24px
|
||||
const NON_VIDEO_WIDGETS = (this.widgets?.filter(
|
||||
(w) => w.name !== "video_preview"
|
||||
).length ?? 0);
|
||||
const headerH = 58 + NON_VIDEO_WIDGETS * 24;
|
||||
|
||||
const totalH = headerH + videoH + labelH;
|
||||
|
||||
this.setSize([nodeWidth, totalH]);
|
||||
this.setDirtyCanvas(true, true);
|
||||
};
|
||||
|
||||
// ── 节点手动缩放时同步更新视频高度 ─────────────────────
|
||||
const origOnResize = nodeType.prototype.onResize;
|
||||
nodeType.prototype.onResize = function (size) {
|
||||
if (origOnResize) origOnResize.apply(this, arguments);
|
||||
if (this._videoAspectRatio && this._videoEl) {
|
||||
// 用新宽度重新计算正确高度，避免拉伸/压缩
|
||||
const innerW = Math.max(size[0] - 16, 10);
|
||||
const videoH = Math.round(innerW * this._videoAspectRatio);
|
||||
this._videoEl.style.height = videoH + "px";
|
||||
}
|
||||
};
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user