publisher/comfyui_o1key
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: publisher/comfyui_o1key:45deea2b064c3a1bb65368fa06a3b905f31a3a50
Branches Tags
publisher/comfyui_o1key:main
publisher/comfyui_o1key:v1.10.0
..
compare: publisher/comfyui_o1key:v1.10.0
Branches Tags
publisher/comfyui_o1key:main
publisher/comfyui_o1key:v1.10.0
1 Commits
45deea2b06 .. v1.10.0
Author SHA1 Message Date
o1key 9ee29e17d0 Initial commit: Comfyui_o1key v1.10.0 2026-02-06 15:56:30 +08:00
62 changed files with 1356 additions and 15775 deletions
Expand all files Collapse all files
.cursorrules
+3 -125
View File
@@ -3,71 +3,6 @@
## 对话原则
始终使用中文进行对话。
## 编码规范 ⚠️ 重要
### 文件编码要求
- **所有文本文件必须使用 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 模型进行图像生成。
@@ -103,13 +38,10 @@ Comfyui_o1key/
│ ├── __init__.py
│ ├── base_client.py # 客户端基类
│ └── gemini_client.py # Gemini API 客户端
├── .config.example # 配置文件模板
├── .config # API 配置文件（不提交）
├── .config.example # 配置示例
├── requirements.txt # 依赖包
└── README.md # 用户文档
├── 设置API密钥(win).bat # Windows 配置脚本
└── 设置API密钥(mac).sh # Mac/Linux 配置脚本
注：.config 文件在本地自动创建，不提交到版本控制
```
---
@@ -434,7 +366,7 @@ 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
from ..utils.config import get_api_key, get_api_key_or_raise, load_config
# 获取 API 密钥（返回 None 如果未找到）
api_key = get_api_key("O1KEY_API_KEY")
@@ -442,64 +374,10 @@ 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 客户端使用
.editorconfig
-38
View File
@@ -1,38 +0,0 @@
# 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
.gitattributes
-31
View File
@@ -1,31 +0,0 @@
# 默认自动处理行结束符
* 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
.gitignore
+3 -3
View File
@@ -1,3 +1,6 @@
# 隐私文件（已弃用配置文件，改用环境变量）
# .config
# Python 缓存
__pycache__/
*.py[cod]
@@ -21,6 +24,3 @@ venv/
# OS
.DS_Store
Thumbs.db
# 用户配置（含 API Key，不提交）
.config
CHANGELOG.md
+36 -49
View File
@@ -6,59 +6,46 @@
---
## [1.10.4] - 2026-04-13
## [1.10.0] - 2026-02-06
### 修复
- 修复香蕉2画草图导致生成2张图片问题
---
## [1.10.3] - 2026-04-13
### 修复
- 香蕉节点修复返回2张图、灰色图片问题
- 修复 SSL 证书报错，统一由【保存图像】节点保存
### 新增
- 支持【立刻取消】生图请求，可立即重新运行
- 支持工作流历史记录恢复
- 并发上限提升，批量生图速度大幅提升
- 全能LLM新增视频/文档分析、流式实时预览
- Seedance 2.0 新增多图及视频URL参考（优化中）
---
## [Unreleased]
### Added ✨
- **快捷配置脚本**
- 新增 `设置API密钥(win).bat` - Windows 一键配置工具
- 新增 `设置API密钥(mac).sh` - Mac/Linux 一键配置工具
- 自动生成 `.config` 配置文件
- 交互式提示引导用户输入 API 密钥
- 自动检测并提示覆盖已存在的配置文件
- 彩色输出和友好的用户提示信息
- **配置模板文件**
- 新增 `.config.example` 作为配置文件示例
### Added ⭐
- **自动更新系统** - 让用户轻松更新插件到最新版本
- 新增 `update.bat` - Windows 自动更新脚本
- 新增 `update.sh` - Linux/Mac 自动更新脚本
- 新增 `version.txt` - 版本号管理文件
- 新增 `utils/update_checker.py` - 启动时自动检查更新
- 新增更新检查功能：每次启动 ComfyUI 时自动检测是否有新版本
- **更新脚本功能**：
- ✅ 自动检查远程更新
- ✅ 自动备份和恢复 `.config` 配置文件
- ✅ 自动拉取最新代码
- ✅ 自动更新 Python 依赖包
- ✅ 显示版本变更信息
- ✅ 显示最近更新日志（前 20 行）
- ✅ 友好的彩色终端输出（Linux/Mac）
- ✅ 完善的错误处理和提示
### Changed
- **502 错误提示优化** (`clients/base_client.py`)
- 当 API 返回 502 时，弹框显示友好文案：「糟糕！请求到上游时遇到超时或过载！别担心，过会儿再次点击运行即可！」
- 在 `request_async` 与 `request_get_async` 中均增加 502 专用分支
- **配置管理策略**
- `.config` 文件现在完全忽略提交（添加到 `.gitignore`）
- 简化配置流程，用户通过快捷脚本自动创建本地配置
- 移除配置文件安全检查机制（不再需要）
- **README 文档**
- 更新配置章节，添加快捷脚本使用说明
- 调整配置方法优先级：快捷脚本 > 环境变量 > 手动配置
- 简化安全提示说明
- **插件启动流程** (`__init__.py`)
- 集成更新检查模块
- 启动时自动检查是否有新版本
- 如有更新，终端显示友好的更新提示
- 静默失败机制，不影响插件正常加载
### Removed
- **安全检查工具**（不再需要）
- 删除 `check_config_safety.py` 配置安全检查脚本
- 删除 `.git-hooks-install.bat` Git Hook 安装脚本
- 彻底杜绝配置文件泄密风险
- **文档更新** (`README.md`)
- 新增"🔄 更新插件"章节
- 提供两种更新方法：自动更新（推荐）和手动更新
- 详细的跨平台更新说明
- 更新提示和注意事项
### Benefits
- 🎯 **用户友好** - 一键更新，无需手动操作 Git
- 🔒 **配置安全** - 自动备份恢复配置，不会丢失设置
- ⚡ **依赖同步** - 自动更新 Python 包，确保兼容性
- 📋 **信息透明** - 显示版本变更和更新日志
- 🌍 **跨平台** - 支持 Windows/Linux/Mac
- 🛡️ **稳定可靠** - 完善的错误处理，不影响插件运行
---
README.md
-255
View File
@@ -10,258 +10,3 @@
- 🎯 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**
__init__.py
+15 -114
View File
@@ -9,128 +9,29 @@ Comfyui_o1key - ComfyUI 自定义节点集合
└── __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, AspectRatioPreset, MultiResPreview, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, KVideoFirstLast, KVideoImage2Video
from .nodes import K3Video, K3VideoFirstLast, K3MotionControl, K3MotionVideoCheck, NanoBananaV2, NanoBananaV2Batch
# 报错弹框友好文案（不修改原节点代码，仅在外层统一处理）
_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(NanoBananaV2)
_wrap_generate_for_error_display(NanoBananaV2Batch)
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini
# 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,
"AspectRatioPreset": AspectRatioPreset,
"MultiResPreview": MultiResPreview,
"BatchImagesO1key": BatchImagesO1key,
"Seedance": Seedance,
"SeedanceMultiModal": SeedanceMultiModal,
"StreamPreview": StreamPreview,
"DoubaoImage": DoubaoImage,
"O1keyGPTImage": O1keyGPTImage,
"KVideoFirstLast": KVideoFirstLast,
"KVideoImage2Video": KVideoImage2Video,
"K3Video": K3Video,
"K3VideoFirstLast": K3VideoFirstLast,
"K3MotionControl": K3MotionControl,
"K3MotionVideoCheck": K3MotionVideoCheck,
"NanoBananaV2": NanoBananaV2,
"NanoBananaV2Batch": NanoBananaV2Batch,
"GoogleGemini": GoogleGemini
}
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": "文/图生视频 自研模型",
"KlingFirstLastFrame": "首尾帧生视频 自研模型",
"KlingMotionControlTest": "动作控制 自研模型",
"AspectRatioPreset": "图片宽高比预设",
"MultiResPreview": "预览图像（v2）",
"BatchImagesO1key": "加载图像（批量）",
"Seedance": "Seedance 视频生成",
"SeedanceMultiModal": "Seedance 多模态参考生视频",
"StreamPreview": "流式文本预览",
"DoubaoImage": "豆包生图",
"O1keyGPTImage": "o1key GPT Image",
"KVideoFirstLast": "K26 图生视频（首尾帧）",
"KVideoImage2Video": "K26 图生视频",
"K3Video": "K3 图生视频 自研",
"K3VideoFirstLast": "首尾帧 K3 自研",
"K3MotionControl": "动作控制 K3 自研",
"K3MotionVideoCheck": "视频时长检测 K3",
"NanoBananaV2": "Nano Banana V2",
"NanoBananaV2Batch": "Nano Banana V2（批量）",
"NanoBananaPro": "Nano Banana Pro",
"BatchNanoBananaPro": "批量 Nano Banana Pro",
"GoogleGemini": "Google Gemini"
}
WEB_DIRECTORY = "./web"
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
# 注册 /o1key/input_dir 接口，供前端文件上传按钮获取 input 目录绝对路径
try:
from aiohttp import web
from server import PromptServer
import folder_paths
@PromptServer.instance.routes.get("/o1key/input_dir")
async def get_input_dir(request):
import os
path = os.path.abspath(folder_paths.get_input_directory())
return web.json_response({"path": path})
except Exception:
pass
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
clients/__init__.py
+1 -5
View File
@@ -6,9 +6,5 @@ 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']
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient']
clients/base_async_provider.py
-183
View File
@@ -1,183 +0,0 @@
"""
异步生图 Provider 抽象基类
定义异步提交+轮询模式的统一接口，支持多种生图模型后端
每个 Provider 封装一种 API 后端的通信协议：
- 如何提交任务（端点、请求体格式）
- 如何轮询状态（端点、状态字段语义）
- 如何解析结果（响应格式、图片提取方式）
新增第三方生图模型时，只需实现此接口即可接入异步节点。
"""
from abc import ABC, abstractmethod
from typing import Dict, List, Optional
from PIL import Image
class BaseAsyncImageProvider(ABC):
"""异步生图 Provider 抽象基类"""
def __init__(self, api_key: str, proxy_url: Optional[str] = None):
self.api_key = api_key
self.proxy_url = proxy_url
# ========================================================================
# 必须实现的抽象方法
# ========================================================================
@property
@abstractmethod
def api_base_url(self) -> str:
"""异步 API 的基础 URL，如 https://cf-api.o1key.com"""
...
@abstractmethod
def get_submit_endpoint(self, model: str, resolution: str) -> str:
"""获取提交任务的 API 端点路径（不含 base_url）"""
...
@abstractmethod
def build_submit_body(
self,
prompt: str,
model: str,
resolution: str,
aspect_ratio: str,
images: Optional[List[Image.Image]] = None,
**kwargs
) -> dict:
"""构建提交任务的请求体"""
...
@abstractmethod
def extract_task_id(self, response: dict) -> str:
"""从提交响应中提取 task_id"""
...
@abstractmethod
def extract_status(self, response: dict) -> str:
"""从轮询响应中提取任务状态（如 SUBMITTED / IN_PROGRESS / SUCCESS / FAILURE）"""
...
@abstractmethod
async def parse_result(
self,
result_data: dict,
session
) -> List[Image.Image]:
"""从任务完成后的 result data 中解析生成的图像列表"""
...
@abstractmethod
def get_models(self) -> List[str]:
"""获取此 Provider 支持的模型 ID 列表"""
...
@abstractmethod
def get_model_aspect_ratios(self, model_id: str) -> List[str]:
"""获取指定模型支持的宽高比"""
...
@abstractmethod
def get_model_resolutions(self, model_id: str) -> List[str]:
"""获取指定模型支持的分辨率"""
...
# ========================================================================
# 可选的覆盖方法
# ========================================================================
def get_poll_endpoint(self, task_id: str) -> str:
"""获取轮询任务状态的 API 端点路径（默认实现适用于 o1key 异步 API）"""
return f"/async/v1/tasks/{task_id}"
def get_headers(self) -> dict:
"""获取 HTTP 请求头"""
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
def get_all_aspect_ratios(self) -> List[str]:
"""获取所有模型支持的宽高比（去重合并）"""
seen = set()
result = []
for model_id in self.get_models():
for ratio in self.get_model_aspect_ratios(model_id):
if ratio not in seen:
seen.add(ratio)
result.append(ratio)
return result
def get_all_resolutions(self) -> List[str]:
"""获取所有模型支持的分辨率（去重，按固定顺序排列）"""
_ORDER = ["512px", "1K", "2K", "4K"]
seen = set()
for model_id in self.get_models():
for res in self.get_model_resolutions(model_id):
seen.add(res)
return [r for r in _ORDER if r in seen]
def get_extra_inputs(self) -> dict:
"""
返回此 Provider 特有的额外 ComfyUI 输入参数。
子类重写以声明 Provider 专有的选项（如 Google Search Grounding）。
Returns:
dict，格式与 ComfyUI INPUT_TYPES 的 optional 字段一致
"""
return {}
def get_extra_kwargs(self, **kwargs) -> dict:
"""
从 ComfyUI kwargs 中提取此 Provider 特有的参数，
转换为 build_submit_body 可接收的 kwargs。
子类重写以处理 Provider 专有参数。
"""
return {}
def extract_progress(self, response: dict) -> Optional[float]:
"""
从轮询响应中提取生成进度。
Args:
response: 轮询接口返回的完整响应字典
Returns:
0.0-1.0 之间的进度值，或 None 表示该响应不含进度信息
"""
return None
def query_balance_sync(self) -> Optional[dict]:
"""
同步查询账户余额（可选）。
返回 None 表示不支持。
"""
return None
def format_balance_info(self, balance_data: dict) -> str:
"""格式化余额信息为展示文本"""
return ""
# ========================================================================
# 工具方法
# ========================================================================
@staticmethod
def build_proxy_url(port: str) -> Optional[str]:
"""
将端口号字符串转为 aiohttp 可用的 HTTP 代理 URL。
兼容 v2rayN (10808)、Clash Verge (7897) 等。
Args:
port: 用户填写的端口号，如 "7897"，空字符串返回 None
Returns:
代理 URL 或 None
"""
port = (port or "").strip()
if not port or not port.isdigit():
return None
return f"http://127.0.0.1:{port}"
clients/base_client.py
+85 -253
View File
@@ -9,13 +9,9 @@ import threading
from abc import ABC, abstractmethod
from typing import Any, Callable, Dict, List, Optional
import os
import time
import aiohttp
class BaseAPIClient(ABC):
"""
API 客户端抽象基类
@@ -30,20 +26,19 @@ class BaseAPIClient(ABC):
self,
base_url: str,
api_key: str,
max_request_size: int = 100 * 1024 * 1024
max_request_size: int = 20 * 1024 * 1024
):
"""
初始化客户端
Args:
base_url: API 基础 URL
api_key: API 密钥
max_request_size: 最大请求体大小（字节），默认 100MB
max_request_size: 最大请求体大小（字节），默认 20MB
"""
self.base_url = base_url
self.api_key = api_key
self.max_request_size = max_request_size
self.proxy_url: Optional[str] = None # 由节点在调用前注入，如 "http://127.0.0.1:7897"
@abstractmethod
def get_endpoint(self, **kwargs) -> str:
@@ -84,15 +79,6 @@ class BaseAPIClient(ABC):
"""
pass
def _make_session(self) -> aiohttp.ClientSession:
"""
创建统一的 aiohttp ClientSession，全局禁用 SSL 验证。
所有需要独立创建 session 的地方都应调用此方法，
避免因客户端系统缺少根证书导致 SSLCertVerificationError。
"""
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
return aiohttp.ClientSession(connector=connector, trust_env=False)
def get_headers(self, use_bearer_token: bool = False) -> Dict[str, str]:
"""
获取请求头
@@ -128,24 +114,13 @@ class BaseAPIClient(ABC):
request_size = len(request_json.encode('utf-8'))
if request_size > self.max_request_size:
size_mb = request_size / 1024 / 1024
limit_mb = self.max_request_size / 1024 / 1024
raise ValueError(
"请求体积超过100MB限制，请调整分辨率或减少图片数量"
f"请求体大小 {size_mb:.2f}MB 超过限制 {limit_mb:.0f}MB，"
"请降低分辨率或减少图片数量"
)
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,
@@ -155,133 +130,84 @@ class BaseAPIClient(ABC):
timeout: Optional[int] = None
) -> Dict[str, Any]:
"""
发送异步 HTTP 请求（带详细计时）
发送异步 HTTP 请求
Args:
endpoint: API 端点
request_body: 请求体
session: aiohttp 会话（可选）
use_bearer_token: 是否使用 Bearer Token 认证
timeout: 超时时间（秒），默认 900 秒
Returns:
响应 JSON
Raises:
RuntimeError: 请求失败时
InterruptProcessingException: 用户点击终止按钮时
"""
import time
# 尝试导入 ComfyUI 中断机制
try:
from comfy.model_management import processing_interrupted, InterruptProcessingException
_interrupt_available = True
except ImportError:
_interrupt_available = False
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 = self._make_session()
session = aiohttp.ClientSession()
close_session = True
# 设置请求超时：连接超时 30s，读取超时 900s（防止服务器出图后卡住）
_timeout_seconds = timeout if timeout is not None else 900
_aiohttp_timeout = aiohttp.ClientTimeout(
total=_timeout_seconds,
connect=30,
sock_read=_timeout_seconds
)
async def _do_request():
connect_start = time.time()
async with session.post(url, json=request_body, headers=headers, timeout=_aiohttp_timeout, proxy=self.proxy_url) as response:
connect_time = time.time() - connect_start
try:
# 设置超时
timeout_obj = aiohttp.ClientTimeout(total=timeout) if timeout else None
async with session.post(url, json=request_body, headers=headers, timeout=timeout_obj) as response:
if response.status != 200:
error_text = await response.text()
raise RuntimeError(error_text)
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
async def _poll_interrupt():
"""每 0.5s 轮询一次中断标志"""
while True:
await asyncio.sleep(0.5)
if processing_interrupted():
return
try:
if _interrupt_available:
request_task = asyncio.ensure_future(_do_request())
interrupt_task = asyncio.ensure_future(_poll_interrupt())
done, pending = await asyncio.wait(
[request_task, interrupt_task],
return_when=asyncio.FIRST_COMPLETED
)
# 取消未完成的任务
for t in pending:
t.cancel()
try:
await t
except (asyncio.CancelledError, Exception):
pass
# 判断是哪个先完成
if interrupt_task in done and request_task not in done:
raise InterruptProcessingException()
# 请求完成，取出结果（可能含异常）
return request_task.result()
else:
return await _do_request()
except InterruptProcessingException:
raise
except aiohttp.ServerTimeoutError as e:
raise RuntimeError(
f"请求超时！等待服务器响应超过 {_timeout_seconds} 秒。\n"
f"服务器可能仍在生成图片，请稍后重试，或检查网络连接。"
) from e
except aiohttp.ClientConnectorError as e:
raise RuntimeError(
f"无法连接到服务器：{str(e)}\n"
f"请检查网络连接是否正常。"
) from e
except asyncio.TimeoutError as e:
raise RuntimeError(
f"请求超时！等待服务器响应超过 {_timeout_seconds} 秒。\n"
f"服务器可能仍在生成图片，请稍后重试，或检查网络连接。"
) from e
# 针对常见错误状态码提供友好提示
if response.status == 504:
raise RuntimeError(
f"API 请求超时 (504 Gateway Timeout)\n"
f"原因：服务器响应超时或该端点暂时不可用\n"
f"建议：\n"
f" - 尝试使用其他模型\n"
f" - 稍后重试\n"
f" - 降低分辨率或减少输入图像数量\n"
f"详细错误: {error_text[:200]}"
)
elif response.status == 503:
raise RuntimeError(
f"服务暂时不可用 (503 Service Unavailable)\n"
f"原因：模型服务过载或维护中\n"
f"建议：\n"
f" - 稍后重试\n"
f" - 尝试使用其他模型"
)
elif response.status == 429:
raise RuntimeError(
f"请求频率超限 (429 Too Many Requests)\n"
f"原因：API 配额用尽或请求过于频繁\n"
f"建议：\n"
f" - 等待一段时间后重试\n"
f" - 检查 API 配额是否充足"
)
elif response.status == 404:
raise RuntimeError(
f"端点不存在 (404 Not Found)\n"
f"原因：API 端点路径错误或模型不存在\n"
f"建议：\n"
f" - 检查模型名称是否正确\n"
f" - 使用其他可用模型"
)
else:
raise RuntimeError(
f"API 请求失败 (状态码: {response.status}): {error_text}"
)
return await response.json()
finally:
if close_session:
await session.close()
async def request_get_async(
self,
endpoint: str,
@@ -296,7 +222,6 @@ class BaseAPIClient(ABC):
endpoint: API 端点
session: aiohttp 会话（可选）
use_bearer_token: 是否使用 Bearer Token 认证（默认为 True）
timeout: 超时时间（秒）- 已废弃，由服务器端控制
Returns:
响应 JSON
@@ -309,81 +234,41 @@ class BaseAPIClient(ABC):
close_session = False
if session is None:
session = self._make_session()
session = aiohttp.ClientSession()
close_session = True
try:
_get_start = time.time()
async with session.get(url, headers=headers) as response:
_get_elapsed = time.time() - _get_start
# 设置超时
timeout_obj = aiohttp.ClientTimeout(total=timeout) if timeout else None
async with session.get(url, headers=headers, timeout=timeout_obj) 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:
if response.status == 504:
raise RuntimeError(
f"API 请求超时 (504 Gateway Timeout)\n"
f"API 返回错误：{error_message}\n"
f"原因：服务器响应超时或该端点暂时不可用\n"
f"建议：稍后重试"
)
elif response.status == 502:
elif response.status == 503:
raise RuntimeError(
"糟糕！请求到上游时遇到超时或过载！别担心，过会儿再次点击运行即可！"
f"服务暂时不可用 (503 Service Unavailable)\n"
f"原因：服务过载或维护中\n"
f"建议：稍后重试"
)
elif response.status == 429:
raise RuntimeError(
f"请求频率超限 (429 Too Many Requests)\n"
f"原因：API 配额用尽或请求过于频繁\n"
f"建议：等待一段时间后重试"
)
else:
raise RuntimeError(
f"API 请求失败 (状态码: {response.status})\n"
f"API 返回错误：{error_message}"
f"API 请求失败 (状态码: {response.status}): {error_text}"
)
_resp_data = await response.json()
return _resp_data
return await response.json()
finally:
if close_session:
@@ -409,7 +294,9 @@ class BaseAPIClient(ABC):
total = len(requests)
# 创建无限制的连接器
async with self._make_session() as session:
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
async with aiohttp.ClientSession(connector=connector) as session:
tasks = []
for req in requests:
@@ -479,58 +366,3 @@ class BaseAPIClient(ABC):
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}"
clients/doubao_image_client.py
-341
View File
@@ -1,341 +0,0 @@
"""
豆包生图 API 客户端
端点：POST /v1/images/generations/
兼容 new-api 透传格式（OpenAI images/generations 兼容）
设计原则：
- 发送完整正确的请求体，new-api 丢弃字段是其侧问题
- 响应永远是同步 JSON（new-api 强制 stream=false）
- 图像输入以 data:image/png;base64,... 格式内联传递
"""
import asyncio
import json
import time
from io import BytesIO
from concurrent.futures import ThreadPoolExecutor
from typing import List, Optional, Union
import aiohttp
from PIL import Image
from ..utils.config import get_api_key_or_raise, get_api_base_url
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
# ── 固定端点 ──────────────────────────────────────────────────────────────────
_ENDPOINT = "/v1/images/generations/"
# ── 轮询 / 请求超时 ───────────────────────────────────────────────────────────
_REQUEST_TIMEOUT = 300 # 单次请求超时秒数（豆包图像生成最长约 60s）
class DoubaoImageClient:
"""
豆包生图客户端（new-api 原生 OpenAI 兼容格式）
new-api 兼容性说明（基于源码分析）：
✅ 透传：model / prompt / size / response_format / watermark / image
❌ 丢弃：seed / sequential_image_generation / sequential_image_generation_options
（进入 Extra map，但 MarshalJSON 中合并代码被注释）
❌ 强制：stream 硬编码 false，图像接口无流式处理
❌ 未实现：/v1/files 文件上传（501）
节点仍发送完整字段，待 new-api 修复后自动生效。
"""
def __init__(self):
self.api_key = get_api_key_or_raise("O1KEY_API_KEY")
self.base_url = get_api_base_url()
# ── 认证头 ────────────────────────────────────────────────────────────────
def _headers(self) -> dict:
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
# ── 图像字段构建 ──────────────────────────────────────────────────────────
def _tensor_to_image_field(self, tensor) -> Union[str, List[str]]:
"""
ComfyUI IMAGE tensor → API image 字段值
单张返回字符串，多张返回字符串列表，格式：
data:image/png;base64,<base64数据>
"""
pil_images = tensor_to_pil(tensor)
data_urls = []
for img in pil_images:
b64 = encode_image_to_base64(img, format="PNG")
data_urls.append(f"data:image/png;base64,{b64}")
return data_urls[0] if len(data_urls) == 1 else data_urls
# ── 请求体构建 ────────────────────────────────────────────────────────────
def _build_body(
self,
model: str,
prompt: str,
size: str,
seed: int,
sequential_image_generation: str,
max_images: int,
image_field=None, # str | list[str] | None
) -> dict:
"""
构建完整请求体。
字段说明（对照官方示例）：
- response_format: 固定 "url"（new-api 原样透传给豆包）
- watermark: 固定 False（UI 已移除该参数）
- stream: 固定 False（new-api 强制非流式，此字段不被读取，仅显式注明）
- sequential_image_generation_options: 仅 sequential=auto 时发送
"""
body = {
"model": model,
"prompt": prompt,
"size": size,
"response_format": "url",
"watermark": False,
"seed": seed,
"sequential_image_generation": sequential_image_generation,
}
# 仅 auto 模式才发送 max_images 选项
if sequential_image_generation == "auto":
body["sequential_image_generation_options"] = {
"max_images": max_images
}
# 图像输入（图生图）
if image_field is not None:
body["image"] = image_field
return body
# ── 图像下载 ──────────────────────────────────────────────────────────────
async def _download_image(
self,
url: str,
session: aiohttp.ClientSession,
) -> Image.Image:
"""从 URL 下载图像，返回 PIL.Image。"""
async with session.get(url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(
f"图像下载失败，HTTP {resp.status}，URL: {url}"
)
data = await resp.read()
try:
img = Image.open(BytesIO(data)).convert("RGB")
except Exception as e:
raise RuntimeError(f"图像解码失败: {e}")
return img
# ── 响应解析 ──────────────────────────────────────────────────────────────
async def _parse_response(
self,
resp_json: dict,
session: aiohttp.ClientSession,
) -> List[Image.Image]:
"""
解析 /v1/images/generations 响应，返回 PIL.Image 列表。
期望格式（new-api 原样透传豆包响应）：
{
"created": 1234567890,
"data": [
{"url": "https://..."},
{"url": "https://..."}
]
}
兼容 b64_json 字段（豆包理论上也支持）。
"""
# 检查 API 层级错误
if "error" in resp_json:
err = resp_json["error"]
if isinstance(err, dict):
msg = err.get("message") or err.get("msg") or json.dumps(err, ensure_ascii=False)
else:
msg = str(err)
raise RuntimeError(f"API 返回错误: {msg}")
data_list = resp_json.get("data")
if not data_list:
raise RuntimeError(
f"API 响应中未找到 data 字段，完整响应：\n"
f"{json.dumps(resp_json, ensure_ascii=False, indent=2)}"
)
images: List[Image.Image] = []
for idx, item in enumerate(data_list):
url = item.get("url", "")
b64 = item.get("b64_json", "")
if url and url.startswith("http"):
# 优先使用 URL 模式
img = await self._download_image(url, session)
images.append(img)
print(f"[豆包生图] 第 {idx + 1} 张下载完成 ({img.size[0]}×{img.size[1]})")
elif b64:
# 回退到 base64 模式
import base64 as _b64
try:
img_data = _b64.b64decode(b64)
img = Image.open(BytesIO(img_data)).convert("RGB")
images.append(img)
print(f"[豆包生图] 第 {idx + 1} 张 base64 解码完成 ({img.size[0]}×{img.size[1]})")
except Exception as e:
raise RuntimeError(f"第 {idx + 1} 张 base64 解码失败: {e}")
else:
print(f"[豆包生图] 警告：第 {idx + 1} 条数据既无 url 也无 b64_json，已跳过")
return images
# ── 核心异步生成方法 ──────────────────────────────────────────────────────
async def _generate_async(
self,
model: str,
prompt: str,
size: str,
seed: int,
sequential_image_generation: str,
max_images: int,
image_tensor=None,
) -> List[Image.Image]:
"""
异步完整流程：构建请求 → POST → 解析 → 下载图像。
"""
# 1. 构建 image 字段
image_field = None
if image_tensor is not None:
image_field = self._tensor_to_image_field(image_tensor)
n_imgs = len(image_field) if isinstance(image_field, list) else 1
print(f"[豆包生图] 图生图模式，参考图 {n_imgs} 张")
else:
print(f"[豆包生图] 文生图模式")
# 2. 构建请求体
body = self._build_body(
model=model,
prompt=prompt,
size=size,
seed=seed,
sequential_image_generation=sequential_image_generation,
max_images=max_images,
image_field=image_field,
)
url = f"{self.base_url}{_ENDPOINT}"
print(f"[豆包生图] 提交请求 → {model} | {size}")
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
# 3. 发送 POST 请求
t0 = time.time()
async with session.post(
url,
json=body,
headers=self._headers(),
) as resp:
elapsed_req = time.time() - t0
text = await resp.text()
if resp.status != 200:
# 尝试解析错误信息
try:
err_json = json.loads(text)
err_obj = err_json.get("error", {})
if isinstance(err_obj, dict):
msg = (
err_obj.get("message")
or err_obj.get("msg")
or text
)
else:
msg = str(err_obj) or text
except Exception:
msg = text
raise RuntimeError(
f"请求失败 HTTP {resp.status}: {msg}"
)
try:
resp_json = json.loads(text)
except Exception:
raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
print(f"[豆包生图] API 响应耗时 {elapsed_req:.1f}s，开始下载图像...")
# 4. 解析响应 & 下载图像（session 复用）
images = await self._parse_response(resp_json, session)
return images
# ── 同步入口（供 ComfyUI 节点调用）──────────────────────────────────────
def generate_sync(
self,
model: str,
prompt: str,
size: str,
seed: int,
sequential_image_generation: str,
max_images: int,
image_tensor=None,
) -> List[Image.Image]:
"""
同步生成接口（在独立线程中运行事件循环，避免与 ComfyUI 主循环冲突）。
Args:
model: 模型 ID
prompt: 提示词
size: 尺寸字符串，如 "2048x2048"
seed: 随机种子
sequential_image_generation: "disabled" | "auto"
max_images: 最大图片数（auto 模式生效）
image_tensor: ComfyUI IMAGE tensor（可选，图生图用）
Returns:
List[PIL.Image]
"""
coro = self._generate_async(
model=model,
prompt=prompt,
size=size,
seed=seed,
sequential_image_generation=sequential_image_generation,
max_images=max_images,
image_tensor=image_tensor,
)
def _run():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
return loop.run_until_complete(coro)
finally:
loop.close()
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(_run)
try:
return future.result(timeout=_REQUEST_TIMEOUT + 30)
except TimeoutError:
raise RuntimeError(
f"豆包生图超时（>{_REQUEST_TIMEOUT}s），请检查网络或稍后重试"
)
clients/flux_edit_client.py
-196
View File
@@ -1,196 +0,0 @@
"""
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)
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}
clients/gemini_async_provider.py
-150
View File
@@ -1,150 +0,0 @@
"""
Gemini 异步生图 Provider
通过 cf-api.o1key.com 的异步提交+轮询接口调用 Gemini 图像生成模型
协议说明：
- 提交：POST {base}/async{gemini_endpoint}?image_format=url
- 轮询：GET {base}/async/v1/tasks/{task_id}
- 结果：可能直接返回 image_url，也可能返回 Gemini 标准 candidates 格式
"""
from io import BytesIO
from typing import Dict, List, Optional
from PIL import Image
from .base_async_provider import BaseAsyncImageProvider
from .gemini_client import GeminiAPIClient
from ..utils.config import get_async_api_base_url, get_api_key_or_raise
from ..models_config import (
get_enabled_models,
get_model_supported_aspect_ratios,
get_model_supported_resolutions,
)
class GeminiAsyncImageProvider(BaseAsyncImageProvider):
"""
Gemini 异步生图 Provider
委托 GeminiAPIClient 处理：
- 端点构造（get_endpoint）
- 请求体构建（build_request_body）
- 响应解析（parse_response_async）
"""
def __init__(self, api_key: str = None, proxy_url: str = None):
if api_key is None:
api_key = get_api_key_or_raise("O1KEY_API_KEY")
super().__init__(api_key=api_key, proxy_url=proxy_url)
self._client = GeminiAPIClient(api_key=api_key)
# ========================================================================
# 抽象方法实现
# ========================================================================
@property
def api_base_url(self) -> str:
return get_async_api_base_url()
def get_submit_endpoint(self, model: str, resolution: str) -> str:
gemini_endpoint = self._client.get_endpoint(
model=model, resolution=resolution, image_format="url"
)
base = gemini_endpoint.split("?")[0]
async_endpoint = f"/async{base}"
if "?" in gemini_endpoint:
async_endpoint += "?" + gemini_endpoint.split("?", 1)[1]
return async_endpoint
def build_submit_body(
self,
prompt: str,
model: str,
resolution: str,
aspect_ratio: str,
images: Optional[List[Image.Image]] = None,
**kwargs
) -> dict:
return self._client.build_request_body(
prompt=prompt,
images=images,
aspect_ratio=aspect_ratio,
resolution=resolution,
enable_grounding=kwargs.get("enable_grounding", False),
enable_image_search=kwargs.get("enable_image_search", False),
image_compression=getattr(self, 'image_compression', None),
)
def extract_task_id(self, response: dict) -> str:
task_id = response.get("task_id")
if not task_id:
raise RuntimeError(f"提交响应中未找到 task_id: {response}")
return task_id
def extract_status(self, response: dict) -> str:
return response.get("status", "UNKNOWN")
async def parse_result(self, result_data: dict, session) -> List[Image.Image]:
# 异步接口可能直接返回 image_url
image_url = result_data.get("image_url", "") if isinstance(result_data, dict) else ""
if image_url:
async with session.get(image_url) as img_resp:
if img_resp.status == 200:
img_bytes = await img_resp.read()
return [Image.open(BytesIO(img_bytes))]
raise RuntimeError(f"下载图片失败 ({img_resp.status}): {image_url}")
# 否则按 Gemini 标准格式解析
images_list, _ = await self._client.parse_response_async(result_data, session=session)
return images_list
def get_models(self) -> List[str]:
return get_enabled_models()
def get_model_aspect_ratios(self, model_id: str) -> List[str]:
return get_model_supported_aspect_ratios(model_id)
def get_model_resolutions(self, model_id: str) -> List[str]:
return get_model_supported_resolutions(model_id)
# ========================================================================
# 可选方法覆盖
# ========================================================================
def get_extra_inputs(self) -> dict:
"""Gemini 专有：Google Search Grounding"""
return {
"联网功能": (["关闭", "打开"], {"default": "关闭"}),
}
def get_extra_kwargs(self, **kwargs) -> dict:
return {
"enable_grounding": kwargs.pop("联网功能", "关闭") == "打开",
}
def extract_progress(self, response: dict) -> Optional[float]:
"""从轮询响应中提取进度（0.0-1.0）"""
# 直接字段：progress / percentage
for field in ("progress", "percentage"):
val = response.get(field)
if val is not None and isinstance(val, (int, float)):
return val / 100.0 if val > 1 else float(val)
# 嵌套字段：progressInfo / progress_info
progress_info = response.get("progressInfo") or response.get("progress_info")
if isinstance(progress_info, dict):
for field in ("progress", "percentage"):
val = progress_info.get(field)
if val is not None and isinstance(val, (int, float)):
return val / 100.0 if val > 1 else float(val)
return None
def query_balance_sync(self) -> Optional[dict]:
try:
return self._client.query_balance_sync()
except Exception:
return None
def format_balance_info(self, balance_data: dict) -> str:
return self._client.format_balance_info(balance_data)
clients/gemini_client.py
+297 -663
View File
File diff suppressed because it is too large Load Diff
clients/gemini_flash_client.py
+46 -76
View File
@@ -8,15 +8,15 @@ 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 ..utils.config import get_api_key_or_raise
from ..models_config import get_flash_model_endpoint, get_enabled_flash_models
from .base_client import BaseAPIClient
# API 基础配置
API_BASE_URL = "https://api.o1key.com"
class GeminiFlashClient(BaseAPIClient):
"""
Gemini Flash API 客户端
@@ -24,7 +24,8 @@ class GeminiFlashClient(BaseAPIClient):
特点：
- 支持图片和视频输入
- 支持动态思考等级端点（不思考/低/中/高）
- 支持系统指令
- 支持不同思考深度
"""
def __init__(self, api_key: Optional[str] = None):
@@ -38,65 +39,46 @@ class GeminiFlashClient(BaseAPIClient):
api_key = get_api_key_or_raise("O1KEY_API_KEY")
super().__init__(
base_url=get_api_base_url(),
base_url=API_BASE_URL,
api_key=api_key,
max_request_size=100 * 1024 * 1024 # 100MB
max_request_size=20 * 1024 * 1024 # 20MB
)
def get_endpoint(
self,
self,
model: str = "gemini-3-flash-preview",
thinking_depth: str = "不思考",
**kwargs
) -> str:
"""
获取模型的 API 端点
根据模型和思考深度获取 API 端点
Args:
model: 模型名称
thinking_depth: 思考深度 ("不思考" 或 "高")
Returns:
API 端点路径
"""
endpoint = get_flash_model_endpoint(model)
endpoint = get_flash_model_endpoint(model, thinking_depth)
if endpoint is None:
# 回退到第一个启用的模型端点
# 回退到默认端点
default_models = get_enabled_flash_models()
if default_models:
endpoint = get_flash_model_endpoint(default_models[0])
endpoint = get_flash_model_endpoint(default_models[0], thinking_depth)
if endpoint is None:
raise ValueError(f"无法获取模型 '{model}' 的端点")
raise ValueError(f"无法获取模型 '{model}' 的端点 (思考深度: {thinking_depth})")
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 = "不思考",
system_instruction: Optional[str] = None,
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]:
"""
@@ -104,11 +86,9 @@ class GeminiFlashClient(BaseAPIClient):
Args:
prompt: 用户提示词
model: 模型名称
thinking_level: 思考等级（不思考/低/中/高）- 通过动态端点控制，不需要在请求体中传递
system_instruction: 系统指令（可选）
image_data: 图片数据列表，每个元素包含 mime_type 和 data
video_data: 视频数据，包含 mime_type 和 data
document_data: 文档数据，包含 mime_type 和 data
Returns:
请求体字典
@@ -138,15 +118,6 @@ class GeminiFlashClient(BaseAPIClient):
}
})
# 添加文档部分（如果有）
if document_data:
parts.append({
"inline_data": {
"mime_type": document_data["mime_type"],
"data": document_data["data"]
}
})
# 构建请求体
request_body = {
"contents": [
@@ -156,15 +127,12 @@ class GeminiFlashClient(BaseAPIClient):
]
}
# 对于支持 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
}
# 添加系统指令（如果有）
if system_instruction and system_instruction.strip():
request_body["system_instruction"] = {
"parts": [
{"text": system_instruction}
]
}
return request_body
@@ -229,10 +197,10 @@ class GeminiFlashClient(BaseAPIClient):
self,
prompt: str,
model: str = "gemini-3-flash-preview",
thinking_level: str = "不思考",
thinking_depth: str = "不思考",
system_instruction: Optional[str] = None,
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:
"""
@@ -241,29 +209,31 @@ class GeminiFlashClient(BaseAPIClient):
Args:
prompt: 用户提示词
model: 模型名称
thinking_level: 思考等级（不思考/低/中/高）
thinking_depth: 思考深度
system_instruction: 系统指令
image_data: 图片数据列表
video_data: 视频数据
document_data: 文档数据
session: aiohttp 会话
Returns:
生成的文本内容
"""
endpoint = self.get_endpoint(model=model)
endpoint = self.get_endpoint(model=model, thinking_depth=thinking_depth)
request_body = self.build_request_body(
prompt=prompt,
model=model,
thinking_level=thinking_level,
system_instruction=system_instruction,
image_data=image_data,
video_data=video_data,
document_data=document_data
video_data=video_data
)
# 根据是否有视频设置超时（视频处理需要更长时间）
timeout = 300 if video_data else 180
response = await self.request_async(
endpoint,
request_body,
session
session,
timeout=timeout
)
return self.parse_response(response)
@@ -272,10 +242,10 @@ class GeminiFlashClient(BaseAPIClient):
self,
prompt: str,
model: str = "gemini-3-flash-preview",
thinking_level: str = "不思考",
thinking_depth: str = "不思考",
system_instruction: Optional[str] = None,
image_data: Optional[List[Dict[str, str]]] = None,
video_data: Optional[Dict[str, str]] = None,
document_data: Optional[Dict[str, str]] = None
video_data: Optional[Dict[str, str]] = None
) -> str:
"""
同步生成文本（用于 ComfyUI 节点）
@@ -283,10 +253,10 @@ class GeminiFlashClient(BaseAPIClient):
Args:
prompt: 用户提示词
model: 模型名称
thinking_level: 思考等级（不思考/低/中/高）
thinking_depth: 思考深度
system_instruction: 系统指令
image_data: 图片数据列表
video_data: 视频数据
document_data: 文档数据
Returns:
生成的文本内容
@@ -294,10 +264,10 @@ class GeminiFlashClient(BaseAPIClient):
coro = self.generate_async(
prompt=prompt,
model=model,
thinking_level=thinking_level,
thinking_depth=thinking_depth,
system_instruction=system_instruction,
image_data=image_data,
video_data=video_data,
document_data=document_data
video_data=video_data
)
return self.run_async_in_thread(coro)
clients/gpt_image_client.py
-553
View File
@@ -1,553 +0,0 @@
"""
GPT Image API 客户端
支持两个接口：
- POST /v1/images/generations/ 文生图 / 图生图（gpt-image-1 / gpt-image-1.5）
- POST /v1/images/edits/ 图像编辑（带蒙版 inpainting）
设计原则：
- 与 doubao_image_client.py 保持相同的异步 + 同步双入口模式
- 图像以 multipart/form-data 方式上传（edits 接口）
- generations 接口使用 JSON 请求体，图像以 data URI base64 内联传递
- 响应支持 url 和 b64_json 两种格式，优先处理 b64_json（避免二次下载）
"""
import asyncio
import base64
import json
import time
from concurrent.futures import ThreadPoolExecutor
from io import BytesIO
from typing import List, Optional
import aiohttp
import numpy as np
import torch
from PIL import Image
from ..utils.config import get_api_key_or_raise, get_api_base_url
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
try:
from comfy.model_management import processing_interrupted, InterruptProcessingException
_INTERRUPT_AVAILABLE = True
except ImportError:
_INTERRUPT_AVAILABLE = False
InterruptProcessingException = RuntimeError
processing_interrupted = lambda: False
# ── 接口端点 ──────────────────────────────────────────────────────────────────
_ENDPOINT_GENERATIONS = "/v1/images/generations/"
_ENDPOINT_EDITS = "/v1/images/edits/"
# ── 模型名映射（UI 显示名 → API 实际参数名）─────────────────────────────────
_MODEL_NAME_MAP = {
"gpt-image-2-按量": "gpt-image-2",
"gpt-image-2-次卡": "gpt-image-2-special",
}
# ── 超时 ──────────────────────────────────────────────────────────────────────
_REQUEST_TIMEOUT = 900 # 秒
class GptImageClient:
"""
GPT Image API 客户端
接口说明：
generations：JSON body，支持 quality / size / n / model
edits：multipart/form-data，必须包含 image（PNG），可选 mask（PNG）
两个接口的响应格式相同：
{ "data": [ {"url": "..."} | {"b64_json": "..."} ] }
"""
def __init__(self):
self.api_key = get_api_key_or_raise("O1KEY_API_KEY")
self.base_url = get_api_base_url()
# ── 认证头 ────────────────────────────────────────────────────────────────
def _auth_headers(self) -> dict:
return {"Authorization": f"Bearer {self.api_key}"}
def _json_headers(self) -> dict:
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
# ── 图像转换工具 ──────────────────────────────────────────────────────────
# ── 请求体大小限制 ────────────────────────────────────────────────────────
_MAX_BODY_BYTES = 20 * 1024 * 1024 # 20 MB
@staticmethod
def _shrink_png_to_limit(png_bytes: bytes, max_bytes: int, label: str = "") -> bytes:
"""
若 PNG bytes 超过 max_bytes，按等比缩放反复压缩直到满足限制。
每次将面积缩小至约 80%（线性尺寸缩小至约 89.4%）。
"""
if len(png_bytes) <= max_bytes:
return png_bytes
img = Image.open(BytesIO(png_bytes))
w, h = img.size
original_size = len(png_bytes)
step = 0
while len(png_bytes) > max_bytes:
scale = 0.894 # sqrt(0.8)，面积缩小 20%
w = max(1, int(w * scale))
h = max(1, int(h * scale))
img = img.resize((w, h), Image.LANCZOS)
buf = BytesIO()
img.save(buf, format="PNG")
png_bytes = buf.getvalue()
step += 1
tag = f" ({label})" if label else ""
print(
f"[o1key GPT Image] 图像{tag}超出 {max_bytes // (1024*1024)}MB 限制，"
f"已等比缩放 {step} 次：{original_size // 1024}KB → {len(png_bytes) // 1024}KB "
f"（{w}×{h}）"
)
return png_bytes
@staticmethod
def _tensor_to_png_bytes(tensor: torch.Tensor) -> bytes:
"""
单张 ComfyUI IMAGE tensor [1, H, W, C] 或 [H, W, C] → PNG bytes
"""
if tensor.dim() == 4:
tensor = tensor.squeeze(0) # [H, W, C]
arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
img = Image.fromarray(arr)
buf = BytesIO()
img.save(buf, format="PNG")
return buf.getvalue()
@staticmethod
def _mask_tensor_to_rgba_png_bytes(mask: torch.Tensor, image_size: tuple) -> bytes:
"""
ComfyUI MASK tensor [1, H, W] 或 [H, W] → RGBA PNG bytes
白色区域（mask=1）→ 透明（alpha=0），即 API 将在此处生成新内容。
"""
if mask.dim() == 3:
mask = mask.squeeze(0) # [H, W]
h, w = mask.shape
ih, iw = image_size
# 尺寸不一致时给出提示（API 侧也会报错）
if (h, w) != (ih, iw):
raise ValueError(
f"蒙版尺寸 ({h}×{w}) 与图像尺寸 ({ih}×{iw}) 不一致，请保持相同尺寸"
)
alpha = ((1.0 - mask.cpu().numpy()) * 255).clip(0, 255).astype(np.uint8)
rgba = np.zeros((h, w, 4), dtype=np.uint8)
rgba[:, :, 3] = alpha # 只设 alpha，RGB 全 0
buf = BytesIO()
Image.fromarray(rgba, mode="RGBA").save(buf, format="PNG")
return buf.getvalue()
@staticmethod
def _pil_list_to_tensor(images: List[Image.Image]) -> torch.Tensor:
"""
PIL Image 列表 → ComfyUI IMAGE tensor [B, H, W, C]，值域 [0, 1]
RGBA 自动转换为 RGBA（保留透明通道）
"""
if not images:
placeholder = Image.new("RGBA", (512, 512), (128, 128, 128, 255))
images = [placeholder]
tensors = []
for img in images:
arr = np.array(img.convert("RGBA")).astype(np.float32) / 255.0
tensors.append(torch.from_numpy(arr))
return torch.stack(tensors, dim=0) # [B, H, W, 4]
# ── 响应解析（通用） ─────────────────────────────────────────────────────
async def _parse_response(
self,
resp_json: dict,
session: aiohttp.ClientSession,
) -> List[Image.Image]:
"""
解析 data 列表，优先取 b64_json，回退到 url 下载
"""
if "error" in resp_json:
err = resp_json["error"]
msg = (
err.get("message") or err.get("msg") or json.dumps(err, ensure_ascii=False)
if isinstance(err, dict)
else str(err)
)
raise RuntimeError(f"API 返回错误: {msg}")
data_list = resp_json.get("data")
if not data_list:
raise RuntimeError(
f"API 响应中未找到 data 字段，完整响应：\n"
f"{json.dumps(resp_json, ensure_ascii=False, indent=2)}"
)
images: List[Image.Image] = []
for idx, item in enumerate(data_list):
b64 = item.get("b64_json", "")
url = item.get("url", "")
if b64:
# 优先 base64（无需二次下载）
try:
img_bytes = base64.b64decode(b64)
img = Image.open(BytesIO(img_bytes))
images.append(img)
print(f"[o1key GPT Image] 第 {idx + 1} 张 base64 解码完成 "
f"({img.size[0]}×{img.size[1]})")
except Exception as e:
raise RuntimeError(f"第 {idx + 1} 张 base64 解码失败: {e}")
elif url and url.startswith("http"):
# 回退：下载 URL
async with session.get(url, allow_redirects=True) as r:
if r.status != 200:
raise RuntimeError(
f"图像下载失败 HTTP {r.status}，URL: {url}"
)
img_bytes = await r.read()
img = Image.open(BytesIO(img_bytes))
images.append(img)
print(f"[o1key GPT Image] 第 {idx + 1} 张下载完成 "
f"({img.size[0]}×{img.size[1]})")
else:
print(f"[o1key GPT Image] 警告：第 {idx + 1} 条数据既无 b64_json 也无 url，已跳过")
return images
# ── 中断轮询 ──────────────────────────────────────────────────────────────
@staticmethod
async def _poll_interrupt():
"""每 0.5s 轮询一次 ComfyUI 中断标志"""
while True:
await asyncio.sleep(0.5)
if _INTERRUPT_AVAILABLE and processing_interrupted():
return
@staticmethod
async def _run_with_interrupt(coro):
"""
将异步任务与中断轮询并发执行。
如果用户点击取消，cancel 掉 coro 并抛出 InterruptProcessingException。
"""
if not _INTERRUPT_AVAILABLE:
return await coro
request_task = asyncio.ensure_future(coro)
interrupt_task = asyncio.ensure_future(GptImageClient._poll_interrupt())
done, pending = await asyncio.wait(
[request_task, interrupt_task],
return_when=asyncio.FIRST_COMPLETED,
)
for t in pending:
t.cancel()
try:
await t
except (asyncio.CancelledError, Exception):
pass
if interrupt_task in done and request_task not in done:
raise InterruptProcessingException()
return request_task.result()
# ── 文生图 / 图生图（generations 接口）───────────────────────────────────
async def _generate_async(
self,
prompt: str,
model: str,
quality: str,
size: str,
n: int,
seed: int,
image_list: Optional[List[torch.Tensor]] = None,
) -> List[Image.Image]:
"""
调用 /v1/images/generations/ 接口。
当传入 image_list 时，以 data URI 格式内联图像（图生图）。
"""
# 模型名映射：UI 显示名 → API 参数名
api_model = _MODEL_NAME_MAP.get(model, model)
body: dict = {
"model": api_model,
"prompt": prompt,
"quality": quality,
"n": n,
"moderation": "low",
}
body["size"] = size if size else "auto"
# 图生图：将 tensor 列表转成 data URI 内联
if image_list is not None:
data_urls = []
for idx_img, img_tensor in enumerate(image_list):
pil_images = tensor_to_pil(img_tensor)
img = pil_images[0]
buf = BytesIO()
img.save(buf, format="PNG")
png_bytes = buf.getvalue()
# 单张图像预算：20MB 按图数平摊，至少保留 1MB 给其他字段
per_image_budget = max(
1024 * 1024,
(self._MAX_BODY_BYTES - 1024 * 1024) // len(image_list),
)
# base64 膨胀约 4/3，所以 PNG 目标上限 = budget * 3/4
png_budget = int(per_image_budget * 3 / 4)
label = f"第{idx_img + 1}张" if len(image_list) > 1 else ""
png_bytes = self._shrink_png_to_limit(png_bytes, png_budget, label)
b64 = base64.b64encode(png_bytes).decode("utf-8")
data_urls.append(f"data:image/png;base64,{b64}")
body["image"] = data_urls[0] if len(data_urls) == 1 else data_urls
mode = f"图生图（参考图 {len(data_urls)} 张）"
else:
mode = "文生图"
url = f"{self.base_url}{_ENDPOINT_GENERATIONS}"
print(f"[o1key GPT Image] {mode} | 模型={model} | quality={quality} | "
f"size={size} | n={n}")
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
async def _do_request():
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
t0 = time.time()
async with session.post(url, json=body, headers=self._json_headers()) as resp:
elapsed = time.time() - t0
text = await resp.text()
if resp.status != 200:
try:
err_json = json.loads(text)
err_obj = err_json.get("error", {})
msg = (
err_obj.get("message") or err_obj.get("msg") or text
if isinstance(err_obj, dict)
else str(err_obj) or text
)
except Exception:
msg = text
raise RuntimeError(f"请求失败 HTTP {resp.status}: {msg}")
try:
resp_json = json.loads(text)
except Exception:
raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
return await self._parse_response(resp_json, session)
return await self._run_with_interrupt(_do_request())
# ── 图像编辑（edits 接口，multipart/form-data）──────────────────────────
async def _edit_async(
self,
prompt: str,
model: str,
quality: str,
size: str,
n: int,
seed: int,
image_list: List[torch.Tensor],
mask_tensor: Optional[torch.Tensor] = None,
) -> List[Image.Image]:
"""
调用 /v1/images/edits/ 接口（multipart/form-data）。
"""
# 模型名映射：UI 显示名 → API 参数名
api_model = _MODEL_NAME_MAP.get(model, model)
# 统一 tensors 为 [1,H,W,C] 格式，支持不同尺寸
normalized_tensors = []
for t in image_list:
if t.dim() == 3:
t = t.unsqueeze(0) # [H,W,C] → [1,H,W,C]
normalized_tensors.append(t)
num_images = len(normalized_tensors)
form = aiohttp.FormData()
form.add_field("model", api_model)
form.add_field("prompt", prompt)
form.add_field("n", str(n))
form.add_field("quality", quality)
form.add_field("size", size if size else "auto")
# 多图：用 image[] 数组字段逐张附加，支持 gpt-image-1.5 最多 16 张
# 预算：20MB 按图数平摊，蒙版预留 1MB
mask_reserve = 1024 * 1024 if mask_tensor is not None else 0
per_image_budget = max(
1024 * 1024,
(self._MAX_BODY_BYTES - mask_reserve) // num_images,
)
for i, frame in enumerate(normalized_tensors):
img_bytes = self._tensor_to_png_bytes(frame)
label = f"第{i + 1}张" if num_images > 1 else ""
img_bytes = self._shrink_png_to_limit(img_bytes, per_image_budget, label)
form.add_field(
"image[]",
img_bytes,
filename=f"image_{i}.png",
content_type="image/png",
)
# 蒙版尺寸校验以第一张图为基准
first_tensor = normalized_tensors[0]
ih, iw = first_tensor.shape[1], first_tensor.shape[2]
if mask_tensor is not None:
mask_png = self._mask_tensor_to_rgba_png_bytes(mask_tensor, (ih, iw))
form.add_field(
"mask",
mask_png,
filename="mask.png",
content_type="image/png",
)
mode = "图像编辑（带蒙版）"
else:
mode = "图像编辑（无蒙版）"
url = f"{self.base_url}{_ENDPOINT_EDITS}"
print(f"[o1key GPT Image] {mode} | 模型={model} | 参考图={num_images}张 | "
f"quality={quality} | size={size} | n={n}")
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
async def _do_request():
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
t0 = time.time()
async with session.post(
url,
data=form,
headers=self._auth_headers(),
) as resp:
elapsed = time.time() - t0
text = await resp.text()
if resp.status != 200:
try:
err_json = json.loads(text)
err_obj = err_json.get("error", {})
msg = (
err_obj.get("message") or err_obj.get("msg") or text
if isinstance(err_obj, dict)
else str(err_obj) or text
)
except Exception:
msg = text
raise RuntimeError(f"请求失败 HTTP {resp.status}: {msg}")
try:
resp_json = json.loads(text)
except Exception:
raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
return await self._parse_response(resp_json, session)
return await self._run_with_interrupt(_do_request())
# ── 同步统一入口（供节点调用）────────────────────────────────────────────
def run_sync(
self,
prompt: str,
model: str,
quality: str,
size: str,
n: int,
seed: int,
image_tensor: Optional[List[torch.Tensor]] = None,
mask_tensor: Optional[torch.Tensor] = None,
) -> List[Image.Image]:
"""
同步入口，在独立线程中运行事件循环，避免与 ComfyUI 主循环冲突。
路由逻辑：
- 无 image_tensor → generations 接口（文生图，JSON body）
- 有 image_tensor → edits 接口（图生图/编辑，multipart/form-data）
"""
use_edits = (image_tensor is not None)
if use_edits:
coro = self._edit_async(
prompt=prompt, model=model, quality=quality,
size=size, n=n, seed=seed,
image_list=image_tensor, mask_tensor=mask_tensor,
)
else:
coro = self._generate_async(
prompt=prompt, model=model, quality=quality,
size=size, n=n, seed=seed,
image_list=image_tensor,
)
def _run():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
return loop.run_until_complete(coro)
finally:
loop.close()
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(_run)
try:
return future.result(timeout=_REQUEST_TIMEOUT + 30)
except TimeoutError:
raise RuntimeError(
f"o1key GPT Image 请求超时（>{_REQUEST_TIMEOUT}s），请检查网络或稍后重试"
)
# ── 余额查询 ──────────────────────────────────────────────────────────────
async def _query_balance_async(self) -> dict:
url = f"{self.base_url}/api/usage/token"
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
timeout = aiohttp.ClientTimeout(total=10)
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
async with session.get(url, headers=self._auth_headers()) as resp:
if resp.status != 200:
raise RuntimeError(f"余额查询失败 HTTP {resp.status}")
return await resp.json()
def query_balance_sync(self) -> dict:
def _run():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
return loop.run_until_complete(self._query_balance_async())
finally:
loop.close()
with ThreadPoolExecutor(max_workers=1) as executor:
return executor.submit(_run).result(timeout=15)
@staticmethod
def format_balance_info(balance_data: dict) -> str:
data = balance_data.get("data", {})
api_name = data.get("name", "未知")
total_available = data.get("total_available", 0)
balance_in_dollars = total_available / 500000
return f"当前余额：{balance_in_dollars:.2f} | API：{api_name}"
clients/kling_client.py
-290
View File
@@ -1,290 +0,0 @@
"""
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",
}
# new API 三段式端点（动作控制走这里）
NEW_API_CREATE = "/v1/videos"
NEW_API_STATUS = "/v1/videos/{video_id}"
NEW_API_CONTENT = "/v1/videos/{video_id}/content"
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(ssl=False, 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
# ── 动作控制：走 new API 三段式流程 ──────────────────────────────
async def motion_control_async(
self,
body: Dict[str, Any],
save_path: str,
on_stage: Optional[Callable[[str], None]] = None,
on_progress: Optional[Callable[[int], None]] = None,
) -> str:
"""
动作控制专用入口：
POST /v1/videos → GET /v1/videos/{id} → GET /v1/videos/{id}/content
body 字段与 Kling 官方动作控制接口一致（image_url/video_url/prompt/...）。
"""
headers = {"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"}
interval = self.POLL_INITIAL_INTERVAL
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
# 1. 提交
if on_stage:
on_stage("submitting")
create_url = f"{self.base_url}{self.NEW_API_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
# 尝试提取友好错误信息
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"动作控制提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
video_id = create_resp.get("id")
if not video_id:
raise RuntimeError(f"API 未返回视频 ID，响应：{create_resp}")
if on_stage:
on_stage(f"submitted:{video_id}")
# 2. 轮询
status_url = f"{self.base_url}{self.NEW_API_STATUS.format(video_id=video_id)}"
while True:
async with session.get(status_url, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
status_resp = json.loads(text)
status = status_resp.get("status", "").lower()
progress_raw = status_resp.get("progress", 0)
try:
progress_pct = int(str(progress_raw).rstrip("%").strip())
except (ValueError, AttributeError):
progress_pct = 0
print(f"[动作控制] 生成中 {progress_pct}%")
if on_progress:
on_progress(progress_pct)
if status == "completed":
break
if status == "failed":
error_info = status_resp.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)
# 3. 下载
if on_stage:
on_stage("downloading")
content_url = f"{self.base_url}{self.NEW_API_CONTENT.format(video_id=video_id)}"
async with session.get(content_url, headers=headers,
allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"视频下载失败 ({resp.status})")
content_type = resp.headers.get("Content-Type", "")
if "application/json" in content_type:
data = await resp.json()
download_url = data.get("url") or data.get("download_url")
if not download_url:
raise RuntimeError("视频下载失败：响应中未找到下载链接")
async with session.get(download_url) as dl_resp:
if dl_resp.status != 200:
raise RuntimeError(f"从下载链接获取视频失败 ({dl_resp.status})")
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
async for chunk in dl_resp.content.iter_chunked(8192):
f.write(chunk)
else:
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)
if on_stage:
on_stage("done")
return save_path
clients/openai_client.py
-782
View File
@@ -1,782 +0,0 @@
"""
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 = self._make_session()
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 图片链接：![alt](url)
urls = re.findall(r'!\[.*?\]\((https?://[^\)]+)\)', text)
if not urls:
urls = re.findall(r'https?://[^\s<>"{}|\\^`\[\]]+', text)
for url in urls:
try:
dl_start = time.time()
async with session.get(url) as img_resp:
if img_resp.status == 200:
img_data = await img_resp.read()
dl_time = time.time() - dl_start
speed = len(img_data) / dl_time if dl_time > 0 else 0
img = Image.open(BytesIO(img_data))
images.append(img)
if format_info["type"] is None:
format_info["type"] = "url"
format_info["size"] = len(img_data)
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
format_info["download_speed"] = speed
except Exception:
pass
except RuntimeError:
raise
except Exception as e:
raise RuntimeError(f"解析 API 响应失败: {str(e)}")
finally:
if close_session:
await session.close()
# ── 3. 无图像但有文本 → API 拒绝说明 ─────────────────────────────
if not images and text_responses:
combined = "\n".join(text_responses)
raise RuntimeError(
f"API 拒绝响应\n\n"
f"API 返回说明：\n{combined}\n\n"
f"建议：\n"
f" - 根据上述说明调整请求内容\n"
f" - 确保提示词和参考图符合使用规范"
)
if not images:
raise RuntimeError("API 响应中未找到生成的图像")
return images, format_info
# ------------------------------------------------------------------ #
# 核心生成方法（接口与 GeminiAPIClient 保持一致，节点可无缝切换） #
# ------------------------------------------------------------------ #
async def generate_single_async(
self,
prompt: str,
model: str,
resolution: str,
aspect_ratio: str,
images: Optional[List[Image.Image]] = None,
session: Optional[aiohttp.ClientSession] = None,
task_index: Optional[int] = None,
total_tasks: Optional[int] = None,
debug: bool = False,
debug_request: bool = False,
enable_grounding: bool = False, # 保留签名兼容，OpenAI 格式暂不使用
enable_image_search: bool = False # 保留签名兼容，OpenAI 格式暂不使用
) -> tuple[List[Image.Image], Dict[str, Any]]:
"""
单次异步生成请求（OpenAI /v1/chat/completions 格式）。
Args:
prompt: 提示词
model: 节点选中的模型 ID（将自动解析为实际模型名）
resolution: 分辨率
aspect_ratio: 宽高比
images: 参考图列表（图生图时传入）
session: 复用的 aiohttp 会话
task_index: 任务序号（批量时用于日志）
total_tasks: 总任务数（批量时用于日志）
debug: 打印完整 API 响应
debug_request: 打印请求体（base64 自动截断）
Returns:
(生成的图像列表, 计时信息字典)
"""
import json
total_start = time.time()
task_prefix = f"[{task_index}/{total_tasks}]" if task_index is not None and total_tasks else ""
# ── 1. 解析模型名 & 构建请求体 ────────────────────────────────────
build_start = time.time()
resolved_model = self.resolve_model_name(model, resolution)
endpoint = self.get_endpoint()
request_body = self.build_request_body(
prompt=prompt,
images=images,
aspect_ratio=aspect_ratio,
resolution=resolution,
model=resolved_model
)
build_time = time.time() - build_start
# ── 调试：打印请求体 ───────────────────────────────────────────────
if debug_request:
import json as _json
def _shorten_b64(obj):
if isinstance(obj, dict):
return {k: _shorten_b64(v) for k, v in obj.items()}
if isinstance(obj, list):
return [_shorten_b64(i) for i in obj]
if isinstance(obj, str):
if obj.startswith("data:"):
header, _, data = obj.partition(",")
return f"{header},<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
timing = response.get("_timing", {})
net_connect = timing.get("connect_time")
net_download = timing.get("download_time")
if net_connect is not None and net_download is not None:
net_str = f" | 连接 {net_connect:.2f}s | 下载 {net_download:.2f}s"
else:
net_str = ""
print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → {download_info} ✓{net_str}")
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(ssl=False, 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)
clients/seedance_client.py
-193
View File
@@ -1,193 +0,0 @@
"""
Seedance 视频生成客户端
使用 new-api 原生格式：POST /v1/video/generations → GET /v1/video/generations/{task_id}
"""
import asyncio
import json
import os
from typing import Any, Callable, Dict, Optional
import aiohttp
from ..utils.config import get_api_key_or_raise
class SeedanceClient:
"""Seedance 视频生成客户端（new-api 原生三段式）"""
# 提交任务
CREATE_ENDPOINT = "/v1/video/generations"
# 查询任务状态：{task_id} 占位
STATUS_ENDPOINT = "/v1/video/generations/{task_id}"
POLL_INITIAL_INTERVAL = 4 # 首次轮询等待秒数
POLL_MAX_INTERVAL = 15 # 最大轮询间隔秒数
# new-api 返回的成功状态值
SUCCESS_STATUSES = {"succeeded", "success", "completed", "done", "finished"}
FAILURE_STATUSES = {"failed", "fail", "error", "expired"}
def __init__(self):
self.api_key = get_api_key_or_raise()
self.base_url = "https://api.o1key.com"
def _headers(self) -> Dict[str, str]:
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
# ── 1. 提交任务 ────────────────────────────────────────────────────
async def submit_async(
self,
body: Dict[str, Any],
session: aiohttp.ClientSession,
) -> str:
"""提交视频生成任务，返回 task_id"""
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
async with session.post(url, json=body, headers=self._headers()) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = (err.get("error", {}).get("message")
or err.get("message")
or text)
except Exception:
msg = text
raise RuntimeError(f"提交失败 ({resp.status}): {msg}")
data = json.loads(text)
# new-api 返回字段：id / task_id
task_id = data.get("id") or data.get("task_id")
if not task_id:
raise RuntimeError(f"API 未返回任务 ID，响应：{data}")
return task_id
# ── 2. 轮询状态 ────────────────────────────────────────────────────
async def poll_async(
self,
task_id: str,
session: aiohttp.ClientSession,
on_progress: Optional[Callable[[int], None]] = None,
) -> str:
"""轮询任务状态，成功后返回视频 URL"""
url = f"{self.base_url}{self.STATUS_ENDPOINT.format(task_id=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:
try:
err = json.loads(text)
msg = (err.get("error", {}).get("message")
or err.get("message")
or text)
except Exception:
msg = text
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
result = json.loads(text)
# new-api 包装格式：真实数据在 result["data"] 里
inner = result.get("data") or result
status = (inner.get("status") or "").lower()
# 解析进度
progress_raw = inner.get("progress", "0")
try:
progress_pct = int(str(progress_raw).rstrip("%").strip())
except (ValueError, AttributeError):
progress_pct = 0
print(f"[Seedance] 生成中 {progress_pct}%")
if on_progress:
on_progress(progress_pct)
if status in self.SUCCESS_STATUSES:
# 响应结构：result["data"] = inner，inner["data"] = platform_data
# 视频 URL 在 inner["result_url"] 或 inner["data"]["content"]["video_url"]
platform_data = inner.get("data") or {}
content = platform_data.get("content") or {}
video_url = (
inner.get("result_url")
or content.get("video_url")
or platform_data.get("video_url")
or inner.get("url")
)
if not video_url:
raise RuntimeError(f"任务成功但未找到视频 URL，响应：{result}")
# 末帧图片 URL 在 inner["data"]["content"]["last_frame_url"]
last_frame_url = (
content.get("last_frame_url")
or platform_data.get("last_frame_url")
or inner.get("last_frame_url")
)
return video_url, last_frame_url
if status in self.FAILURE_STATUSES:
reason = (
inner.get("fail_reason")
or (inner.get("error") or {}).get("message")
or "未知错误"
)
raise RuntimeError(f"视频生成失败：{reason}")
await asyncio.sleep(interval)
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
# ── 3. 下载视频 ────────────────────────────────────────────────────
async def download_async(
self,
video_url: str,
save_path: str,
session: aiohttp.ClientSession,
) -> str:
"""下载视频到本地，返回本地路径"""
print(f"[Seedance] 下载视频...")
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,
body: Dict[str, Any],
save_path: str,
on_stage: Optional[Callable[[str], None]] = None,
on_progress: Optional[Callable[[int], None]] = None,
) -> tuple:
"""提交 → 轮询 → 下载，返回 (本地视频路径, 末帧图片URL或None)"""
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
# 提交
if on_stage:
on_stage("submitting")
task_id = await self.submit_async(body, session)
print(f"[Seedance] 任务已提交 → {task_id}")
if on_stage:
on_stage(f"submitted:{task_id}")
# 轮询
video_url, last_frame_url = await self.poll_async(task_id, session, on_progress=on_progress)
# 下载
if on_stage:
on_stage("downloading")
path = await self.download_async(video_url, save_path, session)
if on_stage:
on_stage("done")
return path, last_frame_url
clients/sora_client.py
-530
View File
@@ -1,530 +0,0 @@
"""
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 = self._make_session()
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 = self._make_session()
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 = self._make_session()
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(ssl=False, 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(ssl=False, 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}"
clients/veo_client.py
-510
View File
@@ -1,510 +0,0 @@
"""
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 = self._make_session()
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 = self._make_session()
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 = self._make_session()
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(ssl=False, 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(ssl=False, 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}"
models_config.py
+62 -531
View File
@@ -8,7 +8,7 @@
3. 重新启用模型: 将模型的 enabled 字段改回 True
模型类型:
- GEMINI_MODELS: Nano Banana 图像生成模型
- GEMINI_MODELS: Nano Banana Pro 图像生成模型
- GEMINI_FLASH_MODELS: Google Gemini Flash 文本生成模型
示例:
@@ -17,21 +17,14 @@
"id": "gemini-新模型名称",
"description": "模型说明和特点",
"enabled": True,
"endpoint_type": "standard",
"endpoint": "/v1beta/models/gemini-新模型名称:generateContent",
"thinking_config": {
"不思考": None,
"低": "low",
"中": None,
"高": "high"
}
"endpoint_type": "standard" # 端点类型: "dynamic", "standard", "flatfee"
}
临时关闭模型:
将对应模型的 "enabled": True 改为 "enabled": False
"""
from typing import List, Dict, Optional, Tuple
from typing import List, Dict, Optional
# ============================================================
@@ -39,72 +32,45 @@ from typing import List, Dict, Optional, Tuple
# ============================================================
# ============================================================
# Nano Banana 图像生成模型
# Nano Banana Pro 图像生成模型
# ============================================================
GEMINI_MODELS = [
{
"id": "nano-banana-pro-次卡",
"description": "Nano Banana Pro 次卡,根据分辨率自动选择端点 (1K/2K/4K),高性能图像生成模型",
"id": "nano-banana-pro",
"description": "Nano Banana Pro,根据分辨率自动选择端点 (1K/2K/4K),高性能图像生成模型",
"enabled": True,
"provider": "gemini_async",
"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"]
"endpoint": None # 动态端点，由代码根据分辨率选择
},
{
"id": "nano-banana-pro-官方计费",
"description": "Nano Banana Pro 官方计费，按分辨率路由 (1K/2K/4K)，使用官方计费通道",
"enabled": True,
"provider": "gemini_async",
"id": "gemini-3-pro-image-preview-url",
"description": "URL 模式,根据分辨率自动选择端点 (1K/2K/4K),推荐用于需要不同分辨率的场景",
"enabled": False,
"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"]
"endpoint": None # 动态端点，由代码根据分辨率选择
},
{
"id": "nano-banana-2-次卡",
"description": "Nano Banana 2 次卡，根据分辨率自动选择端点 (512px/1K/2K/4K)，图像生成模型",
"id": "gemini-3-pro-image-preview",
"description": "标准模式,固定端点,适用于常规图像生成",
"enabled": True,
"provider": "gemini_async",
"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": ["512px", "1K", "2K", "4K"]
},
{
"id": "nano-banana-2-官方计费",
"description": "Nano Banana 2 官方计费，按分辨率路由 (512/1K/2K/4K)，使用官方计费通道",
"enabled": True,
"provider": "gemini_async",
"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": ["512px", "1K", "2K", "4K"]
},
{
"id": "nano-banana-次卡",
"description": "Nano Banana 次卡，固定端点，图像生成模型",
"enabled": True,
"provider": "gemini_async",
"endpoint_type": "standard",
"endpoint": "/v1beta/models/nano-banana: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"]
"endpoint": "/v1beta/models/gemini-3-pro-image-preview:generateContent"
},
{
"id": "gemini-3-pro-image-preview-flatfee",
"description": "固定费用模式,固定端点,按固定价格计费 (暂时不可用-504错误)",
"enabled": False, # 暂时禁用：端点返回 504 错误
"endpoint_type": "flatfee",
"endpoint": "/v1beta/models/gemini-3-pro-image-preview-flatfee:generateContent"
},
{
"id": "nano-banana-2",
"description": "Nano Banana 2 模型,固定端点,适用于高质量图像生成",
"enabled": False,
"endpoint_type": "standard",
"endpoint": "/v1beta/models/nano-banana-2:generateContent"
}
]
@@ -115,39 +81,11 @@ GEMINI_MODELS = [
GEMINI_FLASH_MODELS = [
{
"id": "gemini-3-flash-preview",
"description": "Gemini 3 Flash,快速多模态文本生成,通过 thinkingConfig 控制思考等级",
"description": "Gemini 3 Flash,快速多模态文本生成,支持图片和视频输入",
"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"
"endpoints": {
"不思考": "/v1beta/models/gemini-3-flash-preview-nothinking:generateContent",
"高": "/v1beta/models/gemini-3-flash-preview-high:generateContent"
}
}
]
@@ -246,123 +184,6 @@ def get_model_description(model_id: str) -> str:
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 = ["512px", "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_model_provider(model_id: str) -> Optional[str]:
"""
获取模型的异步 Provider 名称
Args:
model_id: 模型 ID
Returns:
Provider 名称（如 "gemini_async"），如果模型未配置 provider 则返回 None
"""
config = get_model_config(model_id)
if config is None:
return None
return config.get("provider")
def get_enabled_async_models() -> List[str]:
"""
获取所有启用的、支持异步模式的模型 ID 列表
Returns:
模型 ID 列表（仅包含配置了 provider 且 enabled 的模型）
"""
return [
model["id"] for model in GEMINI_MODELS
if model.get("enabled", False) and model.get("provider")
]
def get_endpoint_type(model_id: str) -> Optional[str]:
"""
获取模型的端点类型
@@ -405,233 +226,6 @@ def get_model_endpoint(model_id: str) -> Optional[str]:
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 模型工具函数
# ============================================================
@@ -697,24 +291,29 @@ def is_flash_model_enabled(model_id: str) -> bool:
return config.get("enabled", False)
def get_flash_model_endpoint(model_id: str) -> Optional[str]:
def get_flash_model_endpoint(model_id: str, thinking_depth: str = "不思考") -> Optional[str]:
"""
获取 Flash 模型的 API 端点
Args:
model_id: 模型 ID
thinking_depth: 思考深度 ("不思考" 或 "高")
Returns:
API 端点路径，如果未找到则返回 None
Example:
>>> get_flash_model_endpoint("gemini-3-flash-preview")
'/v1beta/models/gemini-3-flash-preview:generateContent'
>>> get_flash_model_endpoint("gemini-3-flash-preview", "不思考")
'/v1beta/models/gemini-3-flash-preview-nothinking:generateContent'
>>> get_flash_model_endpoint("gemini-3-flash-preview", "高")
'/v1beta/models/gemini-3-flash-preview-high:generateContent'
"""
config = get_flash_model_config(model_id)
if config is None:
return None
return config.get("endpoint")
endpoints = config.get("endpoints", {})
return endpoints.get(thinking_depth)
def get_flash_model_description(model_id: str) -> str:
@@ -733,83 +332,6 @@ def get_flash_model_description(model_id: str) -> str:
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)
# ============================================================
# 向后兼容性检查
# ============================================================
@@ -870,8 +392,8 @@ def validate_flash_models_config() -> None:
验证 Flash 模型配置的完整性
检查:
- 每个模型必须有 id, description, enabled 字段
- 每个模型必须有 endpoint 字段且格式正确
- 每个模型必须有 id, description, enabled, endpoints 字段
- endpoints 必须包含所有思考深度选项
- 至少有一个模型是启用的
Raises:
@@ -880,7 +402,8 @@ def validate_flash_models_config() -> None:
if not GEMINI_FLASH_MODELS:
raise ValueError("GEMINI_FLASH_MODELS 列表不能为空")
required_fields = ["id", "description", "enabled"]
required_fields = ["id", "description", "enabled", "endpoints"]
required_thinking_depths = ["不思考", "高"]
for i, model in enumerate(GEMINI_FLASH_MODELS):
# 检查必需字段
@@ -888,16 +411,24 @@ def validate_flash_models_config() -> None:
if field not in model:
raise ValueError(f"Flash 模型 #{i} 缺少必需字段: {field}")
# 检查端点配置
if "endpoint" not in model:
raise ValueError(f"Flash 模型 {model['id']} 缺少 'endpoint' 字段")
# 检查 endpoints 字典
endpoints = model.get("endpoints", {})
if not isinstance(endpoints, dict):
raise ValueError(f"Flash 模型 {model['id']} 的 endpoints 必须是字典")
endpoint = model.get("endpoint", "")
if not endpoint or not endpoint.startswith("/v1beta/models/"):
raise ValueError(
f"Flash 模型 {model['id']} 的 endpoint '{endpoint}' 格式不正确。"
f"应以 '/v1beta/models/' 开头"
)
# 检查所有思考深度选项都有对应端点
for depth in required_thinking_depths:
if depth not in endpoints:
raise ValueError(
f"Flash 模型 {model['id']} 的 endpoints 缺少 '{depth}' 思考深度"
)
endpoint = endpoints[depth]
if not endpoint or not endpoint.startswith("/v1beta/models/"):
raise ValueError(
f"Flash 模型 {model['id']} 的端点 '{endpoint}' 格式不正确。"
f"应以 '/v1beta/models/' 开头"
)
# 检查至少有一个启用的模型
if not get_enabled_flash_models():
nodes/K3_motion_control.py
-364
View File
@@ -1,364 +0,0 @@
"""
K3 动作控制 自研节点
用参考视频驱动参考图中人物动作，生成视频。
视频通过 R2 上传后传 URL，图片转 base64 直传。
"""
import asyncio
import io
import json
import os
import struct
import tempfile
import aiohttp
from ..utils.config import get_api_key_or_raise, get_async_api_base_url
from ..utils.r2_uploader import upload_video, upload_image
from ..utils.image_utils import tensor_to_pil
try:
from comfy_api.latest import InputImpl
import folder_paths
_FOLDER_PATHS_OK = True
except Exception:
_FOLDER_PATHS_OK = False
# ── 常量 ──────────────────────────────────────────────────────────────────────
_ENDPOINT_CREATE = "/kling/v1/videos/motion-control"
_ENDPOINT_STATUS = "/kling/v1/videos/motion-control/{task_id}"
_POLL_INIT = 5
_POLL_MAX = 15
# ── 工具函数 ───────────────────────────────────────────────────────────────────
# ── 视频时长检测（纯标准库，跨平台） ──────────────────────────────────────────
def _parse_video_duration(data: bytes) -> float | None:
"""从 MP4/MOV 原始字节解析时长（秒）。读取 mvhd box。"""
idx = data.find(b"mvhd")
if idx == -1:
return None
box = data[idx + 4:]
if len(box) < 32:
return None
version = box[0]
try:
if version == 0:
timescale = struct.unpack(">I", box[12:16])[0]
duration = struct.unpack(">I", box[16:20])[0]
else: # version == 1
timescale = struct.unpack(">I", box[20:24])[0]
duration = struct.unpack(">Q", box[24:32])[0]
except struct.error:
return None
return (duration / timescale) if timescale > 0 else None
def _get_video_duration(reference_video) -> float | None:
"""从 ComfyUI VIDEO 对象获取视频时长（秒），失败返回 None。"""
try:
source = reference_video.get_stream_source()
if isinstance(source, str) and os.path.isfile(source):
with open(source, "rb") as f:
data = f.read()
elif isinstance(source, io.BytesIO):
source.seek(0)
data = source.read()
else:
return None
return _parse_video_duration(data)
except Exception:
return None
def _validate_video_duration(reference_video, character_orientation: str):
"""校验视频时长，超限时抛出 ValueError。解析失败时静默跳过。"""
duration = _get_video_duration(reference_video)
if duration is None:
print("[K3 动作控制] 无法解析视频时长，跳过校验。")
return
limit = 10 if character_orientation == "image" else 30
print(f"[K3 动作控制] 检测到视频时长: {duration:.2f}s（限制: 3~{limit}s）")
if not (3 <= duration <= limit):
orientation_label = "图片" if character_orientation == "image" else "视频"
raise ValueError(
f"参考视频时长 {duration:.1f}s 不符合要求。\n"
f"角色朝向为「{orientation_label}」时，时长须在 3~{limit}s 之间。"
)
# ── 节点 ──────────────────────────────────────────────────────────────────────
class K3MotionControl:
"""K3 动作控制 自研 —— 用参考视频驱动参考图人物动作"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"参考图片": ("IMAGE",),
"参考视频": ("VIDEO",),
"提示词": ("STRING", {"multiline": True, "default": ""}),
"模型": (["v3", "v2-6"], {"default": "v3"}),
"模式": (["720p", "1080p"], {"default": "1080p"}),
"时长": ([5, 10, 15, 20, 25, 30], {"default": 5}),
"角色朝向": (["图片", "视频"], {"default": "图片"}),
"保留原声": (["打开", "关闭"], {"default": "打开"}),
"seed": ("INT", {
"default": 0, "min": 0, "max": 2147483647,
"tooltip": "seed 仅控制节点是否重新运行，结果本身不可复现。",
}),
},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("视频",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/KVideo"
async def generate(self, 参考图片, 参考视频, 提示词, 保留原声, 角色朝向, 模式, 模型, 时长, seed, **kwargs):
api_key = get_api_key_or_raise()
base_url = get_async_api_base_url()
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
# ── 参数映射 ──────────────────────────────────────────────────
mode_api = "std" if 模式 == "720p" else "pro"
model_name = f"kling-{模型}-motion-{mode_api}-{时长}s"
character_orientation = "image" if 角色朝向 == "图片" else "video"
keep_sound = "yes" if 保留原声 == "打开" else "no"
prompt = 提示词.strip()
if len(prompt) > 2500:
raise ValueError(f"提示词长度（{len(prompt)}）超过上限 2500 个字符，请缩短后重试。")
# ── 进度条 ────────────────────────────────────────────────────
try:
from comfy.utils import ProgressBar
pbar = ProgressBar(100)
except Exception:
pbar = None
def _stage(s: str):
if s == "uploading":
print("[K3 动作控制] 上传视频到 R2...")
if pbar: pbar.update_absolute(0, 100)
elif s == "submitting":
print("[K3 动作控制] 提交任务...")
if pbar: pbar.update_absolute(10, 100)
elif s.startswith("submitted:"):
print(f"[K3 动作控制] 任务已提交 → {s.split(':', 1)[1]}")
if pbar: pbar.update_absolute(15, 100)
elif s == "downloading":
print("[K3 动作控制] 下载视频...")
if pbar: pbar.update_absolute(99, 100)
elif s == "done":
print("[K3 动作控制] 完成")
if pbar: pbar.update_absolute(100, 100)
def _progress(pct: int):
if pbar: pbar.update_absolute(15 + int(pct * 0.84), 100)
# ── 视频时长校验 ──────────────────────────────────────────────
_validate_video_duration(参考视频, character_orientation)
# 参考视频时长不得超过所选时长（防止用长视频生成短计费）
_dur = _get_video_duration(参考视频)
if _dur is not None and _dur > 时长 + 0.5:
raise ValueError(
f"参考视频时长 {_dur:.1f}s 超过所选时长 {时长}s。\n"
f"请将时长调整为 ≥{_dur:.0f}s 的档位，或更换更短的参考视频。"
)
# ── 图片 & 视频上传 R2 → 获取公网 URL ────────────────────────
_stage("uploading")
pil_list = tensor_to_pil(参考图片)
image_url = await upload_image(pil_list[0].convert("RGB"))
video_url = await upload_video(参考视频)
# ── 构建请求体 ────────────────────────────────────────────────
body: dict = {
"model_name": model_name,
"model": model_name,
"image_url": image_url,
"video_url": video_url,
"character_orientation": character_orientation,
"mode": mode_api,
"keep_original_sound": keep_sound,
}
if prompt:
body["prompt"] = prompt
# ── 保存路径（临时文件，避免与下游保存节点重复落盘）──────────────────
tmp_fd, save_path = tempfile.mkstemp(suffix=".mp4", prefix="k3_motion_")
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
# 1. 提交任务
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(
create_url,
data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
headers=headers,
) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"K3 动作控制提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
# task_id 兼容扁平结构和 data 嵌套结构
task_id = (
create_resp.get("task_id")
or create_resp.get("id")
or create_resp.get("data", {}).get("task_id")
)
if not task_id:
raise RuntimeError(f"API 未返回任务 ID，响应：{create_resp}")
_stage(f"submitted:{task_id}")
# 2. 轮询
status_url = f"{base_url}{_ENDPOINT_STATUS.format(task_id=task_id)}"
interval = _POLL_INIT
video_result_url = None
while True:
await asyncio.sleep(interval)
async with session.get(status_url, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
sr = json.loads(text)
# 兼容扁平结构和 data 嵌套结构
data = sr.get("data", sr)
status = (data.get("status") or sr.get("status") or "").lower()
pct_raw = data.get("progress", 0)
try:
pct = int(str(pct_raw).rstrip("%").strip())
except (ValueError, AttributeError):
pct = 0
print(f"[K3 动作控制] 生成中 {pct}%")
_progress(pct)
if status in ("success", "completed", "done", "finished", "succeed"):
video_result_url = (
data.get("video_url")
or data.get("result_url")
or data.get("url")
or (data.get("result", {}) or {}).get("url")
or sr.get("video_url")
or sr.get("url")
)
break
elif status in ("failed", "fail"):
err_info = data.get("error") or sr.get("error") or {}
err_msg = (err_info.get("message", "未知错误")
if isinstance(err_info, dict) else str(err_info))
raise RuntimeError(f"K3 动作控制生成失败：{err_msg}")
interval = min(interval * 1.3, _POLL_MAX)
if not video_result_url:
raise RuntimeError(f"API 未返回视频 URL，响应：{sr}")
# 3. 下载视频
_stage("downloading")
async with session.get(video_result_url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"视频下载失败 ({resp.status})")
os.close(tmp_fd)
with open(save_path, "wb") as f:
async for chunk in resp.content.iter_chunked(8192):
f.write(chunk)
_stage("done")
if _FOLDER_PATHS_OK:
return (InputImpl.VideoFromFile(save_path),)
return (save_path,)
# ── 视频时长检测测试节点 ──────────────────────────────────────────────────────
class K3MotionVideoCheck:
"""检测视频时长并校验是否满足动作控制的限制，不调用 API。"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"参考视频": ("VIDEO",),
"角色朝向": (["图片", "视频"], {"default": "图片"}),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("检测结果",)
FUNCTION = "check"
CATEGORY = "comfyui_o1key/KVideo"
OUTPUT_NODE = True
def check(self, 参考视频, 角色朝向):
character_orientation = "image" if 角色朝向 == "图片" else "video"
duration = _get_video_duration(参考视频)
if duration is None:
result = "❌ 无法解析视频时长（格式不支持或文件损坏）"
print(f"[K3 视频检测] {result}")
return (result,)
limit = 10 if character_orientation == "image" else 30
orientation_label = 角色朝向
ok = 3 <= duration <= limit
if ok:
result = (
f"✅ 时长检测通过\n"
f"视频时长: {duration:.2f}s\n"
f"角色朝向: {orientation_label}（限制 3~{limit}s）"
)
else:
result = (
f"❌ 时长检测不通过\n"
f"视频时长: {duration:.2f}s\n"
f"角色朝向: {orientation_label}（限制 3~{limit}s）\n"
f"请更换时长在 3~{limit}s 之间的视频。"
)
print(f"[K3 视频检测] {result}")
return (result,)
# ── 节点注册 ──────────────────────────────────────────────────────────────────
NODE_CLASS_MAPPINGS = {
"K3MotionControl": K3MotionControl,
"K3MotionVideoCheck": K3MotionVideoCheck,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"K3MotionControl": "动作控制 K3 自研",
"K3MotionVideoCheck": "视频时长检测 K3",
}
nodes/K3_video.py
-344
View File
@@ -1,344 +0,0 @@
"""
K3 图生视频 自研节点（图生视频 / 多镜头）
模型名根据 模式/时长/音频 动态拼接，不暴露在前端。
起始帧为必填，仅作图生视频；多镜头功能待实现。
"""
import asyncio
import json
import os
import tempfile
import aiohttp
from ..utils.config import get_api_key_or_raise, get_async_api_base_url
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
try:
from comfy_api.latest import InputImpl
import folder_paths
_FOLDER_PATHS_OK = True
except Exception:
_FOLDER_PATHS_OK = False
# ── 常量 ──────────────────────────────────────────────────────────────────────
_MODEL_BASE = "kling-v3" # 动态拼接为 kling-v3-{模式}-{时长}s-{voice}
_MODES = ["720p", "1080p", "4K"]
_MODE_MAP = {"720p": "std", "1080p": "pro", "4K": "4k"}
_MULTI_SHOT_OPTIONS = [
"禁用",
"1个故事板",
"2个故事板",
"3个故事板",
"4个故事板",
"5个故事板",
"6个故事板",
]
_ENDPOINT_CREATE = "/v1/video/generations"
_ENDPOINT_STATUS = "/v1/video/generations/{task_id}"
_POLL_INIT = 3
_POLL_MAX = 15
# ── 工具函数 ───────────────────────────────────────────────────────────────────
def _image_to_base64(tensor) -> str:
pil = tensor_to_pil(tensor)
return encode_image_to_base64(pil[0], format="PNG")
def _prepare_image_base64(tensor) -> str:
"""转换并校验图片，不符合约束时自动等比缩放后返回 base64。"""
import io
import base64
pil_list = tensor_to_pil(tensor)
img = pil_list[0].convert("RGB")
w, h = img.size
# 1. 宽高比校验（无法通过等比缩放修复，直接报错）
ratio = w / h
if ratio < 1 / 2.5 or ratio > 2.5:
raise RuntimeError(
f"图片宽高比 {w}:{h}（{ratio:.2f}）超出允许范围 1:2.5 ~ 2.5:1，请裁剪后重试。"
)
# 2. 最小尺寸：任意边 < 300px 时等比放大
if w < 300 or h < 300:
scale = max(300 / w, 300 / h)
img = img.resize((int(w * scale), int(h * scale)), resample=1) # LANCZOS=1
# 3. 文件大小：循环等比缩小直到 ≤ 10MB
MAX_BYTES = 10 * 1024 * 1024
for _ in range(20): # 最多迭代 20 次，防止死循环
buf = io.BytesIO()
img.save(buf, format="PNG")
if buf.tell() <= MAX_BYTES:
break
scale = (MAX_BYTES / buf.tell()) ** 0.5 * 0.95 # 留 5% 余量
new_w = int(img.width * scale)
new_h = int(img.height * scale)
if new_w < 300 or new_h < 300:
raise RuntimeError(
f"图片压缩至 10MB 以内后尺寸（{new_w}x{new_h}）低于最小限制 300px，无法同时满足两项约束。"
)
img = img.resize((new_w, new_h), resample=1)
else:
raise RuntimeError("图片经过 20 次缩放仍超过 10MB，请检查原始图片。")
buf.seek(0)
return base64.b64encode(buf.read()).decode("utf-8")
# ── 节点 ──────────────────────────────────────────────────────────────────────
class K3Video:
"""K3 图生视频 自研"""
@classmethod
def INPUT_TYPES(cls):
required = {
"多镜头": (_MULTI_SHOT_OPTIONS, {
"default": "禁用",
"tooltip": "禁用：单段模式；N个故事板：启用 N 段分镜。",
}),
"起始帧": ("IMAGE",),
"提示词": ("STRING", {"multiline": True, "default": ""}),
"负向提示词": ("STRING", {"multiline": True, "default": ""}),
"时长": ([5, 10, 15], {"default": 5}),
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
"模式": (_MODES, {"default": "720p"}),
"seed": ("INT", {
"default": 0, "min": 0, "max": 2147483647,
"tooltip": "seed 仅控制节点是否重新运行，结果本身不可复现。",
}),
}
for i in range(1, 7):
required[f"分镜{i}_提示词"] = ("STRING", {
"multiline": True, "default": "",
"tooltip": f"第 {i} 段分镜提示词，最多 512 字符。",
})
required[f"分镜{i}_时长"] = ("INT", {
"default": 4, "min": 1, "max": 15,
"display": "slider",
"tooltip": f"第 {i} 段分镜时长（秒）。",
})
return {"required": required}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("视频",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/KVideo"
async def generate(self, 多镜头, 起始帧, 提示词, 负向提示词, 时长, 生成音频, 模式, seed, **kwargs):
api_key = get_api_key_or_raise()
base_url = get_async_api_base_url()
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
is_multi = 多镜头 != "禁用"
voice = "voice" if 生成音频 == "打开" else "novoice"
mode_api = _MODE_MAP[模式]
if mode_api == "4k":
model_name = f"{_MODEL_BASE}-4k-{时长}s"
else:
model_name = f"{_MODEL_BASE}-{mode_api}-{时长}s-{voice}"
# ── 多镜头模式 ────────────────────────────────────────────────
if is_multi:
shot_count = int(多镜头[0]) # "3个故事板" → 3
# 收集分镜参数
multi_prompt = []
for i in range(1, shot_count + 1):
p = kwargs.get(f"分镜{i}_提示词", "").strip()
d = kwargs.get(f"分镜{i}_时长", 0)
if not p:
raise RuntimeError(f"多镜头模式错误：第 {i} 段分镜提示词不能为空。")
if d < 1:
raise RuntimeError(f"多镜头模式错误：第 {i} 段分镜时长不能小于 1 秒。")
multi_prompt.append({"index": i, "prompt": p, "duration": str(d)})
# 校验时长总和
total = sum(int(s["duration"]) for s in multi_prompt)
if total != 时长:
raise RuntimeError(
f"多镜头模式错误：各分镜时长之和（{total}s）必须等于总时长（{时长}s）。"
)
metadata: dict = {
"multi_shot": "true",
"shot_type": "customize",
"multi_prompt": multi_prompt,
}
if 生成音频 == "打开":
metadata["sound"] = "on"
body: dict = {
"model": model_name,
"prompt": 提示词.strip() or " ",
"mode": mode_api,
"duration": 时长,
"image": _prepare_image_base64(起始帧),
"metadata": metadata,
}
if 负向提示词.strip():
body["negative_prompt"] = 负向提示词.strip()
# ── 单段图生视频模式 ──────────────────────────────────────────
else:
if not 提示词.strip():
raise RuntimeError("单段模式错误：提示词不能为空。")
body = {
"model": model_name,
"prompt": 提示词.strip(),
"mode": mode_api,
"duration": 时长,
"image": _prepare_image_base64(起始帧),
}
if 负向提示词.strip():
body["negative_prompt"] = 负向提示词.strip()
if 生成音频 == "打开":
body["metadata"] = {"sound": "on"}
# ── 进度条 ────────────────────────────────────────────────────
try:
from comfy.utils import ProgressBar
pbar = ProgressBar(100)
except Exception:
pbar = None
tag = "多镜头" if is_multi else "图生视频"
def _stage(s: str):
if s == "submitting":
print(f"[K3 {tag}] 提交中...")
if pbar: pbar.update_absolute(0, 100)
elif s.startswith("submitted:"):
print(f"[K3 {tag}] 任务已提交 → {s.split(':', 1)[1]}")
if pbar: pbar.update_absolute(5, 100)
elif s == "downloading":
print(f"[K3 {tag}] 下载视频...")
if pbar: pbar.update_absolute(99, 100)
elif s == "done":
print(f"[K3 {tag}] 完成")
if pbar: pbar.update_absolute(100, 100)
def _progress(pct: int):
if pbar: pbar.update_absolute(5 + int(pct * 0.94), 100)
# ── 保存路径（临时文件，避免与下游保存节点重复落盘）──────────────────
tmp_fd, save_path = tempfile.mkstemp(suffix=".mp4", prefix="k3_")
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
# 1. 提交
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"K3 提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
task_id = (
create_resp.get("task_id")
or create_resp.get("id")
or create_resp.get("data", {}).get("task_id")
)
if not task_id:
raise RuntimeError(f"API 未返回任务 ID，响应：{create_resp}")
_stage(f"submitted:{task_id}")
# 2. 轮询
status_url = f"{base_url}{_ENDPOINT_STATUS.format(task_id=task_id)}"
interval = _POLL_INIT
video_url = None
while True:
async with session.get(status_url, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
sr = json.loads(text)
data = sr.get("data", sr)
status = (data.get("status") or sr.get("status") or "").lower()
pct_raw = data.get("progress", 0)
try:
pct = int(str(pct_raw).rstrip("%").strip())
except (ValueError, AttributeError):
pct = 0
print(f"[K3 {tag}] 生成中 {pct}%")
_progress(pct)
if status in ("success", "completed", "done", "finished", "succeed"):
video_url = (
data.get("video_url")
or data.get("result_url")
or data.get("url")
or (data.get("result", {}) or {}).get("url")
or sr.get("video_url")
or sr.get("url")
)
break
if status in ("failed", "fail"):
err_info = data.get("error") or sr.get("error") or {}
err_msg = (err_info.get("message", "未知错误")
if isinstance(err_info, dict) else str(err_info))
raise RuntimeError(f"K3 生成失败：{err_msg}")
await asyncio.sleep(interval)
interval = min(interval * 1.5, _POLL_MAX)
if not video_url:
raise RuntimeError(f"API 未返回视频 URL，响应：{sr}")
# 3. 下载
_stage("downloading")
async with session.get(video_url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"视频下载失败 ({resp.status})")
os.close(tmp_fd)
with open(save_path, "wb") as f:
async for chunk in resp.content.iter_chunked(8192):
f.write(chunk)
_stage("done")
if _FOLDER_PATHS_OK:
return (InputImpl.VideoFromFile(save_path),)
return (save_path,)
# ── 节点注册 ──────────────────────────────────────────────────────────────────
NODE_CLASS_MAPPINGS = {
"K3Video": K3Video,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"K3Video": "K3 图生视频 自研",
}
nodes/K3_video_firstlast.py
-283
View File
@@ -1,283 +0,0 @@
"""
首尾帧 K3 自研节点
基于 K3 图生视频 自研，去掉分镜功能，新增尾帧可选输入。
"""
import asyncio
import json
import os
import tempfile
import aiohttp
from ..utils.config import get_api_key_or_raise, get_async_api_base_url
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
try:
from comfy_api.latest import InputImpl
import folder_paths
_FOLDER_PATHS_OK = True
except Exception:
_FOLDER_PATHS_OK = False
# ── 常量 ──────────────────────────────────────────────────────────────────────
_MODEL_BASE = "kling-v3"
_MODES = ["720p", "1080p", "4K"]
_MODE_MAP = {"720p": "std", "1080p": "pro", "4K": "4k"}
_ENDPOINT_CREATE = "/v1/video/generations"
_ENDPOINT_STATUS = "/v1/video/generations/{task_id}"
_POLL_INIT = 3
_POLL_MAX = 15
# ── 工具函数 ───────────────────────────────────────────────────────────────────
def _prepare_image_base64(tensor) -> str:
"""转换并校验图片，不符合约束时自动等比缩放后返回 base64。"""
import io
import base64
pil_list = tensor_to_pil(tensor)
img = pil_list[0].convert("RGB")
w, h = img.size
# 1. 宽高比校验
ratio = w / h
if ratio < 1 / 2.5 or ratio > 2.5:
raise RuntimeError(
f"图片宽高比 {w}:{h}（{ratio:.2f}）超出允许范围 1:2.5 ~ 2.5:1，请裁剪后重试。"
)
# 2. 最小尺寸：任意边 < 300px 时等比放大
if w < 300 or h < 300:
scale = max(300 / w, 300 / h)
img = img.resize((int(w * scale), int(h * scale)), resample=1)
# 3. 文件大小：循环等比缩小直到 ≤ 10MB
MAX_BYTES = 10 * 1024 * 1024
for _ in range(20):
buf = io.BytesIO()
img.save(buf, format="PNG")
if buf.tell() <= MAX_BYTES:
break
scale = (MAX_BYTES / buf.tell()) ** 0.5 * 0.95
new_w = int(img.width * scale)
new_h = int(img.height * scale)
if new_w < 300 or new_h < 300:
raise RuntimeError(
f"图片压缩至 10MB 以内后尺寸（{new_w}x{new_h}）低于最小限制 300px，无法同时满足两项约束。"
)
img = img.resize((new_w, new_h), resample=1)
else:
raise RuntimeError("图片经过 20 次缩放仍超过 10MB，请检查原始图片。")
buf.seek(0)
return base64.b64encode(buf.read()).decode("utf-8")
# ── 节点 ──────────────────────────────────────────────────────────────────────
class K3VideoFirstLast:
"""首尾帧 K3 自研"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"起始帧": ("IMAGE",),
"提示词": ("STRING", {"multiline": True, "default": ""}),
"负向提示词": ("STRING", {"multiline": True, "default": ""}),
"时长": ([5, 10, 15], {"default": 5}),
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
"模式": (_MODES, {"default": "720p"}),
"seed": ("INT", {
"default": 0, "min": 0, "max": 2147483647,
"tooltip": "seed 仅控制节点是否重新运行，结果本身不可复现。",
}),
},
"optional": {
"尾帧": ("IMAGE", {"tooltip": "可选。传入后将作为视频尾帧参考。"}),
},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("视频",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/KVideo"
async def generate(self, 起始帧, 提示词, 负向提示词, 时长, 生成音频, 模式, seed, 尾帧=None):
api_key = get_api_key_or_raise()
base_url = get_async_api_base_url()
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
voice = "voice" if 生成音频 == "打开" else "novoice"
mode_api = _MODE_MAP[模式]
if mode_api == "4k":
model_name = f"{_MODEL_BASE}-4k-{时长}s"
else:
model_name = f"{_MODEL_BASE}-{mode_api}-{时长}s-{voice}"
if not 提示词.strip():
raise RuntimeError("提示词不能为空。")
# ── 构建请求体 ────────────────────────────────────────────────
body: dict = {
"model": model_name,
"prompt": 提示词.strip(),
"mode": mode_api,
"duration": 时长,
"image": _prepare_image_base64(起始帧),
}
if 负向提示词.strip():
body["negative_prompt"] = 负向提示词.strip()
# metadata：尾帧 + 音频
metadata: dict = {}
if 尾帧 is not None:
metadata["image_tail"] = _prepare_image_base64(尾帧)
if 生成音频 == "打开":
metadata["sound"] = "on"
if metadata:
body["metadata"] = metadata
# generate_audio 字段（非 metadata 路径）
if 生成音频 == "打开" and not metadata.get("sound"):
body["generate_audio"] = True
# ── 进度条 ────────────────────────────────────────────────────
try:
from comfy.utils import ProgressBar
pbar = ProgressBar(100)
except Exception:
pbar = None
def _stage(s: str):
if s == "submitting":
print("[K3 首尾帧] 提交中...")
if pbar: pbar.update_absolute(0, 100)
elif s.startswith("submitted:"):
print(f"[K3 首尾帧] 任务已提交 → {s.split(':', 1)[1]}")
if pbar: pbar.update_absolute(5, 100)
elif s == "downloading":
print("[K3 首尾帧] 下载视频...")
if pbar: pbar.update_absolute(99, 100)
elif s == "done":
print("[K3 首尾帧] 完成")
if pbar: pbar.update_absolute(100, 100)
def _progress(pct: int):
if pbar: pbar.update_absolute(5 + int(pct * 0.94), 100)
# ── 保存路径（临时文件，避免与下游保存节点重复落盘）──────────────────
tmp_fd, save_path = tempfile.mkstemp(suffix=".mp4", prefix="k3fl_")
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
# 1. 提交
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"K3 首尾帧提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
task_id = (
create_resp.get("task_id")
or create_resp.get("id")
or create_resp.get("data", {}).get("task_id")
)
if not task_id:
raise RuntimeError(f"API 未返回任务 ID，响应：{create_resp}")
_stage(f"submitted:{task_id}")
# 2. 轮询
status_url = f"{base_url}{_ENDPOINT_STATUS.format(task_id=task_id)}"
interval = _POLL_INIT
video_url = None
while True:
async with session.get(status_url, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
sr = json.loads(text)
data = sr.get("data", sr)
status = (data.get("status") or sr.get("status") or "").lower()
pct_raw = data.get("progress", 0)
try:
pct = int(str(pct_raw).rstrip("%").strip())
except (ValueError, AttributeError):
pct = 0
print(f"[K3 首尾帧] 生成中 {pct}%")
_progress(pct)
if status in ("success", "completed", "done", "finished", "succeed"):
video_url = (
data.get("video_url")
or data.get("result_url")
or data.get("url")
or (data.get("result", {}) or {}).get("url")
or sr.get("video_url")
or sr.get("url")
)
break
if status in ("failed", "fail"):
err_info = data.get("error") or sr.get("error") or {}
err_msg = (err_info.get("message", "未知错误")
if isinstance(err_info, dict) else str(err_info))
raise RuntimeError(f"K3 首尾帧生成失败：{err_msg}")
await asyncio.sleep(interval)
interval = min(interval * 1.5, _POLL_MAX)
if not video_url:
raise RuntimeError(f"API 未返回视频 URL，响应：{sr}")
# 3. 下载
_stage("downloading")
async with session.get(video_url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"视频下载失败 ({resp.status})")
os.close(tmp_fd)
with open(save_path, "wb") as f:
async for chunk in resp.content.iter_chunked(8192):
f.write(chunk)
_stage("done")
if _FOLDER_PATHS_OK:
return (InputImpl.VideoFromFile(save_path),)
return (save_path,)
# ── 节点注册 ──────────────────────────────────────────────────────────────────
NODE_CLASS_MAPPINGS = {
"K3VideoFirstLast": K3VideoFirstLast,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"K3VideoFirstLast": "首尾帧 K3 自研",
}
nodes/K_video_firstlast.py
-251
View File
@@ -1,251 +0,0 @@
"""
K26 图生视频节点
"""
import asyncio
import json
import math
import os
import tempfile
import aiohttp
from ..utils.config import get_api_key_or_raise, get_api_base_url
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
try:
from comfy_api.latest import InputImpl
import folder_paths
_FOLDER_PATHS_OK = True
except ImportError:
_FOLDER_PATHS_OK = False
# 模型基础名，运行时动态拼接完整名称
_MODEL_BASE = "kling-v2-6"
# API 端点
_ENDPOINT_CREATE = "/v1/video/generations"
_ENDPOINT_STATUS = "/v1/video/generations/{task_id}"
_POLL_INIT = 3
_POLL_MAX = 15
def _image_to_base64(tensor, scale=1.0) -> str:
from PIL import Image
pil = tensor_to_pil(tensor)
img = pil[0]
if scale < 1.0:
w, h = img.size
new_w = max(1, int(w * scale))
new_h = max(1, int(h * scale))
img = img.resize((new_w, new_h), Image.LANCZOS)
return encode_image_to_base64(img, format="PNG")
class KVideoFirstLast:
"""K26 图生视频节点（首尾帧）"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"起始帧": ("IMAGE",),
"提示词": ("STRING", {"multiline": True, "default": ""}),
"模式": (["1080p"],),
"时长": ([5, 10],),
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
"seed": ("INT", {
"default": 0, "min": 0, "max": 2147483647,
"tooltip": "seed 仅控制节点是否重新运行，结果本身不可复现。",
}),
},
"optional": {
"尾帧": ("IMAGE",),
},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("视频",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/KVideo"
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 尾帧=None, seed=0):
api_key = get_api_key_or_raise()
base_url = get_api_base_url()
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
# ── 动态拼接模型名 ────────────────────────────────────────────
mode_api = "pro" # 1080p 映射为 pro
voice = "voice" if 生成音频 == "打开" else "novoice"
model_name = f"{_MODEL_BASE}-{mode_api}-{时长}s-{voice}"
# ── 构建请求体（超过 10MB 自动缩放图片）────────────────────────
MAX_BODY = 10 * 1024 * 1024
scale = 1.0
print(f"[K26 图生视频] 请求体大小限制: 10MB，超出将自动缩放图片")
while True:
body = {
"model": model_name,
"prompt": 提示词.strip(),
"image": _image_to_base64(起始帧, scale),
"mode": mode_api,
"duration": 时长,
}
if 生成音频 == "打开":
body["generate_audio"] = True
if 尾帧 is not None:
body["metadata"] = {"image_tail": _image_to_base64(尾帧, scale)}
body_str = json.dumps(body, ensure_ascii=False)
body_size = len(body_str.encode("utf-8"))
if body_size <= MAX_BODY:
print(f"[K26 图生视频] 请求体大小: {body_size / 1024 / 1024:.2f}MB"
+ (f"（已缩放至 {scale:.1%}）" if scale < 1.0 else ""))
break
# 等比缩放：图片像素面积与 base64 长度近似线性
target_ratio = MAX_BODY / body_size
scale = scale * math.sqrt(target_ratio) * 0.95 # 5% 安全余量
if scale < 0.01:
raise RuntimeError("图片缩放后仍超过10MB限制，请使用更小的参考图")
w, h = tensor_to_pil(起始帧)[0].size
print(f"[K26 图生视频] 请求体 {body_size / 1024 / 1024:.2f}MB 超限，"
f"自动缩放至 {scale:.1%}（{int(w * scale)}x{int(h * scale)}）")
# ── 进度条 ────────────────────────────────────────────────────
try:
from comfy.utils import ProgressBar
pbar = ProgressBar(100)
except Exception:
pbar = None
def _stage(s: str):
if s == "submitting":
print("[K26 图生视频] 提交中...")
if pbar: pbar.update_absolute(0, 100)
elif s.startswith("submitted:"):
print(f"[K26 图生视频] 任务已提交 → {s.split(':', 1)[1]}")
if pbar: pbar.update_absolute(5, 100)
elif s == "downloading":
print("[K26 图生视频] 下载视频...")
if pbar: pbar.update_absolute(99, 100)
elif s == "done":
print("[K26 图生视频] 完成")
if pbar: pbar.update_absolute(100, 100)
def _progress(pct: int):
if pbar: pbar.update_absolute(5 + int(pct * 0.94), 100)
# ── 保存路径（临时文件，避免与下游保存节点重复落盘）──────────────────
tmp_fd, save_path = tempfile.mkstemp(suffix=".mp4", prefix="k26_")
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
# 1. 提交
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"K26 提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
task_id = (
create_resp.get("task_id")
or create_resp.get("id")
or create_resp.get("data", {}).get("task_id")
)
if not task_id:
raise RuntimeError(f"API 未返回任务 ID，响应：{create_resp}")
_stage(f"submitted:{task_id}")
# 2. 轮询
status_url = f"{base_url}{_ENDPOINT_STATUS.format(task_id=task_id)}"
interval = _POLL_INIT
video_url = None
while True:
async with session.get(status_url, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
sr = json.loads(text)
data = sr.get("data", sr)
status = (data.get("status") or sr.get("status") or "").lower()
pct_raw = data.get("progress", 0)
try:
pct = int(str(pct_raw).rstrip("%").strip())
except (ValueError, AttributeError):
pct = 0
print(f"[K26 图生视频] 生成中 {pct}%")
_progress(pct)
if status in ("success", "completed", "done", "finished", "succeed"):
# 提取视频 URL
video_url = (
data.get("video_url")
or data.get("result_url")
or data.get("url")
or (data.get("result", {}) or {}).get("url")
or sr.get("video_url")
or sr.get("url")
)
break
if status in ("failed", "fail"):
err_info = data.get("error") or sr.get("error") or {}
err_msg = (err_info.get("message", "未知错误")
if isinstance(err_info, dict) else str(err_info))
raise RuntimeError(f"K26 生成失败：{err_msg}")
await asyncio.sleep(interval)
interval = min(interval * 1.5, _POLL_MAX)
if not video_url:
raise RuntimeError(f"API 未返回视频 URL，响应：{sr}")
# 3. 下载
_stage("downloading")
async with session.get(video_url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"视频下载失败 ({resp.status})")
os.close(tmp_fd)
with open(save_path, "wb") as f:
async for chunk in resp.content.iter_chunked(8192):
f.write(chunk)
_stage("done")
if _FOLDER_PATHS_OK:
return (InputImpl.VideoFromFile(save_path),)
return (save_path,)
NODE_CLASS_MAPPINGS = {
"KVideoFirstLast": KVideoFirstLast,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"KVideoFirstLast": "K26 图生视频（首尾帧）",
}
nodes/K_video_image2video.py
-236
View File
@@ -1,236 +0,0 @@
"""
K26 图生视频节点
支持 720p 和 1080p 模式
"""
import asyncio
import json
import math
import os
import tempfile
import aiohttp
from ..utils.config import get_api_key_or_raise, get_api_base_url
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
try:
from comfy_api.latest import InputImpl
import folder_paths
_FOLDER_PATHS_OK = True
except ImportError:
_FOLDER_PATHS_OK = False
# 模型基础名，运行时动态拼接完整名称
_MODEL_BASE = "kling-v2-6"
# API 端点
_ENDPOINT_CREATE = "/v1/video/generations"
_ENDPOINT_STATUS = "/v1/video/generations/{task_id}"
_POLL_INIT = 3
_POLL_MAX = 15
def _image_to_base64(tensor, scale=1.0) -> str:
from PIL import Image
pil = tensor_to_pil(tensor)
img = pil[0]
if scale < 1.0:
w, h = img.size
new_w = max(1, int(w * scale))
new_h = max(1, int(h * scale))
img = img.resize((new_w, new_h), Image.LANCZOS)
return encode_image_to_base64(img, format="PNG")
class KVideoImage2Video:
"""K26 图生视频节点"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"起始帧": ("IMAGE",),
"提示词": ("STRING", {"multiline": True, "default": ""}),
"模式": (["720p", "1080p"], {"default": "720p"}),
"时长": ([5, 10], {"default": 5}),
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
"seed": ("INT", {
"default": 0, "min": 0, "max": 2147483647,
"tooltip": "seed 仅控制节点是否重新运行，结果本身不可复现。",
}),
},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("视频",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/KVideo"
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", seed=0):
api_key = get_api_key_or_raise()
base_url = get_api_base_url()
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
# ── 动态拼接模型名 ────────────────────────────────────────────
mode_api = "std" if 模式 == "720p" else "pro"
voice = "voice" if 生成音频 == "打开" else "novoice"
model_name = f"{_MODEL_BASE}-{mode_api}-{时长}s-{voice}"
# ── 构建请求体（超过 10MB 自动缩放图片）────────────────────────
MAX_BODY = 10 * 1024 * 1024
scale = 1.0
print(f"[K26 图生视频] 请求体大小限制: 10MB，超出将自动缩放图片")
while True:
body = {
"model": model_name,
"prompt": 提示词.strip(),
"image": _image_to_base64(起始帧, scale),
"mode": mode_api,
"duration": 时长,
}
if 生成音频 == "打开":
body["generate_audio"] = True
body_str = json.dumps(body, ensure_ascii=False)
body_size = len(body_str.encode("utf-8"))
if body_size <= MAX_BODY:
print(f"[K26 图生视频] 请求体大小: {body_size / 1024 / 1024:.2f}MB"
+ (f"（已缩放至 {scale:.1%}）" if scale < 1.0 else ""))
break
# 等比缩放：图片像素面积与 base64 长度近似线性
target_ratio = MAX_BODY / body_size
scale = scale * math.sqrt(target_ratio) * 0.95 # 5% 安全余量
if scale < 0.01:
raise RuntimeError("图片缩放后仍超过10MB限制，请使用更小的参考图")
w, h = tensor_to_pil(起始帧)[0].size
print(f"[K26 图生视频] 请求体 {body_size / 1024 / 1024:.2f}MB 超限，"
f"自动缩放至 {scale:.1%}（{int(w * scale)}x{int(h * scale)}）")
# ── 进度条 ────────────────────────────────────────────────────
try:
from comfy.utils import ProgressBar
pbar = ProgressBar(100)
except Exception:
pbar = None
def _stage(s: str):
if s == "submitting":
print("[K26 图生视频] 提交中...")
if pbar: pbar.update_absolute(0, 100)
elif s.startswith("submitted:"):
print(f"[K26 图生视频] 任务已提交 → {s.split(':', 1)[1]}")
if pbar: pbar.update_absolute(5, 100)
elif s == "downloading":
print("[K26 图生视频] 下载视频...")
if pbar: pbar.update_absolute(99, 100)
elif s == "done":
print("[K26 图生视频] 完成")
if pbar: pbar.update_absolute(100, 100)
def _progress(pct: int):
if pbar: pbar.update_absolute(5 + int(pct * 0.94), 100)
# ── 保存路径（临时文件，避免与下游保存节点重复落盘）──────────────────
tmp_fd, save_path = tempfile.mkstemp(suffix=".mp4", prefix="k26_")
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
# 1. 提交
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
if resp.status != 200:
err_text = await resp.text()
raise RuntimeError(f"提交失败 ({resp.status}): {err_text}")
sr = await resp.json()
task_id = sr.get("task_id") or sr.get("id")
if not task_id:
raise RuntimeError(f"API 未返回 task_id，响应：{sr}")
_stage(f"submitted:{task_id}")
# 2. 轮询
status_url = f"{base_url}{_ENDPOINT_STATUS.format(task_id=task_id)}"
interval = _POLL_INIT
video_url = None
while True:
await asyncio.sleep(interval)
async with session.get(status_url, headers=headers) as resp:
if resp.status != 200:
err_text = await resp.text()
raise RuntimeError(f"查询失败 ({resp.status}): {err_text}")
sr = await resp.json()
data = sr.get("data", {}) or {}
status = (sr.get("status") or data.get("status") or "").lower()
pct_raw = str(data.get("progress", 0)).strip().rstrip('%')
try:
pct = max(0, min(100, int(float(pct_raw))))
except (ValueError, TypeError):
pct = 0
print(f"[K26 图生视频] 生成中 {pct}%")
_progress(pct)
if status in ("success", "completed", "done", "finished", "succeed"):
# 提取视频 URL
video_url = (
data.get("video_url")
or data.get("result_url")
or data.get("url")
or (data.get("result", {}) or {}).get("url")
or sr.get("video_url")
or sr.get("url")
)
break
if status in ("failed", "fail"):
err_info = data.get("error") or sr.get("error") or {}
err_msg = (err_info.get("message", "未知错误")
if isinstance(err_info, dict) else str(err_info))
raise RuntimeError(f"K26 生成失败：{err_msg}")
await asyncio.sleep(interval)
interval = min(interval * 1.5, _POLL_MAX)
if not video_url:
raise RuntimeError(f"API 未返回视频 URL，响应：{sr}")
# 3. 下载
_stage("downloading")
async with session.get(video_url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"视频下载失败 ({resp.status})")
os.close(tmp_fd)
with open(save_path, "wb") as f:
async for chunk in resp.content.iter_chunked(8192):
f.write(chunk)
_stage("done")
if _FOLDER_PATHS_OK:
return (InputImpl.VideoFromFile(save_path),)
return (save_path,)
NODE_CLASS_MAPPINGS = {
"KVideoImage2Video": KVideoImage2Video,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"KVideoImage2Video": "K26 图生视频",
}
nodes/__init__.py
+1 -21
View File
@@ -3,28 +3,8 @@
包含所有 ComfyUI 自定义节点的实现
"""
from .stream_preview import StreamPreview
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 .multi_res_preview import MultiResPreview
from .batch_images_o1key import BatchImagesO1key
from .seedance_video import Seedance, SeedanceMultiModal
from .nano_banana_v2 import NanoBananaV2, NanoBananaV2Batch, AsyncImageGenerator, BatchAsyncImageGenerator
from .doubao_image import DoubaoImage
from .gpt_image import O1keyGPTImage
from .K_video_firstlast import KVideoFirstLast
from .K_video_image2video import KVideoImage2Video
from .K3_video import K3Video
from .K3_video_firstlast import K3VideoFirstLast
from .K3_motion_control import K3MotionControl, K3MotionVideoCheck
__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'MultiResPreview', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator']
__all__ = ['NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini']
nodes/batch_images_o1key.py
-80
View File
@@ -1,80 +0,0 @@
"""
批量图像（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)
nodes/batch_nano_banana_pro.py
+178 -492
View File
File diff suppressed because it is too large Load Diff
nodes/doubao_image.py
-420
View File
@@ -1,420 +0,0 @@
"""
豆包生图节点
后端通过 new-api 兼容层调用豆包官方 API
"""
import asyncio
import time
import numpy as np
import torch
from concurrent.futures import ThreadPoolExecutor
from PIL import Image
from typing import List, Optional
from ..clients.doubao_image_client import DoubaoImageClient
from ..utils.image_utils import tensor_to_pil
# ── 模型列表 ──────────────────────────────────────────────────────────────────
_MODELS = [
"doubao-seedream-5-0-260128",
"doubao-seedream-4-5-251128",
]
# ── 宽高比列表 ─────────────────────────────────────────────────────────────────
_ASPECT_RATIOS = ["1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "21:9"]
# ── 分辨率档位（每个模型支持的档位不同）──────────────────────────────────────
# 5.0：2K / 3K
# 4.5：2K / 4K
_RESOLUTIONS = ["2K", "3K", "4K"]
# ── 像素对照表 ─────────────────────────────────────────────────────────────────
# 结构：{ 模型版本key: { 分辨率: { 宽高比: (宽, 高) } } }
_SIZE_TABLE = {
"5-0": {
"2K": {
"1:1": (2048, 2048),
"4:3": (2304, 1728),
"3:4": (1728, 2304),
"16:9": (2848, 1600),
"9:16": (1600, 2848),
"3:2": (2496, 1664),
"2:3": (1664, 2496),
"21:9": (3136, 1344),
},
"3K": {
"1:1": (3072, 3072),
"4:3": (3456, 2592),
"3:4": (2592, 3456),
"16:9": (4096, 2304),
"9:16": (2304, 4096),
"3:2": (3744, 2496),
"2:3": (2496, 3744),
"21:9": (4704, 2016),
},
},
"4-5": {
"2K": {
"1:1": (2048, 2048),
"4:3": (2304, 1728),
"3:4": (1728, 2304),
"16:9": (2848, 1600),
"9:16": (1600, 2848),
"3:2": (2496, 1664),
"2:3": (1664, 2496),
"21:9": (3136, 1344),
},
"4K": {
"1:1": (4096, 4096),
"4:3": (4704, 3520),
"3:4": (3520, 4704),
"16:9": (5504, 3040),
"9:16": (3040, 5504),
"3:2": (4992, 3328),
"2:3": (3328, 4992),
"21:9": (6240, 2656),
},
},
}
# 每个模型版本支持的分辨率档位
_MODEL_RESOLUTIONS = {
"5-0": ["2K", "3K"],
"4-5": ["2K", "4K"],
}
# 并发请求超时（秒）
_CONCURRENT_TIMEOUT = 330
def _model_key(model: str) -> str:
"""从模型 ID 中提取版本 key（'5-0' 或 '4-5'）。"""
for key in _SIZE_TABLE:
if key in model:
return key
raise ValueError(f"无法识别模型版本：{model}，支持的模型：{_MODELS}")
def _pil_list_to_tensor(images: List[Image.Image]) -> torch.Tensor:
"""
PIL Image 列表 → ComfyUI IMAGE tensor [B, H, W, C]，值域 [0, 1]。
多张尺寸不同时，以最大尺寸为准，较小图像丢弃。
"""
if not images:
placeholder = Image.new("RGB", (512, 512), color=(128, 128, 128))
images = [placeholder]
base_size = max(images, key=lambda img: img.size[0] * img.size[1]).size
matched = [img for img in images if img.size == base_size]
skipped = len(images) - len(matched)
if skipped:
print(f"[豆包生图] 丢弃 {skipped} 张非最大尺寸图像，仅输出 {base_size[0]}×{base_size[1]} 的 {len(matched)} 张")
tensors = []
for img in matched:
arr = np.array(img.convert("RGB")).astype(np.float32) / 255.0
tensors.append(torch.from_numpy(arr))
return torch.stack(tensors, dim=0) # [B, H, W, C]
class DoubaoImage:
"""豆包生图 —— 通过宽高比 + 分辨率档位选择尺寸，后端自动换算真实像素"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"模型": (
_MODELS,
{"default": _MODELS[0]},
),
"提示词": (
"STRING",
{
"multiline": True,
"default": "",
"tooltip": "用于创建或编辑图像的文本提示",
},
),
"宽高比": (
_ASPECT_RATIOS,
{
"default": "1:1",
"tooltip": "图像宽高比。所有分辨率档位均支持这些比例",
},
),
"分辨率": (
_RESOLUTIONS,
{
"default": "2K",
"tooltip": (
"图像分辨率档位。\n"
"• Seedream 5.0：支持 2K / 3K\n"
"• Seedream 4.5：支持 2K / 4K\n"
"（3K 与 4.5 或 4K 与 5.0 搭配时将报错）"
),
},
),
"生图数量": (
"INT",
{
"default": 1,
"min": 1,
"max": 10,
"step": 1,
"tooltip": "生成图像的数量。2-10 张时自动并发请求，加快出图速度",
},
),
"种子": (
"INT",
{
"default": 0,
"min": 0,
"max": 2147483647,
"step": 1,
"control_after_generate": True,
"tooltip": "用于生成的随机种子",
},
),
"部分失败时停止": (
"BOOLEAN",
{
"default": True,
"tooltip": (
"启用时：任意一张失败即抛出错误并中止。\n"
"禁用时：返回已成功生成的图像，忽略失败项"
),
},
),
},
"optional": {
"图像": (
"IMAGE",
{
"tooltip": (
"用于图生图的输入图像。"
"单参考或多参考生成时，可输入1-10张图像列表"
),
},
),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("图像",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/豆包"
# ── 并发核心：在新 event loop 里 gather N 个 _generate_async ─────────────
async def _run_concurrent(
self,
client: DoubaoImageClient,
生图数量: int,
model: str,
prompt: str,
size: str,
seed: int,
image_tensor,
pbar,
) -> List[dict]:
"""
并发发起 生图数量 个独立请求，每完成一个推进一格进度条。
返回结果列表：[{"index": int, "images": [...], "error": str|None}]
"""
# 固定参数（顺序生成功能暂时隐藏）
seq = "disabled"
max_img = 1
async def _one(idx: int) -> dict:
try:
imgs = await client._generate_async(
model=model,
prompt=prompt,
size=size,
seed=seed,
sequential_image_generation=seq,
max_images=max_img,
image_tensor=image_tensor,
)
return {"index": idx, "images": imgs, "error": None}
except Exception as e:
return {"index": idx, "images": [], "error": str(e)}
# 用 as_completed 方式逐个推进进度条
tasks = [asyncio.create_task(_one(i)) for i in range(生图数量)]
results = [None] * 生图数量
completed = 0
for coro in asyncio.as_completed(tasks):
res = await coro
results[res["index"]] = res
completed += 1
status = "✓" if res["error"] is None else f"✗ {res['error']}"
print(f"[豆包生图] [{completed}/{生图数量}] 第 {res['index'] + 1} 张 → {status}")
if pbar is not None:
pbar.update(1)
return results
# ── 节点主入口 ────────────────────────────────────────────────────────────
def generate(
self,
模型: str,
提示词: str,
宽高比: str,
分辨率: str,
生图数量: int,
种子: int,
部分失败时停止: bool,
图像=None,
):
start_time = time.time()
# 顺序图像生成功能暂时隐藏，固定使用默认值
顺序图像生成 = "disabled"
最大图片数 = 1
# ── 1. 校验提示词 ─────────────────────────────────────────────────────
if not 提示词.strip():
raise ValueError("提示词不能为空，请输入图像描述后重试。")
# ── 2. 解析模型版本并校验分辨率兼容性 ────────────────────────────────
try:
mkey = _model_key(模型)
except ValueError as e:
raise ValueError(str(e)) from None
supported = _MODEL_RESOLUTIONS[mkey]
if 分辨率 not in supported:
raise ValueError(
f"模型 {模型} 不支持 {分辨率} 分辨率。\n"
f"该模型支持：{' / '.join(supported)}"
)
# ── 3. 查表换算真实像素 ───────────────────────────────────────────────
w, h = _SIZE_TABLE[mkey][分辨率][宽高比]
size_str = f"{w}x{h}"
# ── 4. 打印概要 ───────────────────────────────────────────────────────
mode_str = "图生图" if 图像 is not None else "文生图"
print(
f"[豆包生图] {mode_str} | 模型={模型} | {分辨率} {宽高比} → {size_str}"
f" | 数量={生图数量} | 种子={种子}"
)
# ── 5. 初始化客户端 ───────────────────────────────────────────────────
try:
client = DoubaoImageClient()
except ValueError as e:
raise ValueError(str(e)) from None
# ── 6. 进度条（按张数计）──────────────────────────────────────────────
try:
from comfy.utils import ProgressBar
pbar = ProgressBar(生图数量)
except Exception:
pbar = None
# ── 7. 单张 / 多张分支 ────────────────────────────────────────────────
if 生图数量 == 1:
# 单张：走原有同步路径
try:
pil_images: List[Image.Image] = client.generate_sync(
model=模型,
prompt=提示词,
size=size_str,
seed=种子,
sequential_image_generation=顺序图像生成,
max_images=最大图片数,
image_tensor=图像,
)
except RuntimeError as e:
raise RuntimeError(str(e)) from None
except Exception as e:
raise RuntimeError(f"豆包生图请求失败: {e}") from None
if pbar is not None:
pbar.update(1)
else:
# 多张：并发请求
def _run_in_thread():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
return loop.run_until_complete(
self._run_concurrent(
client=client,
生图数量=生图数量,
model=模型,
prompt=提示词,
size=size_str,
seed=种子,
image_tensor=图像,
pbar=pbar,
)
)
finally:
loop.close()
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(_run_in_thread)
try:
results = future.result(timeout=_CONCURRENT_TIMEOUT)
except TimeoutError:
raise RuntimeError(
f"并发生图超时（>{_CONCURRENT_TIMEOUT}s），请检查网络或减少生图数量"
)
# 统计成功 / 失败
success_results = [r for r in results if r and r["error"] is None]
failed_results = [r for r in results if r and r["error"] is not None]
if failed_results:
fail_info = "；".join(
f"第{r['index']+1}张: {r['error']}" for r in failed_results
)
if 部分失败时停止:
raise RuntimeError(
f"{len(failed_results)}/{生图数量} 张生成失败：{fail_info}\n"
"（可将【部分失败时停止】设为 False 以返回已成功的图像）"
)
else:
print(f"[豆包生图] 警告：{len(failed_results)}/{生图数量} 张失败，已忽略：{fail_info}")
if not success_results:
raise RuntimeError("所有图像均生成失败，请检查网络或 API 配置。")
# 按原始 index 排序，展平为 PIL 列表
success_results.sort(key=lambda r: r["index"])
pil_images = []
for r in success_results:
pil_images.extend(r["images"])
# ── 8. PIL → tensor ───────────────────────────────────────────────────
output_tensor = _pil_list_to_tensor(pil_images)
# ── 9. 完成日志 ───────────────────────────────────────────────────────
elapsed = time.time() - start_time
print(
f"[豆包生图] 完成！耗时 {elapsed:.1f}s，"
f"输出 {output_tensor.shape[0]} 张 "
f"{output_tensor.shape[2]}×{output_tensor.shape[1]}"
)
return (output_tensor,)
# ── 节点注册 ──────────────────────────────────────────────────────────────────
NODE_CLASS_MAPPINGS = {
"DoubaoImage": DoubaoImage,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DoubaoImage": "豆包生图",
}
nodes/flux_edit.py
-170
View File
@@ -1,170 +0,0 @@
"""
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
nodes/google_gemini.py
+69 -416
View File
@@ -6,49 +6,29 @@ 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",
".mov": "video/mov",
".avi": "video/avi",
".flv": "video/x-flv",
".webm": "video/webm",
".wmv": "video/x-ms-wmv",
".wmv": "video/wmv",
".3gp": "video/3gpp",
".3gpp": "video/3gpp"
}
try:
import subprocess
FFMPEG_AVAILABLE = True
except ImportError:
FFMPEG_AVAILABLE = False
class GoogleGemini:
"""
@@ -56,13 +36,14 @@ class GoogleGemini:
功能：
- 支持多个 Gemini Flash 模型
- 支持图片、视频和文件输入
- 支持不同思考等级（不思考/低/中/高）- 通过 thinkingConfig.thinkingLevel 控制
- 输出生成的文本内容（主要内容 + 思考内容）
- 支持图片和视频输入
- 支持系统指令
- 支持不同思考深度（不思考/高）
- 输出生成的文本内容
"""
# 支持的思考等级选项
THINKING_LEVELS = ["不思考", "低", "中", "高"]
# 支持的思考深度选项
THINKING_DEPTHS = ["不思考", "高"]
def __init__(self):
"""初始化节点"""
@@ -86,20 +67,23 @@ class GoogleGemini:
"default": "",
"multiline": True
}),
"思考等级": (cls.THINKING_LEVELS, {
"思考深度": (cls.THINKING_DEPTHS, {
"default": "不思考"
})
},
"optional": {
"系统指令": ("STRING", {
"default": "",
"multiline": True
}),
"图片": ("IMAGE",),
"视频": ("VIDEO",),
"文件": ("FILE",)
"视频": ("VIDEO",)
}
}
# 返回值类型
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("主要内容",)
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("主要内容", "思考内容")
# 执行函数名
FUNCTION = "generate"
@@ -110,66 +94,6 @@ class GoogleGemini:
# 允许输出到 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]
@@ -177,8 +101,6 @@ class GoogleGemini:
"""
准备图片数据
如果图片超过20MB，会自动进行缩放和压缩
Args:
images: ComfyUI 图片张量 [B, H, W, C]
@@ -188,183 +110,17 @@ class GoogleGemini:
if images is None:
return None
image_data = []
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"))
b64_str = encode_image_to_base64(img)
image_data.append({
"mime_type": "image/png",
"data": b64_str
})
# 多图总体积控制
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(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)
# 所有压缩方法都失败
raise ValueError(
f"视频文件过大 ({original_size / 1024 / 1024:.2f}MB) 且无法压缩到 20MB 以下。"
f"请安装 FFmpeg 以获得更好的压缩效果，或手动压缩视频。"
)
return image_data if image_data else None
def _prepare_video_data(
self,
@@ -375,7 +131,6 @@ class GoogleGemini:
ComfyUI VIDEO 类型包含视频文件路径信息。
读取视频文件并转换为 base64。
如果视频超过 20MB，会自动进行压缩。
Args:
video: ComfyUI VIDEO 类型数据
@@ -386,48 +141,17 @@ class GoogleGemini:
if video is None:
return None
# VIDEO 类型处理：支持多种格式
# VIDEO 类型通常是一个字典，包含 '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
# 尝试获取视频路径
video_path = video.get("video") or video.get("path") or video.get("file")
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
elif hasattr(video, "video"):
video_path = video.video
if not video_path or not os.path.exists(video_path):
print(f"Google Gemini: 视频文件不存在或路径无效: {video_path}")
@@ -439,70 +163,30 @@ class GoogleGemini:
mime_type = VIDEO_MIME_TYPES.get(ext, "video/mp4")
# 检查文件大小（限制 20MB）
file_size = os.path.getsize(video_path)
if file_size > 20 * 1024 * 1024:
raise ValueError(
f"视频文件过大 ({file_size / 1024 / 1024:.2f}MB)，"
f"请使用不超过 20MB 的视频文件"
)
# 读取并编码视频
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)}")
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]:
"""
解析包含思考内容和主要内容的响应
@@ -530,16 +214,16 @@ class GoogleGemini:
# 主要内容
main_text = part.get("text", "")
return main_text
return (main_text, thought_text)
def generate(
self,
模型: str,
提示词: str,
思考等级: str,
思考深度: str,
系统指令: Optional[str] = None,
图片: Optional[torch.Tensor] = None,
视频=None,
文件: Optional[FileData] = None
视频=None
) -> Tuple[str]:
"""
生成文本
@@ -547,13 +231,13 @@ class GoogleGemini:
Args:
模型: 使用的模型名称
提示词: 用户提示词
思考等级: 思考等级选项
思考深度: 思考深度选项
系统指令: 系统级指令
图片: 输入图片
视频: 输入视频
文件: 输入文件（PDF/TXT）
Returns:
(主要内容, 思考内容)
生成的文本 (STRING,)
"""
start_time = time.time()
@@ -575,12 +259,6 @@ class GoogleGemini:
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 提示词:
@@ -589,32 +267,31 @@ class GoogleGemini:
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: 思考等级 = {思考等级}")
print(f"Google Gemini: 思考深度 = {思考深度}")
print(f"Google Gemini: 发送请求...")
# 获取端点和构建请求体
endpoint = self.client.get_endpoint(model=模型)
endpoint = self.client.get_endpoint(model=模型, thinking_depth=思考深度)
request_body = self.client.build_request_body(
prompt=提示词,
model=模型,
thinking_level=思考等级,
system_instruction=系统指令,
image_data=image_data,
video_data=video_data,
document_data=document_data
video_data=video_data
)
print(f"Google Gemini: 发送请求...")
# 根据是否有视频设置超时
timeout = 300 if video_data else 180
# 调用底层 API 获取原始响应
async def get_raw_response():
return await self.client.request_async(
endpoint,
request_body,
session=None
session=None,
timeout=timeout
)
# 在独立线程中执行异步请求
@@ -624,56 +301,32 @@ class GoogleGemini:
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", "")
main_text, thought_text = self._parse_dual_output(raw_response)
# 输出信息
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)} 字符")
print(f"Google Gemini: 思考内容长度: {len(thought_text)} 字符")
# 输出预览
if main_text:
preview = main_text[:100] + "..." if len(main_text) > 100 else main_text
print(f"Google Gemini: 主要内容预览: {preview}")
return (main_text,)
return (main_text, thought_text)
except ValueError as e:
# 检测是否为授权错误
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
print(f"Google Gemini: 输入错误 - {str(e)}")
raise
except RuntimeError as e:
# 日志只打第一行；报错框展示完整多行
error_full = str(e)
print(f"Google Gemini: ❌ {error_full.split('\n')[0]}")
raise RuntimeError(error_full) from None
print(f"Google Gemini: API 错误 - {str(e)}")
raise
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
print(f"Google Gemini: 未知错误 - {str(e)}")
raise
nodes/gpt_image.py
-262
View File
@@ -1,262 +0,0 @@
"""
o1key GPT Image 节点
支持 gpt-image-1 / gpt-image-1.5 模型的文生图、图生图、图像编辑（带蒙版）
"""
import time
from ..clients.gpt_image_client import GptImageClient
from ..utils.image_utils import parse_batch_prompts
try:
from comfy.model_management import processing_interrupted, InterruptProcessingException
_INTERRUPT_AVAILABLE = True
except ImportError:
_INTERRUPT_AVAILABLE = False
processing_interrupted = lambda: False
InterruptProcessingException = RuntimeError
class O1keyGPTImage:
"""
o1key GPT Image 节点
功能：
- 文生图：仅提供 prompt
- 图生图：提供 prompt + 图片（无遮罩）
- 图像编辑：提供 prompt + 图片 + 遮罩（白色区域将被替换）
- 批量模式：prompt 中用单独一行 --- 分隔多条提示词
参数：
- prompt : 文本提示词（多行；用 --- 独占一行分隔批量提示词）
- 模型 : 模型选择
- 分辨率 : 图像尺寸（auto 让 API 自动决定）
- 生图数量 : 每条提示词生成数量 1-8
- 质量 : 生成质量
- seed : 随机种子（0 表示不指定）
- 图片 : 可选参考图（用于图生图或编辑）
- 遮罩 : 可选蒙版（白色区域将被替换）
"""
@classmethod
def INPUT_TYPES(cls):
# 创建9个独立的参考图输入
optional_inputs = {}
for i in range(1, 10):
optional_inputs[f"参考图{i}"] = ("IMAGE", {
"tooltip": f"Optional reference image {i} for image editing.",
})
optional_inputs["模型"] = ([
"gpt-image-2-按量",
"gpt-image-2-次卡",
], {
"default": "gpt-image-2-次卡",
})
optional_inputs["分辨率"] = ([
"智能",
# ── 1K ──
"1024x1024（1K 正方形 1:1）",
"1536x1024（1K 横版 3:2）",
"1024x1536（1K 竖版 2:3）",
"1365x1024（1K 横版 4:3）",
"1024x1365（1K 竖版 3:4）",
"1820x1024（1K 横版 16:9）",
"1024x1820（1K 竖版 9:16）",
# ── 2K ──
"2048x2048（2K 正方形 1:1）",
"3072x2048（2K 横版 3:2）",
"2048x3072（2K 竖版 2:3）",
"2732x2048（2K 横版 4:3）",
"2048x2732（2K 竖版 3:4）",
"3640x2048（2K 横版 16:9）",
"2048x3640（2K 竖版 9:16）",
# ── 4K ──
"3840x3840（4K 正方形 1:1）",
"3840x2560（4K 横版 3:2）",
"2560x3840（4K 竖版 2:3）",
"3840x2880（4K 横版 4:3）",
"2880x3840（4K 竖版 3:4）",
"3840x2160（4K 横版 16:9）",
"2160x3840（4K 竖版 9:16）",
], {
"default": "智能",
"tooltip": "Image size (智能 = API decides)",
})
optional_inputs["生图数量"] = ("INT", {
"default": 1,
"min": 1,
"max": 8,
"step": 1,
"display": "number",
"tooltip": "How many images to generate per prompt",
})
optional_inputs["质量"] = (["高", "中", "低", "自动"], {
"default": "自动",
"tooltip": "Image quality: 高=high, 中=medium, 低=low, 自动=auto",
})
optional_inputs["seed"] = ("INT", {
"default": 0,
"min": 0,
"max": 2**31 - 1,
"step": 1,
"display": "number",
"control_after_generate": True,
"tooltip": "Random seed (0 = not specified)",
})
optional_inputs["遮罩"] = ("MASK", {
"tooltip": "Optional mask for inpainting (white areas will be replaced)",
})
return {
"required": {
"prompt": ("STRING", {
"default": "",
"multiline": True,
"tooltip": "Text prompt for GPT Image. Use --- on its own line to separate batch prompts.",
}),
},
"optional": optional_inputs,
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("IMAGE",)
FUNCTION = "generate"
CATEGORY = "o1key/image"
OUTPUT_NODE = False
def generate(
self,
prompt: str,
模型: str = "gpt-image-2-次卡",
分辨率: str = "auto",
质量: str = "自动",
生图数量: int = 1,
seed: int = 0,
遮罩=None,
**kwargs,
):
"""
生成图像（文生图 / 图生图 / 图像编辑 / 批量提示词）
路由逻辑：
- 无图片 → generations 接口（文生图）
- 有图片，无遮罩 → edits 接口（图生图）
- 有图片，有遮罩 → edits 接口（图像编辑 + 蒙版）
- prompt 含 --- → 批量模式，逐条调用上述接口
"""
start_time = time.time()
# ── 0. 收集多参考图输入 ────────────────────────────────────────────────
reference_tensors = []
for i in range(1, 10):
key = f"参考图{i}"
if key in kwargs and kwargs[key] is not None:
reference_tensors.append(kwargs[key])
图片 = reference_tensors if reference_tensors else None
# ── 1. 参数校验 ───────────────────────────────────────────────────────
if 遮罩 is not None and 图片 is None:
raise ValueError("提供了遮罩但未提供图片，请同时提供图片和遮罩")
# ── 2. 解析分辨率显示值 → API 参数值 ──────────────────────────────────
size = "auto" if 分辨率 == "智能" else 分辨率.split("（")[0].strip()
# ── 2b. 解析质量显示值 → API 参数值 ───────────────────────────────────
_quality_map = {"高": "high", "中": "medium", "低": "low", "自动": "auto"}
quality = _quality_map.get(质量, "auto")
# ── 3. 创建客户端 ─────────────────────────────────────────────────────
try:
client = GptImageClient()
except ValueError as e:
if str(e) == "未授权！":
print("[o1key GPT Image] 请联系作者授权后方可使用！")
raise ValueError("未授权！") from None
raise
try:
# ── 4. 解析批量提示词 ─────────────────────────────────────────────
batch_prompts = parse_batch_prompts(prompt)
# ── 5. 调用 API ───────────────────────────────────────────────────
all_pil_images = []
if batch_prompts:
# 批量模式：逐条提示词调用
total = len(batch_prompts)
print(f"[o1key GPT Image] 批量模式 | {total} 条提示词 | 每条生成 {生图数量} 张")
for idx, p in enumerate(batch_prompts, 1):
if _INTERRUPT_AVAILABLE and processing_interrupted():
print("[o1key GPT Image] 用户取消，已中断批量生成")
raise InterruptProcessingException()
try:
pil_images = client.run_sync(
prompt=p,
model=模型,
quality=quality,
size=size,
n=生图数量,
seed=seed,
image_tensor=图片,
mask_tensor=遮罩,
)
all_pil_images.extend(pil_images)
snippet = p[:30] + ("..." if len(p) >= 30 else "")
print(f"[o1key GPT Image] [{idx}/{total}] ✓ {snippet}")
except InterruptProcessingException:
raise
except Exception as e:
error_msg = str(e).split('\n')[0]
snippet = p[:30] + ("..." if len(p) >= 30 else "")
print(f"[o1key GPT Image] [{idx}/{total}] ❌ {snippet} → {error_msg}")
else:
# 单提示词模式
if not prompt or not prompt.strip():
raise ValueError("提示词不能为空")
try:
pil_images = client.run_sync(
prompt=prompt,
model=模型,
quality=quality,
size=size,
n=生图数量,
seed=seed,
image_tensor=图片,
mask_tensor=遮罩,
)
all_pil_images.extend(pil_images)
except InterruptProcessingException:
raise
except Exception as e:
error_msg = str(e).split('\n')[0]
print(f"[o1key GPT Image] ❌ {error_msg}")
raise RuntimeError(error_msg) from None
# ── 6. 检查是否有可用图像 ─────────────────────────────────────────
if not all_pil_images:
raise RuntimeError("所有提示词均生成失败，无可用图像输出")
# ── 7. PIL → tensor ───────────────────────────────────────────────
output_tensor = GptImageClient._pil_list_to_tensor(all_pil_images)
# ── 8. 完成日志 ───────────────────────────────────────────────────
elapsed = time.time() - start_time
print(
f"[o1key GPT Image] 完成！耗时 {elapsed:.1f}s，"
f"输出 {output_tensor.shape[0]} 张 "
f"{output_tensor.shape[2]}×{output_tensor.shape[1]}"
)
return (output_tensor,)
finally:
self._print_balance(client)
def _print_balance(self, client):
try:
balance_data = client.query_balance_sync()
balance_info = client.format_balance_info(balance_data)
print(f"[o1key GPT Image] {balance_info}")
except Exception:
pass
nodes/image_stitch_pro.py
-241
View File
@@ -1,241 +0,0 @@
"""
高级图像拼接节点
支持最多 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),)
nodes/kling_video.py
-730
View File
@@ -1,730 +0,0 @@
"""
Kling 3.0 Video Nodes
"""
import os
import tempfile
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
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 视频生成节点（支持多镜头）"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"提示词": ("STRING", {"multiline": True, "default": ""}),
"反向提示词": ("STRING", {"multiline": True, "default": ""}),
"模型版本": (["v3", "v2-6"], {"default": "v3"}),
"时长": ([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_时长": ("STRING", {"default": "5"}),
"镜头2_提示词": ("STRING", {"multiline": True, "default": ""}),
"镜头2_时长": ("STRING", {"default": "5"}),
"镜头3_提示词": ("STRING", {"multiline": True, "default": ""}),
"镜头3_时长": ("STRING", {"default": "5"}),
"镜头4_提示词": ("STRING", {"multiline": True, "default": ""}),
"镜头4_时长": ("STRING", {"default": "5"}),
"镜头5_提示词": ("STRING", {"multiline": True, "default": ""}),
"镜头5_时长": ("STRING", {"default": "5"}),
"镜头6_提示词": ("STRING", {"multiline": True, "default": ""}),
"镜头6_时长": ("STRING", {"default": "5"}),
}
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("视频",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/Kling"
async def generate(self, **kwargs):
"""生成视频（支持多镜头）"""
prompt = kwargs["提示词"]
negative_prompt = kwargs["反向提示词"]
model_ver = kwargs.get("模型版本", "v3")
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"
# ── v2-6 模型约束校验 ──────────────────────────────────────────
if model_ver == "v2-6":
if duration == 15:
raise ValueError(
"v2-6 模型不支持 15s 时长，请选择 5s 或 10s。"
)
if mode == "std" and voice == "voice":
raise ValueError(
"v2-6 模型的标准画质（720p）不支持生成音频，请关闭生成音频或切换至 1080p。"
)
# ── 多镜头检测 ────────────────────────────────────────────────
multi_prompt_list = []
for i in range(1, 7):
sb_prompt = kwargs.get(f"镜头{i}_提示词", "").strip()
if sb_prompt:
raw_dur = kwargs.get(f"镜头{i}_时长", "5")
try:
sb_duration = int(str(raw_dur).strip()) if str(raw_dur).strip() else 5
except ValueError:
sb_duration = 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-{model_ver}-{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"
# ── 保存路径（临时文件，避免与下游保存节点重复落盘）──────────────────
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="kling_")
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 首尾帧到视频节点"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"首帧": ("IMAGE",),
"尾帧": ("IMAGE",),
"提示词": ("STRING", {"multiline": True, "default": ""}),
"模型": (["v3", "v2-6"], {"default": "v3"}),
"分辨率": (["1080p", "720p"],),
"时长": ([5, 10, 15],),
"生成音频": (["打开", "关闭"], {"default": "打开"}),
"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/v2-6 → kling-v3/kling-v2-6（后端值还原）
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-{ver}-{mode}-{dur}s-{voice}
mode = "pro" if resolution == "1080p" else "std"
voice = "voice" if generate_audio == "打开" else "novoice"
# ── v2-6 模型约束校验 ──────────────────────────────────────────
model_ver = kwargs["模型"] # "v3" or "v2-6"
if model_ver == "v2-6":
if duration == 15:
raise ValueError(
"v2-6 模型不支持 15s 时长，请选择 5s 或 10s。"
)
if mode == "std" and voice == "voice":
raise ValueError(
"v2-6 模型的标准画质（720p）不支持生成音频，请关闭生成音频或切换至 1080p。"
)
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
# 保存路径（临时文件，避免与下游保存节点重复落盘）
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="kling_")
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": {
"模型": (["v3", "v2-6"], {"default": "v3"}),
"分辨率": (["1080p", "720p"],),
"时长": ([5, 10, 15], {"default": 5}),
"人物朝向": (["video", "image"],),
"保留原声": (["打开", "关闭"], {"default": "打开"}),
"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 类型参考视频 + 图片人物动作迁移（走 new API 三段式）"""
import base64
prompt = kwargs["提示词"]
reference_image = kwargs["参考图片"]
reference_video = kwargs["参考视频"]
keep_original_sound = kwargs.get("保留原声", "打开")
character_orientation = kwargs.get("人物朝向", "video")
mode = kwargs.get("分辨率", "1080p")
duration = kwargs.get("时长", 5)
mode_api = "pro" if mode == "1080p" else "std" # 映射为 API 参数值
model = kwargs.get("模型", "v3")
model_name = f"kling-{model}-motion-{mode_api}-{duration}s"
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 对象获取本地文件路径并读取 ───────────────────────
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:
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")
# ── 构建请求体（new API 格式）─────────────────────────────────
body = {
"model": model_name,
"prompt": prompt,
"image_url": image_b64,
"video_url": video_b64,
"character_orientation": character_orientation,
"mode": mode_api,
"keep_original_sound": "yes" if keep_original_sound == "打开" else "no",
}
# ── 保存路径（临时文件，避免与下游保存节点重复落盘）──────────────────
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="kling_motion_")
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.motion_control_async(
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": "文/图生视频 自研模型",
"KlingFirstLastFrame": "首尾帧生视频 自研模型",
"KlingMotionControlTest": "动作控制 自研模型",
"AspectRatioPreset": "图片宽高比预设",
}
nodes/load_file.py
-125
View File
@@ -1,125 +0,0 @@
"""
LoadFile 节点（增强版）
支持单文件路径和文件夹路径，输出 FILE_LIST 类型供全能LLM等节点使用
"""
import base64
import os
from pathlib import Path
from typing import Tuple, List
from ..utils.file_types import FileData, FileList, DOCUMENT_MIME_TYPES, FILE_SIZE_LIMIT, TOTAL_FILE_SIZE_LIMIT
class LoadFile:
"""
加载文件节点
- 单文件路径：加载指定文件
- 文件夹路径：加载文件夹内所有支持的文件（非递归）
- 两者可同时使用，结果合并输出
- 输出 FILE_LIST 类型，可直接连接到全能LLM对话助手
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {},
"optional": {
"单文件路径": ("STRING", {
"default": "",
"multiline": False,
"placeholder": "文件完整路径，多个文件用英文逗号分隔",
}),
"文件夹路径": ("STRING", {
"default": "",
"multiline": False,
"placeholder": "文件夹路径，自动读取其中所有支持的文件",
}),
},
}
RETURN_TYPES = ("FILE_LIST", "STRING")
RETURN_NAMES = ("文件列表", "文件信息")
FUNCTION = "load_file"
CATEGORY = "file/input"
def load_file(self, 单文件路径: str = "", 文件夹路径: str = "") -> Tuple[FileList, str]:
collected: List[Path] = []
# 1. 单文件路径（逗号分隔，支持多个）
if 单文件路径.strip():
for raw in 单文件路径.split(","):
p = Path(raw.strip().strip('"').strip("'"))
if not p.is_absolute():
p = Path.cwd() / p
if not p.exists():
raise ValueError(f"文件不存在: {p}")
if not p.is_file():
raise ValueError(f"路径不是文件: {p}")
collected.append(p)
# 2. 文件夹路径
if 文件夹路径.strip():
folder = Path(文件夹路径.strip().strip('"').strip("'"))
if not folder.is_absolute():
folder = Path.cwd() / folder
if not folder.exists():
raise ValueError(f"文件夹不存在: {folder}")
if not folder.is_dir():
raise ValueError(f"路径不是文件夹: {folder}")
for p in sorted(folder.iterdir()):
if p.is_file() and p.suffix.lower() in DOCUMENT_MIME_TYPES:
collected.append(p)
if not collected:
raise ValueError(f"文件夹中没有支持的文件: {folder}")
if not collected:
raise ValueError("请至少提供一个文件路径或文件夹路径")
# 去重（保持顺序）
seen = set()
unique: List[Path] = []
for p in collected:
key = str(p.resolve())
if key not in seen:
seen.add(key)
unique.append(p)
# 大小检查 & 读取
total_size = 0
file_list: FileList = []
info_lines = []
for p in unique:
ext = p.suffix.lower()
if ext not in DOCUMENT_MIME_TYPES:
print(f"LoadFile: 跳过不支持的文件类型 {p.name}")
continue
file_size = p.stat().st_size
if file_size > FILE_SIZE_LIMIT:
raise ValueError(
f"文件 {p.name} 大小 {file_size / 1024 / 1024:.1f}MB 超过单文件 50MB 限制"
)
total_size += file_size
if total_size > TOTAL_FILE_SIZE_LIMIT:
raise ValueError(f"所有文件总大小超过 50MB 限制")
mime = DOCUMENT_MIME_TYPES[ext]
with open(p, "rb") as f:
b64 = base64.b64encode(f.read()).decode("utf-8")
file_list.append(FileData(
path=str(p),
filename=p.stem,
extension=ext,
mime_type=mime,
data=b64,
size=file_size,
))
info_lines.append(f" {p.name} ({file_size / 1024:.1f}KB, {mime})")
print(f"LoadFile: 加载 {p.name} ({file_size / 1024:.1f}KB)")
info = f"共 {len(file_list)} 个文件，总大小 {total_size / 1024:.1f}KB\n" + "\n".join(info_lines)
return (file_list, info)
nodes/multi_res_preview.py
-165
View File
@@ -1,165 +0,0 @@
"""
多分辨率图像预览节点
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}}
nodes/nano_banana_pro.py
+96 -546
View File
@@ -1,35 +1,19 @@
"""
Nano Banana Pro 节点
ComfyUI 自定义节点，用于调用 Gemini 模型生成图像
ComfyUI 自定义节点，用于调用 Gemini 3 Pro 模型生成图像
"""
import time
import math
import random
import asyncio
import aiohttp
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Tuple, List
from typing import Optional, Tuple
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
from ..models_config import get_enabled_models, get_model_description
# 导入 ComfyUI 原生进度条
try:
@@ -39,53 +23,6 @@ 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 = True
# 是否启用请求体日志（打印发送给 API 的请求体，base64 图片数据将自动截断）
# 设置为 True 以启用请求体日志，False 以禁用
REQUEST_LOG_ENABLED = True
# ============================================================================
_NODE = "Nano Banana Pro"
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
"""
将 PIL Image 列表转换为 ComfyUI tensor，安全处理多张不同尺寸的情况。
策略：
- 以像素数最大的图尺寸为基准
- 只输出与最大尺寸相同的图，其余较小的图丢弃
"""
if not images:
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
return pil_to_tensor([placeholder])
base_size = max(images, key=lambda img: img.size[0] * img.size[1]).size
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}: 丢弃 {len(skipped)} 张较小尺寸的图 ({sizes_str})，"
f"仅输出最大尺寸 {base_size[0]}×{base_size[1]} 的 {len(matched)} 张"
)
return pil_to_tensor(matched)
class NanoBananaPro:
"""
@@ -104,16 +41,14 @@ class NanoBananaPro:
# 支持的模型列表（从配置文件动态加载）
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"
"2:3", "3:2", "4:5", "5:4", "21:9"
]
# 支持的分辨率列表（全量兜底，实际由 get_all_supported_resolutions() 动态生成）
RESOLUTIONS = ["512px", "1K", "2K", "4K"]
# 支持的分辨率列表
RESOLUTIONS = ["1K", "2K", "4K"]
def __init__(self):
"""初始化节点"""
@@ -135,26 +70,10 @@ class NanoBananaPro:
if not enabled_models:
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
# 动态获取所有启用模型支持的宽高比（去重合并）
all_aspect_ratios = get_all_supported_aspect_ratios()
if not all_aspect_ratios:
all_aspect_ratios = cls.ASPECT_RATIOS
# 动态获取所有启用模型支持的分辨率（去重合并）
all_resolutions = get_all_supported_resolutions()
if not all_resolutions:
all_resolutions = cls.RESOLUTIONS
# 创建9个独立的图像输入
optional_inputs = {}
for i in range(1, 10): # 1-9
optional_inputs[f"参考图{i}"] = ("IMAGE",)
optional_inputs["代理端口（如7897）"] = ("STRING", {
"default": "",
"multiline": False,
"placeholder": "本地代理端口，如 7897（Clash Verge）或 10808（v2rayN），留空不使用"
})
return {
"required": {
@@ -165,10 +84,10 @@ class NanoBananaPro:
"模型": (enabled_models, {
"default": enabled_models[0]
}),
"宽高比": (all_aspect_ratios, {
"宽高比": (cls.ASPECT_RATIOS, {
"default": "1:1"
}),
"分辨率": (all_resolutions, {
"分辨率": (cls.RESOLUTIONS, {
"default": "2K"
}),
"生图数量": ("INT", {
@@ -177,14 +96,15 @@ class NanoBananaPro:
"max": 1000,
"step": 1
}),
"谷歌搜索（联网）": (["关闭", "打开"], {
"default": "关闭"
"像素缩放": ("BOOLEAN", {
"default": False
}),
"图片搜索（联网）": (["关闭", "打开"], {
"default": "关闭"
}),
"返回格式": (["url", "base64"], {
"default": "url"
"分辨率像素": ("FLOAT", {
"default": 1.0,
"min": 0.1,
"max": 100.0,
"step": 0.1,
"display": "number"
}),
"seed": ("INT", {
"default": 0,
@@ -199,13 +119,6 @@ class NanoBananaPro:
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("输出图像",)
# 导入 ComfyUI 的文件夹路径管理
try:
import folder_paths
FOLDER_PATHS_AVAILABLE = True
except ImportError:
FOLDER_PATHS_AVAILABLE = False
# 执行函数名
FUNCTION = "generate"
@@ -283,166 +196,6 @@ class NanoBananaPro:
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,
save_to_disk: bool = True,
image_format: str = "url",
) -> dict:
"""执行单个生成任务"""
result = {
"global_task_index": global_task_index,
"prompt": prompt,
"success": False,
"generated_count": 0,
"saved_files": [],
"output_images": [],
"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,
image_format=image_format,
)
if gen_result:
images_list, _ = gen_result
if save_to_disk:
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
else:
result["output_images"] = images_list
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,
save_to_disk: bool = True,
image_format: str = "url",
) -> 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)
max_concurrent = 50
num_batches = math.ceil(total_tasks / max_concurrent)
all_results = []
completed = 0
success_count = 0
fail_count = 0
connector = aiohttp.TCPConnector(ssl=False, 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,
save_to_disk=save_to_disk,
image_format=image_format,
)
)
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,
@@ -450,36 +203,30 @@ class NanoBananaPro:
宽高比: str,
分辨率: str,
生图数量: int,
像素缩放: bool,
分辨率像素: float,
seed: int,
**kwargs
) -> Tuple[torch.Tensor]:
"""
生成图像
Args:
prompt: 提示词
模型: 模型名称
宽高比: 宽高比
分辨率: 分辨率
生图数量: 批次大小
像素缩放: 是否启用像素缩放
分辨率像素: 目标像素数（百万像素）
seed: 随机种子
**kwargs: 搜索开关（谷歌搜索（联网）/ 图片搜索（联网））及动态参考图输入 (参考图1-9)
注：两个搜索参数名含全角括号，不能作为 Python 形参，从 kwargs 中提取
注意：
调试日志功能已移至文件顶部配置，通过修改 DEBUG_LOG_ENABLED 常量控制
**kwargs: 动态参考图输入 (参考图1-9)
Returns:
生成的图像张量 (IMAGE,)
"""
start_time = time.time()
# 从 kwargs 提取搜索参数（界面显示为「关闭/打开」，转为 bool 供调用）
enable_grounding: bool = (kwargs.pop("谷歌搜索（联网）", "关闭") == "打开")
enable_image_search: bool = (kwargs.pop("图片搜索（联网）", "关闭") == "打开")
proxy_port: str = kwargs.pop("代理端口（如7897）", "")
image_format: str = kwargs.pop("返回格式", "url")
# 创建 ComfyUI 原生进度条
pbar = None
if PROGRESS_BAR_AVAILABLE:
@@ -490,49 +237,12 @@ class NanoBananaPro:
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)}")
# 注入代理设置（每次执行都刷新，支持用户中途修改端口）
self.client.proxy_url = GeminiAPIClient.build_proxy_url(proxy_port)
if self.client.proxy_url:
print(f"Nano Banana Pro: 已启用代理加速 → {self.client.proxy_url}")
# 校验分辨率与模型的兼容性
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 = []
@@ -548,284 +258,124 @@ class NanoBananaPro:
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
print(f"Nano Banana Pro: 已缩放 {len(scaled_images)} 张图像到 {分辨率像素}M 像素")
# 转换为 API 所需的格式
if input_images:
print(f"Nano Banana Pro: 图生图模式 (输入 {len(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"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
# 进度回调 - 打印错误信息并更新进度条，添加内存监控
# 进度回调 - 实时显示每个任务的完成状态，并更新 ComfyUI 进度条
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}] 第 {success_count} 张生成成功")
else:
fail_count += 1
error_brief = error_msg[:50] + "..." if error_msg and len(error_msg) > 50 else error_msg
print(f"Nano Banana Pro: ✗ [{current}/{total}] 生成失败 - {error_brief}")
# 更新 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 * 生图数量
print(f"Nano Banana Pro: 批量提示词模式 ({num_prompts} 个提示词 × {生图数量} 张/提示词 = {total_images} 张图)")
print(f"Nano Banana Pro: 发送请求")
print(f"Nano Banana Pro: 生图中...")
# ===== 批量提示词模式：异步并发，内存输出 =====
# 重新创建进度条以匹配实际总数
if pbar is not None:
pbar = ProgressBar(total_images)
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="",
pbar=pbar,
enable_grounding=enable_grounding,
enable_image_search=enable_image_search,
save_to_disk=False,
image_format=image_format,
)
)
finally:
loop.close()
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(run_async_in_thread)
try:
results = future.result(timeout=900)
except TimeoutError:
raise RuntimeError("任务执行超时（900秒），请减少提示词数量或检查网络连接")
# 统计结果
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)
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}")
# 收集内存中的图像
output_images = []
for r in results:
output_images.extend(r.get("output_images", []))
if not output_images:
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
output_images = [placeholder]
output_tensor = _images_to_tensor_safe(output_images, _NODE)
import gc
gc.collect()
return (output_tensor,)
generated_images = self.client.generate_multi_prompts_sync(
prompts=batch_prompts,
model=模型,
resolution=分辨率,
aspect_ratio=宽高比,
images_per_prompt=生图数量,
images=input_images,
progress_callback=progress_callback
)
if fail_count > 0:
print(f"Nano Banana Pro: 生图完成 (成功: {success_count}, 失败: {fail_count})")
else:
print(f"Nano Banana Pro: 全部生图成功！")
else:
# 单提示词模式
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,
image_format=image_format,
)
else:
# 多张：异步并发，内存输出
if pbar is not None:
pbar = ProgressBar(生图数量)
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="",
pbar=pbar,
enable_grounding=enable_grounding,
enable_image_search=enable_image_search,
save_to_disk=False,
image_format=image_format,
)
)
finally:
loop.close()
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(run_async_in_thread)
try:
results = future.result(timeout=900)
except TimeoutError:
raise RuntimeError("任务执行超时（900秒），请减少生图数量或检查网络连接")
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)
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}")
# 收集内存中的图像
output_images = []
for r in results:
output_images.extend(r.get("output_images", []))
if not output_images:
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
output_images = [placeholder]
output_tensor = _images_to_tensor_safe(output_images, _NODE)
import gc
gc.collect()
return (output_tensor,)
print(f"Nano Banana Pro: {'图生图' if input_images else '文生图'}模式")
print(f"Nano Banana Pro: 发送请求")
print(f"Nano Banana Pro: 生图中...")
# 优化：限制输出图片数量，避免内存爆炸
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
generated_images = self.client.generate_sync(
prompt=prompt,
model=模型,
resolution=分辨率,
aspect_ratio=宽高比,
batch_size=生图数量,
images=input_images,
progress_callback=progress_callback
)
if fail_count > 0:
print(f"Nano Banana Pro: 生图完成 (成功: {success_count}, 失败: {fail_count})")
else:
print(f"Nano Banana Pro: 全部生图成功！")
# 转换输出图像
output_tensor = _images_to_tensor_safe(output_images, _NODE)
output_tensor = pil_to_tensor(generated_images)
# 计算耗时并打印最终统计
# 计算耗时
elapsed = time.time() - start_time
if elapsed < 1:
time_str = f"{elapsed:.3f}s"
else:
time_str = f"{elapsed:.2f}s"
print(f"Nano Banana Pro: 完成生图 (耗时: {elapsed:.2f}s, 成功生成 {len(generated_images)} 张图像)")
# 打印最终汇总
if fail_count > 0:
print(f"完成！总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张")
else:
print(f"完成！总耗时 {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
raise ValueError(str(e)) from None
else:
# 用户输入错误
print(f"Nano Banana Pro: 输入错误 - {str(e)}")
raise
except RuntimeError as e:
raise RuntimeError(str(e)) from None
# API 或网络错误
print(f"Nano Banana Pro: API 错误 - {str(e)}")
raise
except Exception as e:
raise type(e)(str(e)) from None
# 其他未知错误
print(f"Nano Banana Pro: 未知错误 - {str(e)}")
raise
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()
except Exception as e:
print(f"Nano Banana Pro: ⚠️ 余额查询失败 - {str(e)}")
nodes/nano_banana_v2.py
-1238
View File
File diff suppressed because it is too large Load Diff
nodes/remove_metadata.py
-371
View File
@@ -1,371 +0,0 @@
"""
图像元数据去除节点
替代 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,)
nodes/seedance_video.py
-443
View File
@@ -1,443 +0,0 @@
"""
Seedance 视频生成节点
节点列表:
- Seedance: 文生视频 / 图生视频 / 首尾帧生视频（根据图片输入自动切换模式）
"""
import io
import os
import tempfile
import aiohttp
import torch
from ..clients.seedance_client import SeedanceClient
from ..clients.gemini_client import GeminiAPIClient
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64, pil_to_tensor
from ..utils.r2_uploader import upload_video, upload_audio
from comfy_api.latest import InputImpl
# ── 模型列表 ──────────────────────────────────────────────────────────────────
_MODELS = [
"doubao-seedance-2-0-260128",
]
_RESOLUTIONS = ["720p", "1080p", "480p"]
# ── 模型能力判断 ──────────────────────────────────────────────────────────────
def _supports_camera_fixed(model: str) -> bool:
"""2.0 系列不支持固定镜头"""
return False # 当前仅 2.0 模型，均不支持
# ── 工具函数 ──────────────────────────────────────────────────────────────────
def _tensor_to_base64_url(tensor) -> str:
"""ComfyUI IMAGE tensor → data:image/png;base64,xxx"""
pil_images = tensor_to_pil(tensor)
b64 = encode_image_to_base64(pil_images[0], format="PNG")
return f"data:image/png;base64,{b64}"
async def _url_to_tensor(url: str) -> torch.Tensor:
"""从 URL 下载图片并转为 ComfyUI IMAGE tensor，失败时返回 None"""
try:
from PIL import Image
async with aiohttp.ClientSession() as session:
async with session.get(url, allow_redirects=True) as resp:
if resp.status != 200:
return None
data = await resp.read()
img = Image.open(io.BytesIO(data)).convert("RGB")
return pil_to_tensor([img])
except Exception as e:
print(f"[Seedance] 末帧图片下载失败: {e}")
return None
def _show_balance():
"""完成后打印余额（静默失败）"""
try:
client = GeminiAPIClient()
data = client.query_balance_sync()
print(f"Seedance: {client.format_balance_info(data)}")
except Exception:
pass
def _make_pbar():
try:
from comfy.utils import ProgressBar
return ProgressBar(100)
except Exception:
return None
def _make_callbacks(tag: str, pbar):
def on_stage(stage: str):
if stage == "submitting":
print(f"[{tag}] 提交中...")
if pbar: pbar.update_absolute(0, 100)
elif stage.startswith("submitted:"):
print(f"[{tag}] 已提交 → {stage.split(':', 1)[1]}")
if pbar: pbar.update_absolute(5, 100)
elif stage == "downloading":
print(f"[{tag}] 下载视频中...")
if pbar: pbar.update_absolute(99, 100)
elif stage == "done":
print(f"[{tag}] 完成")
if pbar: pbar.update_absolute(100, 100)
def on_progress(pct: int):
if pbar: pbar.update_absolute(5 + int(pct * 0.94), 100)
return on_stage, on_progress
# ── 统一节点 ─────────────────────────────────────────────────────────────────
#
# 模式由图片输入自动判断：
# 首帧 = None → T2V 文生视频 （联网搜索生效）
# 首帧 = 图片，尾帧 = None → I2V 图生视频 （固定镜头生效，当前 2.0 不支持故忽略）
# 首帧 = 图片，尾帧 = 图片 → FlipFlop 首尾帧（联网搜索/固定镜头均忽略）
class Seedance:
"""Seedance 视频生成（文生视频 / 图生视频 / 首尾帧，自动判断模式）"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"提示词": ("STRING", {"multiline": True, "default": ""}),
"模型": (_MODELS, {"default": "doubao-seedance-2-0-260128"}),
"分辨率": (_RESOLUTIONS, {"default": "720p"}),
"宽高比": (["16:9", "adaptive", "9:16", "1:1", "4:3", "3:4", "21:9"],
{"default": "16:9"}),
"时长秒(-1=自动)": ("INT", {"default": 5, "min": -1, "max": 30, "step": 1}),
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
"联网搜索": (["关闭", "打开"], {"default": "关闭"}),
"返回末帧图片": (["关闭", "打开"], {"default": "关闭"}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}),
},
"optional": {
"首帧图片": ("IMAGE",),
"尾帧图片": ("IMAGE",),
},
}
RETURN_TYPES = ("VIDEO", "IMAGE")
RETURN_NAMES = ("视频", "末帧图片")
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/Seedance"
async def generate(self, **kwargs):
prompt = kwargs["提示词"].strip()
model = kwargs["模型"]
resolution = kwargs["分辨率"]
ratio = kwargs["宽高比"]
duration = kwargs["时长秒(-1=自动)"]
gen_audio = kwargs["生成音频"] == "打开"
web_search = kwargs["联网搜索"] == "打开"
return_last = kwargs["返回末帧图片"] == "打开"
seed = kwargs.get("seed", 0)
first_image = kwargs.get("首帧图片", None)
last_image = kwargs.get("尾帧图片", None)
# 模式判断
if first_image is None and last_image is not None:
raise ValueError("请同时接入首帧图片，或仅接入首帧图片。")
if first_image is None:
mode = "t2v"
tag = "Seedance文生视频"
file_prefix = "seedance_t2v"
elif last_image is None:
mode = "i2v"
tag = "Seedance图生视频"
file_prefix = "seedance_i2v"
else:
mode = "flipflop"
tag = "Seedance首尾帧"
file_prefix = "seedance_flip"
if not prompt:
raise ValueError("提示词不能为空。")
if duration == -1 and mode == "t2v":
pass # 2.0 均支持自动时长
elif duration == -1 and mode != "t2v":
pass # 2.0 均支持自动时长
metadata: dict = {
"resolution": resolution,
"watermark": False,
}
if ratio != "adaptive":
metadata["ratio"] = ratio
if duration != -1:
metadata["duration"] = duration
if gen_audio:
metadata["generate_audio"] = True
if return_last:
metadata["return_last_frame"] = True
if seed != 0:
metadata["seed"] = seed
# 模式专属参数
if mode == "t2v":
if web_search:
metadata["tools"] = [{"type": "web_search"}]
body = {
"model": model,
"prompt": prompt,
"metadata": metadata,
}
elif mode == "i2v":
first_url = _tensor_to_base64_url(first_image)
metadata["content"] = [
{
"type": "image_url",
"image_url": {"url": first_url},
"role": "first_frame",
},
{"type": "text", "text": prompt},
]
body = {
"model": model,
"prompt": prompt,
"images": [first_url],
"metadata": metadata,
}
else: # flipflop
first_url = _tensor_to_base64_url(first_image)
last_url = _tensor_to_base64_url(last_image)
metadata["content"] = [
{
"type": "image_url",
"image_url": {"url": first_url},
"role": "first_frame",
},
{
"type": "image_url",
"image_url": {"url": last_url},
"role": "last_frame",
},
{"type": "text", "text": prompt},
]
body = {
"model": model,
"prompt": prompt,
"images": [first_url],
"metadata": metadata,
}
# 保存路径（临时文件，避免与下游保存节点重复落盘）
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix=f"{file_prefix}_")
client = SeedanceClient()
pbar = _make_pbar()
on_stage, on_prog = _make_callbacks(tag, pbar)
try:
result_path, last_frame_url = await client.generate_async(
body=body, save_path=save_path,
on_stage=on_stage, on_progress=on_prog,
)
last_frame_tensor = None
if return_last and last_frame_url:
last_frame_tensor = await _url_to_tensor(last_frame_url)
return (InputImpl.VideoFromFile(result_path), last_frame_tensor)
finally:
_show_balance()
# ── 多模态参考生视频节点 ──────────────────────────────────────────────────────
class SeedanceMultiModal:
"""Seedance 2.0 多模态参考生视频（参考图片 + 参考视频 + 参考音频 + 文本）"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"提示词": ("STRING", {"multiline": True, "default": ""}),
"模型": (_MODELS, {"default": "doubao-seedance-2-0-260128"}),
"分辨率": (_RESOLUTIONS, {"default": "720p"}),
"宽高比": (["adaptive", "16:9", "9:16", "1:1", "4:3", "3:4", "21:9"],
{"default": "adaptive"}),
"时长秒(-1=自动)": ("INT", {"default": 5, "min": -1, "max": 15, "step": 1}),
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
"联网搜索": (["关闭", "打开"], {"default": "关闭"}),
"返回末帧图片": (["关闭", "打开"], {"default": "关闭"}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}),
},
"optional": {
"参考图片": ("IMAGE",),
"参考视频1": ("VIDEO",),
"参考视频2": ("VIDEO",),
"参考视频3": ("VIDEO",),
"参考音频1": ("AUDIO",),
"参考音频2": ("AUDIO",),
"参考音频3": ("AUDIO",),
},
}
RETURN_TYPES = ("VIDEO", "IMAGE")
RETURN_NAMES = ("视频", "末帧图片")
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/Seedance"
INPUT_IS_LIST = True
async def generate(self, **kwargs):
# INPUT_IS_LIST=True 时所有参数都是列表，取第一个元素
def _first(v, default=None):
if isinstance(v, list):
return v[0] if v else default
return v if v is not None else default
prompt = _first(kwargs.get("提示词"), "").strip()
model = _first(kwargs.get("模型"))
resolution = _first(kwargs.get("分辨率"))
ratio = _first(kwargs.get("宽高比"))
duration = _first(kwargs.get("时长秒(-1=自动)"), 5)
gen_audio = _first(kwargs.get("生成音频"), "关闭") == "打开"
web_search = _first(kwargs.get("联网搜索"), "关闭") == "打开"
return_last = _first(kwargs.get("返回末帧图片"), "关闭") == "打开"
seed = _first(kwargs.get("seed"), 0)
# 参考图片：INPUT_IS_LIST 时是 [tensor, tensor, ...] 列表，直接保留
raw_images = kwargs.get("参考图片", None)
ref_images = [img for img in raw_images if img is not None] if raw_images else None
ref_videos = [_first(kwargs.get(f"参考视频{i}")) for i in range(1, 4)]
ref_audios = [_first(kwargs.get(f"参考音频{i}")) for i in range(1, 4)]
ref_videos = [v for v in ref_videos if v is not None]
ref_audios = [a for a in ref_audios if a is not None]
# ── 校验 ──────────────────────────────────────────────────────────
has_image = bool(ref_images)
has_video = len(ref_videos) > 0
has_audio = len(ref_audios) > 0
if not has_image and not has_video and not has_audio and not prompt:
raise ValueError("至少需要提供参考图片、参考视频或提示词之一。")
if has_audio and not has_image and not has_video:
raise ValueError("不可单独输入音频，请至少连接一张参考图片或一个参考视频。")
# ── 构建 content 列表 ─────────────────────────────────────────────
content = []
# 参考图片（批次，最多9张）
if has_image:
imgs = ref_images[:9]
if len(ref_images) > 9:
print(f"[SeedanceMultiModal] 参考图片超过9张，仅取前9张（共{len(ref_images)}张）")
for img_tensor in imgs:
# 每个 tensor 可能是 [1,H,W,C] 或 [H,W,C]，统一确保有 batch 维
if img_tensor.dim() == 3:
img_tensor = img_tensor.unsqueeze(0)
url = _tensor_to_base64_url(img_tensor)
content.append({
"type": "image_url",
"image_url": {"url": url},
"role": "reference_image",
})
# 参考视频（最多3个）
for v in ref_videos:
url = await upload_video(v)
content.append({
"type": "video_url",
"video_url": {"url": url},
"role": "reference_video",
})
# 参考音频（最多3段）
for a in ref_audios:
url = await upload_audio(a)
content.append({
"type": "audio_url",
"audio_url": {"url": url},
"role": "reference_audio",
})
# 文本提示词（放最后）
if prompt:
content.append({"type": "text", "text": prompt})
if not content:
raise ValueError("content 为空，请至少提供参考图片、参考视频或提示词。")
# ── 构建请求体（new-api 兼容格式）──────────────────────────────────
metadata: dict = {
"resolution": resolution,
"watermark": False,
"content": content,
}
if ratio != "adaptive":
metadata["ratio"] = ratio
if duration != -1:
metadata["duration"] = duration
if gen_audio:
metadata["generate_audio"] = True
if return_last:
metadata["return_last_frame"] = True
if seed != 0:
metadata["seed"] = seed
if web_search:
metadata["tools"] = [{"type": "web_search"}]
# 顶层 image：取第一张参考图的 base64（new-api 单图字段）
first_image_url = next(
(item["image_url"]["url"] for item in content if item["type"] == "image_url"),
None,
)
body = {
"model": model,
"prompt": prompt if prompt else " ",
"metadata": metadata,
}
if first_image_url:
body["image"] = first_image_url
# ── 保存路径（临时文件，避免与下游保存节点重复落盘）──────────────────
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="seedance_mm_")
client = SeedanceClient()
pbar = _make_pbar()
on_stage, on_prog = _make_callbacks("Seedance多模态", pbar)
try:
result_path, last_frame_url = await client.generate_async(
body=body, save_path=save_path,
on_stage=on_stage, on_progress=on_prog,
)
last_frame_tensor = None
if return_last and last_frame_url:
last_frame_tensor = await _url_to_tensor(last_frame_url)
return (InputImpl.VideoFromFile(result_path), last_frame_tensor)
finally:
_show_balance()
# ── 节点注册 ──────────────────────────────────────────────────────────────────
NODE_CLASS_MAPPINGS = {
"Seedance": Seedance,
"SeedanceMultiModal": SeedanceMultiModal,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Seedance": "Seedance 视频生成",
"SeedanceMultiModal": "Seedance 多模态参考生视频",
}
nodes/sora_video.py
-526
View File
@@ -1,526 +0,0 @@
"""
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
nodes/stream_preview.py
-28
View File
@@ -1,28 +0,0 @@
"""
流式文本预览节点
接收文本输入，支持 markdown 渲染，通过 ComfyUI 事件系统实时推送内容
"""
class StreamPreview:
"""
流式 Markdown 预览节点
接收任意文本，在节点面板中实时渲染为 Markdown 格式
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"文本": ("STRING", {"forceInput": True}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("文本",)
FUNCTION = "preview"
CATEGORY = "text/preview"
OUTPUT_NODE = True
def preview(self, 文本: str):
return {"ui": {"text": [文本]}, "result": (文本,)}
nodes/universal_llm.py
-506
View File
@@ -1,506 +0,0 @@
"""
全能LLM对话助手节点
ComfyUI 自定义节点，通过 OpenAI 兼容协议调用市面上主流的 AI 对话大模型
支持多模态（图片输入），单轮对话，非流式输出
API 密钥和地址通过插件统一配置（环境变量或 .config 文件），与 Google Gemini 节点一致
"""
import os
import time
import base64
import json
from io import BytesIO
from typing import Optional, Tuple, List
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
from ..utils.file_types import FileList
# ============================================================================
# 模型配置
# ============================================================================
SUPPORTED_MODELS = [
"gpt-5.5",
"gemini-3.1-flash-lite-preview",
"gemini-3.1-pro-preview",
"deepseek-v4-pro",
"doubao-seed-2-0-pro-260215",
]
# 图片缩放最大尺寸
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",),
"视频": ("VIDEO",),
"文件": ("FILE_LIST",),
"令牌": ("STRING", {
"default": "",
"multiline": False,
"placeholder": "留空则使用默认 API Key",
}),
},
"hidden": {
"node_id": "UNIQUE_ID",
},
}
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}"
# 文件大小限制
MAX_FILE_SIZE = 50 * 1024 * 1024 # 单文件 50MB
MAX_TOTAL_FILE_SIZE = 50 * 1024 * 1024 # 所有文件总计 50MB
# 常见 MIME 类型映射
MIME_MAP = {
".pdf": "application/pdf",
".txt": "text/plain",
".md": "text/markdown",
".csv": "text/csv",
".json": "application/json",
".py": "text/x-python",
".js": "text/javascript",
".html": "text/html",
".xml": "application/xml",
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
".zip": "application/zip",
}
# 纯文本类型，直接读取内容
TEXT_EXTS = {".txt", ".md", ".csv", ".json", ".py", ".js", ".ts", ".html",
".xml", ".yaml", ".yml", ".toml", ".ini", ".cfg", ".log",
".sh", ".bat", ".sql", ".css", ".scss", ".jsx", ".tsx"}
def _load_files(self, file_paths_str: str) -> List[dict]:
"""读取文件列表，返回 content part 数组"""
if not file_paths_str or not file_paths_str.strip():
return []
paths = [p.strip() for p in file_paths_str.split(",") if p.strip()]
parts = []
total_size = 0
for path in paths:
if not os.path.isfile(path):
raise ValueError(f"文件不存在: {path}")
file_size = os.path.getsize(path)
if file_size > self.MAX_FILE_SIZE:
raise ValueError(f"文件 {os.path.basename(path)} 大小 {file_size / 1024 / 1024:.1f}MB 超过单文件 50MB 限制")
total_size += file_size
if total_size > self.MAX_TOTAL_FILE_SIZE:
raise ValueError(f"所有文件总大小超过 50MB 限制")
ext = os.path.splitext(path)[1].lower()
mime = self.MIME_MAP.get(ext, "application/octet-stream")
filename = os.path.basename(path)
if ext in self.TEXT_EXTS:
# 文本文件直接读取内容
with open(path, "r", encoding="utf-8", errors="replace") as f:
text_content = f.read()
parts.append({
"type": "text",
"text": f"[文件: {filename}]\n```\n{text_content}\n```",
})
else:
# 二进制文件转 base64，使用 file 格式（OpenAI 兼容协议）
with open(path, "rb") as f:
file_data = base64.b64encode(f.read()).decode("utf-8")
parts.append({
"type": "file",
"file": {
"filename": filename,
"file_data": f"data:{mime};base64,{file_data}",
},
})
print(f"全能LLM: 加载文件 {filename} ({file_size / 1024:.1f}KB, {mime})")
return parts
def _build_input(
self,
prompt: str,
images: Optional[torch.Tensor] = None,
file_paths: str = "",
file_list: Optional[FileList] = None,
video=None,
) -> list:
"""构建 chat/completions 格式的 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 以内")
# 处理视频输入（ComfyUI VIDEO 类型）
video_url_str = ""
if video is not None:
# 从 VIDEO 对象中提取文件路径
vp = None
if isinstance(video, dict):
vp = video.get("video") or video.get("path") or video.get("file") or video.get("filename")
if not vp:
for val in video.values():
if isinstance(val, str) and os.path.exists(val):
vp = val
break
elif isinstance(video, str):
vp = video
else:
for attr in ("video", "path", "filename"):
if hasattr(video, attr):
vp = getattr(video, attr)
break
if not vp and hasattr(video, "__dict__"):
for attr_val in video.__dict__.values():
if isinstance(attr_val, str) and os.path.isfile(attr_val):
vp = attr_val
break
if not vp or not os.path.isfile(vp):
raise ValueError(f"视频文件不存在或路径无效: {vp}")
mime_map = {
".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", ".mkv": "video/x-matroska",
}
ext = os.path.splitext(vp)[1].lower()
mime = mime_map.get(ext, "video/mp4")
file_size = os.path.getsize(vp)
print(f"全能LLM: 加载视频 {os.path.basename(vp)} ({file_size / 1024 / 1024:.1f}MB, {mime})")
with open(vp, "rb") as f:
b64 = base64.b64encode(f.read()).decode("utf-8")
video_url_str = f"data:{mime};base64,{b64}"
# 加载文件：优先使用 FILE_LIST，其次使用字符串路径
file_parts = []
if file_list:
for fd in file_list:
print(f"全能LLM: 使用文件 {fd.filename}{fd.extension} ({fd.size / 1024:.1f}KB)")
file_parts.append({
"type": "file",
"file": {
"filename": fd.filename + fd.extension,
"file_data": f"data:{fd.mime_type};base64,{fd.data}",
},
})
elif file_paths:
file_parts = self._load_files(file_paths)
# 纯文本，无图片无文件无视频
if not image_data_urls and not file_parts and not video_url_str:
return [{"role": "user", "content": prompt}]
content_parts = []
# 图片
for url in image_data_urls:
content_parts.append({
"type": "image_url",
"image_url": {"url": url},
})
# 视频：用 image_url 类型传 data URL（Gemini OpenAI 兼容层支持此格式）
# 同时保留 video_url 类型作为备用（其他支持 video_url 的模型）
if video_url_str:
content_parts.append({
"type": "image_url",
"image_url": {"url": video_url_str},
})
# 文件
for fp in file_parts:
content_parts.append(fp)
content_parts.append({
"type": "text",
"text": prompt,
})
return [{"role": "user", "content": content_parts}]
@staticmethod
def _send_stream_token(node_id, token, done=False):
"""通过 PromptServer 向前端推送流式 token"""
try:
from server import PromptServer
PromptServer.instance.send_sync(
"o1key.stream_token",
{"node_id": str(node_id), "token": token, "done": done},
)
except Exception:
pass
def generate(
self,
模型: str,
提示词: str,
图片: Optional[torch.Tensor] = None,
视频=None,
文件: Optional[FileList] = None,
令牌: str = "",
node_id: str = "",
) -> Tuple[str]:
start_time = time.time()
try:
self._ensure_config()
# 如果用户传入了自定义令牌，则覆盖默认 API Key
effective_api_key = 令牌.strip() if 令牌 and 令牌.strip() else self._api_key
# 构建 input
input_data = self._build_input(提示词, 图片, "", 文件, 视频)
img_count = len(tensor_to_pil(图片)) if 图片 is not None else 0
file_count = len(文件) if 文件 else 0
input_desc = "文本"
if img_count: input_desc += f" + {img_count}张图片"
if 视频 is not None: input_desc += " + 视频"
if file_count: input_desc += f" + {file_count}个文件"
print(f"全能LLM: 模型 = {模型}")
print(f"全能LLM: 输入 = {input_desc}")
# 构建请求体（chat/completions 格式）
request_body = {
"model": 模型,
"messages": input_data,
"stream": True,
}
# 打印请求体，base64 截断显示
def _truncate_for_log(obj):
if isinstance(obj, dict):
return {k: _truncate_for_log(v) for k, v in obj.items()}
if isinstance(obj, list):
return [_truncate_for_log(i) for i in obj]
if isinstance(obj, str) and (obj.startswith("data:image") or obj.startswith("data:application") or obj.startswith("data:text")):
return obj[:60] + f"...[{len(obj)}chars]"
return obj
print(f"全能LLM: 请求原始内容 = {json.dumps(_truncate_for_log(request_body), ensure_ascii=False)}")
# 发送请求（在独立线程中运行异步请求，避免与 ComfyUI 事件循环冲突）
import aiohttp
import asyncio
from concurrent.futures import ThreadPoolExecutor
async def _do_request():
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {effective_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
if status != 200:
body = await resp.text()
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}")
# 流式读取，拼接 delta content
reply_parts = []
async for raw_line in resp.content:
line = raw_line.decode("utf-8").strip()
if not line or not line.startswith("data:"):
continue
data_str = line[len("data:"):].strip()
if data_str == "[DONE]":
break
try:
chunk = json.loads(data_str)
except Exception:
continue
choices = chunk.get("choices")
if not choices:
continue
delta = choices[0].get("delta", {})
content = delta.get("content")
if content:
reply_parts.append(content)
UniversalLLMChat._send_stream_token(node_id, content)
UniversalLLMChat._send_stream_token(node_id, "", done=True)
return "".join(reply_parts)
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:
reply = pool.submit(_run_in_thread).result()
elapsed = time.time() - start_time
print(f"全能LLM: 生成完成 (耗时: {elapsed:.2f}s)")
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
nodes/veo_video.py
-422
View File
@@ -1,422 +0,0 @@
"""
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",
}
nodes/video_preview.py
-90
View File
@@ -1,90 +0,0 @@
"""
视频预览节点
接收 VIDEO 类型，在前端内嵌播放器预览
"""
import os
import io
try:
import folder_paths
FOLDER_PATHS_AVAILABLE = True
except ImportError:
FOLDER_PATHS_AVAILABLE = False
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:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"视频": ("VIDEO",),
},
}
RETURN_TYPES = ()
OUTPUT_NODE = True
FUNCTION = "preview"
CATEGORY = "comfyui_o1key/Utils"
def preview(self, 视频) -> dict:
# 用官方接口取文件路径
source = 视频.get_stream_source()
if isinstance(source, io.BytesIO):
# BytesIO 情况：写到 output/video/ 临时文件
output_dir = os.path.join(_get_output_dir(), "video")
os.makedirs(output_dir, exist_ok=True)
filename = "preview_tmp.mp4"
tmp_path = os.path.join(output_dir, filename)
source.seek(0)
with open(tmp_path, "wb") as f:
f.write(source.read())
subfolder = "video"
else:
video_path = source
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)
else:
# 文件在 output 目录外，复制一份
target_dir = os.path.join(output_dir, "video")
os.makedirs(target_dir, exist_ok=True)
filename = os.path.basename(abs_video)
target_path = os.path.join(target_dir, filename)
if not os.path.exists(target_path):
import shutil
shutil.copy2(abs_video, target_path)
subfolder = "video"
return {
"ui": {
"videos": [{
"filename": filename,
"subfolder": subfolder,
"type": "output",
}],
}
}
NODE_CLASS_MAPPINGS = {
"VideoPreview": VideoPreview,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VideoPreview": "预览视频",
}
update.bat
+87 -63
View File
@@ -1,73 +1,97 @@
@echo off
setlocal EnableDelayedExpansion
title comfyui_o1key Updater
cd /d "%~dp0"
echo.
echo [ comfyui_o1key Updater ]
chcp 65001 > nul
echo ====================================
echo Comfyui_o1key 插件更新工具
echo ====================================
echo.
:: Check git
where git >nul 2>&1
if errorlevel 1 ( set "ERR=Git not found in PATH." & goto :fail )
:: 检查是否在 Git 仓库中
if not exist ".git" (
echo [错误] 当前目录不是 Git 仓库
echo 请确保插件是通过 git clone 安装的
pause
exit /b 1
)
:: Check repo
if not exist ".git" ( set "ERR=Not a git repo. Place this file in the plugin root." & goto :fail )
:: 保存当前版本
if exist "version.txt" (
set /p OLD_VERSION=<version.txt
echo 当前版本: %OLD_VERSION%
) else (
set OLD_VERSION=未知
echo 当前版本: 未知
)
:: Save old hash
for /f %%i in ('git rev-parse --short HEAD 2^>nul') do set "OLD=%%i"
:: Backup .config
if exist ".config" copy /y ".config" ".config.bak" >nul 2>&1
:: Fetch
echo Fetching...
git fetch origin >nul 2>&1
if errorlevel 1 ( set "ERR=Network error. Check GitHub access." & goto :fail )
:: Already up to date?
for /f %%i in ('git rev-parse HEAD 2^>nul') do set "LOCAL=%%i"
for /f %%i in ('git rev-parse origin/main 2^>nul') do set "REMOTE=%%i"
if "%LOCAL%"=="%REMOTE%" ( goto :uptodate )
:: Switch branch & force reset
git branch --list main | findstr "main" >nul 2>&1
if errorlevel 1 ( git checkout -b main origin/main >nul 2>&1 ) else ( git checkout main >nul 2>&1 )
echo Updating...
git reset --hard origin/main >nul 2>&1
if errorlevel 1 ( set "ERR=git reset failed." & goto :fail )
git clean -fd -e ".config" -e ".config.bak" >nul 2>&1
:: Restore .config
if exist ".config.bak" ( copy /y ".config.bak" ".config" >nul 2>&1 & del /f /q ".config.bak" >nul 2>&1 )
for /f %%i in ('git rev-parse --short HEAD 2^>nul') do set "NEW=%%i"
echo.
echo +---------------------------+
echo ^| SUCCESS ^|
echo ^| %OLD% -> %NEW% ^|
echo ^| Restart ComfyUI ^|
echo +---------------------------+
echo.
pause & exit /b 0
echo [1/4] 检查远程更新...
git fetch origin
:uptodate
if exist ".config.bak" ( copy /y ".config.bak" ".config" >nul 2>&1 & del /f /q ".config.bak" >nul 2>&1 )
echo.
echo +---------------------------+
echo ^| Already up to date ^|
echo ^| %LOCAL:~0,7% (no change) ^|
echo +---------------------------+
echo.
pause & exit /b 0
:: 检查是否有更新
git status -uno | findstr "Your branch is behind" > nul
if %errorlevel% equ 0 (
echo 发现新版本！
) else (
echo 已是最新版本
echo.
choice /C YN /M "是否继续检查依赖更新？"
if errorlevel 2 goto :end
)
:fail
if exist ".config.bak" ( copy /y ".config.bak" ".config" >nul 2>&1 & del /f /q ".config.bak" >nul 2>&1 )
echo.
echo +---------------------------+
echo ^| FAILED ^|
echo ^| %ERR%
echo +---------------------------+
echo [2/4] 备份配置文件...
if exist ".config" (
copy /Y ".config" ".config.backup" > nul
echo 已备份 .config 到 .config.backup
)
echo.
pause & exit /b 1
echo [3/4] 拉取最新代码...
git pull origin main
if %errorlevel% neq 0 (
echo [错误] 代码更新失败，请检查网络连接或手动解决冲突
pause
exit /b 1
)
:: 恢复配置文件
if exist ".config.backup" (
copy /Y ".config.backup" ".config" > nul
del ".config.backup"
echo 已恢复配置文件
)
echo.
echo [4/4] 更新依赖包...
python -m pip install -r requirements.txt --upgrade --quiet
if %errorlevel% neq 0 (
echo [警告] 依赖包更新失败，请手动运行: pip install -r requirements.txt
)
echo.
echo ====================================
echo 更新完成！
echo ====================================
:: 显示新版本
if exist "version.txt" (
set /p NEW_VERSION=<version.txt
echo 新版本: %NEW_VERSION%
)
:: 显示最近更新日志
if exist "CHANGELOG.md" (
echo.
echo 最近更新内容:
echo -----------------------------------
powershell -Command "Get-Content CHANGELOG.md -TotalCount 20"
echo -----------------------------------
)
echo.
echo 请重启 ComfyUI 以使更改生效
echo.
pause
goto :end
:end
exit /b 0
update.sh
+101
View File
@@ -0,0 +1,101 @@
#!/bin/bash
# 设置颜色输出
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color
echo "===================================="
echo "Comfyui_o1key 插件更新工具"
echo "===================================="
echo ""
# 检查是否在 Git 仓库中
if [ ! -d ".git" ]; then
echo -e "${RED}[错误] 当前目录不是 Git 仓库${NC}"
echo "请确保插件是通过 git clone 安装的"
exit 1
fi
# 保存当前版本
if [ -f "version.txt" ]; then
OLD_VERSION=$(cat version.txt)
echo "当前版本: $OLD_VERSION"
else
OLD_VERSION="未知"
echo "当前版本: 未知"
fi
echo ""
echo "[1/4] 检查远程更新..."
git fetch origin
# 检查是否有更新
LOCAL=$(git rev-parse @)
REMOTE=$(git rev-parse @{u})
if [ $LOCAL != $REMOTE ]; then
echo -e "${GREEN}发现新版本！${NC}"
else
echo -e "${GREEN}已是最新版本${NC}"
read -p "是否继续检查依赖更新？(y/n) " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
exit 0
fi
fi
echo ""
echo "[2/4] 备份配置文件..."
if [ -f ".config" ]; then
cp .config .config.backup
echo "已备份 .config 到 .config.backup"
fi
echo ""
echo "[3/4] 拉取最新代码..."
if git pull origin main; then
echo -e "${GREEN}代码更新成功${NC}"
else
echo -e "${RED}[错误] 代码更新失败，请检查网络连接或手动解决冲突${NC}"
exit 1
fi
# 恢复配置文件
if [ -f ".config.backup" ]; then
mv .config.backup .config
echo "已恢复配置文件"
fi
echo ""
echo "[4/4] 更新依赖包..."
if python3 -m pip install -r requirements.txt --upgrade --quiet; then
echo -e "${GREEN}依赖包更新成功${NC}"
else
echo -e "${YELLOW}[警告] 依赖包更新失败，请手动运行: pip install -r requirements.txt${NC}"
fi
echo ""
echo "===================================="
echo -e "${GREEN}更新完成！${NC}"
echo "===================================="
# 显示新版本
if [ -f "version.txt" ]; then
NEW_VERSION=$(cat version.txt)
echo "新版本: $NEW_VERSION"
fi
# 显示最近更新日志
if [ -f "CHANGELOG.md" ]; then
echo ""
echo "最近更新内容:"
echo "-----------------------------------"
head -n 20 CHANGELOG.md
echo "-----------------------------------"
fi
echo ""
echo "请重启 ComfyUI 以使更改生效"
echo ""
utils/__init__.py
+4 -2
View File
@@ -15,7 +15,8 @@ from .file_utils import (
load_images_from_folder,
pair_images_indexed,
pair_images_cartesian,
generate_timestamp_filename,
generate_output_filename,
generate_batch_output_filenames,
save_image,
get_folder_image_count
)
@@ -31,7 +32,8 @@ __all__ = [
'load_images_from_folder',
'pair_images_indexed',
'pair_images_cartesian',
'generate_timestamp_filename',
'generate_output_filename',
'generate_batch_output_filenames',
'save_image',
'get_folder_image_count'
]
utils/config.py
+24 -48
View File
@@ -12,16 +12,6 @@ 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://api.o1key.com"
# 异步 API 基础 URL（用于异步提交+轮询模式）
# 可通过环境变量 O1KEY_ASYNC_API_BASE_URL 覆盖
DEFAULT_ASYNC_API_BASE_URL = "https://cf-api.o1key.com"
def load_config(config_path: Optional[str] = None) -> Dict[str, str]:
"""
从配置文件加载所有配置项
@@ -71,16 +61,36 @@ def load_config(config_path: Optional[str] = None) -> Dict[str, str]:
def get_api_key(key_name: str = "O1KEY_API_KEY") -> Optional[str]:
"""
获取 API 密钥
从 .config 文件读取
优先级：环境变量（推荐） > .config 文件（向后兼容）
Args:
key_name: 密钥名称，默认为 O1KEY_API_KEY
Returns:
API 密钥字符串，如果未找到则返回 None
Raises:
ValueError: 如果未找到 API 密钥
Example:
>>> api_key = get_api_key()
>>> if api_key is None:
... raise ValueError("API key not found")
"""
# 1. 优先从环境变量读取（推荐方式）
api_key = os.environ.get(key_name)
if api_key:
return api_key
# 2. 从 .config 文件读取（向后兼容，已弃用）
config = load_config()
return config.get(key_name)
api_key = config.get(key_name)
if api_key:
return api_key
return None
def get_api_key_or_raise(key_name: str = "O1KEY_API_KEY") -> str:
@@ -102,37 +112,3 @@ def get_api_key_or_raise(key_name: str = "O1KEY_API_KEY") -> str:
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
def get_async_api_base_url() -> str:
"""
获取异步 API 基础 URL
从 .config 文件读取，如果未配置则使用默认值
Returns:
异步 API 基础 URL 字符串
"""
config = load_config()
base_url = config.get("O1KEY_ASYNC_API_BASE_URL")
if base_url:
return base_url.rstrip('/')
return DEFAULT_ASYNC_API_BASE_URL
utils/file_types.py
-82
View File
@@ -1,82 +0,0 @@
"""
文件数据类型定义
用于在 ComfyUI 节点间传递文件数据
"""
from typing import NamedTuple, List
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
# FILE_LIST 类型：FileData 的列表，用于多文件传递
# ComfyUI 自定义类型名，节点 RETURN_TYPES / INPUT_TYPES 中使用 "FILE_LIST"
FileList = List[FileData]
# 支持的文件 MIME 类型映射（与 universal_llm.py 的 MIME_MAP 保持一致）
DOCUMENT_MIME_TYPES = {
".pdf": "application/pdf",
".txt": "text/plain",
".md": "text/markdown",
".csv": "text/csv",
".json": "application/json",
".py": "text/x-python",
".js": "text/javascript",
".ts": "text/javascript",
".html": "text/html",
".xml": "application/xml",
".yaml": "text/plain",
".yml": "text/plain",
".toml": "text/plain",
".ini": "text/plain",
".cfg": "text/plain",
".log": "text/plain",
".sh": "text/plain",
".bat": "text/plain",
".sql": "text/plain",
".css": "text/plain",
".scss": "text/plain",
".jsx": "text/javascript",
".tsx": "text/javascript",
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
".zip": "application/zip",
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".webp": "image/webp",
".wav": "audio/wav",
".mp3": "audio/mpeg",
".mp4": "video/mp4",
}
# 单文件大小上限：50MB
FILE_SIZE_LIMIT = 50 * 1024 * 1024
# 所有文件总大小上限：50MB
TOTAL_FILE_SIZE_LIMIT = 50 * 1024 * 1024
# 兼容旧代码
FILE_SIZE_LIMITS = {
".pdf": FILE_SIZE_LIMIT,
".txt": FILE_SIZE_LIMIT,
}
utils/file_utils.py
+90 -137
View File
@@ -4,11 +4,8 @@
"""
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
@@ -16,45 +13,6 @@ 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'}
@@ -175,77 +133,6 @@ def pair_images_indexed(
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, ...]]:
@@ -279,36 +166,100 @@ def pair_images_cartesian(
return list(product(*non_empty_lists))
def generate_timestamp_filename(output_folder: str, prefix: str = "", extension: str = ".png", port_suffix: str = "") -> str:
def generate_output_filename(
source_images: List[ImageInfo],
batch_index: int,
output_folder: str,
extension: str = ".png",
task_id: Optional[str] = None
) -> str:
"""
生成基于时间戳的文件名，确保按文件名排序 = 按生成时间排序。
格式：{prefix}{HHMMSS_YYYYMMDD_mmm}{port_suffix}{extension}
例如：161700_20260322_001.png 或 去除ai_161700_20260322_001.png
生成智能输出文件名
基于源图片文件名生成输出文件名，使用任务ID和时间戳确保并发安全。
Args:
output_folder: 输出目录
prefix: 文件名前缀（如 "去除ai_"）
extension: 文件扩展名（如 ".png"）
port_suffix: 端口后缀（如 "_8189"），为空时自动获取
source_images: 源图片信息列表
batch_index: 批次索引（从 0 开始）
output_folder: 输出文件夹路径
extension: 输出文件扩展名
task_id: 任务唯一标识符（用于并发场景）
Returns:
完整文件路径
完整的输出文件路径
Example:
>>> # 单图片: hello.png -> hello_task0_12345_000.png
>>> # 多图片: hello.png + ref.png -> hello_ref_task0_12345_000.png
>>> # 并发安全：每个任务有唯一的 task_id 和时间戳
"""
Path(output_folder).mkdir(parents=True, exist_ok=True)
if not port_suffix:
port_suffix = _get_port_suffix()
# 构建基础文件名
if len(source_images) == 1:
base_name = source_images[0].filename
else:
# 多个源图片，组合文件名
names = [info.filename for info in source_images]
base_name = "_".join(names)
# 确保输出文件夹存在
output_path = Path(output_folder)
output_path.mkdir(parents=True, exist_ok=True)
# 生成唯一性标识
if task_id is None:
# 如果没有提供 task_id，使用 UUID 前8位
task_id = str(uuid.uuid4())[:8]
# 使用时间戳（毫秒级）增加唯一性
timestamp = int(time.time() * 1000) % 100000 # 精确到毫秒的后5位
# 生成文件名：基础名_任务ID_时间戳_批次索引
filename = f"{base_name}_{task_id}_{timestamp:05d}_{batch_index:03d}{extension}"
full_path = output_path / filename
# 极小概率的冲突处理
counter = 1
while full_path.exists():
filename = f"{base_name}_{task_id}_{timestamp:05d}_{batch_index:03d}_{counter}{extension}"
full_path = output_path / filename
counter += 1
return str(full_path)
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 generate_batch_output_filenames(
source_images: List[ImageInfo],
count: int,
output_folder: str,
extension: str = ".png",
task_id: Optional[str] = None
) -> List[str]:
"""
批量生成输出文件名
Args:
source_images: 源图片信息列表
count: 需要生成的文件名数量
output_folder: 输出文件夹路径
extension: 输出文件扩展名
task_id: 任务唯一标识符（用于并发场景）
Returns:
输出文件路径列表
"""
filenames = []
for i in range(count):
filename = generate_output_filename(
source_images=source_images,
batch_index=i,
output_folder=output_folder,
extension=extension,
task_id=task_id
)
filenames.append(filename)
return filenames
def save_image(
@@ -339,6 +290,8 @@ def save_image(
if image.mode != 'RGB':
image = image.convert('RGB')
image.save(output_path, quality=quality)
elif ext == '.png':
image.save(output_path)
elif ext == '.webp':
image.save(output_path, quality=quality)
else:
utils/image_utils.py
+9 -9
View File
@@ -49,36 +49,36 @@ def tensor_to_pil(tensor: torch.Tensor) -> List[Image.Image]:
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)
utils/r2_uploader.py
-145
View File
@@ -1,145 +0,0 @@
"""
R2 文件上传工具（通过 o1key 后端预签名接口）
- 插件内零 R2 凭证，仅使用用户的 O1KEY_API_KEY
- 流程：请求预签名 URL → PUT 直传 R2 → 返回公网 URL
"""
import io
import os
import uuid
import aiohttp
from .config import get_api_key_or_raise, get_api_base_url
async def _presign(filename: str, content_type: str) -> tuple:
"""向 o1key 后端请求预签名 URL，返回 (upload_url, public_url)"""
api_key = get_api_key_or_raise()
base_url = get_api_base_url()
connector = aiohttp.TCPConnector(ssl=False)
async with aiohttp.ClientSession(connector=connector) as session:
async with session.post(
f"{base_url}/v1/storage/presign",
headers={"Authorization": f"Bearer {api_key}"},
json={"filename": filename, "content_type": content_type},
timeout=aiohttp.ClientTimeout(total=10),
) as resp:
if resp.status != 200:
text = await resp.text()
raise RuntimeError(f"预签名请求失败 ({resp.status}): {text}")
data = await resp.json()
return data["upload_url"], data["public_url"]
async def _put_upload(upload_url: str, data: bytes, content_type: str):
"""用预签名 URL 直传文件到 R2（不带 Authorization）"""
connector = aiohttp.TCPConnector(ssl=False)
async with aiohttp.ClientSession(connector=connector) as session:
async with session.put(
upload_url,
data=data,
headers={"Content-Type": content_type},
timeout=aiohttp.ClientTimeout(total=120),
) as resp:
if resp.status not in (200, 204):
text = await resp.text()
raise RuntimeError(f"文件上传失败 ({resp.status}): {text}")
async def upload_image(pil_image) -> str:
"""
接受 PIL Image 对象，编码为 PNG 上传到 R2，返回公网 URL。
"""
import io as _io
buf = _io.BytesIO()
pil_image.save(buf, format="PNG")
data = buf.getvalue()
filename = f"{uuid.uuid4()}.png"
upload_url, public_url = await _presign(filename, "image/png")
await _put_upload(upload_url, data, "image/png")
print(f"[R2] 图片已上传: {public_url}")
return public_url
async def upload_video(video) -> str:
"""
接受 ComfyUI VIDEO 对象，上传到 R2，返回公网 URL。
支持 mp4 / mov 格式。
"""
source = video.get_stream_source()
if isinstance(source, io.BytesIO):
source.seek(0)
data = source.read()
ext = "mp4"
else:
video_path = source
if not video_path or not os.path.isfile(video_path):
raise ValueError(f"无法获取参考视频文件路径（当前路径：{video_path}）")
ext = os.path.splitext(video_path)[1].lower().lstrip(".")
if ext not in ("mp4", "mov"):
raise ValueError(f"参考视频格式须为 mp4 或 mov，当前为 .{ext}")
with open(video_path, "rb") as f:
data = f.read()
content_type = "video/mp4" if ext == "mp4" else "video/quicktime"
filename = f"{uuid.uuid4()}.{ext}"
upload_url, public_url = await _presign(filename, content_type)
await _put_upload(upload_url, data, content_type)
print(f"[R2] 视频已上传: {public_url}")
return public_url
async def upload_audio(audio) -> str:
"""
接受 ComfyUI AUDIO dict（waveform tensor + sample_rate），
编码为 WAV 后上传到 R2，返回公网 URL。
"""
import struct
import numpy as np
waveform = audio["waveform"] # shape: [B, C, N] or [C, N]
sample_rate = int(audio["sample_rate"])
if waveform.dim() == 3:
waveform = waveform[0]
wav_np = waveform.cpu().numpy()
if wav_np.ndim == 2:
wav_np = wav_np.mean(axis=0)
wav_np = np.clip(wav_np, -1.0, 1.0)
pcm = (wav_np * 32767).astype(np.int16)
num_samples = len(pcm)
num_channels = 1
bits_per_sample = 16
byte_rate = sample_rate * num_channels * bits_per_sample // 8
block_align = num_channels * bits_per_sample // 8
data_size = num_samples * block_align
buf = io.BytesIO()
buf.write(b"RIFF")
buf.write(struct.pack("<I", 36 + data_size))
buf.write(b"WAVE")
buf.write(b"fmt ")
buf.write(struct.pack("<IHHIIHH", 16, 1, num_channels, sample_rate,
byte_rate, block_align, bits_per_sample))
buf.write(b"data")
buf.write(struct.pack("<I", data_size))
buf.write(pcm.tobytes())
data = buf.getvalue()
filename = f"{uuid.uuid4()}.wav"
upload_url, public_url = await _presign(filename, "audio/wav")
await _put_upload(upload_url, data, "audio/wav")
print(f"[R2] 音频已上传: {public_url}")
return public_url
utils/update_checker.py
+14 -79
View File
@@ -39,15 +39,12 @@ def check_for_updates() -> bool:
if not os.path.exists(git_dir):
return False
# 执行 git fetch（禁止弹出认证弹框，失败时静默处理）
env = os.environ.copy()
env['GIT_TERMINAL_PROMPT'] = '0'
# 执行 git fetch
subprocess.run(
['git', 'fetch', 'origin'],
cwd=plugin_dir,
capture_output=True,
timeout=10,
env=env
timeout=10
)
# 检查本地和远程版本
@@ -71,78 +68,16 @@ def check_for_updates() -> bool:
return False
def get_update_changelog() -> list:
"""从远程 CHANGELOG.md 最新版本块中提取更新内容（最多5条）"""
try:
plugin_dir = os.path.dirname(os.path.dirname(__file__))
result = subprocess.run(
['git', 'show', 'origin/main:CHANGELOG.md'],
cwd=plugin_dir,
capture_output=True,
text=True,
encoding='utf-8'
)
lines = result.stdout.splitlines()
in_block = False
items = []
for line in lines:
if line.startswith('## [') and not line.startswith('## [Unreleased]'):
if in_block:
break
in_block = True
continue
if in_block:
stripped = line.strip()
if stripped and not stripped.startswith('#') and not stripped.startswith('---'):
text = stripped.lstrip('- ').replace('**', '').strip()
if text and len(text) > 3:
items.append(text)
if len(items) >= 5:
break
return items
except Exception:
return []
def notify_new_version():
"""检测到新版本时，推送蓝色更新通知弹框"""
changelog = get_update_changelog()
try:
import threading
from server import PromptServer
def _send():
try:
PromptServer.instance.send_sync(
"o1key.new_version",
{"changelog": changelog}
)
except Exception:
pass
threading.Timer(3.0, _send).start()
except Exception:
pass
def notify_update_available():
"""通知用户有更新可用（前端弹窗）"""
try:
import threading
from server import PromptServer
def _send():
try:
PromptServer.instance.send_sync(
"o1key.update_available",
{"message": "欢迎使用o1key工作流，祝您马年，马上有福，马上有钱，马到成功！！！"}
)
except Exception:
pass
threading.Timer(3.0, _send).start()
except Exception:
pass
"""通知用户有更新可用"""
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")
version.txt
+1 -1
View File
@@ -1 +1 @@
v1.10.3
v1.10.0
web/js/fileUpload.js
-122
View File
@@ -1,122 +0,0 @@
import { app } from "../../../scripts/app.js";
import { api } from "../../../scripts/api.js";
// 上传单个文件到 ComfyUI input 目录，返回服务端绝对路径
async function uploadFile(file) {
const formData = new FormData();
formData.append("image", file, file.name);
const resp = await api.fetchApi("/upload/image", { method: "POST", body: formData });
if (!resp.ok) throw new Error(`上传失败: ${file.name}`);
const data = await resp.json();
const inputDir = await getInputDir();
// 拼成绝对路径（Windows 用反斜杠也可以，用正斜杠 Python 也认）
return inputDir ? inputDir.replace(/\\/g, "/") + "/" + data.name : data.name;
}
// 获取 ComfyUI input 目录绝对路径（缓存）
let _inputDir = null;
async function getInputDir() {
if (_inputDir !== null) return _inputDir;
try {
const resp = await api.fetchApi("/o1key/input_dir");
if (resp.ok) _inputDir = (await resp.json()).path;
else _inputDir = "";
} catch { _inputDir = ""; }
return _inputDir;
}
// 创建一个"选择文件"按钮，点击后弹出文件选择框
// onPaths(paths: string[]) 回调拿到上传后的路径列表
function makeUploadButton(label, accept, multiple, onPaths) {
const btn = document.createElement("button");
btn.textContent = label;
btn.style.cssText =
"width:100%;padding:4px 8px;cursor:pointer;margin-top:2px;" +
"background:#3a5a3a;color:#ddd;border:1px solid #666;" +
"border-radius:4px;font-size:12px;";
const fileInput = document.createElement("input");
fileInput.type = "file";
fileInput.multiple = multiple;
fileInput.accept = accept;
fileInput.style.display = "none";
document.body.appendChild(fileInput);
btn.addEventListener("click", () => fileInput.click());
fileInput.addEventListener("change", async () => {
const files = Array.from(fileInput.files);
if (!files.length) return;
btn.textContent = "⏳ 上传中...";
btn.disabled = true;
try {
const paths = [];
for (const f of files) paths.push(await uploadFile(f));
onPaths(paths);
btn.textContent = `✅ 已上传 ${files.length} 个`;
setTimeout(() => { btn.textContent = label; }, 2000);
} catch (e) {
console.error("[o1key fileUpload]", e);
btn.textContent = "❌ 上传失败";
setTimeout(() => { btn.textContent = label; }, 2000);
} finally {
btn.disabled = false;
fileInput.value = "";
}
});
return btn;
}
const ACCEPT = ".pdf,.txt,.md,.csv,.json,.py,.js,.ts,.html,.xml,.docx,.xlsx,.pptx,.zip,.wav,.mp3,.png,.jpg,.jpeg,.webp";
app.registerExtension({
name: "o1key.fileUpload",
async beforeRegisterNodeDef(nodeType, nodeData) {
if (nodeData.name !== "LoadFile") return;
const origCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = function () {
origCreated?.call(this);
const singleWidget = this.widgets?.find(w => w.name === "单文件路径");
const folderWidget = this.widgets?.find(w => w.name === "文件夹路径");
// "单文件路径"下方加按钮（支持多选，追加路径）
if (singleWidget) {
const btn = makeUploadButton("📂 选择文件（可多选）", ACCEPT, true, (paths) => {
const existing = singleWidget.value?.trim();
singleWidget.value = existing
? existing + ", " + paths.join(", ")
: paths.join(", ");
singleWidget.callback?.(singleWidget.value);
app.graph.setDirtyCanvas(true);
});
this.addDOMWidget("upload_single_btn", "btn", btn, {
getValue() { return null; },
setValue() {},
});
}
// 清空按钮：同时清空单文件路径和文件夹路径
if (singleWidget || folderWidget) {
const clearBtn = document.createElement("button");
clearBtn.textContent = "🗑 清空文件路径";
clearBtn.style.cssText =
"width:100%;padding:4px 8px;cursor:pointer;margin-top:2px;" +
"background:#5a3a3a;color:#ddd;border:1px solid #666;" +
"border-radius:4px;font-size:12px;";
clearBtn.addEventListener("click", () => {
if (singleWidget) { singleWidget.value = ""; singleWidget.callback?.(""); }
if (folderWidget) { folderWidget.value = ""; folderWidget.callback?.(""); }
app.graph.setDirtyCanvas(true);
});
this.addDOMWidget("clear_paths_btn", "btn", clearBtn, {
getValue() { return null; },
setValue() {},
});
}
};
},
});
web/js/streamPreview.js
-210
View File
@@ -1,210 +0,0 @@
import { app } from "../../../scripts/app.js";
import { api } from "../../../scripts/api.js";
// ── marked.js 懒加载 ──────────────────────────────────────────────────────────
let markedReady = null;
function loadMarked() {
if (markedReady) return markedReady;
markedReady = new Promise((resolve) => {
if (window.marked) { resolve(window.marked); return; }
const s = document.createElement("script");
s.src = "https://cdn.jsdelivr.net/npm/marked/marked.min.js";
s.onload = () => resolve(window.marked);
s.onerror = () => resolve(null);
document.head.appendChild(s);
});
return markedReady;
}
// ── 节点 UI 构建 ──────────────────────────────────────────────────────────────
function buildUI(node) {
if (node._spContainer) return;
const container = document.createElement("div");
container.style.cssText =
"width:100%;height:100%;box-sizing:border-box;padding:6px;" +
"display:flex;flex-direction:column;gap:4px;";
const toolbar = document.createElement("div");
toolbar.style.cssText =
"display:flex;justify-content:flex-end;gap:6px;align-items:center;";
const mdToggle = document.createElement("button");
mdToggle.textContent = "MD";
mdToggle.title = "切换 Markdown / 纯文本";
mdToggle.style.cssText =
"font-size:10px;padding:2px 6px;border-radius:3px;cursor:pointer;" +
"background:#2a5a2a;color:#ccc;border:1px solid #666;";
const copyBtn = document.createElement("button");
copyBtn.textContent = "复制";
copyBtn.style.cssText =
"font-size:10px;padding:2px 6px;border-radius:3px;cursor:pointer;" +
"background:#444;color:#ccc;border:1px solid #666;";
toolbar.appendChild(mdToggle);
toolbar.appendChild(copyBtn);
const isDedicatedPreview = node.comfyClass === "StreamPreview" || node.type === "StreamPreview";
const content = document.createElement("div");
if (isDedicatedPreview) {
content.style.cssText =
"flex:1;min-height:0;overflow:hidden;" +
"background:#1a1a1a;border:1px solid #444;border-radius:4px;" +
"padding:8px;box-sizing:border-box;font-size:13px;line-height:1.6;" +
"color:#ddd;white-space:pre-wrap;word-break:break-word;";
} else {
content.style.cssText =
"width:100%;min-height:60px;max-height:480px;overflow-y:auto;" +
"background:#1a1a1a;border:1px solid #444;border-radius:4px;" +
"padding:8px;box-sizing:border-box;font-size:13px;line-height:1.6;" +
"color:#ddd;white-space:pre-wrap;word-break:break-word;";
}
const status = document.createElement("div");
status.style.cssText =
"font-size:10px;color:#888;text-align:right;min-height:14px;";
container.appendChild(toolbar);
container.appendChild(content);
container.appendChild(status);
node._spContainer = container;
node._spContent = content;
node._spStatus = status;
node._spMdToggle = mdToggle;
node._spRawText = "";
node._spMarkdown = true;
mdToggle.addEventListener("click", () => {
node._spMarkdown = !node._spMarkdown;
mdToggle.style.background = node._spMarkdown ? "#2a5a2a" : "#444";
renderContent(node);
});
copyBtn.addEventListener("click", () => {
navigator.clipboard.writeText(node._spRawText).then(() => {
copyBtn.textContent = "已复制";
setTimeout(() => { copyBtn.textContent = "复制"; }, 1500);
});
});
const widget = node.addDOMWidget("stream_preview_widget", "preview", container, {
getValue() { return node._spRawText; },
setValue(v) { },
});
widget.computeSize = (width) => {
const isDedicatedPreview = node.comfyClass === "StreamPreview" || node.type === "StreamPreview";
if (isDedicatedPreview) {
const nodeHeight = node.size?.[1] ?? 320;
const overhead = 60;
return [width, Math.max(120, nodeHeight - overhead)];
}
return [width, 320];
};
loadMarked();
}
async function renderContent(node) {
const text = node._spRawText;
const el = node._spContent;
if (!text) { el.innerHTML = ""; return; }
if (node._spMarkdown) {
const marked = await loadMarked();
if (marked) {
el.style.whiteSpace = "normal";
el.innerHTML = marked.parse(text);
} else {
el.style.whiteSpace = "pre-wrap";
el.textContent = text;
}
} else {
el.style.whiteSpace = "pre-wrap";
el.textContent = text;
}
const isPreview = node.comfyClass === "StreamPreview" || node.type === "StreamPreview";
if (!isPreview) el.scrollTop = el.scrollHeight;
}
// ── 流式事件监听 ──────────────────────────────────────────────────────────────
api.addEventListener("o1key.stream_token", (event) => {
const { node_id, token, done } = event.detail;
const node = app.graph.getNodeById(parseInt(node_id));
if (!node) return;
buildUI(node);
if (done) {
node._spStreaming = false;
node._spStatus.textContent = "生成完成";
node._spStatus.style.color = "#4a4";
return;
}
// 第一个 token 到来时清空上一次内容
if (!node._spStreaming) {
node._spStreaming = true;
node._spRawText = "";
}
node._spRawText += token;
node._spStatus.textContent = "生成中…";
node._spStatus.style.color = "#a84";
renderContent(node);
});
// ── 节点注册 ──────────────────────────────────────────────────────────────────
app.registerExtension({
name: "comfyui_o1key.streamPreview",
async beforeRegisterNodeDef(nodeType, nodeData) {
if (nodeData.name !== "StreamPreview") return;
const origOnNodeCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = function () {
if (origOnNodeCreated) origOnNodeCreated.apply(this, arguments);
buildUI(this);
};
nodeType.prototype.onResize = function () {
this.setDirtyCanvas(true, false);
};
const origOnExecuted = nodeType.prototype.onExecuted;
nodeType.prototype.onExecuted = function (message) {
if (origOnExecuted) origOnExecuted.apply(this, arguments);
buildUI(this);
const texts = message?.text;
if (!texts || texts.length === 0) return;
this._spRawText = texts[0];
this._spStatus.textContent = "完成";
this._spStatus.style.color = "#4a4";
renderContent(this);
this.setDirtyCanvas(true, true);
};
const origOnSerialize = nodeType.prototype.onSerialize;
nodeType.prototype.onSerialize = function (o) {
if (origOnSerialize) origOnSerialize.apply(this, arguments);
o.sp_text = this._spRawText || "";
o.sp_markdown = this._spMarkdown !== false;
};
const origOnConfigure = nodeType.prototype.onConfigure;
nodeType.prototype.onConfigure = function (o) {
if (origOnConfigure) origOnConfigure.apply(this, arguments);
buildUI(this);
if (o.sp_text) {
this._spRawText = o.sp_text;
this._spMarkdown = o.sp_markdown !== false;
this._spMdToggle.style.background = this._spMarkdown ? "#2a5a2a" : "#444";
renderContent(this);
}
};
},
});
web/js/updateNotifier.js
-90
View File
@@ -1,90 +0,0 @@
import { api } from "../../../scripts/api.js";
function startCountdown(toast, closeBtn, seconds, accentColor) {
let remaining = seconds;
closeBtn.textContent = `× ${remaining}s`;
const interval = setInterval(() => {
remaining--;
if (remaining <= 0) {
clearInterval(interval);
toast.style.transition = "opacity 0.4s ease";
toast.style.opacity = "0";
setTimeout(() => toast.remove(), 400);
} else {
closeBtn.textContent = `× ${remaining}s`;
}
}, 1000);
closeBtn.onclick = () => {
clearInterval(interval);
toast.remove();
};
}
api.addEventListener("o1key.new_version", (event) => {
const changelog = event.detail?.changelog || [];
const style = document.createElement("style");
style.textContent = `
@keyframes o1key-slide-in {
from { opacity: 0; transform: translateY(16px) scale(0.97); }
to { opacity: 1; transform: translateY(0) scale(1); }
}
`;
document.head.appendChild(style);
const toast = document.createElement("div");
toast.style.cssText = `
position: fixed;
bottom: 28px;
left: 28px;
background: linear-gradient(135deg, #0a1628 0%, #0d2b4e 60%, #1a4a7a 100%);
color: #d6eaf8;
border: 1px solid #2e86c1;
border-radius: 12px;
padding: 12px 16px;
font-size: 14px;
z-index: 100000;
box-shadow: 0 6px 24px rgba(46,134,193,0.35), 0 2px 8px rgba(0,0,0,0.5);
max-width: 340px;
animation: o1key-slide-in 0.4s cubic-bezier(.22,.68,0,1.2);
`;
const header = document.createElement("div");
header.style.cssText = `display: flex; align-items: center; justify-content: space-between;`;
const title = document.createElement("span");
title.textContent = "🔔 检测到有新版本发布！";
title.style.cssText = `font-weight: bold; font-size: 13px; color: #7fb3d3; letter-spacing: 0.5px;`;
const closeBtn = document.createElement("button");
closeBtn.style.cssText = `background: none; border: none; color: #7fb3d3; font-size: 13px; cursor: pointer; padding: 0; line-height: 1;`;
header.appendChild(title);
header.appendChild(closeBtn);
const divider = document.createElement("div");
divider.style.cssText = `height: 1px; background: rgba(46,134,193,0.3); margin: 8px 0;`;
toast.appendChild(header);
toast.appendChild(divider);
const body = document.createElement("div");
const items = changelog.length > 0 ? changelog : ["暂无更新说明"];
items.forEach(item => {
const line = document.createElement("div");
line.textContent = `• ${item}`;
line.style.cssText = `margin-bottom: 4px; font-size: 12px; line-height: 1.6; color: #d6eaf8;`;
body.appendChild(line);
});
const more = document.createElement("div");
more.textContent = "...";
more.style.cssText = `color: #7fb3d3; font-size: 12px; margin-top: 2px;`;
body.appendChild(more);
toast.appendChild(body);
document.body.appendChild(toast);
startCountdown(toast, closeBtn, 5, "#7fb3d3");
});
web/js/videoPreview.js
-146
View File
@@ -1,146 +0,0 @@
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_widget"
).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._resizeToVideo();
}
};
},
});
更新说明.md
+134
View File
@@ -0,0 +1,134 @@
# Comfyui_o1key 更新说明
## 🎉 插件已支持一键自动更新！
从 v1.10.0 版本开始，插件支持自动更新功能。你只需运行更新脚本，即可轻松获取最新版本。
---
## 📦 如何更新
### Windows 用户
1. 打开文件资源管理器
2. 进入插件目录：`ComfyUI\custom_nodes\Comfyui_o1key`
3. 双击运行 `update.bat` 文件
4. 等待更新完成（通常只需几秒钟）
5. 重启 ComfyUI
### Linux/Mac 用户
打开终端，执行以下命令：
```bash
cd ComfyUI/custom_nodes/Comfyui_o1key
chmod +x update.sh # 首次运行需要添加执行权限
./update.sh
```
---
## ✨ 更新脚本功能
✅ **自动检查更新** - 自动检测是否有新版本
✅ **备份配置** - 自动备份和恢复 `.config` 配置文件
✅ **拉取代码** - 自动从 GitHub 拉取最新代码
✅ **更新依赖** - 自动更新 Python 依赖包
✅ **显示日志** - 显示最近的更新内容
✅ **完善提示** - 友好的中文提示和错误处理
---
## 🔔 更新检查
插件会在每次启动 ComfyUI 时自动检查是否有新版本：
- 如果发现新版本，终端会显示更新提示
- 不会影响插件加载速度
- 不会弹窗打断工作流程
- 检查失败不影响插件正常使用
**终端提示示例：**
```
============================================================
🎉 Comfyui_o1key 有新版本可用 (当前版本: v1.9.1)
============================================================
更新方法：
Windows: 双击运行 update.bat
Linux/Mac: 运行 ./update.sh
或手动执行: git pull origin main
============================================================
```
---
## 🛡️ 安全保障
- **配置安全**：更新前自动备份 `.config` 文件，更新后自动恢复
- **环境变量**：存储在系统级别的 API 密钥不受影响
- **错误处理**：更新失败不会破坏现有安装
- **回退方案**：如有问题可用 `git reset` 回退
---
## 🔧 手动更新（备用方案）
如果自动更新脚本无法使用，可以手动执行：
```bash
cd ComfyUI/custom_nodes/Comfyui_o1key
git pull origin main
pip install -r requirements.txt --upgrade
```
---
## ❓ 常见问题
### Q1: 更新会覆盖我的配置吗？
**不会。** 更新脚本会自动备份和恢复你的 `.config` 文件。环境变量中的 API 密钥也不受影响。
### Q2: 更新失败怎么办？
1. 检查网络连接是否正常
2. 确认 Git 已正确安装
3. 尝试手动更新（见上方"手动更新"部分）
4. 如有未提交的修改，先备份后执行 `git reset --hard origin/main`
### Q3: 更新后插件无法启动怎么办？
1. 检查终端错误信息
2. 重新运行依赖安装：`pip install -r requirements.txt --upgrade`
3. 确认 Python 版本 ≥ 3.7
4. 查看 [GitHub Issues](https://github.com/你的用户名/Comfyui_o1key/issues) 寻求帮助
### Q4: 可以禁用启动时的更新检查吗？
暂不支持配置禁用，但更新检查：
- 速度极快（<1 秒）
- 完全静默（无更新时不显示任何信息）
- 失败不影响插件加载
### Q5: 如何查看当前版本？
查看插件目录下的 `version.txt` 文件，或在更新时会显示当前版本号。
### Q6: 多久检查一次更新？
仅在 ComfyUI 启动时检查一次，不会在运行过程中反复检查。
---
## 📮 反馈与支持
如果在更新过程中遇到问题，请：
1. 查看终端输出的错误信息
2. 查阅 [GitHub Issues](https://github.com/你的用户名/Comfyui_o1key/issues)
3. 提交新 Issue 并附上错误信息
---
**享受自动更新带来的便利！** 🎉