Compare commits
63
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5fccb3e8eb | ||
|
|
30e9603f77 | ||
|
|
5d9aff9ca7 | ||
|
|
3f0f4099fb | ||
|
|
c974df1b5e | ||
|
|
815c2c598e | ||
|
|
228c4c5141 | ||
|
|
8eae3da298 | ||
|
|
69279c654d | ||
|
|
c491731c99 | ||
|
|
45deea2b06 | ||
|
|
d8ee3f1b97 | ||
|
|
9f68d7dda8 | ||
|
|
aae98c0f89 | ||
|
|
e9b79669e9 | ||
|
|
a11478df70 | ||
|
|
1a813bfd1d | ||
|
|
844401dbb2 | ||
|
|
9eaf785425 | ||
|
|
a3541cdcec | ||
|
|
3ca7581080 | ||
|
|
40b10209a4 | ||
|
|
2e93a34434 | ||
|
|
b6de4e49ab | ||
|
|
ad49a3c886 | ||
|
|
d4c887d440 | ||
|
|
b40212f826 | ||
|
|
85f9228220 | ||
|
|
c646b0d1d7 | ||
|
|
8299594646 | ||
|
|
35333b296b | ||
|
|
caec23b5cc | ||
|
|
1941357ae4 | ||
|
|
53384f3820 | ||
|
|
dbf6bdfcc8 | ||
|
|
1de500a6a9 | ||
|
|
5cf5c6b6d6 | ||
|
|
fe3cc65b71 | ||
|
|
949f7bb180 | ||
|
|
afa732b93a | ||
|
|
a4edf56503 | ||
|
|
beadf0e365 | ||
|
|
ba468f5ca0 | ||
|
|
07f0c5ed5f | ||
|
|
c4bb8d9724 | ||
|
|
a2665b4010 | ||
|
|
659f94656c | ||
|
|
0b9d7583c7 | ||
|
|
a5dfebb1eb | ||
|
|
dafc4cf0f4 | ||
|
|
b4e82fecf7 | ||
|
|
a977522564 | ||
|
|
00295d65c7 | ||
|
|
82752d34fc | ||
|
|
92bcf65d14 | ||
|
|
bbc5f4a2c4 | ||
|
|
0ddc571f20 | ||
|
|
03d477648a | ||
|
|
9abd175316 | ||
|
|
2b64a45b8e | ||
|
|
0b9a1ebf7a | ||
|
|
34d5c43cad | ||
|
|
9ab209b2b7 |
+125
-3
@@ -3,6 +3,71 @@
|
||||
## 对话原则
|
||||
始终使用中文进行对话。
|
||||
|
||||
## 编码规范 ⚠️ 重要
|
||||
|
||||
### 文件编码要求
|
||||
- **所有文本文件必须使用 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 模型进行图像生成。
|
||||
@@ -38,10 +103,13 @@ Comfyui_o1key/
|
||||
│ ├── __init__.py
|
||||
│ ├── base_client.py # 客户端基类
|
||||
│ └── gemini_client.py # Gemini API 客户端
|
||||
├── .config # API 配置文件(不提交)
|
||||
├── .config.example # 配置示例
|
||||
├── .config.example # 配置文件模板
|
||||
├── requirements.txt # 依赖包
|
||||
└── README.md # 用户文档
|
||||
├── 设置API密钥(win).bat # Windows 配置脚本
|
||||
└── 设置API密钥(mac).sh # Mac/Linux 配置脚本
|
||||
|
||||
注:.config 文件在本地自动创建,不提交到版本控制
|
||||
```
|
||||
|
||||
---
|
||||
@@ -366,7 +434,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
|
||||
from ..utils.config import get_api_key, get_api_key_or_raise, load_config, get_api_base_url
|
||||
|
||||
# 获取 API 密钥(返回 None 如果未找到)
|
||||
api_key = get_api_key("O1KEY_API_KEY")
|
||||
@@ -374,10 +442,64 @@ 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 客户端使用
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
# EditorConfig 配置文件
|
||||
# https://editorconfig.org
|
||||
|
||||
root = true
|
||||
|
||||
# 默认配置
|
||||
[*]
|
||||
charset = utf-8
|
||||
end_of_line = lf
|
||||
insert_final_newline = true
|
||||
trim_trailing_whitespace = true
|
||||
indent_style = space
|
||||
indent_size = 4
|
||||
|
||||
# Python 文件
|
||||
[*.py]
|
||||
indent_size = 4
|
||||
|
||||
# Shell 脚本
|
||||
[*.sh]
|
||||
indent_size = 4
|
||||
|
||||
# Windows 批处理文件
|
||||
[*.{bat,cmd}]
|
||||
end_of_line = crlf
|
||||
indent_size = 4
|
||||
|
||||
# Markdown 文件
|
||||
[*.md]
|
||||
trim_trailing_whitespace = false
|
||||
|
||||
# YAML 文件
|
||||
[*.{yml,yaml}]
|
||||
indent_size = 2
|
||||
|
||||
# JSON 文件
|
||||
[*.json]
|
||||
indent_size = 2
|
||||
@@ -0,0 +1,31 @@
|
||||
# 默认自动处理行结束符
|
||||
* text=auto
|
||||
|
||||
# Python 文件使用 LF
|
||||
*.py text eol=lf
|
||||
|
||||
# Shell 脚本使用 LF
|
||||
*.sh text eol=lf
|
||||
|
||||
# Windows 批处理文件使用 CRLF
|
||||
*.bat text eol=crlf
|
||||
*.cmd text eol=crlf
|
||||
|
||||
# 配置文件使用 LF
|
||||
.config text eol=lf
|
||||
.config.* text eol=lf
|
||||
|
||||
# Markdown 文档使用 LF
|
||||
*.md text eol=lf
|
||||
|
||||
# 二进制文件
|
||||
*.png binary
|
||||
*.jpg binary
|
||||
*.jpeg binary
|
||||
*.gif binary
|
||||
*.ico binary
|
||||
*.mov binary
|
||||
*.mp4 binary
|
||||
*.mp3 binary
|
||||
*.zip binary
|
||||
*.psd binary
|
||||
+3
-3
@@ -1,6 +1,3 @@
|
||||
# 隐私文件(已弃用配置文件,改用环境变量)
|
||||
# .config
|
||||
|
||||
# Python 缓存
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
@@ -24,3 +21,6 @@ venv/
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# 用户配置(含 API Key,不提交)
|
||||
.config
|
||||
|
||||
+49
-36
@@ -6,46 +6,59 @@
|
||||
|
||||
---
|
||||
|
||||
## [1.10.0] - 2026-02-06
|
||||
## [1.10.4] - 2026-04-13
|
||||
|
||||
### Added ⭐
|
||||
- **自动更新系统** - 让用户轻松更新插件到最新版本
|
||||
- 新增 `update.bat` - Windows 自动更新脚本
|
||||
- 新增 `update.sh` - Linux/Mac 自动更新脚本
|
||||
- 新增 `version.txt` - 版本号管理文件
|
||||
- 新增 `utils/update_checker.py` - 启动时自动检查更新
|
||||
- 新增更新检查功能:每次启动 ComfyUI 时自动检测是否有新版本
|
||||
|
||||
- **更新脚本功能**:
|
||||
- ✅ 自动检查远程更新
|
||||
- ✅ 自动备份和恢复 `.config` 配置文件
|
||||
- ✅ 自动拉取最新代码
|
||||
- ✅ 自动更新 Python 依赖包
|
||||
- ✅ 显示版本变更信息
|
||||
- ✅ 显示最近更新日志(前 20 行)
|
||||
- ✅ 友好的彩色终端输出(Linux/Mac)
|
||||
- ✅ 完善的错误处理和提示
|
||||
### 修复
|
||||
- 修复香蕉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` 作为配置文件示例
|
||||
|
||||
### Changed
|
||||
- **插件启动流程** (`__init__.py`)
|
||||
- 集成更新检查模块
|
||||
- 启动时自动检查是否有新版本
|
||||
- 如有更新,终端显示友好的更新提示
|
||||
- 静默失败机制,不影响插件正常加载
|
||||
- **502 错误提示优化** (`clients/base_client.py`)
|
||||
- 当 API 返回 502 时,弹框显示友好文案:「糟糕!请求到上游时遇到超时或过载!别担心,过会儿再次点击运行即可!」
|
||||
- 在 `request_async` 与 `request_get_async` 中均增加 502 专用分支
|
||||
- **配置管理策略**
|
||||
- `.config` 文件现在完全忽略提交(添加到 `.gitignore`)
|
||||
- 简化配置流程,用户通过快捷脚本自动创建本地配置
|
||||
- 移除配置文件安全检查机制(不再需要)
|
||||
- **README 文档**
|
||||
- 更新配置章节,添加快捷脚本使用说明
|
||||
- 调整配置方法优先级:快捷脚本 > 环境变量 > 手动配置
|
||||
- 简化安全提示说明
|
||||
|
||||
- **文档更新** (`README.md`)
|
||||
- 新增"🔄 更新插件"章节
|
||||
- 提供两种更新方法:自动更新(推荐)和手动更新
|
||||
- 详细的跨平台更新说明
|
||||
- 更新提示和注意事项
|
||||
|
||||
### Benefits
|
||||
- 🎯 **用户友好** - 一键更新,无需手动操作 Git
|
||||
- 🔒 **配置安全** - 自动备份恢复配置,不会丢失设置
|
||||
- ⚡ **依赖同步** - 自动更新 Python 包,确保兼容性
|
||||
- 📋 **信息透明** - 显示版本变更和更新日志
|
||||
- 🌍 **跨平台** - 支持 Windows/Linux/Mac
|
||||
- 🛡️ **稳定可靠** - 完善的错误处理,不影响插件运行
|
||||
### Removed
|
||||
- **安全检查工具**(不再需要)
|
||||
- 删除 `check_config_safety.py` 配置安全检查脚本
|
||||
- 删除 `.git-hooks-install.bat` Git Hook 安装脚本
|
||||
- 彻底杜绝配置文件泄密风险
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -10,3 +10,241 @@
|
||||
- 🎯 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` 常量
|
||||
|
||||
---
|
||||
|
||||
## 🔄 更新插件
|
||||
|
||||
### 界面更新
|
||||
|
||||
在 ComfyUI 左侧功能栏点击「更新」(位于「重启」下方)。按钮会从当前 Git 仓库的 `origin/main` 拉取最新版本。完成后点击「重启」使新版本生效。
|
||||
|
||||
界面更新需要通过 Git 安装、处于 `main` 分支,且节点包文件没有本地修改。更新仅允许快进,不会覆盖本地修改或删除配置。ZIP 安装、分支分叉或网络连接失败时,界面会显示原因,需要手动处理。
|
||||
|
||||
如果提示依赖列表已变化,请在 ComfyUI 使用的 Python 环境中执行:
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes/comfyui_o1key
|
||||
python -m pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 手动更新
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes/comfyui_o1key
|
||||
git pull --ff-only origin main
|
||||
python -m pip install -r requirements.txt
|
||||
```
|
||||
|
||||
更新保留环境变量中配置的 API 密钥。启动时仍会检查是否有新版本。
|
||||
|
||||
---
|
||||
|
||||
## 📚 节点说明
|
||||
|
||||
### 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**
|
||||
|
||||
+148
@@ -0,0 +1,148 @@
|
||||
"""
|
||||
ComfyUI V3 节点开发参考
|
||||
========================
|
||||
|
||||
本文件记录了将 V1 节点迁移到 V3 的关键经验,供后续节点开发快速参考。
|
||||
基于 nano_banana.py 的实际迁移总结。
|
||||
|
||||
核心发现:V3 节点可以直接放入 V1 的 NODE_CLASS_MAPPINGS 中注册,
|
||||
ComfyUI 通过 issubclass(obj_class, _ComfyNodeInternal) 自动识别并
|
||||
调用 GET_NODE_INFO_V1() 生成前端所需的节点信息。无需 comfy_entrypoint。
|
||||
|
||||
=== 最小 V3 节点模板 ===
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
class MyNode(io.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="MyNode", # 必须与 NODE_CLASS_MAPPINGS 的 key 一致
|
||||
display_name="我的节点",
|
||||
category="image/generation",
|
||||
inputs=[...],
|
||||
outputs=[io.Image.Output(display_name="输出")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, input1, input2, ...) -> io.NodeOutput:
|
||||
# 业务逻辑
|
||||
return io.NodeOutput(result)
|
||||
|
||||
=== V1 → V3 对照表 ===
|
||||
|
||||
V1 V3
|
||||
─────────────────────────────────────────────────────
|
||||
INPUT_TYPES() classmethod define_schema() → io.Schema
|
||||
RETURN_TYPES = ("IMAGE",) outputs=[io.Image.Output()]
|
||||
RETURN_NAMES = ("输出",) io.Image.Output(display_name="输出")
|
||||
FUNCTION = "generate" 固定为 execute
|
||||
CATEGORY = "xxx" Schema(category="xxx")
|
||||
generate(self, ...) execute(cls, ...) classmethod
|
||||
self.xxx 实例状态 模块级单例函数
|
||||
|
||||
=== DynamicCombo(动态联动下拉框)===
|
||||
|
||||
场景:一个 combo 的选项决定其他 combo 显示哪些值。
|
||||
|
||||
io.DynamicCombo.Input("模型", options=[
|
||||
io.DynamicCombo.Option("选项A", [
|
||||
io.Combo.Input("子参数1", options=["x", "y"]),
|
||||
io.Combo.Input("子参数2", options=["1K", "2K"]),
|
||||
]),
|
||||
io.DynamicCombo.Option("选项B", [
|
||||
io.Combo.Input("子参数1", options=["x", "y", "z", "w"]),
|
||||
io.Combo.Input("子参数2", options=["512px", "1K", "2K", "4K"]),
|
||||
]),
|
||||
])
|
||||
|
||||
execute 中接收为 dict:
|
||||
def execute(cls, 模型, ...):
|
||||
selected = 模型["模型"] # "选项A" 或 "选项B"
|
||||
sub1 = 模型["子参数1"] # 对应选项下的子输入值
|
||||
sub2 = 模型["子参数2"]
|
||||
|
||||
注意:dict 的 key 是 DynamicCombo.Input 的 id("模型"),
|
||||
子输入的 key 是各 Combo.Input 的 id。
|
||||
|
||||
=== Autogrow(自动增长输入槽)===
|
||||
|
||||
场景:用户连接一个槽后自动出现下一个,最多 N 个。
|
||||
|
||||
io.Autogrow.Input("参考图",
|
||||
template=io.Autogrow.TemplatePrefix(
|
||||
input=io.Image.Input("img"),
|
||||
prefix="参考图", # 生成 参考图0, 参考图1, ...
|
||||
min=0, # 最少显示几个槽
|
||||
max=9, # 最多几个槽
|
||||
),
|
||||
)
|
||||
|
||||
execute 中接收为 dict(或 io.Autogrow.Type):
|
||||
def execute(cls, 参考图=None, ...):
|
||||
if 参考图:
|
||||
for key, tensor in 参考图.items():
|
||||
# key = "参考图0", "参考图1", ...
|
||||
# tensor = IMAGE tensor 或 None
|
||||
|
||||
=== 实例状态处理 ===
|
||||
|
||||
V3 的 execute 是 classmethod,无法用 self。
|
||||
用模块级单例替代:
|
||||
|
||||
_client = None
|
||||
|
||||
def _get_client():
|
||||
global _client
|
||||
if _client is None:
|
||||
_client = MyAPIClient()
|
||||
return _client
|
||||
|
||||
=== 注册方式(与 V1 共存)===
|
||||
|
||||
在 __init__.py 中照常注册,无需任何特殊处理:
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MyV1Node": MyV1Node, # V1 节点
|
||||
"MyV3Node": MyV3Node, # V3 节点,自动识别
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MyV1Node": "V1 节点",
|
||||
"MyV3Node": "V3 节点", # 也可省略,V3 用 Schema.display_name
|
||||
}
|
||||
|
||||
=== 注意事项 ===
|
||||
|
||||
1. node_id 必须与 NODE_CLASS_MAPPINGS 的 key 完全一致
|
||||
2. V3 execute 返回 io.NodeOutput(tensor),不是 tuple
|
||||
3. _wrap_generate_for_error_display 等 V1 包装器对 V3 无效
|
||||
(找不到 generate 方法会安全跳过)
|
||||
4. V3 支持 async execute(直接加 async 即可)
|
||||
5. 输入参数名必须与 Schema inputs 的 id 一致
|
||||
6. DynamicCombo 的子输入在前端会随选项切换动态显示/隐藏
|
||||
7. Autogrow 的 widget 输入会被强制为 force_input(仅连接,无控件)
|
||||
|
||||
=== 可用输入类型速查 ===
|
||||
|
||||
io.String.Input(id, default="", multiline=False)
|
||||
io.Int.Input(id, default=0, min=0, max=N, step=1)
|
||||
io.Float.Input(id, default=0.0, min=0.0, max=N, step=0.01)
|
||||
io.Combo.Input(id, options=[...], default="...")
|
||||
io.Boolean.Input(id, default=False)
|
||||
io.Image.Input(id)
|
||||
io.Mask.Input(id)
|
||||
io.Latent.Input(id)
|
||||
io.DynamicCombo.Input(id, options=[DynamicCombo.Option(...)])
|
||||
io.Autogrow.Input(id, template=TemplatePrefix/TemplateNames)
|
||||
|
||||
=== 可用输出类型速查 ===
|
||||
|
||||
io.Image.Output(display_name="...")
|
||||
io.String.Output(display_name="...")
|
||||
io.Int.Output()
|
||||
io.Float.Output()
|
||||
io.Latent.Output()
|
||||
io.Mask.Output()
|
||||
"""
|
||||
+732
-14
@@ -9,29 +9,747 @@ Comfyui_o1key - ComfyUI 自定义节点集合
|
||||
└── __init__.py # 节点注册入口
|
||||
"""
|
||||
|
||||
# 检查更新(仅在启动时检查一次)
|
||||
|
||||
import ssl
|
||||
import logging
|
||||
import asyncio
|
||||
|
||||
# 屏蔽 ComfyUI 资产扫描的终端日志输出
|
||||
_seeder_filter = lambda record: not any(
|
||||
kw in record.getMessage()
|
||||
for kw in ("Seeder start", "Asset scan", "Scan(", "Fast scan")
|
||||
)
|
||||
logging.getLogger().addFilter(_seeder_filter)
|
||||
|
||||
|
||||
def _is_ignored_asyncio_win10054(context):
|
||||
exc = context.get("exception")
|
||||
if not (
|
||||
isinstance(exc, ConnectionResetError)
|
||||
and getattr(exc, "winerror", None) == 10054
|
||||
):
|
||||
return False
|
||||
|
||||
handle = str(context.get("handle", ""))
|
||||
message = str(context.get("message", ""))
|
||||
marker = "_ProactorBasePipeTransport._call_connection_lost"
|
||||
return marker in handle or marker in message
|
||||
|
||||
|
||||
def _install_asyncio_win10054_filter(loop):
|
||||
if getattr(loop, "_o1key_win10054_filter_installed", False):
|
||||
return loop
|
||||
|
||||
previous_handler = loop.get_exception_handler()
|
||||
|
||||
def _o1key_asyncio_exception_handler(loop, context):
|
||||
if _is_ignored_asyncio_win10054(context):
|
||||
return
|
||||
if previous_handler is not None:
|
||||
previous_handler(loop, context)
|
||||
else:
|
||||
loop.default_exception_handler(context)
|
||||
|
||||
loop.set_exception_handler(_o1key_asyncio_exception_handler)
|
||||
setattr(loop, "_o1key_win10054_filter_installed", True)
|
||||
return loop
|
||||
|
||||
|
||||
try:
|
||||
from .utils.update_checker import check_for_updates, notify_update_available
|
||||
|
||||
if check_for_updates():
|
||||
notify_update_available()
|
||||
except Exception:
|
||||
# 静默失败,不影响插件加载
|
||||
_install_asyncio_win10054_filter(asyncio.get_event_loop())
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini
|
||||
if not getattr(asyncio, "_o1key_new_event_loop_patched", False):
|
||||
_o1key_original_new_event_loop = asyncio.new_event_loop
|
||||
|
||||
def _o1key_new_event_loop(*args, **kwargs):
|
||||
return _install_asyncio_win10054_filter(
|
||||
_o1key_original_new_event_loop(*args, **kwargs)
|
||||
)
|
||||
|
||||
asyncio.new_event_loop = _o1key_new_event_loop
|
||||
asyncio._o1key_new_event_loop_patched = True
|
||||
|
||||
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, BatchCleanMetadata, VideoPreview, GoogleVeo, Google31Video, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, O1keyGPTImageBatch, O1keyGrokImage, O1keyGrokVideo, KVideoFirstLast, KVideoImage2Video
|
||||
from .nodes import K3Video, K3VideoFirstLast, K3MotionControl, K3MotionVideoCheck, NanoBananaV2, NanoBananaV2Batch, SaveImageFormat
|
||||
from .nodes import O1keySavePSD
|
||||
from .nodes import O1keyRemoveBackground
|
||||
from .nodes import O1keyColorRemoveBG
|
||||
from .nodes import O1keyGridSplitter
|
||||
|
||||
# 报错弹框友好文案(不修改原节点代码,仅在外层统一处理)
|
||||
_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)
|
||||
|
||||
# ComfyUI 节点注册
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NanoBananaPro": NanoBananaPro,
|
||||
"NanoBanana": NanoBananaPro,
|
||||
"BatchNanoBananaPro": BatchNanoBananaPro,
|
||||
"GoogleGemini": GoogleGemini
|
||||
"GoogleGemini": GoogleGemini,
|
||||
"LoadFile": LoadFile,
|
||||
"ImageStitchPro": ImageStitchPro,
|
||||
|
||||
"BatchCleanMetadata": BatchCleanMetadata,
|
||||
"VideoPreview": VideoPreview,
|
||||
"GoogleVeo": GoogleVeo,
|
||||
"Google31Video": Google31Video,
|
||||
"FluxImageEdit": FluxImageEdit,
|
||||
"UniversalLLMChat": UniversalLLMChat,
|
||||
"KlingVideo": KlingVideo,
|
||||
"KlingFirstLastFrame": KlingFirstLastFrame,
|
||||
"KlingMotionControlTest": KlingMotionControlTest,
|
||||
"AspectRatioPreset": AspectRatioPreset,
|
||||
|
||||
"BatchImagesO1key": BatchImagesO1key,
|
||||
"Seedance": Seedance,
|
||||
"SeedanceMultiModal": SeedanceMultiModal,
|
||||
"StreamPreview": StreamPreview,
|
||||
"DoubaoImage": DoubaoImage,
|
||||
"O1keyGPTImage": O1keyGPTImage,
|
||||
"O1keyGPTImageBatch": O1keyGPTImageBatch,
|
||||
"O1keyGrokImage": O1keyGrokImage,
|
||||
"O1keyGrokVideo": O1keyGrokVideo,
|
||||
"KVideoFirstLast": KVideoFirstLast,
|
||||
"KVideoImage2Video": KVideoImage2Video,
|
||||
"K3Video": K3Video,
|
||||
"K3VideoFirstLast": K3VideoFirstLast,
|
||||
"K3MotionControl": K3MotionControl,
|
||||
"K3MotionVideoCheck": K3MotionVideoCheck,
|
||||
"NanoBananaV2": NanoBananaV2,
|
||||
"NanoBananaV2Batch": NanoBananaV2Batch,
|
||||
"SaveImageFormat": SaveImageFormat,
|
||||
"O1keySavePSD": O1keySavePSD,
|
||||
"O1keyRemoveBackground": O1keyRemoveBackground,
|
||||
"O1keyColorRemoveBG": O1keyColorRemoveBG,
|
||||
"O1keyGridSplitter": O1keyGridSplitter,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NanoBananaPro": "Nano Banana Pro",
|
||||
"BatchNanoBananaPro": "批量 Nano Banana Pro",
|
||||
"GoogleGemini": "Google Gemini"
|
||||
"NanoBanana": "Nano Banana",
|
||||
"BatchNanoBananaPro": "批量 Nano Banana",
|
||||
"GoogleGemini": "Google Gemini",
|
||||
"LoadFile": "加载文件",
|
||||
"ImageStitchPro": "图像拼接 Pro",
|
||||
|
||||
"BatchCleanMetadata": "批量任务(防AI识别)",
|
||||
"VideoPreview": "预览视频",
|
||||
"GoogleVeo": "Google Veo - ab",
|
||||
"Google31Video": "Google 3.1 Video",
|
||||
"FluxImageEdit": "Flux2 图像编辑",
|
||||
"UniversalLLMChat": "全能LLM对话助手",
|
||||
"KlingVideo": "文/图生视频 自研模型",
|
||||
"KlingFirstLastFrame": "首尾帧生视频 自研模型",
|
||||
"KlingMotionControlTest": "动作控制 自研模型",
|
||||
"AspectRatioPreset": "图片宽高比预设",
|
||||
|
||||
"BatchImagesO1key": "加载图像(批量)",
|
||||
"Seedance": "Seedance 视频生成",
|
||||
"SeedanceMultiModal": "Seedance 多模态参考生视频",
|
||||
"StreamPreview": "流式文本预览",
|
||||
"DoubaoImage": "豆包生图",
|
||||
"O1keyGPTImage": "o1key GPT Image",
|
||||
"O1keyGPTImageBatch": "o1key GPT Image(批量)",
|
||||
"O1keyGrokImage": "Grok Image",
|
||||
"O1keyGrokVideo": "Grok Video",
|
||||
"KVideoFirstLast": "K26 图生视频(首尾帧)",
|
||||
"KVideoImage2Video": "K26 图生视频",
|
||||
"K3Video": "K3 图生视频 自研",
|
||||
"K3VideoFirstLast": "首尾帧 K3 自研",
|
||||
"K3MotionControl": "动作控制 K3 自研",
|
||||
"K3MotionVideoCheck": "视频时长检测 K3",
|
||||
"NanoBananaV2": "Nano Banana V2",
|
||||
"NanoBananaV2Batch": "Nano Banana V2(批量)",
|
||||
"SaveImageFormat": "保存图像(格式转换)",
|
||||
"O1keySavePSD": "保存 PSD(分层)",
|
||||
"O1keyRemoveBackground": "去背景(rembg)",
|
||||
"O1keyColorRemoveBG": "颜色去背景",
|
||||
"O1keyGridSplitter": "合并图智能切割",
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
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
|
||||
from .utils.config import CONFIG_FILE, load_config, NETWORK_ROUTES
|
||||
from .utils.updater import UpdateError, update_package
|
||||
import threading as _update_threading
|
||||
|
||||
_update_lock = _update_threading.Lock()
|
||||
|
||||
def _get_o1key_server_port():
|
||||
try:
|
||||
import comfy.cli_args as _cli_args
|
||||
args = getattr(_cli_args, "args", None)
|
||||
port = getattr(args, "port", None) if args else None
|
||||
port = port or getattr(_cli_args, "server_port", None) or getattr(_cli_args, "port", None)
|
||||
if port is not None:
|
||||
return str(int(port))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
import sys as _sys
|
||||
for idx, arg in enumerate(_sys.argv):
|
||||
if arg in ("--port", "--listen-port") and idx + 1 < len(_sys.argv):
|
||||
return str(int(_sys.argv[idx + 1]))
|
||||
for prefix in ("--port=", "--listen-port="):
|
||||
if arg.startswith(prefix):
|
||||
return str(int(arg.split("=", 1)[1]))
|
||||
except Exception:
|
||||
pass
|
||||
return "8188"
|
||||
|
||||
def _get_o1key_history_meta_file(output_dir):
|
||||
import os as _os_history
|
||||
return _os_history.path.join(
|
||||
output_dir,
|
||||
f".o1key_history_{_get_o1key_server_port()}.json",
|
||||
)
|
||||
|
||||
def _get_o1key_notes_file():
|
||||
import os as _os_notes
|
||||
input_dir = _os_notes.path.abspath(folder_paths.get_input_directory())
|
||||
_os_notes.makedirs(input_dir, exist_ok=True)
|
||||
return _os_notes.path.join(input_dir, "o1key-notes.json")
|
||||
|
||||
def _extract_o1key_notes(payload):
|
||||
if isinstance(payload, list):
|
||||
return payload
|
||||
if isinstance(payload, dict) and isinstance(payload.get("notes"), list):
|
||||
return payload["notes"]
|
||||
return None
|
||||
|
||||
@PromptServer.instance.routes.get("/o1key/notes")
|
||||
async def get_o1key_notes(request):
|
||||
import os as _os_notes
|
||||
import json as _json_notes
|
||||
|
||||
notes_file = _get_o1key_notes_file()
|
||||
exists = _os_notes.path.isfile(notes_file)
|
||||
notes = []
|
||||
|
||||
if exists:
|
||||
try:
|
||||
with open(notes_file, "r", encoding="utf-8") as nf:
|
||||
loaded = _json_notes.load(nf)
|
||||
notes = _extract_o1key_notes(loaded)
|
||||
if notes is None:
|
||||
return web.json_response(
|
||||
{"error": "invalid notes file", "path": notes_file},
|
||||
status=500,
|
||||
)
|
||||
except Exception as e:
|
||||
return web.json_response(
|
||||
{"error": str(e), "path": notes_file},
|
||||
status=500,
|
||||
)
|
||||
|
||||
return web.json_response({"notes": notes, "path": notes_file, "exists": exists})
|
||||
|
||||
@PromptServer.instance.routes.post("/o1key/notes")
|
||||
async def save_o1key_notes(request):
|
||||
import os as _os_notes
|
||||
import json as _json_notes
|
||||
|
||||
try:
|
||||
payload = await request.json()
|
||||
notes = _extract_o1key_notes(payload)
|
||||
if notes is None:
|
||||
return web.json_response({"error": "notes must be a list"}, status=400)
|
||||
except Exception as e:
|
||||
return web.json_response({"error": f"invalid notes payload: {str(e)}"}, status=400)
|
||||
|
||||
notes_file = _get_o1key_notes_file()
|
||||
temp_file = notes_file + ".tmp"
|
||||
try:
|
||||
with open(temp_file, "w", encoding="utf-8") as nf:
|
||||
_json_notes.dump(notes, nf, ensure_ascii=False, indent=2)
|
||||
nf.write("\n")
|
||||
_os_notes.replace(temp_file, notes_file)
|
||||
except Exception as e:
|
||||
return web.json_response({"error": f"save notes failed: {str(e)}"}, status=500)
|
||||
|
||||
return web.json_response({
|
||||
"success": True,
|
||||
"path": notes_file,
|
||||
"count": len(notes),
|
||||
})
|
||||
|
||||
@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})
|
||||
|
||||
@PromptServer.instance.routes.get("/o1key/api_key")
|
||||
async def get_api_key_route(request):
|
||||
config = load_config()
|
||||
key = config.get("O1KEY_API_KEY", "")
|
||||
masked = ""
|
||||
if key:
|
||||
if len(key) > 8:
|
||||
masked = key[:3] + "****" + key[-4:]
|
||||
else:
|
||||
masked = "****"
|
||||
return web.json_response({"has_key": bool(key), "masked": masked})
|
||||
|
||||
@PromptServer.instance.routes.post("/o1key/api_key")
|
||||
async def set_api_key_route(request):
|
||||
import os
|
||||
data = await request.json()
|
||||
new_key = data.get("api_key", "").strip()
|
||||
if not new_key:
|
||||
return web.json_response({"error": "API Key 不能为空"}, status=400)
|
||||
config = load_config()
|
||||
config["O1KEY_API_KEY"] = new_key
|
||||
lines = []
|
||||
for k, v in config.items():
|
||||
lines.append(f"{k}={v}")
|
||||
with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
|
||||
f.write("\n".join(lines) + "\n")
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@PromptServer.instance.routes.post("/o1key/test_key")
|
||||
async def test_api_key_route(request):
|
||||
import aiohttp as _aiohttp
|
||||
data = await request.json()
|
||||
test_key = data.get("api_key", "").strip()
|
||||
if not test_key:
|
||||
return web.json_response({"valid": False, "error": "密钥不能为空"})
|
||||
base_url = NETWORK_ROUTES.get("CF加速", "https://cf-api.o1key.com")
|
||||
url = f"{base_url}/v1/models"
|
||||
headers = {"Authorization": f"Bearer {test_key}"}
|
||||
try:
|
||||
async with _aiohttp.ClientSession() as session:
|
||||
async with session.get(url, headers=headers, timeout=_aiohttp.ClientTimeout(total=10)) as resp:
|
||||
if resp.status == 200:
|
||||
return web.json_response({"valid": True})
|
||||
elif resp.status == 401:
|
||||
return web.json_response({"valid": False, "error": "密钥无效或已过期"})
|
||||
else:
|
||||
text = await resp.text()
|
||||
return web.json_response({"valid": False, "error": f"验证失败 ({resp.status})"})
|
||||
except Exception as e:
|
||||
return web.json_response({"valid": False, "error": f"网络错误: {str(e)}"})
|
||||
|
||||
@PromptServer.instance.routes.delete("/o1key/api_key")
|
||||
async def delete_api_key_route(request):
|
||||
config = load_config()
|
||||
config.pop("O1KEY_API_KEY", None)
|
||||
lines = []
|
||||
for k, v in config.items():
|
||||
lines.append(f"{k}={v}")
|
||||
with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
|
||||
f.write("\n".join(lines) + "\n")
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@PromptServer.instance.routes.get("/o1key/output_history")
|
||||
async def get_output_history(request):
|
||||
"""读取 output 目录文件,按执行分组返回 /api/jobs 兼容格式"""
|
||||
import os, json as _json
|
||||
limit = int(request.query.get("limit", "200"))
|
||||
offset = int(request.query.get("offset", "0"))
|
||||
output_dir = os.path.abspath(folder_paths.get_output_directory())
|
||||
meta_file = _get_o1key_history_meta_file(output_dir)
|
||||
meta = {}
|
||||
if os.path.isfile(meta_file):
|
||||
try:
|
||||
with open(meta_file, "r", encoding="utf-8") as mf:
|
||||
meta = _json.load(mf)
|
||||
except Exception:
|
||||
pass
|
||||
supported_ext = {'.png', '.jpg', '.jpeg', '.webp', '.gif', '.mp4', '.webm'}
|
||||
# 收集所有文件并按 workflow_id 分组
|
||||
all_files = []
|
||||
for fname in meta.keys():
|
||||
ext = os.path.splitext(fname)[1].lower()
|
||||
if ext not in supported_ext:
|
||||
continue
|
||||
fpath = os.path.join(output_dir, fname)
|
||||
if not os.path.isfile(fpath):
|
||||
continue
|
||||
mtime = os.path.getmtime(fpath)
|
||||
media = "images" if ext in {'.png','.jpg','.jpeg','.webp','.gif'} else "video"
|
||||
all_files.append({"name": fname, "mtime": mtime, "media": media})
|
||||
# 按 workflow_id 分组(同一次执行合并为一个 job)
|
||||
groups = {}
|
||||
for f in all_files:
|
||||
m = meta.get(f["name"], {})
|
||||
wid = m.get("workflow_id")
|
||||
if wid:
|
||||
groups.setdefault(wid, []).append((f, m))
|
||||
# 构建 job 列表
|
||||
jobs = []
|
||||
for wid, items in groups.items():
|
||||
items.sort(key=lambda x: x[0]["mtime"], reverse=True)
|
||||
latest = items[0]
|
||||
f, m = latest
|
||||
start_ms = int(m.get("start_time", f["mtime"]) * 1000)
|
||||
end_ms = int(m.get("end_time", f["mtime"]) * 1000)
|
||||
jobs.append({
|
||||
"id": wid,
|
||||
"status": "completed",
|
||||
"create_time": start_ms,
|
||||
"execution_start_time": start_ms,
|
||||
"execution_end_time": end_ms,
|
||||
"preview_output": {
|
||||
"filename": f["name"],
|
||||
"subfolder": "",
|
||||
"type": "output",
|
||||
"nodeId": "0",
|
||||
"mediaType": f["media"],
|
||||
},
|
||||
"outputs_count": len(items),
|
||||
"execution_error": None,
|
||||
"workflow_id": wid,
|
||||
})
|
||||
# 按时间倒序排列,分页
|
||||
jobs.sort(key=lambda x: x["create_time"], reverse=True)
|
||||
total = len(jobs)
|
||||
page = jobs[offset:offset+limit]
|
||||
return web.json_response({
|
||||
"jobs": page,
|
||||
"pagination": {"offset": offset, "limit": limit, "total": total, "has_more": offset + limit < total}
|
||||
})
|
||||
|
||||
@PromptServer.instance.routes.get("/o1key/output_workflow")
|
||||
async def get_output_workflow(request):
|
||||
"""从 PNG 元数据中读取工作流,供前端恢复使用"""
|
||||
import os, struct, json as _json
|
||||
filename = request.query.get("filename", "")
|
||||
if not filename:
|
||||
return web.json_response({"error": "missing filename"}, status=400)
|
||||
output_dir = os.path.abspath(folder_paths.get_output_directory())
|
||||
fpath = os.path.join(output_dir, filename)
|
||||
if not os.path.isfile(fpath) or not fpath.lower().endswith(".png"):
|
||||
return web.json_response({"error": "file not found"}, status=404)
|
||||
workflow = None
|
||||
prompt_data = None
|
||||
try:
|
||||
with open(fpath, "rb") as pf:
|
||||
pf.read(8) # PNG signature
|
||||
while True:
|
||||
raw = pf.read(8)
|
||||
if len(raw) < 8:
|
||||
break
|
||||
length = struct.unpack(">I", raw[:4])[0]
|
||||
chunk_type = raw[4:8]
|
||||
data = pf.read(length)
|
||||
pf.read(4) # CRC
|
||||
if chunk_type == b"tEXt":
|
||||
key, val = data.split(b"\x00", 1)
|
||||
k = key.decode("ascii", errors="replace")
|
||||
if k == "workflow":
|
||||
workflow = _json.loads(val)
|
||||
elif k == "prompt":
|
||||
prompt_data = _json.loads(val)
|
||||
elif chunk_type == b"IEND":
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
return web.json_response({"workflow": workflow, "prompt": prompt_data})
|
||||
|
||||
@PromptServer.instance.routes.get("/o1key/job_detail/{job_id}")
|
||||
async def get_job_detail(request):
|
||||
"""根据 job_id 返回当前端口持久化历史中的 job 详情"""
|
||||
import os, struct, json as _json
|
||||
job_id = request.match_info["job_id"]
|
||||
output_dir = os.path.abspath(folder_paths.get_output_directory())
|
||||
meta_file = _get_o1key_history_meta_file(output_dir)
|
||||
meta = {}
|
||||
if os.path.isfile(meta_file):
|
||||
try:
|
||||
with open(meta_file, "r", encoding="utf-8") as mf:
|
||||
meta = _json.load(mf)
|
||||
except Exception:
|
||||
pass
|
||||
supported_ext = {'.png', '.jpg', '.jpeg', '.webp', '.gif', '.mp4', '.webm'}
|
||||
# 只在当前端口的持久化记录中查找该 job 的文件
|
||||
matched_files = []
|
||||
for fname in meta.keys():
|
||||
ext = os.path.splitext(fname)[1].lower()
|
||||
if ext not in supported_ext:
|
||||
continue
|
||||
m = meta.get(fname, {})
|
||||
if m.get("workflow_id") == job_id:
|
||||
matched_files.append(fname)
|
||||
if not matched_files:
|
||||
return web.json_response({"error": "not found"}, status=404)
|
||||
# 用最新文件作为代表
|
||||
matched_files.sort(key=lambda f: os.path.getmtime(os.path.join(output_dir, f)), reverse=True)
|
||||
target_file = matched_files[0]
|
||||
fpath = os.path.join(output_dir, target_file)
|
||||
m = meta.get(target_file, {})
|
||||
mtime = os.path.getmtime(fpath)
|
||||
start_ms = int(m.get("start_time", mtime) * 1000)
|
||||
end_ms = int(m.get("end_time", mtime) * 1000)
|
||||
ext = os.path.splitext(target_file)[1].lower()
|
||||
media = "images" if ext in {'.png','.jpg','.jpeg','.webp','.gif'} else "video"
|
||||
# 读取 PNG 工作流元数据
|
||||
workflow = None
|
||||
if ext == ".png":
|
||||
try:
|
||||
with open(fpath, "rb") as pf:
|
||||
pf.read(8)
|
||||
while True:
|
||||
raw = pf.read(8)
|
||||
if len(raw) < 8:
|
||||
break
|
||||
length = struct.unpack(">I", raw[:4])[0]
|
||||
chunk_type = raw[4:8]
|
||||
data = pf.read(length)
|
||||
pf.read(4)
|
||||
if chunk_type == b"tEXt":
|
||||
key, val = data.split(b"\x00", 1)
|
||||
k = key.decode("ascii", errors="replace")
|
||||
if k == "workflow":
|
||||
workflow = _json.loads(val)
|
||||
elif chunk_type == b"IEND":
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
# 构建 outputs:包含该执行的所有文件
|
||||
outputs = {}
|
||||
for i, fname in enumerate(matched_files):
|
||||
e = os.path.splitext(fname)[1].lower()
|
||||
mt = "images" if e in {'.png','.jpg','.jpeg','.webp','.gif'} else "gifs"
|
||||
outputs.setdefault(str(i), {}).setdefault(mt, []).append(
|
||||
{"filename": fname, "subfolder": "", "type": "output"}
|
||||
)
|
||||
job_detail = {
|
||||
"id": job_id,
|
||||
"status": "completed",
|
||||
"create_time": start_ms,
|
||||
"execution_start_time": start_ms,
|
||||
"execution_end_time": end_ms,
|
||||
"preview_output": {
|
||||
"filename": target_file,
|
||||
"subfolder": "",
|
||||
"type": "output",
|
||||
"nodeId": "0",
|
||||
"mediaType": media,
|
||||
},
|
||||
"outputs_count": len(matched_files),
|
||||
"execution_error": None,
|
||||
"workflow_id": job_id,
|
||||
"workflow": {
|
||||
"extra_data": {
|
||||
"extra_pnginfo": {"workflow": workflow}
|
||||
}
|
||||
} if workflow else None,
|
||||
"outputs": outputs,
|
||||
}
|
||||
return web.json_response(job_detail)
|
||||
|
||||
@PromptServer.instance.routes.post("/o1key/delete_history")
|
||||
async def delete_history_item(request):
|
||||
"""删除持久化历史记录及对应的输出文件"""
|
||||
import os, json as _json
|
||||
body = await request.json()
|
||||
job_ids = body.get("delete", [])
|
||||
if not job_ids:
|
||||
return web.json_response({"success": False, "error": "missing ids"}, status=400)
|
||||
output_dir = os.path.abspath(folder_paths.get_output_directory())
|
||||
meta_file = _get_o1key_history_meta_file(output_dir)
|
||||
meta = {}
|
||||
if os.path.isfile(meta_file):
|
||||
try:
|
||||
with open(meta_file, "r", encoding="utf-8") as mf:
|
||||
meta = _json.load(mf)
|
||||
except Exception:
|
||||
pass
|
||||
deleted_files = []
|
||||
for job_id in job_ids:
|
||||
files_to_remove = []
|
||||
for fname, m in list(meta.items()):
|
||||
if m.get("workflow_id") == job_id:
|
||||
files_to_remove.append(fname)
|
||||
for fname in files_to_remove:
|
||||
meta.pop(fname, None)
|
||||
fpath = os.path.join(output_dir, fname)
|
||||
if os.path.isfile(fpath):
|
||||
try:
|
||||
os.remove(fpath)
|
||||
deleted_files.append(fname)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
with open(meta_file, "w", encoding="utf-8") as mf:
|
||||
_json.dump(meta, mf, ensure_ascii=False)
|
||||
except Exception:
|
||||
pass
|
||||
return web.json_response({"success": True, "deleted": deleted_files})
|
||||
|
||||
@PromptServer.instance.routes.post("/o1key/update")
|
||||
async def update_node_package(request):
|
||||
if request.headers.get("X-O1Key-Update") != "1":
|
||||
return web.json_response({"error": "无效的更新请求。"}, status=403)
|
||||
if not _update_lock.acquire(blocking=False):
|
||||
return web.json_response({"error": "更新正在进行,请稍候。"}, status=409)
|
||||
try:
|
||||
result = await asyncio.to_thread(update_package)
|
||||
return web.json_response(result)
|
||||
except UpdateError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=409)
|
||||
except Exception:
|
||||
logging.exception("o1key update failed")
|
||||
return web.json_response({"error": "更新失败,请查看 ComfyUI 日志。"}, status=500)
|
||||
finally:
|
||||
_update_lock.release()
|
||||
|
||||
# === AI 聊天代理(流式 SSE 透传) ===
|
||||
@PromptServer.instance.routes.post("/o1key/restart")
|
||||
async def restart_server(request):
|
||||
import sys, os as _ros, subprocess, threading
|
||||
def _do_restart():
|
||||
import time
|
||||
time.sleep(1.5)
|
||||
skip = {"--auto-launch", "--auto_launch", "--launch", "--windows-standalone-build"}
|
||||
args = [a for a in sys.argv if a not in skip]
|
||||
args.append("--disable-auto-launch")
|
||||
subprocess.Popen([sys.executable] + args, cwd=_ros.getcwd())
|
||||
_ros._exit(0)
|
||||
threading.Thread(target=_do_restart, daemon=True).start()
|
||||
return web.json_response({"success": True, "message": "正在重启..."})
|
||||
|
||||
# === AI 聊天代理(流式 SSE 透传) ===
|
||||
@PromptServer.instance.routes.post("/o1key/chat/completions")
|
||||
async def chat_completions_proxy(request):
|
||||
import aiohttp as _aiohttp
|
||||
import json as _cjson
|
||||
|
||||
data = await request.json()
|
||||
config = load_config()
|
||||
api_key = config.get("O1KEY_API_KEY", "")
|
||||
if not api_key:
|
||||
return web.json_response({"error": "未配置 API Key"}, status=401)
|
||||
|
||||
route = data.get("route", "CF加速")
|
||||
base_url = NETWORK_ROUTES.get(route, "https://cf-api.o1key.com")
|
||||
model = data.get("model", "gpt-5.5")
|
||||
messages = data.get("messages", [])
|
||||
|
||||
url = f"{base_url}/v1/chat/completions"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
body = {"model": model, "messages": messages, "stream": True}
|
||||
|
||||
resp = web.StreamResponse(
|
||||
status=200, reason="OK",
|
||||
headers={
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
}
|
||||
)
|
||||
await resp.prepare(request)
|
||||
|
||||
try:
|
||||
timeout = _aiohttp.ClientTimeout(total=120)
|
||||
async with _aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.post(url, headers=headers, json=body) as upstream:
|
||||
if upstream.status != 200:
|
||||
err = await upstream.text()
|
||||
await resp.write(f"data: {_cjson.dumps({'error': err})}\n\n".encode())
|
||||
await resp.write(b"data: [DONE]\n\n")
|
||||
return resp
|
||||
async for chunk in upstream.content.iter_any():
|
||||
await resp.write(chunk)
|
||||
except Exception as e:
|
||||
await resp.write(f"data: {_cjson.dumps({'error': str(e)})}\n\n".encode())
|
||||
await resp.write(b"data: [DONE]\n\n")
|
||||
|
||||
return resp
|
||||
|
||||
# === 执行事件 Hook:持久化耗时元数据 ===
|
||||
import time as _time, json as _json2, os as _os
|
||||
_execution_tracker = {}
|
||||
|
||||
_orig_send_sync = PromptServer.instance.send_sync
|
||||
|
||||
def _patched_send_sync(event, data, *args, **kwargs):
|
||||
try:
|
||||
if event == "execution_start":
|
||||
pid = data.get("prompt_id", "")
|
||||
if pid:
|
||||
_execution_tracker[pid] = {"start": _time.time(), "outputs": []}
|
||||
elif event == "executed":
|
||||
pid = data.get("prompt_id", "")
|
||||
output = data.get("output") or {}
|
||||
if pid and pid in _execution_tracker:
|
||||
for img in output.get("images", []) + output.get("gifs", []):
|
||||
if img.get("type") == "output" and img.get("filename"):
|
||||
_execution_tracker[pid]["outputs"].append(img["filename"])
|
||||
elif event == "executing" and data.get("node") is None:
|
||||
pid = data.get("prompt_id", "")
|
||||
tracker = _execution_tracker.pop(pid, None)
|
||||
if tracker and tracker["outputs"]:
|
||||
end_time = _time.time()
|
||||
start_time = tracker["start"]
|
||||
output_dir = _os.path.abspath(folder_paths.get_output_directory())
|
||||
meta_file = _get_o1key_history_meta_file(output_dir)
|
||||
meta = {}
|
||||
if _os.path.isfile(meta_file):
|
||||
try:
|
||||
with open(meta_file, "r", encoding="utf-8") as mf:
|
||||
meta = _json2.load(mf)
|
||||
except Exception:
|
||||
pass
|
||||
for fname in tracker["outputs"]:
|
||||
meta[fname] = {
|
||||
"start_time": start_time,
|
||||
"end_time": end_time,
|
||||
"outputs_count": len(tracker["outputs"]),
|
||||
"workflow_id": pid,
|
||||
}
|
||||
try:
|
||||
with open(meta_file, "w", encoding="utf-8") as mf:
|
||||
_json2.dump(meta, mf, ensure_ascii=False)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
return _orig_send_sync(event, data, *args, **kwargs)
|
||||
|
||||
PromptServer.instance.send_sync = _patched_send_sync
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
+7
-1
@@ -6,5 +6,11 @@ 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 .newapi_veo_client import NewAPIVeoClient
|
||||
from .grok_video_client import GrokVideoClient
|
||||
from .openai_client import OpenAIAPIClient
|
||||
|
||||
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient']
|
||||
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'NewAPIVeoClient', 'GrokVideoClient', 'OpenAIAPIClient']
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
"""
|
||||
异步生图 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}"
|
||||
+260
-99
@@ -9,8 +9,14 @@ import threading
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..utils.http_error import HTTP_ERROR_MESSAGES, RETRYABLE_STATUS_CODES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR, get_friendly_message
|
||||
|
||||
|
||||
|
||||
class BaseAPIClient(ABC):
|
||||
"""
|
||||
@@ -26,19 +32,21 @@ class BaseAPIClient(ABC):
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
max_request_size: int = 20 * 1024 * 1024
|
||||
max_request_size: int = 100 * 1024 * 1024
|
||||
):
|
||||
"""
|
||||
初始化客户端
|
||||
|
||||
|
||||
Args:
|
||||
base_url: API 基础 URL
|
||||
api_key: API 密钥
|
||||
max_request_size: 最大请求体大小(字节),默认 20MB
|
||||
max_request_size: 兼容参数;基类不再用它限制 JSON 请求体,
|
||||
部分子类仍用它作为上传文件大小限制
|
||||
"""
|
||||
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:
|
||||
@@ -79,6 +87,15 @@ 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]:
|
||||
"""
|
||||
获取请求头
|
||||
@@ -100,26 +117,19 @@ class BaseAPIClient(ABC):
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
def check_request_size(self, request_body: Dict[str, Any]) -> None:
|
||||
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
|
||||
"""
|
||||
检查请求体大小是否超过限制
|
||||
子类可重写:为指定 HTTP 状态码返回自定义错误文案。
|
||||
若返回 None,则使用基类默认拼接文案。
|
||||
|
||||
Args:
|
||||
request_body: 请求体字典
|
||||
status_code: HTTP 状态码(如 429、503)
|
||||
error_message: API 返回的原始错误信息
|
||||
|
||||
Raises:
|
||||
ValueError: 如果请求体超过限制
|
||||
Returns:
|
||||
自定义完整错误文案,或 None 表示使用默认
|
||||
"""
|
||||
request_json = json.dumps(request_body)
|
||||
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(
|
||||
f"请求体大小 {size_mb:.2f}MB 超过限制 {limit_mb:.0f}MB,"
|
||||
"请降低分辨率或减少图片数量"
|
||||
)
|
||||
return None
|
||||
|
||||
async def request_async(
|
||||
self,
|
||||
@@ -130,84 +140,155 @@ 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 = aiohttp.ClientSession()
|
||||
session = self._make_session()
|
||||
close_session = True
|
||||
|
||||
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:
|
||||
|
||||
# 设置请求超时:连接超时 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
|
||||
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
|
||||
# 针对常见错误状态码提供友好提示
|
||||
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()
|
||||
|
||||
# 返回状态码和错误文本,由外层处理重试
|
||||
return {"_error": True, "_status": response.status, "_text": 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:
|
||||
last_error_status = None
|
||||
last_error_text = ""
|
||||
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
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()
|
||||
|
||||
result = request_task.result()
|
||||
else:
|
||||
result = await _do_request()
|
||||
|
||||
if isinstance(result, dict) and result.get("_error"):
|
||||
status = result["_status"]
|
||||
error_text = result["_text"]
|
||||
last_error_status = status
|
||||
last_error_text = error_text
|
||||
|
||||
if status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
|
||||
friendly = HTTP_ERROR_MESSAGES.get(status, f"请求失败 ({status})")
|
||||
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
|
||||
print(f"{friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
|
||||
if status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[status])
|
||||
raise RuntimeError(get_friendly_message(status, error_text))
|
||||
|
||||
return result
|
||||
|
||||
if last_error_status and last_error_status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[last_error_status])
|
||||
raise RuntimeError(get_friendly_message(last_error_status or 0, last_error_text))
|
||||
|
||||
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
|
||||
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
|
||||
async def request_get_async(
|
||||
self,
|
||||
endpoint: str,
|
||||
@@ -222,6 +303,7 @@ class BaseAPIClient(ABC):
|
||||
endpoint: API 端点
|
||||
session: aiohttp 会话(可选)
|
||||
use_bearer_token: 是否使用 Bearer Token 认证(默认为 True)
|
||||
timeout: 超时时间(秒)- 已废弃,由服务器端控制
|
||||
|
||||
Returns:
|
||||
响应 JSON
|
||||
@@ -234,41 +316,67 @@ class BaseAPIClient(ABC):
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = aiohttp.ClientSession()
|
||||
session = self._make_session()
|
||||
close_session = True
|
||||
|
||||
|
||||
try:
|
||||
# 设置超时
|
||||
timeout_obj = aiohttp.ClientTimeout(total=timeout) if timeout else None
|
||||
async with session.get(url, headers=headers, timeout=timeout_obj) as response:
|
||||
_get_start = time.time()
|
||||
async with session.get(url, headers=headers) as response:
|
||||
_get_elapsed = time.time() - _get_start
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
|
||||
# 尝试解析 JSON 错误信息,提取关键内容
|
||||
error_message = error_text
|
||||
try:
|
||||
error_json = json.loads(error_text)
|
||||
# 尝试从多个常见位置提取错误信息
|
||||
if "error" in error_json:
|
||||
if isinstance(error_json["error"], dict):
|
||||
error_message = error_json["error"].get("message", error_text)
|
||||
else:
|
||||
error_message = str(error_json["error"])
|
||||
elif "message" in error_json:
|
||||
error_message = error_json["message"]
|
||||
except:
|
||||
# 如果不是 JSON,使用原始文本
|
||||
pass
|
||||
|
||||
# 针对常见错误状态码提供友好提示
|
||||
if response.status == 504:
|
||||
if response.status == 400:
|
||||
raise RuntimeError(
|
||||
f"API 请求超时 (504 Gateway Timeout)\n"
|
||||
f"原因:服务器响应超时或该端点暂时不可用\n"
|
||||
f"建议:稍后重试"
|
||||
f"请求参数错误 (400 Bad Request)\n"
|
||||
f"API 返回错误:{error_message}\n"
|
||||
f"建议:检查请求参数"
|
||||
)
|
||||
elif response.status == 503:
|
||||
elif response.status == 401:
|
||||
raise RuntimeError(
|
||||
f"服务暂时不可用 (503 Service Unavailable)\n"
|
||||
f"原因:服务过载或维护中\n"
|
||||
f"建议:稍后重试"
|
||||
f"认证失败 (401 Unauthorized)\n"
|
||||
f"API 返回错误:{error_message}\n"
|
||||
f"建议:检查 API 密钥"
|
||||
)
|
||||
elif response.status == 429:
|
||||
raise RuntimeError(
|
||||
f"请求频率超限 (429 Too Many Requests)\n"
|
||||
f"原因:API 配额用尽或请求过于频繁\n"
|
||||
f"建议:等待一段时间后重试"
|
||||
)
|
||||
custom = self.get_http_error_message(429, error_message)
|
||||
if custom is not None:
|
||||
raise RuntimeError(custom)
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[429])
|
||||
elif response.status == 503:
|
||||
custom = self.get_http_error_message(503, error_message)
|
||||
if custom is not None:
|
||||
raise RuntimeError(custom)
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[503])
|
||||
elif response.status == 504:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[504])
|
||||
elif response.status == 502:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[502])
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"API 请求失败 (状态码: {response.status}): {error_text}"
|
||||
f"API 请求失败 (状态码: {response.status})\n"
|
||||
f"API 返回错误:{error_message}"
|
||||
)
|
||||
|
||||
return await response.json()
|
||||
_resp_data = await response.json()
|
||||
return _resp_data
|
||||
|
||||
finally:
|
||||
if close_session:
|
||||
@@ -294,9 +402,7 @@ class BaseAPIClient(ABC):
|
||||
total = len(requests)
|
||||
|
||||
# 创建无限制的连接器
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
async with self._make_session() as session:
|
||||
tasks = []
|
||||
|
||||
for req in requests:
|
||||
@@ -366,3 +472,58 @@ 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}"
|
||||
|
||||
@@ -0,0 +1,358 @@
|
||||
"""
|
||||
豆包生图 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
|
||||
from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
|
||||
|
||||
|
||||
# ── 固定端点 ──────────────────────────────────────────────────────────────────
|
||||
_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 请求(带退避重试)
|
||||
last_status = None
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
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:
|
||||
last_status = resp.status
|
||||
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
|
||||
friendly = HTTP_ERROR_MESSAGES.get(resp.status)
|
||||
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
|
||||
print(f"[豆包生图] {friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
if resp.status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
|
||||
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]}")
|
||||
|
||||
break
|
||||
else:
|
||||
if last_status and last_status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
|
||||
raise RuntimeError(f"请求失败: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
|
||||
|
||||
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),请检查网络或稍后重试"
|
||||
)
|
||||
@@ -0,0 +1,199 @@
|
||||
"""
|
||||
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
|
||||
from ..utils.http_error import HTTP_ERROR_MESSAGES
|
||||
|
||||
|
||||
# 显示名 → 实际请求值的映射
|
||||
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:
|
||||
if resp.status_code in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status_code])
|
||||
raise RuntimeError(
|
||||
f"提交任务失败 (HTTP {resp.status_code})\n"
|
||||
f"响应: {resp.text[:500]}"
|
||||
)
|
||||
|
||||
result = resp.json()
|
||||
task_id = result.get("id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"服务器返回异常: 未获取到任务ID\n{result}")
|
||||
|
||||
return task_id
|
||||
|
||||
def _poll_result_sync(
|
||||
self,
|
||||
task_id: str,
|
||||
poll_interval: int,
|
||||
progress_callback=None,
|
||||
) -> bytes:
|
||||
"""同步轮询任务状态(直连容器),返回结果图像二进制"""
|
||||
url = f"{POLL_BASE_URL}{self.STATUS_ENDPOINT.format(task_id=task_id)}"
|
||||
|
||||
start_time = time.time()
|
||||
last_status = None
|
||||
|
||||
while True:
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
try:
|
||||
resp = requests.get(url, timeout=30)
|
||||
except requests.exceptions.ConnectionError:
|
||||
raise RuntimeError("轮询时无法连接到服务器,请检查网络")
|
||||
|
||||
if resp.status_code != 200:
|
||||
raise RuntimeError(
|
||||
f"查询任务状态失败 (HTTP {resp.status_code})\n"
|
||||
f"响应: {resp.text[:500]}"
|
||||
)
|
||||
|
||||
result = resp.json()
|
||||
status = result.get("status", "unknown")
|
||||
|
||||
# 状态变化时打印日志
|
||||
if status != last_status:
|
||||
elapsed_str = f"{elapsed:.0f}s"
|
||||
print(f"Flux Edit: [{elapsed_str}] 任务 {task_id[:8]}... → {status}")
|
||||
last_status = status
|
||||
|
||||
if progress_callback:
|
||||
elapsed_str = f"{elapsed:.0f}s"
|
||||
status_desc = {
|
||||
"pending": "排队中",
|
||||
"processing": "处理中",
|
||||
"generating": "生图中,请耐心等待,预计耗时140s左右",
|
||||
}.get(status, status)
|
||||
progress_callback(f"{status_desc} (当前进度:{elapsed_str})")
|
||||
|
||||
if status == "completed":
|
||||
# 解码 base64 图像
|
||||
b64_data = result.get("result")
|
||||
if not b64_data:
|
||||
raise RuntimeError("任务完成但未返回图像数据")
|
||||
return base64.b64decode(b64_data)
|
||||
|
||||
elif status == "failed":
|
||||
error_msg = result.get("error", "未知错误")
|
||||
raise RuntimeError(
|
||||
f"图像编辑任务失败\n"
|
||||
f"错误: {error_msg}"
|
||||
)
|
||||
|
||||
elif status in ("not_found",):
|
||||
raise RuntimeError(
|
||||
f"任务未找到: {task_id}\n"
|
||||
f"可能已被清理或 ID 无效"
|
||||
)
|
||||
|
||||
# 继续等待
|
||||
time.sleep(poll_interval)
|
||||
|
||||
def query_balance_sync(self) -> dict:
|
||||
"""查询余额(兼容现有节点的 finally 块调用)"""
|
||||
return {"name": "flux-edit", "total_available": 0}
|
||||
@@ -0,0 +1,194 @@
|
||||
"""
|
||||
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 getattr(self, '_route_base_url', None) or 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),
|
||||
thinking_level=kwargs.get("thinking_level"),
|
||||
request_log_enabled=False,
|
||||
)
|
||||
|
||||
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]:
|
||||
images = result_data.get("images") if isinstance(result_data, dict) else None
|
||||
if isinstance(images, list) and images:
|
||||
parsed = []
|
||||
for item in images:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
image_url = item.get("url") or item.get("image_url")
|
||||
if image_url:
|
||||
async with session.get(image_url) as img_resp:
|
||||
if img_resp.status == 200:
|
||||
img_bytes = await img_resp.read()
|
||||
parsed.append(Image.open(BytesIO(img_bytes)).convert("RGB"))
|
||||
else:
|
||||
raise RuntimeError(f"下载图片失败 ({img_resp.status}): {image_url}")
|
||||
if parsed:
|
||||
return parsed
|
||||
|
||||
# 异步接口可能直接返回 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)"""
|
||||
|
||||
def _coerce(val) -> Optional[float]:
|
||||
if val is None or isinstance(val, bool):
|
||||
return None
|
||||
if isinstance(val, (int, float)):
|
||||
progress = float(val)
|
||||
elif isinstance(val, str):
|
||||
text = val.strip()
|
||||
if not text:
|
||||
return None
|
||||
has_percent_suffix = text.endswith("%")
|
||||
if has_percent_suffix:
|
||||
text = text[:-1].strip()
|
||||
try:
|
||||
progress = float(text)
|
||||
except ValueError:
|
||||
return None
|
||||
if has_percent_suffix:
|
||||
progress /= 100.0
|
||||
else:
|
||||
return None
|
||||
if progress > 1.0:
|
||||
progress /= 100.0
|
||||
return max(0.0, min(progress, 1.0))
|
||||
|
||||
# 直接字段:progress / percentage
|
||||
for field in ("progress", "percentage", "percent"):
|
||||
progress = _coerce(response.get(field))
|
||||
if progress is not None:
|
||||
return progress
|
||||
|
||||
# 嵌套字段:progressInfo / progress_info
|
||||
progress_info = response.get("progressInfo") or response.get("progress_info")
|
||||
if isinstance(progress_info, dict):
|
||||
for field in ("progress", "percentage", "percent"):
|
||||
progress = _coerce(progress_info.get(field))
|
||||
if progress is not None:
|
||||
return progress
|
||||
|
||||
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)
|
||||
+719
-315
File diff suppressed because it is too large
Load Diff
@@ -8,15 +8,15 @@ from typing import Any, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..utils.config import get_api_key_or_raise
|
||||
from ..models_config import get_flash_model_endpoint, get_enabled_flash_models
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
from ..models_config import (
|
||||
get_flash_model_endpoint,
|
||||
get_enabled_flash_models,
|
||||
get_flash_model_thinking_level_value,
|
||||
)
|
||||
from .base_client import BaseAPIClient
|
||||
|
||||
|
||||
# API 基础配置
|
||||
API_BASE_URL = "https://api.o1key.com"
|
||||
|
||||
|
||||
class GeminiFlashClient(BaseAPIClient):
|
||||
"""
|
||||
Gemini Flash API 客户端
|
||||
@@ -24,8 +24,7 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
|
||||
特点:
|
||||
- 支持图片和视频输入
|
||||
- 支持系统指令
|
||||
- 支持不同思考深度
|
||||
- 支持动态思考等级端点(不思考/低/中/高)
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
@@ -39,46 +38,65 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
|
||||
super().__init__(
|
||||
base_url=API_BASE_URL,
|
||||
base_url=get_api_base_url(),
|
||||
api_key=api_key,
|
||||
max_request_size=20 * 1024 * 1024 # 20MB
|
||||
max_request_size=100 * 1024 * 1024 # 100MB
|
||||
)
|
||||
|
||||
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, thinking_depth)
|
||||
endpoint = get_flash_model_endpoint(model)
|
||||
|
||||
if endpoint is None:
|
||||
# 回退到默认端点
|
||||
# 回退到第一个启用的模型端点
|
||||
default_models = get_enabled_flash_models()
|
||||
if default_models:
|
||||
endpoint = get_flash_model_endpoint(default_models[0], thinking_depth)
|
||||
endpoint = get_flash_model_endpoint(default_models[0])
|
||||
|
||||
if endpoint is None:
|
||||
raise ValueError(f"无法获取模型 '{model}' 的端点 (思考深度: {thinking_depth})")
|
||||
raise ValueError(f"无法获取模型 '{model}' 的端点")
|
||||
|
||||
return endpoint
|
||||
|
||||
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
|
||||
"""Gemini 请求 429/503 时返回图中约定的多行错误框文案。"""
|
||||
if status_code == 429:
|
||||
return (
|
||||
"莫慌!该模型暂时超出速率限制啦\n"
|
||||
"解决方案如下(任意一种):\n"
|
||||
"1.切换当前模型\n"
|
||||
"2.前往后台,修改令牌分组"
|
||||
)
|
||||
if status_code == 503:
|
||||
return (
|
||||
"警报!谷歌服务器当前过载!\n"
|
||||
"解决方案如下:\n"
|
||||
"1.摸会儿鱼吧,我也没办法,谷歌会尽快恢复,嘿嘿~\n"
|
||||
"2.切换其他模型\n"
|
||||
"3.前往后台,修改令牌分组"
|
||||
)
|
||||
return None
|
||||
|
||||
def build_request_body(
|
||||
self,
|
||||
prompt: str = "",
|
||||
system_instruction: Optional[str] = None,
|
||||
model: str = "gemini-3-flash-preview",
|
||||
thinking_level: str = "不思考",
|
||||
image_data: Optional[List[Dict[str, str]]] = None,
|
||||
video_data: Optional[Dict[str, str]] = None,
|
||||
document_data: Optional[Dict[str, str]] = None,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
@@ -86,9 +104,11 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
|
||||
Args:
|
||||
prompt: 用户提示词
|
||||
system_instruction: 系统指令(可选)
|
||||
model: 模型名称
|
||||
thinking_level: 思考等级(不思考/低/中/高)- 通过动态端点控制,不需要在请求体中传递
|
||||
image_data: 图片数据列表,每个元素包含 mime_type 和 data
|
||||
video_data: 视频数据,包含 mime_type 和 data
|
||||
document_data: 文档数据,包含 mime_type 和 data
|
||||
|
||||
Returns:
|
||||
请求体字典
|
||||
@@ -118,6 +138,15 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
}
|
||||
})
|
||||
|
||||
# 添加文档部分(如果有)
|
||||
if document_data:
|
||||
parts.append({
|
||||
"inline_data": {
|
||||
"mime_type": document_data["mime_type"],
|
||||
"data": document_data["data"]
|
||||
}
|
||||
})
|
||||
|
||||
# 构建请求体
|
||||
request_body = {
|
||||
"contents": [
|
||||
@@ -127,12 +156,15 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
]
|
||||
}
|
||||
|
||||
# 添加系统指令(如果有)
|
||||
if system_instruction and system_instruction.strip():
|
||||
request_body["system_instruction"] = {
|
||||
"parts": [
|
||||
{"text": system_instruction}
|
||||
]
|
||||
# 对于支持 thinkingConfig 的固定端点模型(如 gemini-3-pro-preview)
|
||||
# 通过请求体传递思考等级;动态端点模型(如 gemini-3-flash-preview)
|
||||
# 通过不同 URL 端点控制,无需此字段
|
||||
thinking_level_value = get_flash_model_thinking_level_value(model, thinking_level)
|
||||
if thinking_level_value is not None:
|
||||
request_body["generationConfig"] = {
|
||||
"thinkingConfig": {
|
||||
"thinkingLevel": thinking_level_value
|
||||
}
|
||||
}
|
||||
|
||||
return request_body
|
||||
@@ -197,10 +229,10 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
self,
|
||||
prompt: str,
|
||||
model: str = "gemini-3-flash-preview",
|
||||
thinking_depth: str = "不思考",
|
||||
system_instruction: Optional[str] = None,
|
||||
thinking_level: str = "不思考",
|
||||
image_data: Optional[List[Dict[str, str]]] = None,
|
||||
video_data: Optional[Dict[str, str]] = None,
|
||||
document_data: Optional[Dict[str, str]] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None
|
||||
) -> str:
|
||||
"""
|
||||
@@ -209,31 +241,29 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
Args:
|
||||
prompt: 用户提示词
|
||||
model: 模型名称
|
||||
thinking_depth: 思考深度
|
||||
system_instruction: 系统指令
|
||||
thinking_level: 思考等级(不思考/低/中/高)
|
||||
image_data: 图片数据列表
|
||||
video_data: 视频数据
|
||||
document_data: 文档数据
|
||||
session: aiohttp 会话
|
||||
|
||||
Returns:
|
||||
生成的文本内容
|
||||
"""
|
||||
endpoint = self.get_endpoint(model=model, thinking_depth=thinking_depth)
|
||||
endpoint = self.get_endpoint(model=model)
|
||||
request_body = self.build_request_body(
|
||||
prompt=prompt,
|
||||
system_instruction=system_instruction,
|
||||
model=model,
|
||||
thinking_level=thinking_level,
|
||||
image_data=image_data,
|
||||
video_data=video_data
|
||||
video_data=video_data,
|
||||
document_data=document_data
|
||||
)
|
||||
|
||||
# 根据是否有视频设置超时(视频处理需要更长时间)
|
||||
timeout = 300 if video_data else 180
|
||||
|
||||
response = await self.request_async(
|
||||
endpoint,
|
||||
request_body,
|
||||
session,
|
||||
timeout=timeout
|
||||
session
|
||||
)
|
||||
|
||||
return self.parse_response(response)
|
||||
@@ -242,10 +272,10 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
self,
|
||||
prompt: str,
|
||||
model: str = "gemini-3-flash-preview",
|
||||
thinking_depth: str = "不思考",
|
||||
system_instruction: Optional[str] = None,
|
||||
thinking_level: str = "不思考",
|
||||
image_data: Optional[List[Dict[str, str]]] = None,
|
||||
video_data: Optional[Dict[str, str]] = None
|
||||
video_data: Optional[Dict[str, str]] = None,
|
||||
document_data: Optional[Dict[str, str]] = None
|
||||
) -> str:
|
||||
"""
|
||||
同步生成文本(用于 ComfyUI 节点)
|
||||
@@ -253,10 +283,10 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
Args:
|
||||
prompt: 用户提示词
|
||||
model: 模型名称
|
||||
thinking_depth: 思考深度
|
||||
system_instruction: 系统指令
|
||||
thinking_level: 思考等级(不思考/低/中/高)
|
||||
image_data: 图片数据列表
|
||||
video_data: 视频数据
|
||||
document_data: 文档数据
|
||||
|
||||
Returns:
|
||||
生成的文本内容
|
||||
@@ -264,10 +294,10 @@ class GeminiFlashClient(BaseAPIClient):
|
||||
coro = self.generate_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
thinking_depth=thinking_depth,
|
||||
system_instruction=system_instruction,
|
||||
thinking_level=thinking_level,
|
||||
image_data=image_data,
|
||||
video_data=video_data
|
||||
video_data=video_data,
|
||||
document_data=document_data
|
||||
)
|
||||
|
||||
return self.run_async_in_thread(coro)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,355 @@
|
||||
"""
|
||||
Grok Image API 客户端
|
||||
支持两个接口:
|
||||
- POST /v1/images/generations 文生图
|
||||
- POST /v1/images/edits 图生图(带参考图)
|
||||
|
||||
上游 API 格式与 OpenAI Images API 兼容。
|
||||
"""
|
||||
|
||||
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_base_url_by_route
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
|
||||
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"
|
||||
|
||||
_MODEL_NAME_MAP = {
|
||||
"Grok Image": "grok-imagine-image",
|
||||
"Grok Image Pro": "grok-imagine-image-quality",
|
||||
}
|
||||
|
||||
_REQUEST_TIMEOUT = 900
|
||||
_MAX_BODY_BYTES = 20 * 1024 * 1024
|
||||
|
||||
_MAX_RETRIES = 3
|
||||
_RETRY_DELAY = 5
|
||||
|
||||
|
||||
class GrokImageClient:
|
||||
|
||||
def __init__(self, route: str = "全球加速"):
|
||||
self.api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
self.base_url = get_base_url_by_route(route)
|
||||
|
||||
def _json_headers(self) -> dict:
|
||||
return {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def _auth_headers(self) -> dict:
|
||||
return {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
# ── 图像工具 ──────────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _shrink_png_to_limit(png_bytes: bytes, max_bytes: int, label: str = "") -> bytes:
|
||||
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
|
||||
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 Grok 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 _pil_list_to_tensor(images: List[Image.Image]) -> torch.Tensor:
|
||||
if not images:
|
||||
placeholder = Image.new("RGB", (512, 512), (128, 128, 128))
|
||||
images = [placeholder]
|
||||
tensors = []
|
||||
for img in images:
|
||||
arr = np.array(img.convert("RGB")).astype(np.float32) / 255.0
|
||||
tensors.append(torch.from_numpy(arr))
|
||||
return torch.stack(tensors, dim=0)
|
||||
|
||||
# ── 中断轮询 ──────────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
async def _poll_interrupt():
|
||||
while True:
|
||||
await asyncio.sleep(0.5)
|
||||
if _INTERRUPT_AVAILABLE and processing_interrupted():
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
async def _run_with_interrupt(coro):
|
||||
if not _INTERRUPT_AVAILABLE:
|
||||
return await coro
|
||||
request_task = asyncio.ensure_future(coro)
|
||||
interrupt_task = asyncio.ensure_future(GrokImageClient._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()
|
||||
|
||||
# ── 响应解析 ──────────────────────────────────────────────────────────────
|
||||
|
||||
async def _parse_response(self, resp_json: dict, session: aiohttp.ClientSession) -> List[Image.Image]:
|
||||
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 字段")
|
||||
images: List[Image.Image] = []
|
||||
for idx, item in enumerate(data_list):
|
||||
b64 = item.get("b64_json", "")
|
||||
url = item.get("url", "")
|
||||
if b64:
|
||||
img_bytes = base64.b64decode(b64)
|
||||
img = Image.open(BytesIO(img_bytes))
|
||||
images.append(img)
|
||||
elif url and url.startswith("http"):
|
||||
async with session.get(url, allow_redirects=True) as r:
|
||||
if r.status != 200:
|
||||
raise RuntimeError(f"图像下载失败 HTTP {r.status}")
|
||||
img_bytes = await r.read()
|
||||
images.append(Image.open(BytesIO(img_bytes)))
|
||||
else:
|
||||
print(f"[o1key Grok Image] 警告:第 {idx + 1} 条数据无有效图像,已跳过")
|
||||
return images
|
||||
|
||||
# ── 文生图(generations 接口)─────────────────────────────────────────────
|
||||
|
||||
async def _generate_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
n: int,
|
||||
) -> List[Image.Image]:
|
||||
api_model = _MODEL_NAME_MAP.get(model, model)
|
||||
body: dict = {
|
||||
"model": api_model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio if aspect_ratio else "auto",
|
||||
"resolution": resolution if resolution else "1k",
|
||||
"response_format": "b64_json",
|
||||
}
|
||||
|
||||
url = f"{self.base_url}{_ENDPOINT_GENERATIONS}"
|
||||
log_body = {k: v for k, v in body.items()}
|
||||
print(f"[o1key Grok Image] 请求 URL: {url}")
|
||||
print(f"[o1key Grok Image] 请求体: {json.dumps(log_body, ensure_ascii=False)}")
|
||||
|
||||
results = []
|
||||
for i in range(n):
|
||||
images = await self._do_request_with_retry(url, body)
|
||||
results.extend(images)
|
||||
if n > 1:
|
||||
print(f"[o1key Grok Image] 第 {i+1}/{n} 张完成")
|
||||
return results
|
||||
|
||||
# ── 图生图(edits 接口)───────────────────────────────────────────────────
|
||||
|
||||
async def _edit_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
n: int,
|
||||
image_list: List[torch.Tensor],
|
||||
) -> List[Image.Image]:
|
||||
api_model = _MODEL_NAME_MAP.get(model, model)
|
||||
body: dict = {
|
||||
"model": api_model,
|
||||
"prompt": prompt,
|
||||
"response_format": "b64_json",
|
||||
}
|
||||
if aspect_ratio and aspect_ratio != "auto":
|
||||
body["aspect_ratio"] = aspect_ratio
|
||||
if resolution:
|
||||
body["resolution"] = resolution
|
||||
|
||||
# 参考图转 base64 字符串
|
||||
pil_images = tensor_to_pil(image_list[0])
|
||||
img = pil_images[0]
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
png_bytes = buf.getvalue()
|
||||
png_bytes = self._shrink_png_to_limit(png_bytes, _MAX_BODY_BYTES // 2)
|
||||
body["image"] = base64.b64encode(png_bytes).decode("utf-8")
|
||||
|
||||
url = f"{self.base_url}{_ENDPOINT_EDITS}"
|
||||
log_body = {k: (v[:50] + "..." if k == "image" and len(v) > 50 else v) for k, v in body.items()}
|
||||
print(f"[o1key Grok Image] 请求 URL: {url}")
|
||||
print(f"[o1key Grok Image] 请求体: {json.dumps(log_body, ensure_ascii=False)}")
|
||||
|
||||
results = []
|
||||
for i in range(n):
|
||||
images = await self._do_request_with_retry(url, body)
|
||||
results.extend(images)
|
||||
if n > 1:
|
||||
print(f"[o1key Grok Image] 第 {i+1}/{n} 张完成")
|
||||
return results
|
||||
|
||||
# ── 带重试的请求 ────────────────────────────────────────────────────────
|
||||
|
||||
async def _do_request_with_retry(self, url: str, body: dict) -> List[Image.Image]:
|
||||
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:
|
||||
last_error = None
|
||||
for attempt in range(1, _MAX_RETRIES + 1):
|
||||
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 == 429 or resp.status in (502, 503, 504):
|
||||
last_error = f"HTTP {resp.status}"
|
||||
print(f"[o1key Grok Image] 重试 {attempt}/{_MAX_RETRIES}({last_error})")
|
||||
await asyncio.sleep(_RETRY_DELAY * attempt)
|
||||
continue
|
||||
|
||||
if resp.status == 400 and "high load" in text.lower():
|
||||
last_error = "high load"
|
||||
print(f"[o1key Grok Image] 重试 {attempt}/{_MAX_RETRIES}(服务繁忙)")
|
||||
await asyncio.sleep(_RETRY_DELAY * attempt)
|
||||
continue
|
||||
|
||||
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 Grok Image] API 响应耗时 {elapsed:.1f}s")
|
||||
return await self._parse_response(resp_json, session)
|
||||
|
||||
raise RuntimeError(f"重试 {_MAX_RETRIES} 次后仍失败: {last_error}")
|
||||
|
||||
return await self._run_with_interrupt(_do_request())
|
||||
|
||||
# ── 同步入口 ──────────────────────────────────────────────────────────────
|
||||
|
||||
def run_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
n: int,
|
||||
image_list: Optional[List[torch.Tensor]] = None,
|
||||
) -> List[Image.Image]:
|
||||
if image_list:
|
||||
coro = self._edit_async(
|
||||
prompt=prompt, model=model, aspect_ratio=aspect_ratio,
|
||||
resolution=resolution, n=n, image_list=image_list,
|
||||
)
|
||||
else:
|
||||
coro = self._generate_async(
|
||||
prompt=prompt, model=model, aspect_ratio=aspect_ratio,
|
||||
resolution=resolution, n=n,
|
||||
)
|
||||
|
||||
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("Grok Image 请求超时,请检查网络或稍后重试")
|
||||
|
||||
# ── 余额查询 ──────────────────────────────────────────────────────────────
|
||||
|
||||
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}"
|
||||
@@ -0,0 +1,483 @@
|
||||
"""
|
||||
Grok Video API client.
|
||||
|
||||
Flow:
|
||||
1. POST /v1/videos
|
||||
2. GET /v1/videos/{task_id}
|
||||
3. GET /v1/videos/{task_id}/content, or download a URL from the status body
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from .base_client import BaseAPIClient
|
||||
from ..utils.config import get_api_base_url, get_api_key_or_raise
|
||||
from ..utils.http_error import RETRYABLE_STATUS_CODES, get_friendly_message
|
||||
from ..utils.video_task import (
|
||||
check_interrupt,
|
||||
extract_error_message,
|
||||
extract_progress,
|
||||
extract_status,
|
||||
extract_video_url,
|
||||
interruptible_sleep,
|
||||
is_failure_status,
|
||||
is_success_status,
|
||||
run_with_interrupt,
|
||||
)
|
||||
|
||||
|
||||
class GrokVideoClient(BaseAPIClient):
|
||||
CREATE_ENDPOINT = "/v1/videos"
|
||||
STATUS_ENDPOINT = "/v1/videos/{task_id}"
|
||||
CONTENT_ENDPOINT = "/v1/videos/{task_id}/content"
|
||||
|
||||
MODEL_OPTIONS = ["grok-imagine-video-1.5-preview", "grok-imagine-1.0-video"]
|
||||
ASPECT_RATIO_OPTIONS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"]
|
||||
QUALITY_OPTIONS = ["720p"]
|
||||
MODEL_SECONDS_OPTIONS = {
|
||||
"grok-imagine-1.0-video": [6, 10, 12, 16, 20],
|
||||
}
|
||||
QUALITY_API_MAP = {
|
||||
"720p": "high",
|
||||
"high": "high",
|
||||
}
|
||||
|
||||
SUCCESS_STATUSES = {"complete", "completed", "succeed", "succeeded", "success", "done", "finished"}
|
||||
FAILURE_STATUSES = {"fail", "failed", "failure", "error", "expired", "timeout", "cancelled", "canceled"}
|
||||
|
||||
def __init__(self, base_url: Optional[str] = None):
|
||||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
resolved_base_url = (base_url or "").strip() or get_api_base_url()
|
||||
super().__init__(base_url=resolved_base_url.rstrip("/"), api_key=api_key)
|
||||
|
||||
def get_endpoint(self, **kwargs) -> str:
|
||||
return self.CREATE_ENDPOINT
|
||||
|
||||
def build_request_body(self, **kwargs) -> Dict[str, Any]:
|
||||
return self.build_video_body(**kwargs)
|
||||
|
||||
def parse_response(self, response: Dict[str, Any]) -> Any:
|
||||
return response
|
||||
|
||||
@classmethod
|
||||
def build_video_body(
|
||||
cls,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
seconds: int,
|
||||
quality: str = "720p",
|
||||
images: Optional[List[str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
prompt = (prompt or "").strip()
|
||||
if not prompt:
|
||||
raise ValueError("提示词不能为空。")
|
||||
|
||||
if model not in cls.MODEL_OPTIONS:
|
||||
raise ValueError(f"模型仅支持: {', '.join(cls.MODEL_OPTIONS)}。")
|
||||
|
||||
if aspect_ratio not in cls.ASPECT_RATIO_OPTIONS:
|
||||
raise ValueError(f"宽高比仅支持: {', '.join(cls.ASPECT_RATIO_OPTIONS)}。")
|
||||
|
||||
try:
|
||||
seconds_value = int(seconds)
|
||||
except (TypeError, ValueError):
|
||||
raise ValueError("秒数必须是整数。") from None
|
||||
|
||||
allowed_seconds = cls.MODEL_SECONDS_OPTIONS.get(model)
|
||||
if allowed_seconds is not None:
|
||||
if seconds_value not in allowed_seconds:
|
||||
raise ValueError(
|
||||
f"模型 {model} 仅支持秒数: "
|
||||
f"{', '.join(str(s) for s in allowed_seconds)}。"
|
||||
"请修改为正确的秒数后再发起请求。"
|
||||
)
|
||||
elif seconds_value < 5 or seconds_value > 15:
|
||||
raise ValueError("秒数仅支持 5 到 15。")
|
||||
|
||||
api_quality = cls.QUALITY_API_MAP.get(str(quality), str(quality))
|
||||
if api_quality != "high":
|
||||
raise ValueError("画质仅支持 720p。")
|
||||
|
||||
body: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"seconds": str(seconds_value),
|
||||
"quality": api_quality,
|
||||
}
|
||||
|
||||
image_list = [img for img in (images or []) if img]
|
||||
if image_list:
|
||||
body["images"] = image_list[:3]
|
||||
|
||||
return body
|
||||
|
||||
@staticmethod
|
||||
def _safe_task_filename(task_id: str) -> str:
|
||||
safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", task_id).strip("._")
|
||||
return safe or "grok_video"
|
||||
|
||||
@staticmethod
|
||||
def _mask_body_for_log(body: Dict[str, Any]) -> Dict[str, Any]:
|
||||
log_body = dict(body)
|
||||
images = log_body.get("images")
|
||||
if isinstance(images, list):
|
||||
log_body["images"] = [f"<data-url chars={len(item)}>" for item in images]
|
||||
return log_body
|
||||
|
||||
@staticmethod
|
||||
def _extract_task_id(payload: Dict[str, Any]) -> Optional[str]:
|
||||
sources = [payload]
|
||||
data = payload.get("data")
|
||||
if isinstance(data, dict):
|
||||
sources.append(data)
|
||||
|
||||
for source in sources:
|
||||
for key in ("id", "task_id", "video_id"):
|
||||
value = source.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _format_http_error(endpoint: str, status: int, error_text: str, task_id: Optional[str] = None) -> str:
|
||||
message = get_friendly_message(status, error_text)
|
||||
parts = [
|
||||
"Grok Video 请求失败。",
|
||||
f"endpoint: {endpoint}",
|
||||
f"http_status: {status}",
|
||||
]
|
||||
if task_id:
|
||||
parts.append(f"task_id: {task_id}")
|
||||
if message:
|
||||
parts.append(f"message: {message}")
|
||||
return "\n".join(parts)
|
||||
|
||||
@classmethod
|
||||
def _format_task_failure(cls, task_id: str, payload: Dict[str, Any]) -> str:
|
||||
return "\n".join(
|
||||
[
|
||||
"Grok Video 任务失败。",
|
||||
f"endpoint: {cls.STATUS_ENDPOINT.format(task_id=task_id)}",
|
||||
f"task_id: {task_id}",
|
||||
f"message: {extract_error_message(payload)}",
|
||||
]
|
||||
)
|
||||
|
||||
async def _request_json_with_retry(
|
||||
self,
|
||||
method: str,
|
||||
endpoint: str,
|
||||
session: aiohttp.ClientSession,
|
||||
task_id: Optional[str] = None,
|
||||
json_body: Optional[Dict[str, Any]] = None,
|
||||
max_retries: int = 3,
|
||||
timeout_seconds: int = 120,
|
||||
) -> Dict[str, Any]:
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
timeout = aiohttp.ClientTimeout(total=timeout_seconds, connect=30, sock_read=timeout_seconds)
|
||||
|
||||
last_status = 0
|
||||
last_text = ""
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
check_interrupt()
|
||||
response = None
|
||||
try:
|
||||
response = await run_with_interrupt(
|
||||
session.request(method, url, json=json_body, headers=headers, timeout=timeout)
|
||||
)
|
||||
text = await run_with_interrupt(response.text())
|
||||
last_status = response.status
|
||||
last_text = text
|
||||
|
||||
if 200 <= response.status < 300:
|
||||
if not text.strip():
|
||||
return {}
|
||||
try:
|
||||
return json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError(f"Grok Video 响应 JSON 解析失败,原始内容:{text[:500]}") from None
|
||||
|
||||
if response.status in RETRYABLE_STATUS_CODES and attempt < max_retries:
|
||||
delay = min(2 ** attempt, 8)
|
||||
print(
|
||||
f"Grok Video:{get_friendly_message(response.status)} "
|
||||
f"{delay}s 后重试 ({attempt + 1}/{max_retries})..."
|
||||
)
|
||||
await interruptible_sleep(delay)
|
||||
continue
|
||||
|
||||
break
|
||||
|
||||
except (aiohttp.ClientError, asyncio.TimeoutError) as e:
|
||||
if attempt < max_retries:
|
||||
delay = min(2 ** attempt, 8)
|
||||
print(f"Grok Video:网络错误,{delay}s 后重试 ({attempt + 1}/{max_retries})...")
|
||||
await interruptible_sleep(delay)
|
||||
continue
|
||||
raise RuntimeError(f"Grok Video 网络错误: {e}") from None
|
||||
|
||||
finally:
|
||||
if response is not None:
|
||||
response.release()
|
||||
|
||||
raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id))
|
||||
|
||||
async def create_video_async(
|
||||
self,
|
||||
body: Dict[str, Any],
|
||||
session: aiohttp.ClientSession,
|
||||
) -> Dict[str, Any]:
|
||||
print("Grok Video:正在提交任务...")
|
||||
return await self._request_json_with_retry(
|
||||
"POST",
|
||||
self.CREATE_ENDPOINT,
|
||||
session=session,
|
||||
json_body=body,
|
||||
timeout_seconds=180,
|
||||
)
|
||||
|
||||
async def poll_video_status_async(
|
||||
self,
|
||||
task_id: str,
|
||||
session: aiohttp.ClientSession,
|
||||
poll_interval: int = 5,
|
||||
timeout: int = 900,
|
||||
progress_callback: Optional[Callable[[int, str, float], None]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
endpoint = self.STATUS_ENDPOINT.format(task_id=task_id)
|
||||
start = time.time()
|
||||
interval = max(1, int(poll_interval))
|
||||
|
||||
await interruptible_sleep(interval)
|
||||
|
||||
while True:
|
||||
data = await self._request_json_with_retry(
|
||||
"GET",
|
||||
endpoint,
|
||||
session=session,
|
||||
task_id=task_id,
|
||||
timeout_seconds=60,
|
||||
)
|
||||
|
||||
status = extract_status(data)
|
||||
progress = extract_progress(data)
|
||||
elapsed = time.time() - start
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(progress, status, elapsed)
|
||||
|
||||
if status in self.SUCCESS_STATUSES or is_success_status(status):
|
||||
return data
|
||||
|
||||
if status in self.FAILURE_STATUSES or is_failure_status(status, data):
|
||||
raise RuntimeError(self._format_task_failure(task_id, data))
|
||||
|
||||
if elapsed >= timeout:
|
||||
raise TimeoutError(
|
||||
"Grok Video 任务轮询超时;任务未被标记为失败,可用 task_id 继续查询。\n"
|
||||
f"endpoint: {endpoint}\n"
|
||||
f"task_id: {task_id}\n"
|
||||
f"status: {status or 'unknown'}\n"
|
||||
f"timeout: {timeout}s"
|
||||
)
|
||||
|
||||
await interruptible_sleep(min(interval, max(0.0, timeout - elapsed)))
|
||||
|
||||
async def _download_url_to_file(
|
||||
self,
|
||||
url: str,
|
||||
save_path: str,
|
||||
session: aiohttp.ClientSession,
|
||||
max_retries: int = 3,
|
||||
) -> str:
|
||||
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
|
||||
last_status = 0
|
||||
last_text = ""
|
||||
headers = None
|
||||
resolved_url = url
|
||||
|
||||
if url.startswith("data:"):
|
||||
if "," not in url:
|
||||
raise RuntimeError("Grok Video 下载失败:data URL 格式无效。")
|
||||
_, b64_data = url.split(",", 1)
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
with open(save_path, "wb") as f:
|
||||
f.write(base64.b64decode(b64_data))
|
||||
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
|
||||
raise RuntimeError("Grok Video 下载失败:保存后的文件为空。")
|
||||
return save_path
|
||||
|
||||
if url.startswith("/"):
|
||||
resolved_url = f"{self.base_url}{url}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
check_interrupt()
|
||||
async with session.get(
|
||||
resolved_url,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
allow_redirects=True,
|
||||
) as response:
|
||||
if 200 <= response.status < 300:
|
||||
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(1024 * 1024):
|
||||
check_interrupt()
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
|
||||
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
|
||||
raise RuntimeError("Grok Video 下载失败:保存后的文件为空。")
|
||||
return save_path
|
||||
|
||||
last_status = response.status
|
||||
last_text = await response.text()
|
||||
if response.status not in RETRYABLE_STATUS_CODES or attempt >= max_retries:
|
||||
break
|
||||
|
||||
delay = min(2 ** attempt, 8)
|
||||
print(f"Grok Video:下载重试 {attempt + 1}/{max_retries},{delay}s 后继续...")
|
||||
await interruptible_sleep(delay)
|
||||
|
||||
raise RuntimeError(self._format_http_error("download_url", last_status, last_text))
|
||||
|
||||
async def download_video_async(
|
||||
self,
|
||||
task_id: str,
|
||||
save_path: str,
|
||||
session: aiohttp.ClientSession,
|
||||
) -> str:
|
||||
endpoint = self.CONTENT_ENDPOINT.format(task_id=task_id)
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
|
||||
|
||||
last_status = 0
|
||||
last_text = ""
|
||||
|
||||
for attempt in range(4):
|
||||
check_interrupt()
|
||||
async with session.get(url, headers=headers, timeout=timeout, allow_redirects=True) as response:
|
||||
if 200 <= response.status < 300:
|
||||
content_type = response.headers.get("Content-Type", "").lower()
|
||||
if "application/json" in content_type:
|
||||
data = await response.json(content_type=None)
|
||||
download_url = extract_video_url(data)
|
||||
if not download_url:
|
||||
raise RuntimeError(
|
||||
"Grok Video 下载失败:content 响应为 JSON,但未包含视频 URL。\n"
|
||||
f"endpoint: {endpoint}\n"
|
||||
f"task_id: {task_id}"
|
||||
)
|
||||
return await self._download_url_to_file(download_url, save_path, session)
|
||||
|
||||
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(1024 * 1024):
|
||||
check_interrupt()
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
|
||||
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
|
||||
raise RuntimeError(
|
||||
"Grok Video 下载失败:保存后的文件为空。\n"
|
||||
f"endpoint: {endpoint}\n"
|
||||
f"task_id: {task_id}"
|
||||
)
|
||||
return save_path
|
||||
|
||||
last_status = response.status
|
||||
last_text = await response.text()
|
||||
if response.status not in RETRYABLE_STATUS_CODES or attempt >= 3:
|
||||
break
|
||||
|
||||
delay = min(2 ** attempt, 8)
|
||||
print(f"Grok Video:content 下载重试 {attempt + 1}/3,{delay}s 后继续...")
|
||||
await interruptible_sleep(delay)
|
||||
|
||||
raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id))
|
||||
|
||||
def generate_video_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
seconds: int,
|
||||
quality: str,
|
||||
images: Optional[List[str]],
|
||||
output_dir: Optional[str] = None,
|
||||
save_path: Optional[str] = None,
|
||||
poll_interval: int = 5,
|
||||
timeout: int = 900,
|
||||
progress_callback: Optional[Callable[[int, str, float], None]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
async def _run():
|
||||
async with self._make_session() as session:
|
||||
body = self.build_video_body(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seconds=seconds,
|
||||
quality=quality,
|
||||
images=images,
|
||||
)
|
||||
|
||||
create_response = await self.create_video_async(body, session)
|
||||
task_id = self._extract_task_id(create_response) or ""
|
||||
if not task_id:
|
||||
raise RuntimeError(
|
||||
"Grok Video 未返回任务 ID。\n"
|
||||
f"endpoint: {self.CREATE_ENDPOINT}\n"
|
||||
f"response: {json.dumps(create_response, ensure_ascii=False)[:1200]}"
|
||||
)
|
||||
|
||||
print(f"Grok Video:任务已提交,任务ID:{task_id}")
|
||||
print("Grok Video:视频生成中...")
|
||||
status_response = await self.poll_video_status_async(
|
||||
task_id=task_id,
|
||||
session=session,
|
||||
poll_interval=poll_interval,
|
||||
timeout=timeout,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
video_url = extract_video_url(status_response)
|
||||
print("Grok Video:视频生成完成,正在下载...")
|
||||
if save_path is None:
|
||||
resolved_output_dir = output_dir or os.getcwd()
|
||||
os.makedirs(resolved_output_dir, exist_ok=True)
|
||||
target_path = os.path.join(
|
||||
resolved_output_dir,
|
||||
f"{self._safe_task_filename(task_id)}.mp4",
|
||||
)
|
||||
else:
|
||||
target_path = save_path
|
||||
|
||||
if video_url:
|
||||
video_path = await self._download_url_to_file(video_url, target_path, session)
|
||||
else:
|
||||
video_path = await self.download_video_async(task_id, target_path, session)
|
||||
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": extract_status(status_response),
|
||||
"video_path": video_path,
|
||||
"raw_json": {
|
||||
"create": create_response,
|
||||
"status": status_response,
|
||||
},
|
||||
}
|
||||
|
||||
return self.run_async_in_thread(_run())
|
||||
@@ -0,0 +1,285 @@
|
||||
"""
|
||||
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
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
from ..utils.video_task import (
|
||||
check_interrupt,
|
||||
extract_error_message,
|
||||
extract_progress,
|
||||
extract_status,
|
||||
interruptible_sleep,
|
||||
is_failure_status,
|
||||
is_success_status,
|
||||
run_with_interrupt,
|
||||
)
|
||||
|
||||
|
||||
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]}"
|
||||
|
||||
check_interrupt()
|
||||
resp = await run_with_interrupt(async_request_with_retry(
|
||||
session, "POST", url, json=body, headers=self._headers(), prefix="Kling 提交: "
|
||||
))
|
||||
check_interrupt()
|
||||
text = await resp.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:
|
||||
check_interrupt()
|
||||
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 = extract_status(result)
|
||||
|
||||
progress_pct = extract_progress(result)
|
||||
|
||||
print(f"[视频生成] 生成中 {progress_pct}%")
|
||||
|
||||
if on_progress:
|
||||
on_progress(progress_pct)
|
||||
|
||||
if is_success_status(status):
|
||||
return result
|
||||
elif is_failure_status(status, result):
|
||||
error_msg = extract_error_message(result)
|
||||
raise RuntimeError(f"生成失败:{error_msg}")
|
||||
|
||||
await interruptible_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("[视频生成] 下载视频...")
|
||||
check_interrupt()
|
||||
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):
|
||||
check_interrupt()
|
||||
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. 提交
|
||||
check_interrupt()
|
||||
if on_stage:
|
||||
on_stage("submitting")
|
||||
create_url = f"{self.base_url}{self.NEW_API_CREATE}"
|
||||
resp = await run_with_interrupt(async_request_with_retry(
|
||||
session, "POST", create_url, json=body, headers=headers, prefix="Kling 动作控制提交: "
|
||||
))
|
||||
check_interrupt()
|
||||
text = await resp.text()
|
||||
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:
|
||||
check_interrupt()
|
||||
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 = extract_status(status_resp)
|
||||
progress_pct = extract_progress(status_resp)
|
||||
|
||||
print(f"[动作控制] 生成中 {progress_pct}%")
|
||||
if on_progress:
|
||||
on_progress(progress_pct)
|
||||
|
||||
if is_success_status(status):
|
||||
break
|
||||
if is_failure_status(status, status_resp):
|
||||
error_msg = extract_error_message(status_resp)
|
||||
raise RuntimeError(f"动作控制生成失败:{error_msg}")
|
||||
|
||||
await interruptible_sleep(interval)
|
||||
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
|
||||
|
||||
# 3. 下载
|
||||
check_interrupt()
|
||||
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):
|
||||
check_interrupt()
|
||||
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):
|
||||
check_interrupt()
|
||||
f.write(chunk)
|
||||
|
||||
if on_stage:
|
||||
on_stage("done")
|
||||
return save_path
|
||||
@@ -0,0 +1,542 @@
|
||||
"""
|
||||
new-api Veo 3.1 video client.
|
||||
|
||||
Implements the OpenAI-compatible /v1/videos task flow:
|
||||
submit, poll, and stream-download video content.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from .base_client import BaseAPIClient
|
||||
from ..utils.config import get_api_base_url, get_api_key_or_raise
|
||||
|
||||
|
||||
class NewAPIVeoClient(BaseAPIClient):
|
||||
CREATE_ENDPOINT = "/v1/videos"
|
||||
STATUS_ENDPOINT = "/v1/videos/{task_id}"
|
||||
CONTENT_ENDPOINT = "/v1/videos/{task_id}/content"
|
||||
|
||||
RETRYABLE_STATUS_CODES = {408, 409, 425, 429, 500, 502, 503, 504}
|
||||
COMPLETED_STATUSES = {"completed", "succeeded", "success", "done"}
|
||||
FAILED_STATUSES = {"failed", "error", "cancelled", "canceled"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: Optional[str] = None,
|
||||
):
|
||||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
resolved_base_url = (base_url or "").strip() or get_api_base_url()
|
||||
super().__init__(base_url=resolved_base_url.rstrip("/"), api_key=api_key)
|
||||
|
||||
def get_endpoint(self, **kwargs) -> str:
|
||||
return self.CREATE_ENDPOINT
|
||||
|
||||
def build_request_body(self, **kwargs) -> Dict[str, Any]:
|
||||
return self._build_video_body(**kwargs)
|
||||
|
||||
def parse_response(self, response: Dict[str, Any]) -> Any:
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def _build_video_body(
|
||||
prompt: str,
|
||||
model: str,
|
||||
duration: int,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
negative_prompt: str = "",
|
||||
generate_audio: bool = True,
|
||||
) -> Dict[str, Any]:
|
||||
metadata: Dict[str, Any] = {
|
||||
"aspectRatio": aspect_ratio,
|
||||
"resolution": resolution,
|
||||
"generateAudio": bool(generate_audio),
|
||||
}
|
||||
|
||||
negative_prompt = (negative_prompt or "").strip()
|
||||
if negative_prompt:
|
||||
metadata["negativePrompt"] = negative_prompt
|
||||
|
||||
body: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"duration": int(duration),
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
return body
|
||||
|
||||
@staticmethod
|
||||
def _print_request_body(body: Dict[str, Any], image_bytes: Optional[bytes] = None) -> None:
|
||||
log_body = dict(body)
|
||||
if image_bytes is not None:
|
||||
log_body["input_reference"] = f"<PNG bytes: {len(image_bytes)}>"
|
||||
print(
|
||||
"NewAPI Veo request body:\n"
|
||||
f"{json.dumps(log_body, ensure_ascii=False, indent=2)}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _safe_task_filename(task_id: str) -> str:
|
||||
safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", task_id).strip("._")
|
||||
return safe or "newapi_veo"
|
||||
|
||||
@staticmethod
|
||||
def _extract_task_id(data: Dict[str, Any]) -> Optional[str]:
|
||||
for key in ("id", "task_id", "video_id"):
|
||||
value = data.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
|
||||
nested = data.get("data")
|
||||
if isinstance(nested, dict):
|
||||
for key in ("id", "task_id", "video_id"):
|
||||
value = nested.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _extract_status(data: Dict[str, Any]) -> str:
|
||||
for key in ("status", "state", "task_status"):
|
||||
value = data.get(key)
|
||||
if value:
|
||||
return str(value).lower()
|
||||
|
||||
nested = data.get("data")
|
||||
if isinstance(nested, dict):
|
||||
for key in ("status", "state", "task_status"):
|
||||
value = nested.get(key)
|
||||
if value:
|
||||
return str(value).lower()
|
||||
return "unknown"
|
||||
|
||||
@staticmethod
|
||||
def _extract_progress(data: Dict[str, Any]) -> int:
|
||||
progress = data.get("progress")
|
||||
if progress is None and isinstance(data.get("data"), dict):
|
||||
progress = data["data"].get("progress")
|
||||
|
||||
if isinstance(progress, str):
|
||||
progress = progress.rstrip("%").strip()
|
||||
try:
|
||||
return int(float(progress))
|
||||
except ValueError:
|
||||
return 0
|
||||
if isinstance(progress, (int, float)):
|
||||
return int(progress)
|
||||
return 0
|
||||
|
||||
@classmethod
|
||||
def _format_http_error(
|
||||
cls,
|
||||
endpoint: str,
|
||||
status: int,
|
||||
error_text: str,
|
||||
task_id: Optional[str] = None,
|
||||
) -> str:
|
||||
code = ""
|
||||
message = error_text
|
||||
try:
|
||||
payload = json.loads(error_text)
|
||||
error = payload.get("error", payload)
|
||||
if isinstance(error, dict):
|
||||
code = str(error.get("code") or error.get("type") or "")
|
||||
message = str(error.get("message") or payload.get("message") or error_text)
|
||||
elif error is not None:
|
||||
message = str(error)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
message = (message or "").strip()
|
||||
if len(message) > 1200:
|
||||
message = message[:1200] + "...(truncated)"
|
||||
|
||||
if status in (401, 403):
|
||||
hint = "凭证或分组权限问题,请检查 new-api token、模型分组或渠道权限。"
|
||||
elif status == 429:
|
||||
hint = "频率或额度限制,请稍后重试或检查 new-api 额度。"
|
||||
elif status in (502, 503, 504):
|
||||
hint = "上游服务暂时不可用或超时,请稍后用 task_id 继续查询。"
|
||||
elif status == 400:
|
||||
hint = "请求参数错误,请检查 model、duration、metadata 和图片输入。"
|
||||
else:
|
||||
hint = "new-api 视频请求失败。"
|
||||
|
||||
parts = [
|
||||
hint,
|
||||
f"endpoint: {endpoint}",
|
||||
f"http_status: {status}",
|
||||
]
|
||||
if task_id:
|
||||
parts.append(f"task_id: {task_id}")
|
||||
if code:
|
||||
parts.append(f"error_code: {code}")
|
||||
if message:
|
||||
parts.append(f"message: {message}")
|
||||
return "\n".join(parts)
|
||||
|
||||
@classmethod
|
||||
def _format_task_failure(cls, task_id: str, data: Dict[str, Any]) -> str:
|
||||
error = data.get("error")
|
||||
if error is None and isinstance(data.get("data"), dict):
|
||||
error = data["data"].get("error")
|
||||
|
||||
if isinstance(error, dict):
|
||||
code = error.get("code") or error.get("type") or ""
|
||||
message = error.get("message") or json.dumps(error, ensure_ascii=False)
|
||||
else:
|
||||
code = ""
|
||||
message = str(error or "未知错误")
|
||||
|
||||
return "\n".join(
|
||||
[
|
||||
"Veo 视频任务失败。",
|
||||
f"endpoint: {cls.STATUS_ENDPOINT.format(task_id=task_id)}",
|
||||
f"task_id: {task_id}",
|
||||
f"error_code: {code}",
|
||||
f"message: {message}",
|
||||
]
|
||||
)
|
||||
|
||||
async def create_video_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
duration: int,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
negative_prompt: str = "",
|
||||
generate_audio: bool = True,
|
||||
image_bytes: Optional[bytes] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> Dict[str, Any]:
|
||||
body = self._build_video_body(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
duration=duration,
|
||||
aspect_ratio=aspect_ratio,
|
||||
resolution=resolution,
|
||||
negative_prompt=negative_prompt,
|
||||
generate_audio=generate_audio,
|
||||
)
|
||||
|
||||
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = self._make_session()
|
||||
close_session = True
|
||||
|
||||
try:
|
||||
timeout = aiohttp.ClientTimeout(total=120, connect=30, sock_read=120)
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
if image_bytes is not None:
|
||||
if len(image_bytes) > self.max_request_size:
|
||||
raise ValueError(
|
||||
f"输入图片过大,超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制"
|
||||
)
|
||||
|
||||
self._print_request_body(body, image_bytes=image_bytes)
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("model", body["model"])
|
||||
form.add_field("prompt", body["prompt"])
|
||||
form.add_field("duration", str(body["duration"]))
|
||||
form.add_field("metadata", json.dumps(body["metadata"], ensure_ascii=False))
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_bytes,
|
||||
filename="input_reference.png",
|
||||
content_type="image/png",
|
||||
)
|
||||
request_kwargs = {"data": form, "headers": headers}
|
||||
print(
|
||||
"NewAPI Veo: POST /v1/videos multipart "
|
||||
f"| model={model} | duration={duration}s | {resolution} {aspect_ratio}"
|
||||
)
|
||||
else:
|
||||
self._print_request_body(body)
|
||||
headers["Content-Type"] = "application/json"
|
||||
request_kwargs = {"json": body, "headers": headers}
|
||||
print(
|
||||
"NewAPI Veo: POST /v1/videos json "
|
||||
f"| model={model} | duration={duration}s | {resolution} {aspect_ratio}"
|
||||
)
|
||||
|
||||
async with session.post(url, timeout=timeout, **request_kwargs) as response:
|
||||
if response.status >= 300:
|
||||
error_text = await response.text()
|
||||
raise RuntimeError(
|
||||
self._format_http_error(
|
||||
self.CREATE_ENDPOINT,
|
||||
response.status,
|
||||
error_text,
|
||||
)
|
||||
)
|
||||
return await response.json()
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
async def _get_json_with_retry(
|
||||
self,
|
||||
endpoint: str,
|
||||
session: aiohttp.ClientSession,
|
||||
task_id: Optional[str] = None,
|
||||
max_retries: int = 3,
|
||||
) -> Dict[str, Any]:
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
timeout = aiohttp.ClientTimeout(total=60, connect=30, sock_read=60)
|
||||
|
||||
last_error = ""
|
||||
last_status = 0
|
||||
for attempt in range(max_retries + 1):
|
||||
async with session.get(url, headers=headers, timeout=timeout) as response:
|
||||
if response.status < 300:
|
||||
return await response.json()
|
||||
|
||||
last_status = response.status
|
||||
last_error = await response.text()
|
||||
if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= max_retries:
|
||||
break
|
||||
|
||||
await asyncio.sleep(min(2 ** attempt, 8))
|
||||
|
||||
raise RuntimeError(
|
||||
self._format_http_error(endpoint, last_status, last_error, task_id=task_id)
|
||||
)
|
||||
|
||||
async def poll_video_status_async(
|
||||
self,
|
||||
task_id: str,
|
||||
poll_interval: int = 5,
|
||||
timeout: int = 900,
|
||||
progress_callback: Optional[Callable[[int, str, float], None]] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> Dict[str, Any]:
|
||||
endpoint = self.STATUS_ENDPOINT.format(task_id=task_id)
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = self._make_session()
|
||||
close_session = True
|
||||
|
||||
start = time.time()
|
||||
try:
|
||||
poll_interval = max(1, int(poll_interval))
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
while True:
|
||||
data = await self._get_json_with_retry(endpoint, session, task_id=task_id)
|
||||
status = self._extract_status(data)
|
||||
elapsed = time.time() - start
|
||||
progress = self._extract_progress(data)
|
||||
|
||||
if status == "unknown":
|
||||
print(
|
||||
"NewAPI Veo status response did not include a recognized status field:\n"
|
||||
f"{json.dumps(data, ensure_ascii=False, indent=2)[:1200]}"
|
||||
)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(progress, status, elapsed)
|
||||
|
||||
if status in self.COMPLETED_STATUSES:
|
||||
return data
|
||||
|
||||
if status in self.FAILED_STATUSES:
|
||||
raise RuntimeError(self._format_task_failure(task_id, data))
|
||||
|
||||
if elapsed >= timeout:
|
||||
raise TimeoutError(
|
||||
"Veo 视频任务轮询超时;任务未被标记为失败,可用 task_id 继续查询。\n"
|
||||
f"endpoint: {endpoint}\n"
|
||||
f"task_id: {task_id}\n"
|
||||
f"status: {status}\n"
|
||||
f"timeout: {timeout}s"
|
||||
)
|
||||
|
||||
remaining = max(0.0, timeout - elapsed)
|
||||
await asyncio.sleep(min(poll_interval, remaining))
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
async def _download_url_to_file(
|
||||
self,
|
||||
url: str,
|
||||
save_path: str,
|
||||
session: aiohttp.ClientSession,
|
||||
max_retries: int = 3,
|
||||
) -> None:
|
||||
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
|
||||
last_status = 0
|
||||
last_error = ""
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
async with session.get(url, timeout=timeout, allow_redirects=True) as response:
|
||||
if response.status < 300:
|
||||
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(1024 * 1024):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
return
|
||||
|
||||
last_status = response.status
|
||||
last_error = await response.text()
|
||||
if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= max_retries:
|
||||
break
|
||||
|
||||
await asyncio.sleep(min(2 ** attempt, 8))
|
||||
|
||||
raise RuntimeError(
|
||||
self._format_http_error("download_url", last_status, last_error)
|
||||
)
|
||||
|
||||
async def download_video_async(
|
||||
self,
|
||||
task_id: str,
|
||||
save_path: str,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> str:
|
||||
endpoint = self.CONTENT_ENDPOINT.format(task_id=task_id)
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = self._make_session()
|
||||
close_session = True
|
||||
|
||||
try:
|
||||
last_status = 0
|
||||
last_error = ""
|
||||
for attempt in range(4):
|
||||
async with session.get(url, headers=headers, timeout=timeout, allow_redirects=True) as response:
|
||||
if response.status < 300:
|
||||
content_type = response.headers.get("Content-Type", "")
|
||||
if "application/json" in content_type.lower():
|
||||
data = await response.json()
|
||||
nested = data.get("data") if isinstance(data.get("data"), dict) else {}
|
||||
download_url = (
|
||||
data.get("url")
|
||||
or data.get("download_url")
|
||||
or nested.get("url")
|
||||
or nested.get("download_url")
|
||||
)
|
||||
if not download_url:
|
||||
raise RuntimeError(
|
||||
"视频下载失败: content 响应为 JSON,但未包含 url/download_url。\n"
|
||||
f"endpoint: {endpoint}\n"
|
||||
f"task_id: {task_id}"
|
||||
)
|
||||
await self._download_url_to_file(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(1024 * 1024):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
|
||||
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
|
||||
raise RuntimeError(
|
||||
"视频下载失败: 保存后的文件为空。\n"
|
||||
f"endpoint: {endpoint}\n"
|
||||
f"task_id: {task_id}"
|
||||
)
|
||||
return save_path
|
||||
|
||||
last_status = response.status
|
||||
last_error = await response.text()
|
||||
if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= 3:
|
||||
break
|
||||
|
||||
await asyncio.sleep(min(2 ** attempt, 8))
|
||||
|
||||
raise RuntimeError(
|
||||
self._format_http_error(endpoint, last_status, last_error, task_id=task_id)
|
||||
)
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
def generate_video_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
duration: int,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
output_dir: str,
|
||||
negative_prompt: str = "",
|
||||
generate_audio: bool = True,
|
||||
image_bytes: Optional[bytes] = None,
|
||||
poll_interval: int = 5,
|
||||
timeout: int = 900,
|
||||
reuse_task_id: str = "",
|
||||
progress_callback: Optional[Callable[[int, str, float], None]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
async def _run():
|
||||
async with self._make_session() as session:
|
||||
create_response: Dict[str, Any] = {}
|
||||
task_id = (reuse_task_id or "").strip()
|
||||
if task_id:
|
||||
print(f"NewAPI Veo: reuse task_id={task_id}")
|
||||
else:
|
||||
create_response = await self.create_video_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
duration=duration,
|
||||
aspect_ratio=aspect_ratio,
|
||||
resolution=resolution,
|
||||
negative_prompt=negative_prompt,
|
||||
generate_audio=generate_audio,
|
||||
image_bytes=image_bytes,
|
||||
session=session,
|
||||
)
|
||||
task_id = self._extract_task_id(create_response) or ""
|
||||
if not task_id:
|
||||
raise RuntimeError(
|
||||
"new-api 未返回视频任务 ID。\n"
|
||||
f"endpoint: {self.CREATE_ENDPOINT}\n"
|
||||
f"response: {json.dumps(create_response, ensure_ascii=False)[:1200]}"
|
||||
)
|
||||
|
||||
status_response = await self.poll_video_status_async(
|
||||
task_id=task_id,
|
||||
poll_interval=poll_interval,
|
||||
timeout=timeout,
|
||||
progress_callback=progress_callback,
|
||||
session=session,
|
||||
)
|
||||
status = self._extract_status(status_response)
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
filename = f"{self._safe_task_filename(task_id)}.mp4"
|
||||
save_path = os.path.join(output_dir, filename)
|
||||
video_path = await self.download_video_async(
|
||||
task_id=task_id,
|
||||
save_path=save_path,
|
||||
session=session,
|
||||
)
|
||||
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": status,
|
||||
"video_path": video_path,
|
||||
"raw_json": {
|
||||
"create": create_response,
|
||||
"status": status_response,
|
||||
},
|
||||
}
|
||||
|
||||
return self.run_async_in_thread(_run())
|
||||
@@ -0,0 +1,784 @@
|
||||
"""
|
||||
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": {
|
||||
"image_size": api_image_size
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
request_body["extra_body"]["google"]["image_config"]["aspect_ratio"] = aspect_ratio
|
||||
|
||||
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 图片链接:
|
||||
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)
|
||||
@@ -0,0 +1,195 @@
|
||||
"""
|
||||
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
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
from ..utils.video_task import (
|
||||
check_interrupt,
|
||||
extract_error_message,
|
||||
extract_progress,
|
||||
extract_status,
|
||||
interruptible_sleep,
|
||||
is_failure_status,
|
||||
is_success_status,
|
||||
run_with_interrupt,
|
||||
)
|
||||
|
||||
|
||||
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}"
|
||||
check_interrupt()
|
||||
resp = await run_with_interrupt(async_request_with_retry(
|
||||
session, "POST", url, json=body, headers=self._headers(), prefix="Seedance 提交: "
|
||||
))
|
||||
check_interrupt()
|
||||
text = await resp.text()
|
||||
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:
|
||||
check_interrupt()
|
||||
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 = extract_status(result)
|
||||
|
||||
# 解析进度
|
||||
progress_pct = extract_progress(result)
|
||||
|
||||
print(f"[Seedance] 生成中 {progress_pct}%")
|
||||
if on_progress:
|
||||
on_progress(progress_pct)
|
||||
|
||||
if is_success_status(status):
|
||||
# 响应结构: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 is_failure_status(status, result):
|
||||
reason = extract_error_message(result)
|
||||
raise RuntimeError(f"视频生成失败:{reason}")
|
||||
|
||||
await interruptible_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] 下载视频...")
|
||||
check_interrupt()
|
||||
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):
|
||||
check_interrupt()
|
||||
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:
|
||||
|
||||
# 提交
|
||||
check_interrupt()
|
||||
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
|
||||
@@ -0,0 +1,530 @@
|
||||
"""
|
||||
Sora 视频生成 API 客户端
|
||||
提供视频创建、状态轮询、视频下载功能
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from .base_client import BaseAPIClient
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
from ..utils.image_utils import encode_image_to_base64
|
||||
|
||||
|
||||
def _translate_error_message(msg: str) -> str:
|
||||
"""将 API 返回的已知英文错误信息翻译为中文友好提示"""
|
||||
if "people-in-user-uploads" in msg or (
|
||||
"moderation" in msg and "inputs" in msg
|
||||
):
|
||||
return "上传的参考图片中包含了真实人物【官方风控】,请尝试使用其他办法绕开。"
|
||||
return msg
|
||||
|
||||
|
||||
class SoraClient(BaseAPIClient):
|
||||
"""
|
||||
Sora 视频生成客户端
|
||||
|
||||
工作流程:
|
||||
1. create_video → POST /v1/videos (提交生成任务)
|
||||
2. poll_status → GET /v1/videos/{id} (轮询直到完成/失败)
|
||||
3. download_video→ GET /v1/videos/{id}/content (下载视频文件)
|
||||
"""
|
||||
|
||||
CREATE_ENDPOINT = "/v1/videos"
|
||||
STATUS_ENDPOINT = "/v1/videos/{video_id}"
|
||||
CONTENT_ENDPOINT = "/v1/videos/{video_id}/content"
|
||||
|
||||
POLL_INITIAL_INTERVAL = 3
|
||||
POLL_MAX_INTERVAL = 15
|
||||
|
||||
def __init__(self):
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_api_base_url()
|
||||
super().__init__(base_url=base_url, api_key=api_key)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# BaseAPIClient 抽象方法实现(本客户端主要使用自定义方法)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_endpoint(self, **kwargs) -> str:
|
||||
return self.CREATE_ENDPOINT
|
||||
|
||||
def build_request_body(self, **kwargs) -> Dict[str, Any]:
|
||||
return {}
|
||||
|
||||
def parse_response(self, response: Dict[str, Any]) -> Any:
|
||||
return response
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 核心异步方法
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def create_video_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
seconds: int = 4,
|
||||
size: str = "720x1280",
|
||||
input_reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
提交视频生成任务
|
||||
|
||||
格式策略(根据抓包确认):
|
||||
- 无参考图片:application/json
|
||||
- 有参考图片:multipart/form-data,input_reference 以 PNG 文件上传
|
||||
|
||||
注意:seed 不被上游 API 接受,仅在 ComfyUI 节点侧用于缓存刷新
|
||||
|
||||
Returns:
|
||||
API 响应 JSON,包含 video id 和初始状态
|
||||
"""
|
||||
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
# ============================================================
|
||||
# ⚠️ 已验证可用的标准请求方案,请勿随意修改!(2026-02-28)
|
||||
# ============================================================
|
||||
# 经多轮调试确认:
|
||||
# - 有图片:必须使用 multipart/form-data,input_reference 以 PNG 文件上传
|
||||
# · filename="reference.png", content_type="image/png"(与抓包一致)
|
||||
# · 不可改为 application/json + base64 → 400 "expected a file, got a string"
|
||||
# · 不可改为 application/json + data URI → 500 upstream error
|
||||
# · 不可改为 multipart + image/jpeg → 400 "Inpaint image must match..."(尺寸校验失败)
|
||||
# - 无图片:使用 application/json,已验证成功
|
||||
# ============================================================
|
||||
if input_reference_bytes:
|
||||
if len(input_reference_bytes) > self.max_request_size:
|
||||
raise ValueError(
|
||||
f"参考图片约 {len(input_reference_bytes) / 1024 / 1024:.1f}MB,"
|
||||
f"超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制,请使用较小的图片"
|
||||
)
|
||||
# ⚠️ 有图片:multipart/form-data + PNG 文件上传(唯一验证成功的方案)
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("prompt", prompt)
|
||||
form.add_field("model", model)
|
||||
form.add_field("seconds", str(seconds))
|
||||
form.add_field("size", size)
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
input_reference_bytes,
|
||||
filename="reference.png", # ⚠️ 不可改文件名/扩展名
|
||||
content_type="image/png", # ⚠️ 不可改为 image/jpeg
|
||||
)
|
||||
send_kwargs: Dict[str, Any] = {"data": form, "headers": headers}
|
||||
else:
|
||||
# ⚠️ 无图片:application/json(已验证成功)
|
||||
body: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"seconds": str(seconds),
|
||||
"size": size,
|
||||
}
|
||||
send_kwargs = {"json": body, "headers": headers}
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = 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}"
|
||||
@@ -0,0 +1,510 @@
|
||||
"""
|
||||
Veo 视频生成 API 客户端
|
||||
提供视频创建、状态轮询、视频下载功能
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from .base_client import BaseAPIClient
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
from ..utils.image_utils import encode_image_to_base64
|
||||
|
||||
|
||||
class VeoClient(BaseAPIClient):
|
||||
"""
|
||||
Veo 视频生成客户端
|
||||
|
||||
工作流程:
|
||||
1. create_video → POST /v1/videos (提交生成任务)
|
||||
2. poll_status → GET /v1/videos/{id} (轮询直到完成/失败)
|
||||
3. download_video→ GET /v1/videos/{id}/content (下载视频文件)
|
||||
"""
|
||||
|
||||
CREATE_ENDPOINT = "/v1/videos"
|
||||
STATUS_ENDPOINT = "/v1/videos/{video_id}"
|
||||
CONTENT_ENDPOINT = "/v1/videos/{video_id}/content"
|
||||
|
||||
POLL_INITIAL_INTERVAL = 3
|
||||
POLL_MAX_INTERVAL = 15
|
||||
|
||||
def __init__(self):
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_api_base_url()
|
||||
super().__init__(base_url=base_url, api_key=api_key)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# BaseAPIClient 抽象方法实现
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_endpoint(self, **kwargs) -> str:
|
||||
return self.CREATE_ENDPOINT
|
||||
|
||||
def build_request_body(self, **kwargs) -> Dict[str, Any]:
|
||||
return {}
|
||||
|
||||
def parse_response(self, response: Dict[str, Any]) -> Any:
|
||||
return response
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 核心异步方法
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def create_video_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
seconds: int = 8,
|
||||
size: str = "720x1280",
|
||||
first_frame_bytes: Optional[bytes] = None,
|
||||
last_frame_bytes: Optional[bytes] = None,
|
||||
reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
提交视频生成任务
|
||||
|
||||
格式策略:
|
||||
- 无参考图片:application/json
|
||||
- 有参考图片:multipart/form-data,图片以 PNG 文件上传
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
model: 模型名称
|
||||
seconds: 视频时长(秒)
|
||||
size: 分辨率
|
||||
first_frame_bytes: 首帧图片字节
|
||||
last_frame_bytes: 尾帧图片字节
|
||||
reference_bytes: 参考图片字节
|
||||
seed: 随机种子
|
||||
session: aiohttp 会话
|
||||
|
||||
Returns:
|
||||
API 响应 JSON,包含 video id 和初始状态
|
||||
"""
|
||||
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
# 检查是否有图片
|
||||
has_images = any([first_frame_bytes, last_frame_bytes, reference_bytes])
|
||||
|
||||
if has_images:
|
||||
# 有图片:multipart/form-data + PNG 文件上传
|
||||
if first_frame_bytes and len(first_frame_bytes) > self.max_request_size:
|
||||
raise ValueError(f"首帧图片过大,超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制")
|
||||
if last_frame_bytes and len(last_frame_bytes) > self.max_request_size:
|
||||
raise ValueError(f"尾帧图片过大,超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制")
|
||||
if reference_bytes and len(reference_bytes) > self.max_request_size:
|
||||
raise ValueError(f"参考图片过大,超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制")
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("prompt", prompt)
|
||||
form.add_field("model", model)
|
||||
form.add_field("seconds", str(seconds))
|
||||
form.add_field("size", size)
|
||||
# 注意:seed 不被上游 API 接受,仅在 ComfyUI 节点侧用于缓存刷新
|
||||
# if seed is not None:
|
||||
# form.add_field("seed", str(seed))
|
||||
|
||||
# 使用 input_reference 字段(OpenAI兼容格式)
|
||||
# 尝试支持多张图片:按顺序添加多个 input_reference 字段
|
||||
if first_frame_bytes:
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
first_frame_bytes,
|
||||
filename="first_frame.png",
|
||||
content_type="image/png",
|
||||
)
|
||||
if last_frame_bytes:
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
last_frame_bytes,
|
||||
filename="last_frame.png",
|
||||
content_type="image/png",
|
||||
)
|
||||
if reference_bytes:
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
reference_bytes,
|
||||
filename="reference.png",
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
send_kwargs: Dict[str, Any] = {"data": form, "headers": headers}
|
||||
else:
|
||||
# 无图片:application/json
|
||||
body: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"seconds": str(seconds),
|
||||
"size": size,
|
||||
}
|
||||
# 注意:seed 不被上游 API 接受,仅在 ComfyUI 节点侧用于缓存刷新
|
||||
# if seed is not None:
|
||||
# body["seed"] = str(seed)
|
||||
send_kwargs = {"json": body, "headers": headers}
|
||||
|
||||
# 打印请求调试信息
|
||||
import json
|
||||
if has_images:
|
||||
print(f"Veo: 使用 multipart/form-data 格式上传图片")
|
||||
else:
|
||||
print(f"Veo API 请求体: {json.dumps(body, ensure_ascii=False)}")
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = 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}"
|
||||
@@ -0,0 +1,555 @@
|
||||
<!doctype html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="utf-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1" />
|
||||
<title>O1Key 笔记侧栏草图</title>
|
||||
<style>
|
||||
:root {
|
||||
--bg: #151515;
|
||||
--rail: #1b1b1b;
|
||||
--panel: #202020;
|
||||
--panel-2: #262626;
|
||||
--field: #181818;
|
||||
--line: rgba(255,255,255,.09);
|
||||
--line-2: rgba(255,255,255,.16);
|
||||
--text: #e6e6e6;
|
||||
--soft: #b5b5b5;
|
||||
--muted: #777;
|
||||
--blue: #4f8cff;
|
||||
--blue-2: #7eb8f7;
|
||||
--blue-soft: rgba(79,140,255,.14);
|
||||
--danger: #d76565;
|
||||
}
|
||||
|
||||
* { box-sizing: border-box; }
|
||||
body {
|
||||
margin: 0;
|
||||
min-height: 100vh;
|
||||
background: #101010;
|
||||
color: var(--text);
|
||||
font-family: "Microsoft YaHei", "Segoe UI", sans-serif;
|
||||
display: grid;
|
||||
place-items: center;
|
||||
}
|
||||
|
||||
.frame {
|
||||
width: 1180px;
|
||||
height: 740px;
|
||||
background: var(--bg);
|
||||
border: 1px solid #303030;
|
||||
display: grid;
|
||||
grid-template-columns: 52px 352px 1fr;
|
||||
overflow: hidden;
|
||||
box-shadow: 0 18px 60px rgba(0,0,0,.45);
|
||||
position: relative;
|
||||
}
|
||||
|
||||
.rail {
|
||||
background: var(--rail);
|
||||
border-right: 1px solid var(--line);
|
||||
padding: 10px 7px;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 8px;
|
||||
align-items: center;
|
||||
}
|
||||
|
||||
.rail button {
|
||||
width: 36px;
|
||||
height: 36px;
|
||||
border: 1px solid transparent;
|
||||
border-radius: 8px;
|
||||
background: transparent;
|
||||
color: #8a8a8a;
|
||||
display: grid;
|
||||
place-items: center;
|
||||
cursor: default;
|
||||
font-size: 15px;
|
||||
}
|
||||
|
||||
.rail button.active {
|
||||
color: var(--blue-2);
|
||||
background: rgba(79,140,255,.13);
|
||||
border-color: rgba(79,140,255,.36);
|
||||
}
|
||||
|
||||
.note-panel {
|
||||
background: var(--panel);
|
||||
border-right: 1px solid var(--line);
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
min-width: 0;
|
||||
min-height: 0;
|
||||
}
|
||||
|
||||
.header {
|
||||
padding: 14px 14px 10px;
|
||||
border-bottom: 1px solid rgba(255,255,255,.06);
|
||||
flex: 0 0 auto;
|
||||
}
|
||||
|
||||
.title-row,
|
||||
.actions,
|
||||
.tag-strip,
|
||||
.note-top,
|
||||
.note-meta,
|
||||
.edit-top,
|
||||
.edit-actions,
|
||||
.tag-row {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
}
|
||||
|
||||
.title-row {
|
||||
justify-content: space-between;
|
||||
margin-bottom: 12px;
|
||||
}
|
||||
|
||||
.title {
|
||||
font-size: 14px;
|
||||
font-weight: 700;
|
||||
letter-spacing: .2px;
|
||||
}
|
||||
|
||||
.actions { gap: 5px; }
|
||||
|
||||
.icon-btn {
|
||||
width: 28px;
|
||||
height: 28px;
|
||||
border: 1px solid var(--line);
|
||||
border-radius: 7px;
|
||||
color: #9b9b9b;
|
||||
background: rgba(255,255,255,.035);
|
||||
display: grid;
|
||||
place-items: center;
|
||||
}
|
||||
|
||||
.icon-btn.primary {
|
||||
color: white;
|
||||
background: var(--blue);
|
||||
border-color: transparent;
|
||||
}
|
||||
|
||||
.search {
|
||||
height: 34px;
|
||||
border: 1px solid var(--line);
|
||||
border-radius: 8px;
|
||||
background: rgba(255,255,255,.045);
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
padding: 0 10px;
|
||||
color: var(--muted);
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.tag-strip {
|
||||
gap: 6px;
|
||||
margin-top: 10px;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.tag-filter {
|
||||
height: 26px;
|
||||
padding: 0 9px;
|
||||
border-radius: 7px;
|
||||
border: 1px solid var(--line);
|
||||
background: transparent;
|
||||
color: #969696;
|
||||
font-size: 12px;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.tag-filter.active {
|
||||
color: var(--blue-2);
|
||||
border-color: rgba(79,140,255,.36);
|
||||
background: var(--blue-soft);
|
||||
}
|
||||
|
||||
.list {
|
||||
flex: 1;
|
||||
min-height: 0;
|
||||
overflow: hidden;
|
||||
padding: 8px 10px 12px;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 6px;
|
||||
}
|
||||
|
||||
.scroll-hint {
|
||||
height: 24px;
|
||||
color: #606060;
|
||||
font-size: 11px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
flex: 0 0 auto;
|
||||
border-top: 1px solid rgba(255,255,255,.04);
|
||||
margin: 2px 4px 0;
|
||||
}
|
||||
|
||||
.note {
|
||||
border: 1px solid transparent;
|
||||
border-radius: 8px;
|
||||
padding: 10px;
|
||||
background: transparent;
|
||||
}
|
||||
|
||||
.note.active {
|
||||
background: rgba(255,255,255,.055);
|
||||
border-color: rgba(255,255,255,.11);
|
||||
box-shadow: inset 2px 0 0 var(--blue);
|
||||
}
|
||||
|
||||
.note-title {
|
||||
font-size: 13px;
|
||||
color: #e0e0e0;
|
||||
white-space: nowrap;
|
||||
overflow: hidden;
|
||||
text-overflow: ellipsis;
|
||||
flex: 1;
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.note-action {
|
||||
color: #777;
|
||||
font-size: 11px;
|
||||
margin-left: 8px;
|
||||
}
|
||||
|
||||
.note-text {
|
||||
margin-top: 6px;
|
||||
font-size: 12px;
|
||||
line-height: 1.45;
|
||||
color: #8c8c8c;
|
||||
display: -webkit-box;
|
||||
-webkit-line-clamp: 3;
|
||||
-webkit-box-orient: vertical;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.note-meta {
|
||||
margin-top: 8px;
|
||||
justify-content: space-between;
|
||||
color: #666;
|
||||
font-size: 10px;
|
||||
gap: 8px;
|
||||
}
|
||||
|
||||
.tags {
|
||||
display: flex;
|
||||
gap: 4px;
|
||||
min-width: 0;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.tag {
|
||||
color: var(--blue-2);
|
||||
background: rgba(79,140,255,.1);
|
||||
border: 1px solid rgba(79,140,255,.2);
|
||||
border-radius: 5px;
|
||||
padding: 2px 5px;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.canvas {
|
||||
position: relative;
|
||||
background:
|
||||
linear-gradient(rgba(255,255,255,.025) 1px, transparent 1px),
|
||||
linear-gradient(90deg, rgba(255,255,255,.025) 1px, transparent 1px);
|
||||
background-size: 26px 26px;
|
||||
min-width: 0;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.canvas-label {
|
||||
position: absolute;
|
||||
left: 24px;
|
||||
top: 22px;
|
||||
color: #5f5f5f;
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.node {
|
||||
position: absolute;
|
||||
width: 186px;
|
||||
height: 92px;
|
||||
border-radius: 8px;
|
||||
border: 1px solid rgba(255,255,255,.12);
|
||||
background: #222;
|
||||
box-shadow: 0 12px 30px rgba(0,0,0,.2);
|
||||
}
|
||||
|
||||
.node.one { left: 160px; top: 160px; }
|
||||
.node.two { left: 430px; top: 290px; }
|
||||
.node::before {
|
||||
content: "";
|
||||
display: block;
|
||||
height: 28px;
|
||||
border-bottom: 1px solid rgba(255,255,255,.08);
|
||||
background: rgba(79,140,255,.12);
|
||||
border-radius: 8px 8px 0 0;
|
||||
}
|
||||
|
||||
.edit-popover {
|
||||
position: absolute;
|
||||
right: 24px;
|
||||
top: 76px;
|
||||
width: 420px;
|
||||
height: 588px;
|
||||
border: 1px solid rgba(79,140,255,.28);
|
||||
border-radius: 10px;
|
||||
background: #202020;
|
||||
box-shadow: 0 22px 70px rgba(0,0,0,.46);
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.edit-top {
|
||||
height: 48px;
|
||||
padding: 0 14px;
|
||||
border-bottom: 1px solid rgba(255,255,255,.08);
|
||||
justify-content: space-between;
|
||||
flex: 0 0 auto;
|
||||
}
|
||||
|
||||
.edit-title {
|
||||
font-size: 13px;
|
||||
font-weight: 700;
|
||||
}
|
||||
|
||||
.edit-body {
|
||||
padding: 12px;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 9px;
|
||||
min-height: 0;
|
||||
flex: 1;
|
||||
}
|
||||
|
||||
.input,
|
||||
.textarea {
|
||||
border: 1px solid var(--line);
|
||||
border-radius: 8px;
|
||||
background: rgba(255,255,255,.045);
|
||||
color: #ddd;
|
||||
font-family: inherit;
|
||||
}
|
||||
|
||||
.input {
|
||||
height: 34px;
|
||||
padding: 0 10px;
|
||||
font-size: 13px;
|
||||
font-weight: 700;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
}
|
||||
|
||||
.textarea {
|
||||
flex: 1;
|
||||
min-height: 220px;
|
||||
padding: 12px;
|
||||
font-size: 13px;
|
||||
line-height: 1.6;
|
||||
white-space: pre-wrap;
|
||||
font-family: "JetBrains Mono", "Consolas", monospace;
|
||||
}
|
||||
|
||||
.tag-editor {
|
||||
min-height: 74px;
|
||||
border: 1px solid var(--line);
|
||||
border-radius: 8px;
|
||||
background: var(--field);
|
||||
padding: 8px;
|
||||
}
|
||||
|
||||
.tag-row {
|
||||
gap: 6px;
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
|
||||
.tag-token {
|
||||
height: 24px;
|
||||
display: inline-flex;
|
||||
align-items: center;
|
||||
gap: 6px;
|
||||
border-radius: 6px;
|
||||
padding: 0 7px;
|
||||
border: 1px solid rgba(79,140,255,.25);
|
||||
color: var(--blue-2);
|
||||
background: rgba(79,140,255,.1);
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.tag-token .x {
|
||||
color: #8baee8;
|
||||
font-size: 13px;
|
||||
}
|
||||
|
||||
.tag-input {
|
||||
height: 24px;
|
||||
min-width: 106px;
|
||||
padding: 0 6px;
|
||||
border: 1px dashed rgba(255,255,255,.14);
|
||||
border-radius: 6px;
|
||||
color: #8f8f8f;
|
||||
display: inline-flex;
|
||||
align-items: center;
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.hint {
|
||||
margin-top: 7px;
|
||||
color: #666;
|
||||
font-size: 11px;
|
||||
}
|
||||
|
||||
.edit-actions {
|
||||
gap: 8px;
|
||||
padding: 12px;
|
||||
border-top: 1px solid rgba(255,255,255,.08);
|
||||
flex: 0 0 auto;
|
||||
}
|
||||
|
||||
.btn {
|
||||
height: 32px;
|
||||
border-radius: 7px;
|
||||
border: 1px solid var(--line);
|
||||
background: rgba(255,255,255,.04);
|
||||
color: #aaa;
|
||||
padding: 0 12px;
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.btn.primary {
|
||||
background: var(--blue);
|
||||
border-color: transparent;
|
||||
color: white;
|
||||
font-weight: 700;
|
||||
flex: 1;
|
||||
}
|
||||
|
||||
.btn.cancel {
|
||||
color: #ccc;
|
||||
}
|
||||
|
||||
.btn.danger {
|
||||
color: #ee9f9f;
|
||||
border-color: rgba(215,101,101,.24);
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="frame">
|
||||
<nav class="rail">
|
||||
<button>☰</button>
|
||||
<button>💬</button>
|
||||
<button class="active">✎</button>
|
||||
<button>🖼</button>
|
||||
<button>⚙</button>
|
||||
</nav>
|
||||
|
||||
<aside class="note-panel">
|
||||
<div class="header">
|
||||
<div class="title-row">
|
||||
<div class="title">笔记</div>
|
||||
<div class="actions">
|
||||
<button class="icon-btn">⇩</button>
|
||||
<button class="icon-btn primary">+</button>
|
||||
</div>
|
||||
</div>
|
||||
<div class="search">⌕ 搜索笔记内容或标签</div>
|
||||
<div class="tag-strip">
|
||||
<div class="tag-filter active">全部 18</div>
|
||||
<div class="tag-filter">#产品图</div>
|
||||
<div class="tag-filter">#Nano Banana</div>
|
||||
<div class="tag-filter">#负向词</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="list">
|
||||
<div class="note active">
|
||||
<div class="note-top">
|
||||
<div class="note-title">产品主图:高级玻璃质感</div>
|
||||
<div class="note-action">编辑</div>
|
||||
</div>
|
||||
<div class="note-text">Clean studio lighting, translucent glass material, subtle caustics, soft shadow, premium product photography...</div>
|
||||
<div class="note-meta"><div class="tags"><span class="tag">#产品图</span><span class="tag">#玻璃</span></div><span>今天 14:22</span></div>
|
||||
</div>
|
||||
<div class="note">
|
||||
<div class="note-top">
|
||||
<div class="note-title">Nano Banana 参考图经验</div>
|
||||
<div class="note-action">编辑</div>
|
||||
</div>
|
||||
<div class="note-text">参考图越多越容易跑偏,主体一致性优先用 1-3 张图;复杂场景建议分两步...</div>
|
||||
<div class="note-meta"><div class="tags"><span class="tag">#Nano Banana</span><span class="tag">#参考图</span></div><span>昨天</span></div>
|
||||
</div>
|
||||
<div class="note">
|
||||
<div class="note-top">
|
||||
<div class="note-title">电商模特换装模板</div>
|
||||
<div class="note-action">编辑</div>
|
||||
</div>
|
||||
<div class="note-text">Keep face identity, preserve pose, replace outfit with [服装描述], realistic fabric texture...</div>
|
||||
<div class="note-meta"><div class="tags"><span class="tag">#电商</span><span class="tag">#换装</span></div><span>05/25</span></div>
|
||||
</div>
|
||||
<div class="note">
|
||||
<div class="note-top">
|
||||
<div class="note-title">常用负向词</div>
|
||||
<div class="note-action">编辑</div>
|
||||
</div>
|
||||
<div class="note-text">low quality, blurry, deformed hands, extra fingers, bad anatomy, distorted text, watermark...</div>
|
||||
<div class="note-meta"><div class="tags"><span class="tag">#负向词</span><span class="tag">#通用</span></div><span>05/20</span></div>
|
||||
</div>
|
||||
<div class="note">
|
||||
<div class="note-top">
|
||||
<div class="note-title">批量任务命名经验</div>
|
||||
<div class="note-action">编辑</div>
|
||||
</div>
|
||||
<div class="note-text">小批量先跑 2-3 张确认风格,固定提示词和参考图后再放大批量数量...</div>
|
||||
<div class="note-meta"><div class="tags"><span class="tag">#批量</span><span class="tag">#工作流</span></div><span>05/18</span></div>
|
||||
</div>
|
||||
<div class="scroll-hint">列表默认可滚动,编辑框不常驻</div>
|
||||
</div>
|
||||
</aside>
|
||||
|
||||
<main class="canvas">
|
||||
<div class="canvas-label">ComfyUI 画布区域,笔记列表不再占用右侧空间</div>
|
||||
<div class="node one"></div>
|
||||
<div class="node two"></div>
|
||||
</main>
|
||||
|
||||
<section class="edit-popover">
|
||||
<div class="edit-top">
|
||||
<div class="edit-title">编辑笔记</div>
|
||||
<button class="icon-btn">×</button>
|
||||
</div>
|
||||
<div class="edit-body">
|
||||
<div class="input">产品主图:高级玻璃质感</div>
|
||||
<div class="tag-editor">
|
||||
<div class="tag-row">
|
||||
<span class="tag-token">产品图 <span class="x">×</span></span>
|
||||
<span class="tag-token">玻璃 <span class="x">×</span></span>
|
||||
<span class="tag-token">灯光 <span class="x">×</span></span>
|
||||
<span class="tag-input">+ 添加标签</span>
|
||||
</div>
|
||||
<div class="hint">只保留标签作为组织方式;可搜索、筛选、删除。</div>
|
||||
</div>
|
||||
<div class="textarea">Clean studio lighting, translucent glass material, subtle caustics, soft shadow, premium product photography, 85mm lens, minimal background, high detail.
|
||||
|
||||
使用方式:
|
||||
1. 把产品图作为参考图输入
|
||||
2. 保留主体轮廓,只调整材质和灯光
|
||||
3. 如果玻璃过亮,降低 “caustics” 权重</div>
|
||||
</div>
|
||||
<div class="edit-actions">
|
||||
<button class="btn cancel">取消</button>
|
||||
<button class="btn">复制</button>
|
||||
<button class="btn danger">删除</button>
|
||||
<button class="btn primary">保存</button>
|
||||
</div>
|
||||
</section>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
+520
-64
@@ -8,7 +8,7 @@
|
||||
3. 重新启用模型: 将模型的 enabled 字段改回 True
|
||||
|
||||
模型类型:
|
||||
- GEMINI_MODELS: Nano Banana Pro 图像生成模型
|
||||
- GEMINI_MODELS: Nano Banana 图像生成模型
|
||||
- GEMINI_FLASH_MODELS: Google Gemini Flash 文本生成模型
|
||||
|
||||
示例:
|
||||
@@ -17,14 +17,21 @@
|
||||
"id": "gemini-新模型名称",
|
||||
"description": "模型说明和特点",
|
||||
"enabled": True,
|
||||
"endpoint_type": "standard" # 端点类型: "dynamic", "standard", "flatfee"
|
||||
"endpoint_type": "standard",
|
||||
"endpoint": "/v1beta/models/gemini-新模型名称:generateContent",
|
||||
"thinking_config": {
|
||||
"不思考": None,
|
||||
"低": "low",
|
||||
"中": None,
|
||||
"高": "high"
|
||||
}
|
||||
}
|
||||
|
||||
临时关闭模型:
|
||||
将对应模型的 "enabled": True 改为 "enabled": False
|
||||
"""
|
||||
|
||||
from typing import List, Dict, Optional
|
||||
from typing import List, Dict, Optional, Tuple
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -32,45 +39,72 @@ from typing import List, Dict, Optional
|
||||
# ============================================================
|
||||
|
||||
# ============================================================
|
||||
# Nano Banana Pro 图像生成模型
|
||||
# Nano Banana 图像生成模型
|
||||
# ============================================================
|
||||
|
||||
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 # 动态端点,由代码根据分辨率选择
|
||||
"endpoint": None, # 动态端点,由代码根据分辨率选择
|
||||
"supported_aspect_ratios": [
|
||||
"1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"
|
||||
],
|
||||
"supported_resolutions": ["1K", "2K", "4K"]
|
||||
},
|
||||
{
|
||||
"id": "gemini-3-pro-image-preview-url",
|
||||
"description": "URL 模式,根据分辨率自动选择端点 (1K/2K/4K),推荐用于需要不同分辨率的场景",
|
||||
"enabled": False,
|
||||
"endpoint_type": "dynamic",
|
||||
"endpoint": None # 动态端点,由代码根据分辨率选择
|
||||
},
|
||||
{
|
||||
"id": "gemini-3-pro-image-preview",
|
||||
"description": "标准模式,固定端点,适用于常规图像生成",
|
||||
"id": "nano-banana-pro-官方计费",
|
||||
"description": "Nano Banana Pro 官方计费,按分辨率路由 (1K/2K/4K),使用官方计费通道",
|
||||
"enabled": True,
|
||||
"endpoint_type": "standard",
|
||||
"endpoint": "/v1beta/models/gemini-3-pro-image-preview:generateContent"
|
||||
"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"]
|
||||
},
|
||||
{
|
||||
"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 次卡,根据分辨率自动选择端点 (512px/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-2",
|
||||
"description": "Nano Banana 2 模型,固定端点,适用于高质量图像生成",
|
||||
"enabled": False,
|
||||
"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-2:generateContent"
|
||||
}
|
||||
"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"]
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@@ -80,14 +114,29 @@ GEMINI_MODELS = [
|
||||
|
||||
GEMINI_FLASH_MODELS = [
|
||||
{
|
||||
"id": "gemini-3-flash-preview",
|
||||
"description": "Gemini 3 Flash,快速多模态文本生成,支持图片和视频输入",
|
||||
"id": "gemini-3.5-flash",
|
||||
"description": "Gemini 3.5 Flash,快速多模态文本生成,通过 thinkingConfig 控制思考等级",
|
||||
"enabled": True,
|
||||
"endpoints": {
|
||||
"不思考": "/v1beta/models/gemini-3-flash-preview-nothinking:generateContent",
|
||||
"高": "/v1beta/models/gemini-3-flash-preview-high:generateContent"
|
||||
"endpoint_type": "standard",
|
||||
"endpoint": "/v1beta/models/gemini-3.5-flash: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"
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@@ -184,6 +233,123 @@ 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]:
|
||||
"""
|
||||
获取模型的端点类型
|
||||
@@ -226,6 +392,233 @@ 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 模型工具函数
|
||||
# ============================================================
|
||||
@@ -291,29 +684,24 @@ def is_flash_model_enabled(model_id: str) -> bool:
|
||||
return config.get("enabled", False)
|
||||
|
||||
|
||||
def get_flash_model_endpoint(model_id: str, thinking_depth: str = "不思考") -> Optional[str]:
|
||||
def get_flash_model_endpoint(model_id: 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-nothinking:generateContent'
|
||||
>>> get_flash_model_endpoint("gemini-3-flash-preview", "高")
|
||||
'/v1beta/models/gemini-3-flash-preview-high:generateContent'
|
||||
>>> get_flash_model_endpoint("gemini-3-flash-preview")
|
||||
'/v1beta/models/gemini-3-flash-preview:generateContent'
|
||||
"""
|
||||
config = get_flash_model_config(model_id)
|
||||
if config is None:
|
||||
return None
|
||||
|
||||
endpoints = config.get("endpoints", {})
|
||||
return endpoints.get(thinking_depth)
|
||||
return config.get("endpoint")
|
||||
|
||||
|
||||
def get_flash_model_description(model_id: str) -> str:
|
||||
@@ -332,6 +720,83 @@ 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)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 向后兼容性检查
|
||||
# ============================================================
|
||||
@@ -392,8 +857,8 @@ def validate_flash_models_config() -> None:
|
||||
验证 Flash 模型配置的完整性
|
||||
|
||||
检查:
|
||||
- 每个模型必须有 id, description, enabled, endpoints 字段
|
||||
- endpoints 必须包含所有思考深度选项
|
||||
- 每个模型必须有 id, description, enabled 字段
|
||||
- 每个模型必须有 endpoint 字段且格式正确
|
||||
- 至少有一个模型是启用的
|
||||
|
||||
Raises:
|
||||
@@ -402,8 +867,7 @@ def validate_flash_models_config() -> None:
|
||||
if not GEMINI_FLASH_MODELS:
|
||||
raise ValueError("GEMINI_FLASH_MODELS 列表不能为空")
|
||||
|
||||
required_fields = ["id", "description", "enabled", "endpoints"]
|
||||
required_thinking_depths = ["不思考", "高"]
|
||||
required_fields = ["id", "description", "enabled"]
|
||||
|
||||
for i, model in enumerate(GEMINI_FLASH_MODELS):
|
||||
# 检查必需字段
|
||||
@@ -411,24 +875,16 @@ def validate_flash_models_config() -> None:
|
||||
if field not in model:
|
||||
raise ValueError(f"Flash 模型 #{i} 缺少必需字段: {field}")
|
||||
|
||||
# 检查 endpoints 字典
|
||||
endpoints = model.get("endpoints", {})
|
||||
if not isinstance(endpoints, dict):
|
||||
raise ValueError(f"Flash 模型 {model['id']} 的 endpoints 必须是字典")
|
||||
# 检查端点配置
|
||||
if "endpoint" not in model:
|
||||
raise ValueError(f"Flash 模型 {model['id']} 缺少 'endpoint' 字段")
|
||||
|
||||
# 检查所有思考深度选项都有对应端点
|
||||
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/' 开头"
|
||||
)
|
||||
endpoint = model.get("endpoint", "")
|
||||
if not endpoint or not endpoint.startswith("/v1beta/models/"):
|
||||
raise ValueError(
|
||||
f"Flash 模型 {model['id']} 的 endpoint '{endpoint}' 格式不正确。"
|
||||
f"应以 '/v1beta/models/' 开头"
|
||||
)
|
||||
|
||||
# 检查至少有一个启用的模型
|
||||
if not get_enabled_flash_models():
|
||||
|
||||
@@ -0,0 +1,364 @@
|
||||
"""
|
||||
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, NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.r2_uploader import upload_video, upload_image
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
from ..utils.video_task import (
|
||||
check_interrupt,
|
||||
extract_error_message,
|
||||
extract_progress,
|
||||
extract_status,
|
||||
extract_video_url,
|
||||
interruptible_sleep,
|
||||
is_failure_status,
|
||||
is_success_status,
|
||||
run_with_interrupt,
|
||||
)
|
||||
|
||||
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": ""}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"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_base_url_by_route(网络线路)
|
||||
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")
|
||||
check_interrupt()
|
||||
pil_list = tensor_to_pil(参考图片)
|
||||
image_url = await upload_image(pil_list[0].convert("RGB"))
|
||||
check_interrupt()
|
||||
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. 提交任务
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
resp = await run_with_interrupt(async_request_with_retry(
|
||||
session, "POST", create_url,
|
||||
data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
|
||||
headers=headers, prefix="K3 动作控制提交: "
|
||||
))
|
||||
check_interrupt()
|
||||
text = await resp.text()
|
||||
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 interruptible_sleep(interval)
|
||||
check_interrupt()
|
||||
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 = extract_status(sr)
|
||||
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K3 动作控制] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if is_success_status(status):
|
||||
video_result_url = extract_video_url(sr)
|
||||
break
|
||||
elif is_failure_status(status, sr):
|
||||
err_msg = extract_error_message(sr)
|
||||
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. 下载视频
|
||||
check_interrupt()
|
||||
_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):
|
||||
check_interrupt()
|
||||
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",
|
||||
}
|
||||
@@ -0,0 +1,344 @@
|
||||
"""
|
||||
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, NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
from ..utils.video_task import (
|
||||
check_interrupt,
|
||||
extract_error_message,
|
||||
extract_progress,
|
||||
extract_status,
|
||||
extract_video_url,
|
||||
interruptible_sleep,
|
||||
is_failure_status,
|
||||
is_success_status,
|
||||
run_with_interrupt,
|
||||
)
|
||||
|
||||
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"}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"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_base_url_by_route(网络线路)
|
||||
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. 提交
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
resp = await run_with_interrupt(async_request_with_retry(
|
||||
session, "POST", create_url, json=body, headers=headers, prefix="K3 提交: "
|
||||
))
|
||||
check_interrupt()
|
||||
text = await resp.text()
|
||||
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:
|
||||
check_interrupt()
|
||||
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 = extract_status(sr)
|
||||
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K3 {tag}] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if is_success_status(status):
|
||||
video_url = extract_video_url(sr)
|
||||
break
|
||||
if is_failure_status(status, sr):
|
||||
err_msg = extract_error_message(sr)
|
||||
raise RuntimeError(f"K3 生成失败:{err_msg}")
|
||||
|
||||
await interruptible_sleep(interval)
|
||||
interval = min(interval * 1.5, _POLL_MAX)
|
||||
|
||||
if not video_url:
|
||||
raise RuntimeError(f"API 未返回视频 URL,响应:{sr}")
|
||||
|
||||
# 3. 下载
|
||||
check_interrupt()
|
||||
_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):
|
||||
check_interrupt()
|
||||
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 图生视频 自研",
|
||||
}
|
||||
@@ -0,0 +1,283 @@
|
||||
"""
|
||||
首尾帧 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, NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
from ..utils.video_task import (
|
||||
check_interrupt,
|
||||
extract_error_message,
|
||||
extract_progress,
|
||||
extract_status,
|
||||
extract_video_url,
|
||||
interruptible_sleep,
|
||||
is_failure_status,
|
||||
is_success_status,
|
||||
run_with_interrupt,
|
||||
)
|
||||
|
||||
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"}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"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_base_url_by_route(网络线路)
|
||||
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. 提交
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
resp = await run_with_interrupt(async_request_with_retry(
|
||||
session, "POST", create_url, json=body, headers=headers, prefix="K3 首尾帧提交: "
|
||||
))
|
||||
check_interrupt()
|
||||
text = await resp.text()
|
||||
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:
|
||||
check_interrupt()
|
||||
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 = extract_status(sr)
|
||||
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K3 首尾帧] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if is_success_status(status):
|
||||
video_url = extract_video_url(sr)
|
||||
break
|
||||
if is_failure_status(status, sr):
|
||||
err_msg = extract_error_message(sr)
|
||||
raise RuntimeError(f"K3 首尾帧生成失败:{err_msg}")
|
||||
|
||||
await interruptible_sleep(interval)
|
||||
interval = min(interval * 1.5, _POLL_MAX)
|
||||
|
||||
if not video_url:
|
||||
raise RuntimeError(f"API 未返回视频 URL,响应:{sr}")
|
||||
|
||||
# 3. 下载
|
||||
check_interrupt()
|
||||
_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):
|
||||
check_interrupt()
|
||||
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 自研",
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
"""
|
||||
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, NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
from ..utils.video_task import (
|
||||
check_interrupt,
|
||||
extract_error_message,
|
||||
extract_progress,
|
||||
extract_status,
|
||||
extract_video_url,
|
||||
interruptible_sleep,
|
||||
is_failure_status,
|
||||
is_success_status,
|
||||
run_with_interrupt,
|
||||
)
|
||||
|
||||
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": "关闭"}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"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_base_url_by_route(网络线路)
|
||||
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": 时长,
|
||||
}
|
||||
metadata = {}
|
||||
if 尾帧 is not None:
|
||||
metadata["image_tail"] = _image_to_base64(尾帧, scale)
|
||||
if 生成音频 == "打开":
|
||||
metadata["sound"] = "on"
|
||||
if metadata:
|
||||
body["metadata"] = metadata
|
||||
|
||||
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. 提交
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
resp = await run_with_interrupt(async_request_with_retry(
|
||||
session, "POST", create_url, json=body, headers=headers, prefix="K26 提交: "
|
||||
))
|
||||
check_interrupt()
|
||||
text = await resp.text()
|
||||
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:
|
||||
check_interrupt()
|
||||
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 = extract_status(sr)
|
||||
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K26 图生视频] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if is_success_status(status):
|
||||
# 提取视频 URL
|
||||
video_url = extract_video_url(sr)
|
||||
break
|
||||
if is_failure_status(status, sr):
|
||||
err_msg = extract_error_message(sr)
|
||||
raise RuntimeError(f"K26 生成失败:{err_msg}")
|
||||
|
||||
await interruptible_sleep(interval)
|
||||
interval = min(interval * 1.5, _POLL_MAX)
|
||||
|
||||
if not video_url:
|
||||
raise RuntimeError(f"API 未返回视频 URL,响应:{sr}")
|
||||
|
||||
# 3. 下载
|
||||
check_interrupt()
|
||||
_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):
|
||||
check_interrupt()
|
||||
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 图生视频(首尾帧)",
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
"""
|
||||
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, NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
from ..utils.video_task import (
|
||||
check_interrupt,
|
||||
extract_error_message,
|
||||
extract_progress,
|
||||
extract_status,
|
||||
extract_video_url,
|
||||
interruptible_sleep,
|
||||
is_failure_status,
|
||||
is_success_status,
|
||||
run_with_interrupt,
|
||||
)
|
||||
|
||||
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": "关闭"}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"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):
|
||||
if 模式 == "720p" and 生成音频 == "打开":
|
||||
raise RuntimeError("K26 仅1080p支持音频,请将模式切换为1080p或关闭生成音频。")
|
||||
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_base_url_by_route(网络线路)
|
||||
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["metadata"] = {"sound": "on"}
|
||||
|
||||
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. 提交
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
resp = await run_with_interrupt(async_request_with_retry(
|
||||
session, "POST", create_url, json=body, headers=headers, prefix="K26 图生视频提交: "
|
||||
))
|
||||
check_interrupt()
|
||||
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 interruptible_sleep(interval)
|
||||
|
||||
check_interrupt()
|
||||
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 = extract_status(sr)
|
||||
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K26 图生视频] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if is_success_status(status):
|
||||
# 提取视频 URL
|
||||
video_url = extract_video_url(sr)
|
||||
break
|
||||
if is_failure_status(status, sr):
|
||||
err_msg = extract_error_message(sr)
|
||||
raise RuntimeError(f"K26 生成失败:{err_msg}")
|
||||
|
||||
interval = min(interval * 1.5, _POLL_MAX)
|
||||
|
||||
if not video_url:
|
||||
raise RuntimeError(f"API 未返回视频 URL,响应:{sr}")
|
||||
|
||||
# 3. 下载
|
||||
check_interrupt()
|
||||
_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):
|
||||
check_interrupt()
|
||||
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 图生视频",
|
||||
}
|
||||
+31
-3
@@ -3,8 +3,36 @@
|
||||
包含所有 ComfyUI 自定义节点的实现
|
||||
"""
|
||||
|
||||
from .nano_banana_pro import NanoBananaPro
|
||||
from .batch_nano_banana_pro import BatchNanoBananaPro
|
||||
from .stream_preview import StreamPreview
|
||||
from .nano_banana import NanoBanana
|
||||
NanoBananaPro = NanoBanana
|
||||
from .batch_nano_banana import BatchNanoBananaPro
|
||||
from .google_gemini import GoogleGemini
|
||||
from .load_file import LoadFile
|
||||
from .image_stitch_pro import ImageStitchPro
|
||||
from .remove_metadata import BatchCleanMetadata
|
||||
from .video_preview import VideoPreview
|
||||
from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset
|
||||
from .veo_video import GoogleVeo
|
||||
from .newapi_veo_video import Google31Video
|
||||
from .flux_edit import FluxImageEdit
|
||||
from .universal_llm import UniversalLLMChat
|
||||
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, O1keyGPTImageBatch
|
||||
from .grok_image import O1keyGrokImage
|
||||
from .grok_video import O1keyGrokVideo
|
||||
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
|
||||
from .save_image_format import SaveImageFormat
|
||||
from .save_psd import O1keySavePSD
|
||||
from .remove_bg import O1keyRemoveBackground
|
||||
from .color_remove_bg import O1keyColorRemoveBG
|
||||
from .grid_splitter import O1keyGridSplitter
|
||||
|
||||
__all__ = ['NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini']
|
||||
__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBanana', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'Google31Video', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGPTImageBatch', 'O1keyGrokImage', 'O1keyGrokVideo', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat', 'O1keySavePSD', 'O1keyRemoveBackground', 'O1keyColorRemoveBG', 'O1keyGridSplitter']
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
"""
|
||||
批量图像(o1key)节点
|
||||
复刻 ComfyUI 原生「批量图像」节点的动态输入行为:
|
||||
|
||||
- 默认显示 2 个图像输入端口(图1, 图2)
|
||||
- 当最后一个端口连上图像后,自动追加新端口
|
||||
- 断开连线后,多余的端口自动消失,最少保留 2 个
|
||||
|
||||
与原生节点的区别:
|
||||
原生节点会把所有图像强制 resize 到第一张的分辨率再合并为单一 tensor。
|
||||
本节点保留每张图的原始分辨率,以 list[Tensor] 形式输出(is_output_list)。
|
||||
下游节点(如「多分辨率图像预览」)需开启 INPUT_IS_LIST 才能正确接收。
|
||||
|
||||
实现方式:使用 V3 API 的 io.Autogrow.TemplateNames,
|
||||
框架原生支持动态 slot 增减,无需编写任何 JS 扩展。
|
||||
"""
|
||||
|
||||
import torch
|
||||
from comfy_api.latest import io
|
||||
|
||||
# 预生成 50 个端口名:图1, 图2, ..., 图50
|
||||
_SLOT_NAMES = [f"图{i}" for i in range(1, 51)]
|
||||
|
||||
|
||||
class BatchImagesO1key(io.ComfyNode):
|
||||
"""
|
||||
批量图像(o1key)
|
||||
|
||||
- 动态输入端口(默认 2 个,最多 50 个),端口名为 图1、图2、图3...
|
||||
- 连接最后一个端口时自动增加新端口
|
||||
- 断开后自动减少,保持界面整洁
|
||||
- 保留每张图的原始分辨率,不做任何 resize / 裁剪
|
||||
- 输出为图像列表,可直接接入「多分辨率图像预览」节点
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
autogrow_template = io.Autogrow.TemplateNames(
|
||||
input=io.Image.Input("image"),
|
||||
names=_SLOT_NAMES,
|
||||
min=2,
|
||||
)
|
||||
return io.Schema(
|
||||
node_id="BatchImagesO1key",
|
||||
display_name="加载图像(批量)",
|
||||
category="image",
|
||||
description=(
|
||||
"将多个独立图像收集为图像列表输出,保留每张图的原始分辨率。\n"
|
||||
"• 默认显示 2 个输入端口(图1、图2),连接最后一个后自动追加新端口\n"
|
||||
"• 断开连线后端口自动减少,最少保留 2 个\n"
|
||||
"• 不做任何 resize / 裁剪,原图尺寸原样输出\n"
|
||||
"• 输出为图像列表,可直接接入「多分辨率图像预览」节点"
|
||||
),
|
||||
search_aliases=["批量图像", "batch images", "合并图像", "图像合并", "stack images"],
|
||||
inputs=[
|
||||
io.Autogrow.Input("images", template=autogrow_template)
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="图像", is_output_list=True),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images: io.Autogrow.Type) -> io.NodeOutput:
|
||||
# images 是 dict,key 为 "图1", "图2", ... ;未连接的 slot 值为 None
|
||||
tensors = [v for v in images.values() if v is not None]
|
||||
|
||||
if not tensors:
|
||||
raise ValueError("批量图像(o1key):请至少连接一张图像")
|
||||
|
||||
for i, t in enumerate(tensors):
|
||||
h, w = t.shape[1], t.shape[2]
|
||||
print(f"批量图像(o1key):图{i + 1} → {w}×{h},shape={list(t.shape)}")
|
||||
|
||||
print(f"批量图像(o1key):共收集 {len(tensors)} 张,原始分辨率原样输出")
|
||||
|
||||
# 以 list[Tensor] 形式返回,每张图保持自身分辨率
|
||||
return io.NodeOutput(tensors)
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,751 +0,0 @@
|
||||
"""
|
||||
批量 Nano Banana Pro 节点
|
||||
ComfyUI 自定义节点,用于批量处理图像生成任务
|
||||
支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存
|
||||
"""
|
||||
|
||||
import time
|
||||
import math
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Optional, Tuple, List
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor
|
||||
from ..utils.file_utils import (
|
||||
ImageInfo,
|
||||
load_images_from_folder,
|
||||
pair_images_indexed,
|
||||
pair_images_cartesian,
|
||||
generate_output_filename,
|
||||
save_image
|
||||
)
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
from ..models_config import get_enabled_models
|
||||
|
||||
# 导入 ComfyUI 原生进度条
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ BatchNanoBananaPro: comfy.utils.ProgressBar 不可用,将只使用终端进度显示")
|
||||
|
||||
|
||||
class BatchNanoBananaPro:
|
||||
"""
|
||||
批量 Nano Banana Pro 节点
|
||||
|
||||
功能:
|
||||
- 从多个文件夹加载图片
|
||||
- 支持三种配对模式:
|
||||
* 1:1 - 索引配对(文件夹之间按位置配对)
|
||||
* 1*N - 笛卡尔积配对(所有可能组合)
|
||||
* 不配对 - 固定参考图模式(文件夹图片依次与所有参考图组合)
|
||||
- 批量调用 API 生成图像
|
||||
- 智能命名保存(保留原始文件名)
|
||||
- 并发控制(默认最大 100)
|
||||
|
||||
注意:
|
||||
- 「不配对」模式只支持单个文件夹
|
||||
- 支持的模型列表从 models_config.py 动态加载
|
||||
- 要添加/禁用模型,请编辑 models_config.py 文件
|
||||
"""
|
||||
|
||||
# 支持的模型列表(从配置文件动态加载)
|
||||
MODELS = None # 将在 INPUT_TYPES 中动态获取
|
||||
|
||||
# 支持的宽高比列表
|
||||
ASPECT_RATIOS = [
|
||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
||||
"2:3", "3:2", "4:5", "5:4", "21:9"
|
||||
]
|
||||
|
||||
# 支持的分辨率列表
|
||||
RESOLUTIONS = ["1K", "2K", "4K"]
|
||||
|
||||
# 配对模式
|
||||
PAIRING_MODES = ["1:1", "1*N", "不配对"]
|
||||
|
||||
def __init__(self):
|
||||
"""初始化节点"""
|
||||
self.client = None
|
||||
|
||||
def resize_to_megapixels(
|
||||
self,
|
||||
image: Image.Image,
|
||||
target_megapixels: float
|
||||
) -> Image.Image:
|
||||
"""
|
||||
将图像缩放到指定的总像素数,保持纵横比
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
target_megapixels: 目标像素数(百万像素)
|
||||
|
||||
Returns:
|
||||
缩放后的 PIL Image
|
||||
|
||||
Example:
|
||||
>>> resized = self.resize_to_megapixels(img, 2.0) # 缩放到2百万像素
|
||||
"""
|
||||
# 计算当前像素数
|
||||
current_pixels = image.width * image.height
|
||||
target_pixels = int(target_megapixels * 1_000_000)
|
||||
|
||||
# 如果当前像素数已经接近目标,则不缩放
|
||||
if abs(current_pixels - target_pixels) / target_pixels < 0.05:
|
||||
return image
|
||||
|
||||
# 计算缩放比例
|
||||
scale = (target_pixels / current_pixels) ** 0.5
|
||||
|
||||
# 计算新尺寸
|
||||
new_width = int(image.width * scale)
|
||||
new_height = int(image.height * scale)
|
||||
|
||||
# 确保至少为1像素
|
||||
new_width = max(1, new_width)
|
||||
new_height = max(1, new_height)
|
||||
|
||||
# 使用 Lanczos 重采样
|
||||
resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
return resized_image
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
|
||||
ComfyUI 节点规范:
|
||||
- required: 必选参数
|
||||
- optional: 可选参数
|
||||
"""
|
||||
# 从配置文件动态获取启用的模型列表
|
||||
enabled_models = get_enabled_models()
|
||||
|
||||
# 如果没有启用的模型,使用空列表(会导致节点不可用,提示用户配置)
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
||||
|
||||
# 创建9个独立的图像输入
|
||||
optional_inputs = {}
|
||||
for i in range(1, 10): # 1-9
|
||||
optional_inputs[f"参考图{i}"] = ("IMAGE",)
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "一个中国女子的OOTD",
|
||||
"multiline": True
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0]
|
||||
}),
|
||||
"宽高比": (cls.ASPECT_RATIOS, {
|
||||
"default": "1:1"
|
||||
}),
|
||||
"分辨率": (cls.RESOLUTIONS, {
|
||||
"default": "2K"
|
||||
}),
|
||||
"像素缩放": ("BOOLEAN", {
|
||||
"default": False
|
||||
}),
|
||||
"分辨率像素": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.1,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"display": "number"
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff
|
||||
}),
|
||||
"文件夹1": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹2": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹3": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹4": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"保存路径": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"图片配对模式": (cls.PAIRING_MODES, {
|
||||
"default": "不配对"
|
||||
})
|
||||
},
|
||||
"optional": optional_inputs
|
||||
}
|
||||
|
||||
# 返回值类型
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("输出图像",)
|
||||
|
||||
# 执行函数名
|
||||
FUNCTION = "process_batch"
|
||||
|
||||
# 节点分类
|
||||
CATEGORY = "image/batch"
|
||||
|
||||
def _load_folders(
|
||||
self,
|
||||
folder1: str,
|
||||
folder2: Optional[str],
|
||||
folder3: Optional[str],
|
||||
folder4: Optional[str],
|
||||
enable_scaling: bool,
|
||||
target_megapixels: float
|
||||
) -> List[List[ImageInfo]]:
|
||||
"""
|
||||
加载所有文件夹中的图片
|
||||
|
||||
Args:
|
||||
folder1-4: 文件夹路径
|
||||
enable_scaling: 是否启用像素缩放
|
||||
target_megapixels: 目标像素数(百万像素)
|
||||
|
||||
Returns:
|
||||
图片列表的列表
|
||||
"""
|
||||
folders = [folder1, folder2, folder3, folder4]
|
||||
all_images = []
|
||||
|
||||
for i, folder in enumerate(folders, 1):
|
||||
if folder and folder.strip():
|
||||
try:
|
||||
images = load_images_from_folder(folder)
|
||||
if images:
|
||||
# 应用像素缩放
|
||||
if enable_scaling:
|
||||
scaled_images = []
|
||||
for img_info in images:
|
||||
scaled_img = self.resize_to_megapixels(
|
||||
img_info.image,
|
||||
target_megapixels
|
||||
)
|
||||
# 创建新的 ImageInfo,保留其他元数据
|
||||
scaled_info = ImageInfo(
|
||||
image=scaled_img,
|
||||
filename=img_info.filename,
|
||||
extension=img_info.extension,
|
||||
source_path=img_info.source_path
|
||||
)
|
||||
scaled_images.append(scaled_info)
|
||||
images = scaled_images
|
||||
|
||||
all_images.append(images)
|
||||
print(f"BatchNanoBananaPro: 文件夹{i} 加载了 {len(images)} 张图片")
|
||||
else:
|
||||
print(f"BatchNanoBananaPro: 文件夹{i} 为空或没有有效图片")
|
||||
except ValueError as e:
|
||||
print(f"BatchNanoBananaPro: 文件夹{i} 加载失败 - {e}")
|
||||
|
||||
return all_images
|
||||
|
||||
def _create_pairs(
|
||||
self,
|
||||
image_lists: List[List[ImageInfo]],
|
||||
pairing_mode: str,
|
||||
manual_images: Optional[List[ImageInfo]] = None
|
||||
) -> List[Tuple[ImageInfo, ...]]:
|
||||
"""
|
||||
根据配对模式创建图片组合
|
||||
|
||||
Args:
|
||||
image_lists: 从文件夹加载的图片列表
|
||||
pairing_mode: 配对模式 (1:1, 1*N, 不配对)
|
||||
manual_images: 手动输入的参考图
|
||||
|
||||
Returns:
|
||||
配对后的元组列表
|
||||
|
||||
Raises:
|
||||
ValueError: 不配对模式下填入多个文件夹时
|
||||
"""
|
||||
# === 新模式:不配对 ===
|
||||
if pairing_mode == "不配对":
|
||||
# 验证:只支持单个文件夹
|
||||
if len(image_lists) > 1:
|
||||
raise ValueError("「不配对」模式只支持单个文件夹,请清空其他文件夹路径")
|
||||
|
||||
# 场景1:有文件夹 + 有参考图
|
||||
if image_lists and manual_images:
|
||||
folder_images = image_lists[0]
|
||||
# 每张文件夹图片 + 所有参考图
|
||||
pairs = []
|
||||
for img in folder_images:
|
||||
pair = (img,) + tuple(manual_images)
|
||||
pairs.append(pair)
|
||||
return pairs
|
||||
|
||||
# 场景2:有文件夹 + 无参考图
|
||||
elif image_lists:
|
||||
# 每张图片单独成组
|
||||
return [(img,) for img in image_lists[0]]
|
||||
|
||||
# 场景3:无文件夹 + 有参考图
|
||||
elif manual_images:
|
||||
# 每张参考图单独成组
|
||||
return [(img,) for img in manual_images]
|
||||
|
||||
else:
|
||||
return []
|
||||
|
||||
# === 原有逻辑:1:1 和 1*N ===
|
||||
# 如果有手动参考图,添加到列表中(所有参考图作为一个列表)
|
||||
if manual_images:
|
||||
image_lists.append(manual_images)
|
||||
|
||||
if not image_lists:
|
||||
return []
|
||||
|
||||
# 如果只有一个列表,直接返回每个图片作为单元素元组
|
||||
if len(image_lists) == 1:
|
||||
return [(img,) for img in image_lists[0]]
|
||||
|
||||
# 根据配对模式选择配对函数
|
||||
if pairing_mode == "1:1":
|
||||
pairs = pair_images_indexed(*image_lists)
|
||||
else: # 1*N
|
||||
pairs = pair_images_cartesian(*image_lists)
|
||||
|
||||
return pairs
|
||||
|
||||
async def _generate_single_task(
|
||||
self,
|
||||
client: GeminiAPIClient,
|
||||
session: aiohttp.ClientSession,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: List[ImageInfo],
|
||||
output_folder: str,
|
||||
task_index: int
|
||||
) -> dict:
|
||||
"""
|
||||
执行单个生成任务
|
||||
|
||||
Args:
|
||||
client: API 客户端
|
||||
session: aiohttp 会话
|
||||
prompt: 提示词
|
||||
model: 模型名称
|
||||
resolution: 分辨率
|
||||
aspect_ratio: 宽高比
|
||||
images: 输入图片列表
|
||||
output_folder: 输出文件夹
|
||||
task_index: 任务索引
|
||||
|
||||
Returns:
|
||||
包含结果信息的字典
|
||||
"""
|
||||
result = {
|
||||
"task_index": task_index,
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"saved_files": [],
|
||||
"error": None
|
||||
}
|
||||
|
||||
try:
|
||||
# 准备输入图片
|
||||
input_pil_images = [info.image for info in images]
|
||||
|
||||
# 调用 API 生成图片(固定生成1次)
|
||||
generated_images = []
|
||||
try:
|
||||
gen_result = await client.generate_single_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_pil_images,
|
||||
session=session
|
||||
)
|
||||
if gen_result:
|
||||
generated_images.extend(gen_result)
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"BatchNanoBananaPro: 任务 {task_index + 1} 生成失败 - {error_msg}")
|
||||
result["error"] = error_msg
|
||||
|
||||
# 保存生成的图片
|
||||
for i, gen_img in enumerate(generated_images):
|
||||
# 使用任务索引作为唯一标识,确保并发安全
|
||||
output_path = generate_output_filename(
|
||||
source_images=list(images),
|
||||
batch_index=i,
|
||||
output_folder=output_folder,
|
||||
extension=".png",
|
||||
task_id=f"task{task_index}"
|
||||
)
|
||||
save_image(gen_img, output_path)
|
||||
result["saved_files"].append(output_path)
|
||||
|
||||
# 只有生成了图片才标记为成功
|
||||
if len(generated_images) > 0:
|
||||
result["success"] = True
|
||||
result["generated_count"] = len(generated_images)
|
||||
|
||||
except Exception as e:
|
||||
result["error"] = str(e)
|
||||
|
||||
return result
|
||||
|
||||
async def _process_batch_async(
|
||||
self,
|
||||
pairs: List[Tuple[ImageInfo, ...]],
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
output_folder: str,
|
||||
pbar=None
|
||||
) -> List[dict]:
|
||||
"""
|
||||
异步批量处理所有任务
|
||||
|
||||
Args:
|
||||
pairs: 配对后的图片组合
|
||||
prompt: 提示词
|
||||
model: 模型名称
|
||||
resolution: 分辨率
|
||||
aspect_ratio: 宽高比
|
||||
output_folder: 输出文件夹
|
||||
|
||||
Returns:
|
||||
所有任务的结果列表
|
||||
"""
|
||||
if self.client is None:
|
||||
self.client = GeminiAPIClient()
|
||||
|
||||
# 固定最大并发数为 100
|
||||
max_concurrent = 100
|
||||
|
||||
total_tasks = len(pairs)
|
||||
all_results = []
|
||||
completed = 0
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
# 计算分批数量
|
||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
||||
|
||||
# 进度打印配置:任务数 >= 50 时,额外显示百分比里程碑
|
||||
show_milestone = total_tasks >= 50
|
||||
milestones = [0.2, 0.4, 0.6, 0.8, 1.0] # 20%, 40%, 60%, 80%, 100%
|
||||
milestone_index = 0
|
||||
|
||||
if num_batches > 1:
|
||||
print(f"BatchNanoBananaPro: 任务数 {total_tasks} 超过并发上限 {max_concurrent},将分 {num_batches} 批执行")
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
for batch_idx in range(num_batches):
|
||||
start_idx = batch_idx * max_concurrent
|
||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||||
batch_pairs = pairs[start_idx:end_idx]
|
||||
|
||||
if num_batches > 1:
|
||||
print(f"BatchNanoBananaPro: 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...")
|
||||
|
||||
# 创建当前批次的任务
|
||||
tasks = []
|
||||
for i, pair in enumerate(batch_pairs):
|
||||
task = asyncio.create_task(
|
||||
self._generate_single_task(
|
||||
client=self.client,
|
||||
session=session,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=list(pair),
|
||||
output_folder=output_folder,
|
||||
task_index=start_idx + i
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
# 使用 as_completed 实时获取完成的任务
|
||||
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": []
|
||||
}
|
||||
all_results.append(result_data)
|
||||
else:
|
||||
result_data = result
|
||||
all_results.append(result)
|
||||
except Exception as e:
|
||||
result_data = {
|
||||
"success": False,
|
||||
"error": str(e),
|
||||
"generated_count": 0,
|
||||
"saved_files": []
|
||||
}
|
||||
all_results.append(result_data)
|
||||
|
||||
completed += 1
|
||||
|
||||
# 根据成功/失败状态打印不同信息
|
||||
if result_data and result_data.get("success", False):
|
||||
success_count += 1
|
||||
print(f"BatchNanoBananaPro: 任务 {completed}/{total_tasks} 成功 ✓")
|
||||
else:
|
||||
fail_count += 1
|
||||
# 提取错误信息的第一行
|
||||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
||||
# 截取第一行或前50个字符
|
||||
if '\n' in error_msg:
|
||||
error_msg = error_msg.split('\n')[0]
|
||||
if len(error_msg) > 50:
|
||||
error_msg = error_msg[:50] + "..."
|
||||
print(f"BatchNanoBananaPro: 任务 {completed}/{total_tasks} 失败 ✗ - {error_msg}")
|
||||
|
||||
# 更新 ComfyUI 原生进度条
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
# 大任务额外显示百分比里程碑
|
||||
if show_milestone and milestone_index < len(milestones):
|
||||
progress = completed / total_tasks
|
||||
if progress >= milestones[milestone_index]:
|
||||
percentage = int(milestones[milestone_index] * 100)
|
||||
print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<")
|
||||
milestone_index += 1
|
||||
|
||||
return all_results
|
||||
|
||||
def process_batch(
|
||||
self,
|
||||
prompt: str,
|
||||
文件夹1: str,
|
||||
文件夹2: str,
|
||||
文件夹3: str,
|
||||
文件夹4: str,
|
||||
像素缩放: bool,
|
||||
分辨率像素: float,
|
||||
seed: int,
|
||||
保存路径: str,
|
||||
图片配对模式: str,
|
||||
模型: str,
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""
|
||||
批量处理图像生成任务
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
文件夹1-4: 图片文件夹路径
|
||||
像素缩放: 是否启用像素缩放
|
||||
分辨率像素: 目标像素数(百万像素)
|
||||
seed: 随机种子
|
||||
保存路径: 输出保存路径
|
||||
图片配对模式: 1:1 或 1*N
|
||||
模型: 模型名称
|
||||
宽高比: 输出宽高比
|
||||
分辨率: 输出分辨率
|
||||
**kwargs: 动态参考图输入 (参考图1-9)
|
||||
|
||||
Returns:
|
||||
输出图像张量
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 设置随机种子(用于本地随机操作)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed % (2**32))
|
||||
# 验证保存路径
|
||||
if not 保存路径 or not 保存路径.strip():
|
||||
raise ValueError("请提供保存路径")
|
||||
|
||||
# 加载文件夹图片
|
||||
print("BatchNanoBananaPro: 开始加载图片...")
|
||||
image_lists = self._load_folders(
|
||||
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
||||
像素缩放, 分辨率像素
|
||||
)
|
||||
|
||||
# 处理独立的参考图输入
|
||||
manual_images = []
|
||||
for i in range(1, 10): # 1-9
|
||||
key = f"参考图{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
pil_images = tensor_to_pil(kwargs[key])
|
||||
for j, img in enumerate(pil_images):
|
||||
# 如果启用像素缩放,也对参考图进行缩放
|
||||
if 像素缩放:
|
||||
img = self.resize_to_megapixels(img, 分辨率像素)
|
||||
|
||||
manual_images.append(
|
||||
ImageInfo(
|
||||
image=img,
|
||||
filename=f"manual_{i}_{j}",
|
||||
extension=".png",
|
||||
source_path=""
|
||||
)
|
||||
)
|
||||
|
||||
if manual_images:
|
||||
print(f"BatchNanoBananaPro: 加载了 {len(manual_images)} 张参考图")
|
||||
|
||||
# 验证是否有图片
|
||||
total_folder_images = sum(len(lst) for lst in image_lists)
|
||||
total_manual_images = len(manual_images)
|
||||
|
||||
if total_folder_images == 0 and total_manual_images == 0:
|
||||
raise ValueError("未找到任何图片,请检查文件夹路径或提供参考图")
|
||||
|
||||
# 创建配对
|
||||
print(f"BatchNanoBananaPro: 使用 {图片配对模式} 模式创建配对...")
|
||||
pairs = self._create_pairs(image_lists, 图片配对模式, manual_images if manual_images else None)
|
||||
|
||||
if not pairs:
|
||||
raise ValueError("配对结果为空,请检查输入")
|
||||
|
||||
total_tasks = len(pairs)
|
||||
print(f"BatchNanoBananaPro: 共 {total_tasks} 组配对")
|
||||
|
||||
# 创建 ComfyUI 原生进度条
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(total_tasks)
|
||||
|
||||
# 初始化 API 客户端
|
||||
if self.client is None:
|
||||
try:
|
||||
self.client = GeminiAPIClient()
|
||||
except ValueError as e:
|
||||
raise ValueError(f"初始化 API 客户端失败: {str(e)}")
|
||||
|
||||
# 执行批量生成
|
||||
print("BatchNanoBananaPro: 开始批量生成...")
|
||||
|
||||
# 在新线程中运行异步代码,避免事件循环冲突
|
||||
def run_async_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
self._process_batch_async(
|
||||
pairs=pairs,
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
output_folder=保存路径,
|
||||
pbar=pbar
|
||||
)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
# 使用线程池在新线程中运行事件循环
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async_in_thread)
|
||||
results = future.result()
|
||||
|
||||
# 统计结果
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||
all_saved_files = []
|
||||
for r in results:
|
||||
all_saved_files.extend(r.get("saved_files", []))
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# 精简统计信息
|
||||
print("=" * 50)
|
||||
print(f"BatchNanoBananaPro 处理完成 | 总耗时: {elapsed:.2f}s | 成功: {success_count}/{total_tasks} | 生成: {total_generated}张")
|
||||
print(f"保存路径: {保存路径}")
|
||||
|
||||
# 失败详情(如果有)
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
if failed_results:
|
||||
# 收集失败任务的索引
|
||||
failed_indices = [str(r.get('task_index', '?') + 1) for r in failed_results[:5]]
|
||||
failed_str = ",".join(failed_indices)
|
||||
if len(failed_results) > 5:
|
||||
failed_str += f"... (共{len(failed_results)}个)"
|
||||
# 显示第一个失败原因作为示例
|
||||
first_error = failed_results[0].get('error', '未知错误')
|
||||
print(f"失败 {len(failed_results)}个: 任务{failed_str} - {first_error}")
|
||||
|
||||
# 收集所有生成的图片
|
||||
output_images = []
|
||||
for file_path in all_saved_files:
|
||||
try:
|
||||
img = Image.open(file_path)
|
||||
output_images.append(img)
|
||||
except Exception as e:
|
||||
print(f"BatchNanoBananaPro: 无法加载图片 {file_path} - {e}")
|
||||
|
||||
# 如果没有生成成功的图片,创建一个占位图
|
||||
if not output_images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
output_images = [placeholder]
|
||||
|
||||
# 转换为张量
|
||||
output_tensor = pil_to_tensor(output_images)
|
||||
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as e:
|
||||
# 检测是否为授权错误
|
||||
if str(e) == "未授权!":
|
||||
print("请联系作者授权后方可使用!")
|
||||
else:
|
||||
print(f"BatchNanoBananaPro: 输入错误 - {str(e)}")
|
||||
raise
|
||||
|
||||
except RuntimeError as e:
|
||||
print(f"BatchNanoBananaPro: 运行时错误 - {str(e)}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
print(f"BatchNanoBananaPro: 未知错误 - {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"{balance_info}")
|
||||
print("=" * 50)
|
||||
except Exception as e:
|
||||
print(f"⚠️ 余额查询失败 - {str(e)}")
|
||||
print("=" * 50)
|
||||
@@ -0,0 +1,93 @@
|
||||
"""
|
||||
o1key 颜色去背景节点
|
||||
基于颜色距离计算,精确可控,不依赖 AI 模型
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class O1keyColorRemoveBG:
|
||||
"""
|
||||
颜色去背景 - 精确移除纯色背景
|
||||
|
||||
模式说明:
|
||||
- 白色(white): 移除白色背景,适合大多数场景
|
||||
- 白色保护(white-preserve): 移除白底但保护浅色前景物体
|
||||
- 自动检测(corner): 自动采样四角颜色作为背景色
|
||||
- 指定颜色(color): 手动指定要移除的背景颜色
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"模式": (["白色", "白色保护", "自动检测", "指定颜色"], {
|
||||
"default": "白色",
|
||||
}),
|
||||
"容差": ("FLOAT", {
|
||||
"default": 8.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 1.0,
|
||||
"tooltip": "颜色距离阈值,越大去除范围越广",
|
||||
}),
|
||||
"羽化": ("FLOAT", {
|
||||
"default": 45.0,
|
||||
"min": 0.0,
|
||||
"max": 200.0,
|
||||
"step": 1.0,
|
||||
"tooltip": "边缘过渡范围,越大边缘越柔和",
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"背景色R": ("INT", {"default": 255, "min": 0, "max": 255}),
|
||||
"背景色G": ("INT", {"default": 255, "min": 0, "max": 255}),
|
||||
"背景色B": ("INT", {"default": 255, "min": 0, "max": 255}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("RGBA图像",)
|
||||
FUNCTION = "remove_bg"
|
||||
CATEGORY = "o1key/image"
|
||||
|
||||
_MODE_MAP = {
|
||||
"白色": "white",
|
||||
"白色保护": "white-preserve",
|
||||
"自动检测": "corner",
|
||||
"指定颜色": "color",
|
||||
}
|
||||
|
||||
def remove_bg(self, image, 模式, 容差, 羽化, 背景色R=255, 背景色G=255, 背景色B=255):
|
||||
from ..utils.color_key import remove_background
|
||||
|
||||
mode = self._MODE_MAP.get(模式, "white")
|
||||
bg_color = (背景色R, 背景色G, 背景色B)
|
||||
|
||||
batch_size = image.shape[0]
|
||||
results = []
|
||||
|
||||
for i in range(batch_size):
|
||||
frame = image[i] # [H, W, C]
|
||||
arr = (frame.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||||
|
||||
if arr.shape[2] == 4:
|
||||
pil_img = Image.fromarray(arr, mode="RGBA")
|
||||
else:
|
||||
pil_img = Image.fromarray(arr, mode="RGB")
|
||||
|
||||
result = remove_background(
|
||||
pil_img, mode=mode, bg_color=bg_color,
|
||||
tolerance=容差, feather=羽化,
|
||||
)
|
||||
|
||||
result_arr = np.array(result.convert("RGBA")).astype(np.float32) / 255.0
|
||||
results.append(torch.from_numpy(result_arr))
|
||||
|
||||
output = torch.stack(results, dim=0)
|
||||
print(f"[o1key 颜色去背景] 模式={模式}, 容差={容差}, 羽化={羽化}, "
|
||||
f"处理 {batch_size} 张")
|
||||
return (output,)
|
||||
@@ -0,0 +1,420 @@
|
||||
"""
|
||||
豆包生图节点
|
||||
后端通过 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": "豆包生图",
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
"""
|
||||
Flux2 图像编辑节点
|
||||
通过 vip.o1key.com 调用 Flux2 + SeedVR2 远程服务进行图像编辑和超分辨率
|
||||
|
||||
功能:
|
||||
- 接收主图和参考图
|
||||
- 上传到远程服务器执行图像编辑
|
||||
- 轮询等待 SeedVR2 超分辨率结果
|
||||
- 返回最终放大后的图像
|
||||
"""
|
||||
|
||||
import time
|
||||
from io import BytesIO
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor
|
||||
from ..clients.flux_edit_client import FluxEditClient
|
||||
|
||||
|
||||
class FluxImageEdit:
|
||||
"""
|
||||
Flux2 图像编辑节点
|
||||
|
||||
通过远程 API 将主图与参考图结合,按照提示词进行图像编辑,
|
||||
并经 SeedVR2 超分辨率放大后返回最终结果。
|
||||
"""
|
||||
|
||||
SIZES = ["2K", "4K"]
|
||||
|
||||
def __init__(self):
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"主图": ("IMAGE",),
|
||||
"参考图": ("IMAGE",),
|
||||
"提示词": ("STRING", {
|
||||
"default": "Replace the woman's underwear in Figure 1 with the strapless bra in Figure 2",
|
||||
"multiline": True,
|
||||
}),
|
||||
"分辨率": (cls.SIZES, {
|
||||
"default": "4K",
|
||||
}),
|
||||
"轮询间隔": ("INT", {
|
||||
"default": 15,
|
||||
"min": 5,
|
||||
"max": 60,
|
||||
"step": 5,
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff,
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("输出图像",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "image/edit"
|
||||
|
||||
def _image_to_jpeg_bytes(self, image: Image.Image, quality: int = 92) -> bytes:
|
||||
"""将 PIL Image 转为 JPEG 二进制"""
|
||||
if image.mode in ("RGBA", "P", "LA"):
|
||||
image = image.convert("RGB")
|
||||
buf = BytesIO()
|
||||
image.save(buf, format="JPEG", quality=quality)
|
||||
return buf.getvalue()
|
||||
|
||||
def generate(
|
||||
self,
|
||||
主图: torch.Tensor,
|
||||
参考图: torch.Tensor,
|
||||
提示词: str,
|
||||
分辨率: str,
|
||||
轮询间隔: int,
|
||||
seed: int,
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""
|
||||
执行图像编辑
|
||||
|
||||
Args:
|
||||
主图: 要编辑的原始图像 (ComfyUI tensor, [B, H, W, C])
|
||||
参考图: 参考/风格图像 (ComfyUI tensor, [B, H, W, C])
|
||||
提示词: 编辑指令
|
||||
分辨率: 超分辨率目标 ("2K" 或 "4K",会自动映射为 2048/4096)
|
||||
轮询间隔: 轮询秒数
|
||||
seed: 随机种子
|
||||
|
||||
Returns:
|
||||
输出图像 tensor (IMAGE,)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 初始化客户端
|
||||
if self.client is None:
|
||||
self.client = FluxEditClient()
|
||||
|
||||
# Tensor → PIL(取第一张)
|
||||
main_pils = tensor_to_pil(主图)
|
||||
ref_pils = tensor_to_pil(参考图)
|
||||
|
||||
if not main_pils:
|
||||
raise ValueError("主图不能为空")
|
||||
if not ref_pils:
|
||||
raise ValueError("参考图不能为空")
|
||||
|
||||
main_img = main_pils[0]
|
||||
ref_img = ref_pils[0]
|
||||
|
||||
# PIL → JPEG bytes
|
||||
main_bytes = self._image_to_jpeg_bytes(main_img)
|
||||
ref_bytes = self._image_to_jpeg_bytes(ref_img)
|
||||
|
||||
print(f"Flux Edit: 开始处理 | 主图 {main_img.size} | 参考图 {ref_img.size} | 分辨率 {分辨率} | seed {seed}")
|
||||
|
||||
# 进度回调
|
||||
def progress_callback(status_str: str):
|
||||
print(f"Flux Edit: {status_str}")
|
||||
|
||||
# 提交任务并等待结果
|
||||
result_bytes = self.client.submit_and_wait(
|
||||
image_bytes=main_bytes,
|
||||
mask_bytes=ref_bytes,
|
||||
prompt=提示词,
|
||||
size=分辨率,
|
||||
poll_interval=轮询间隔,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
# 解码结果
|
||||
result_img = Image.open(BytesIO(result_bytes))
|
||||
if result_img.mode != "RGB":
|
||||
result_img = result_img.convert("RGB")
|
||||
|
||||
print(f"Flux Edit: 结果图像尺寸 {result_img.size}")
|
||||
|
||||
# 转为 tensor
|
||||
output_tensor = pil_to_tensor([result_img])
|
||||
|
||||
# 打印耗时
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed < 60:
|
||||
time_str = f"{elapsed:.1f}s"
|
||||
else:
|
||||
minutes = int(elapsed // 60)
|
||||
seconds = elapsed % 60
|
||||
time_str = f"{minutes}m {seconds:.0f}s"
|
||||
print(f"Flux Edit: 完成!总耗时 {time_str}")
|
||||
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as e:
|
||||
if str(e) == "未授权!":
|
||||
print("请联系作者授权后方可使用!")
|
||||
raise ValueError("未授权!") from None
|
||||
print(f"Flux Edit: ❌ {e}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"Flux Edit: ❌ {error_msg}")
|
||||
raise RuntimeError(error_msg) from None
|
||||
+416
-69
@@ -6,29 +6,49 @@ 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/mov",
|
||||
".avi": "video/avi",
|
||||
".mov": "video/quicktime",
|
||||
".avi": "video/x-msvideo",
|
||||
".flv": "video/x-flv",
|
||||
".webm": "video/webm",
|
||||
".wmv": "video/wmv",
|
||||
".wmv": "video/x-ms-wmv",
|
||||
".3gp": "video/3gpp",
|
||||
".3gpp": "video/3gpp"
|
||||
}
|
||||
|
||||
try:
|
||||
import subprocess
|
||||
FFMPEG_AVAILABLE = True
|
||||
except ImportError:
|
||||
FFMPEG_AVAILABLE = False
|
||||
|
||||
|
||||
class GoogleGemini:
|
||||
"""
|
||||
@@ -36,14 +56,13 @@ class GoogleGemini:
|
||||
|
||||
功能:
|
||||
- 支持多个 Gemini Flash 模型
|
||||
- 支持图片和视频输入
|
||||
- 支持系统指令
|
||||
- 支持不同思考深度(不思考/高)
|
||||
- 输出生成的文本内容
|
||||
- 支持图片、视频和文件输入
|
||||
- 支持不同思考等级(不思考/低/中/高)- 通过 thinkingConfig.thinkingLevel 控制
|
||||
- 输出生成的文本内容(主要内容 + 思考内容)
|
||||
"""
|
||||
|
||||
# 支持的思考深度选项
|
||||
THINKING_DEPTHS = ["不思考", "高"]
|
||||
# 支持的思考等级选项
|
||||
THINKING_LEVELS = ["不思考", "低", "中", "高"]
|
||||
|
||||
def __init__(self):
|
||||
"""初始化节点"""
|
||||
@@ -67,23 +86,20 @@ class GoogleGemini:
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"思考深度": (cls.THINKING_DEPTHS, {
|
||||
"思考等级": (cls.THINKING_LEVELS, {
|
||||
"default": "不思考"
|
||||
})
|
||||
},
|
||||
"optional": {
|
||||
"系统指令": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"图片": ("IMAGE",),
|
||||
"视频": ("VIDEO",)
|
||||
"视频": ("VIDEO",),
|
||||
"文件": ("FILE",)
|
||||
}
|
||||
}
|
||||
|
||||
# 返回值类型
|
||||
RETURN_TYPES = ("STRING", "STRING")
|
||||
RETURN_NAMES = ("主要内容", "思考内容")
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("主要内容",)
|
||||
|
||||
# 执行函数名
|
||||
FUNCTION = "generate"
|
||||
@@ -94,6 +110,66 @@ 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]
|
||||
@@ -101,6 +177,8 @@ class GoogleGemini:
|
||||
"""
|
||||
准备图片数据
|
||||
|
||||
如果图片超过20MB,会自动进行缩放和压缩
|
||||
|
||||
Args:
|
||||
images: ComfyUI 图片张量 [B, H, W, C]
|
||||
|
||||
@@ -110,17 +188,183 @@ 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:
|
||||
b64_str = encode_image_to_base64(img)
|
||||
image_data.append({
|
||||
"mime_type": "image/png",
|
||||
"data": b64_str
|
||||
})
|
||||
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"))
|
||||
|
||||
return image_data if image_data else None
|
||||
# 多图总体积控制
|
||||
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 以获得更好的压缩效果,或手动压缩视频。"
|
||||
)
|
||||
|
||||
def _prepare_video_data(
|
||||
self,
|
||||
@@ -131,6 +375,7 @@ class GoogleGemini:
|
||||
|
||||
ComfyUI VIDEO 类型包含视频文件路径信息。
|
||||
读取视频文件并转换为 base64。
|
||||
如果视频超过 20MB,会自动进行压缩。
|
||||
|
||||
Args:
|
||||
video: ComfyUI VIDEO 类型数据
|
||||
@@ -141,17 +386,48 @@ 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")
|
||||
# 字典格式:尝试常见的键名
|
||||
video_path = video.get("video") or video.get("path") or video.get("file") or video.get("filename")
|
||||
# 如果还是找不到,遍历所有键找到有效路径
|
||||
if not video_path:
|
||||
for key, val in video.items():
|
||||
if isinstance(val, str) and os.path.exists(val):
|
||||
video_path = val
|
||||
break
|
||||
elif isinstance(video, str):
|
||||
# 字符串格式:直接作为路径
|
||||
video_path = video
|
||||
elif hasattr(video, "video"):
|
||||
video_path = video.video
|
||||
else:
|
||||
# 对象格式:尝试常见属性
|
||||
# 1. 尝试 __file 属性(VideoFromFile 对象)
|
||||
if hasattr(video, "__file"):
|
||||
video_path = video.__file
|
||||
# 2. 尝试其他常见属性
|
||||
elif hasattr(video, "video"):
|
||||
video_path = video.video
|
||||
elif hasattr(video, "path"):
|
||||
video_path = video.path
|
||||
elif hasattr(video, "filename"):
|
||||
video_path = video.filename
|
||||
# 3. 尝试从 __dict__ 中查找路径(支持私有属性如 _VideoFromFile__file)
|
||||
elif hasattr(video, "__dict__"):
|
||||
for attr_name, attr_value in video.__dict__.items():
|
||||
# 查找字符串类型的属性,且包含 file 或 path 关键字
|
||||
if isinstance(attr_value, str):
|
||||
if "file" in attr_name.lower() or "path" in attr_name.lower():
|
||||
# 验证路径是否有效
|
||||
if os.path.exists(attr_value):
|
||||
video_path = attr_value
|
||||
break
|
||||
# 如果属性值本身看起来像文件路径,也尝试使用
|
||||
elif os.path.exists(attr_value) and os.path.isfile(attr_value):
|
||||
video_path = attr_value
|
||||
break
|
||||
|
||||
if not video_path or not os.path.exists(video_path):
|
||||
print(f"Google Gemini: 视频文件不存在或路径无效: {video_path}")
|
||||
@@ -163,30 +439,70 @@ 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:
|
||||
print(f"Google Gemini: 读取视频文件失败 - {str(e)}")
|
||||
# 清理临时文件
|
||||
if temp_compressed_path and os.path.exists(temp_compressed_path):
|
||||
try:
|
||||
os.remove(temp_compressed_path)
|
||||
except:
|
||||
pass
|
||||
|
||||
print(f"Google Gemini: 处理视频文件失败 - {str(e)}")
|
||||
return None
|
||||
|
||||
def _prepare_file_data(
|
||||
self,
|
||||
file: Optional[FileData]
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
准备文件数据
|
||||
|
||||
从 FILE 类型提取文件数据
|
||||
|
||||
Args:
|
||||
file: FileData 对象(来自 LoadFile 节点)
|
||||
|
||||
Returns:
|
||||
文件数据字典,包含 mime_type 和 data
|
||||
"""
|
||||
if file is None:
|
||||
return None
|
||||
|
||||
return {
|
||||
"mime_type": file.mime_type,
|
||||
"data": file.data
|
||||
}
|
||||
|
||||
def _parse_dual_output(self, raw_response: Dict) -> Tuple[str, str]:
|
||||
"""
|
||||
解析包含思考内容和主要内容的响应
|
||||
@@ -214,16 +530,16 @@ class GoogleGemini:
|
||||
# 主要内容
|
||||
main_text = part.get("text", "")
|
||||
|
||||
return (main_text, thought_text)
|
||||
return main_text
|
||||
|
||||
def generate(
|
||||
self,
|
||||
模型: str,
|
||||
提示词: str,
|
||||
思考深度: str,
|
||||
系统指令: Optional[str] = None,
|
||||
思考等级: str,
|
||||
图片: Optional[torch.Tensor] = None,
|
||||
视频=None
|
||||
视频=None,
|
||||
文件: Optional[FileData] = None
|
||||
) -> Tuple[str]:
|
||||
"""
|
||||
生成文本
|
||||
@@ -231,13 +547,13 @@ class GoogleGemini:
|
||||
Args:
|
||||
模型: 使用的模型名称
|
||||
提示词: 用户提示词
|
||||
思考深度: 思考深度选项
|
||||
系统指令: 系统级指令
|
||||
思考等级: 思考等级选项
|
||||
图片: 输入图片
|
||||
视频: 输入视频
|
||||
文件: 输入文件(PDF/TXT)
|
||||
|
||||
Returns:
|
||||
生成的文本 (STRING,)
|
||||
(主要内容, 思考内容)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
@@ -259,6 +575,12 @@ 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 提示词:
|
||||
@@ -267,31 +589,32 @@ 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=模型, thinking_depth=思考深度)
|
||||
endpoint = self.client.get_endpoint(model=模型)
|
||||
request_body = self.client.build_request_body(
|
||||
prompt=提示词,
|
||||
system_instruction=系统指令,
|
||||
model=模型,
|
||||
thinking_level=思考等级,
|
||||
image_data=image_data,
|
||||
video_data=video_data
|
||||
video_data=video_data,
|
||||
document_data=document_data
|
||||
)
|
||||
|
||||
# 根据是否有视频设置超时
|
||||
timeout = 300 if video_data else 180
|
||||
print(f"Google Gemini: 发送请求...")
|
||||
|
||||
# 调用底层 API 获取原始响应
|
||||
async def get_raw_response():
|
||||
return await self.client.request_async(
|
||||
endpoint,
|
||||
request_body,
|
||||
session=None,
|
||||
timeout=timeout
|
||||
session=None
|
||||
)
|
||||
|
||||
# 在独立线程中执行异步请求
|
||||
@@ -301,32 +624,56 @@ class GoogleGemini:
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# 解析响应,分离主要内容和思考内容
|
||||
main_text, thought_text = self._parse_dual_output(raw_response)
|
||||
main_text = self._parse_dual_output(raw_response)
|
||||
|
||||
# 打印响应 token 用量
|
||||
usage = raw_response.get("usageMetadata", {})
|
||||
prompt_tokens = usage.get("promptTokenCount", 0)
|
||||
candidates_tokens = usage.get("candidatesTokenCount", 0)
|
||||
thoughts_tokens = usage.get("thoughtsTokenCount", 0)
|
||||
total_tokens = usage.get("totalTokenCount", 0)
|
||||
finish_reason = ""
|
||||
candidates = raw_response.get("candidates", [])
|
||||
if candidates:
|
||||
finish_reason = candidates[0].get("finishReason", "")
|
||||
|
||||
# 输出信息
|
||||
print(f"Google Gemini: 生成完成 (耗时: {elapsed:.2f}s)")
|
||||
print(f"Google Gemini: finishReason = {finish_reason}")
|
||||
print(f"Google Gemini: Token 用量 — 输入: {prompt_tokens}, 输出: {candidates_tokens}, 思考: {thoughts_tokens}, 合计: {total_tokens}")
|
||||
print(f"Google Gemini: 主要内容长度: {len(main_text)} 字符")
|
||||
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, thought_text)
|
||||
|
||||
except ValueError as e:
|
||||
# 检测是否为授权错误
|
||||
return (main_text,)
|
||||
if str(e) == "未授权!":
|
||||
print("请联系作者授权后方可使用!")
|
||||
raise ValueError("未授权!") from None
|
||||
else:
|
||||
print(f"Google Gemini: 输入错误 - {str(e)}")
|
||||
raise
|
||||
# 用户输入错误 - 只显示简洁信息
|
||||
error_msg = str(e).split('\n')[0] # 只取第一行
|
||||
print(f"Google Gemini: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
print(f"Google Gemini: API 错误 - {str(e)}")
|
||||
raise
|
||||
# 日志只打第一行;报错框展示完整多行
|
||||
error_full = str(e)
|
||||
print(f"Google Gemini: ❌ {error_full.split('\n')[0]}")
|
||||
raise RuntimeError(error_full) from None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Google Gemini: 未知错误 - {str(e)}")
|
||||
raise
|
||||
# 其他未知错误 - 只显示简洁信息
|
||||
error_msg = str(e).split('\n')[0]
|
||||
print(f"Google Gemini: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
finally:
|
||||
if self.client is not None:
|
||||
try:
|
||||
balance_data = self.client.query_balance_sync()
|
||||
balance_info = self.client.format_balance_info(balance_data)
|
||||
print(f"Google Gemini: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,824 @@
|
||||
"""
|
||||
o1key GPT Image 节点
|
||||
支持 gpt-image-1 / gpt-image-1.5 模型的文生图、图生图、图像编辑(带蒙版)
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from ..clients.gpt_image_client import GptImageClient
|
||||
from ..utils.image_utils import parse_batch_prompts, pil_to_tensor, tensor_to_pil
|
||||
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.file_utils import (
|
||||
ImageInfo,
|
||||
generate_timestamp_filename,
|
||||
load_images_from_folder,
|
||||
pair_images_by_name,
|
||||
pair_images_cartesian,
|
||||
save_image,
|
||||
)
|
||||
|
||||
try:
|
||||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||||
_INTERRUPT_AVAILABLE = True
|
||||
except ImportError:
|
||||
_INTERRUPT_AVAILABLE = False
|
||||
processing_interrupted = lambda: False
|
||||
InterruptProcessingException = RuntimeError
|
||||
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
_PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
_PROGRESS_BAR_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
_FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
_FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
|
||||
def _make_node_progress_callback(progress_bar, task_index: int, total_tasks: int):
|
||||
if progress_bar is None:
|
||||
return None
|
||||
|
||||
total_units = max(1, total_tasks) * 100
|
||||
base_units = max(0, task_index - 1) * 100
|
||||
last_pct = {"value": -1}
|
||||
|
||||
def _callback(pct: int):
|
||||
try:
|
||||
pct_value = int(round(float(pct)))
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
pct_value = max(0, min(100, pct_value))
|
||||
if pct_value < last_pct["value"]:
|
||||
return
|
||||
last_pct["value"] = pct_value
|
||||
progress_bar.update_absolute(
|
||||
min(total_units, base_units + pct_value),
|
||||
total_units,
|
||||
)
|
||||
|
||||
return _callback
|
||||
|
||||
|
||||
def _resolve_async_size(value: str) -> str:
|
||||
value = (value or "").strip()
|
||||
if not value or value == "智能" or value.lower() == "auto":
|
||||
return "auto"
|
||||
|
||||
first_part = value.split("(")[0].strip()
|
||||
normalized_size = first_part.lower().replace("*", "x").replace("×", "x")
|
||||
size_parts = [part.strip() for part in normalized_size.split("x")]
|
||||
if len(size_parts) == 2 and all(part.isdigit() for part in size_parts):
|
||||
return f"{int(size_parts[0])}x{int(size_parts[1])}"
|
||||
|
||||
allowed = {"auto", "1024x1024", "1K", "2K", "4K"}
|
||||
if first_part in allowed:
|
||||
return first_part
|
||||
|
||||
if "4K" in value:
|
||||
return "4K"
|
||||
if "2K" in value:
|
||||
return "2K"
|
||||
if "1K" in value:
|
||||
return "1K"
|
||||
|
||||
return "auto"
|
||||
|
||||
|
||||
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["网络"] = (NETWORK_ROUTE_OPTIONS, {
|
||||
"default": "全球加速",
|
||||
})
|
||||
optional_inputs["分辨率"] = ([
|
||||
"智能",
|
||||
# ── 1K ──
|
||||
"1024x1024(1K 正方形 1:1)",
|
||||
"1536x1024(1K 横版 3:2)",
|
||||
"1024x1536(1K 竖版 2:3)",
|
||||
"1360x1024(1K 横版 4:3)",
|
||||
"1024x1360(1K 竖版 3:4)",
|
||||
"1824x1024(1K 横版 16:9)",
|
||||
"1024x1824(1K 竖版 9:16)",
|
||||
# ── 2K ──
|
||||
"2048x2048(2K 正方形 1:1)",
|
||||
"3072x2048(2K 横版 3:2)",
|
||||
"2048x3072(2K 竖版 2:3)",
|
||||
"2736x2048(2K 横版 4:3)",
|
||||
"2048x2736(2K 竖版 3:4)",
|
||||
"3648x2048(2K 横版 16:9)",
|
||||
"2048x3648(2K 竖版 9:16)",
|
||||
# ── 4K ──
|
||||
"2880x2880(4K 正方形 1:1)",
|
||||
"3504x2336(4K 横版 3:2)",
|
||||
"2336x3504(4K 竖版 2:3)",
|
||||
"3264x2448(4K 横版 4:3)",
|
||||
"2448x3264(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["输出格式"] = (["png", "jpeg", "webp"], {
|
||||
"default": "jpeg",
|
||||
"tooltip": "Generated image output format",
|
||||
})
|
||||
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 = "全球加速",
|
||||
分辨率: str = "auto",
|
||||
质量: str = "自动",
|
||||
输出格式: str = "jpeg",
|
||||
生图数量: 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 = _resolve_async_size(分辨率)
|
||||
|
||||
# ── 2b. 解析模型显示值 → API 参数值 ───────────────────────────────────
|
||||
_model_map = {"gpt-image-2-次卡": "gpt-image-2-c", "gpt-image-2-按量": "gpt-image-2"}
|
||||
model = _model_map.get(模型, 模型)
|
||||
|
||||
# ── 2c. 解析质量显示值 → API 参数值 ───────────────────────────────────
|
||||
_quality_map = {"高": "high", "中": "medium", "低": "low", "自动": "auto"}
|
||||
quality = _quality_map.get(质量, "auto")
|
||||
|
||||
# ── 3. 创建客户端 ─────────────────────────────────────────────────────
|
||||
try:
|
||||
client = GptImageClient()
|
||||
client.base_url = get_base_url_by_route(网络)
|
||||
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 = []
|
||||
progress_total = len(batch_prompts) if batch_prompts else 1
|
||||
progress_bar = ProgressBar(progress_total * 100) if _PROGRESS_BAR_AVAILABLE else None
|
||||
|
||||
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.generate_image_async_sync(
|
||||
prompt=p,
|
||||
model=model,
|
||||
quality=quality,
|
||||
size=size,
|
||||
n=生图数量,
|
||||
seed=seed,
|
||||
image_tensor=图片,
|
||||
mask_tensor=遮罩,
|
||||
output_format=输出格式,
|
||||
progress_callback=_make_node_progress_callback(progress_bar, idx, total),
|
||||
)
|
||||
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}")
|
||||
if progress_bar is not None:
|
||||
progress_bar.update_absolute(idx * 100, total * 100)
|
||||
else:
|
||||
# 单提示词模式
|
||||
if not prompt or not prompt.strip():
|
||||
raise ValueError("提示词不能为空")
|
||||
try:
|
||||
pil_images = client.generate_image_async_sync(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
quality=quality,
|
||||
size=size,
|
||||
n=生图数量,
|
||||
seed=seed,
|
||||
image_tensor=图片,
|
||||
mask_tensor=遮罩,
|
||||
output_format=输出格式,
|
||||
progress_callback=_make_node_progress_callback(progress_bar, 1, 1),
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
class O1keyGPTImageBatch:
|
||||
"""
|
||||
o1key GPT Image 批量节点
|
||||
|
||||
复用 BatchNanoBananaPro 的批量思路:
|
||||
- 从文件夹批量加载图片
|
||||
- 按文件名同名 / 1*N / 不配对 三种模式创建任务
|
||||
- 可追加节点手动输入参考图
|
||||
- prompt 支持用独占一行 --- 展开为多提示词任务
|
||||
- 每个任务调用 GPT Image 客户端并保存到磁盘
|
||||
"""
|
||||
|
||||
PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"]
|
||||
IMAGE_FORMATS = ["原始", "JPEG", "PNG", "WebP"]
|
||||
MODEL_OPTIONS = ["gpt-image-2-按量", "gpt-image-2-次卡"]
|
||||
QUALITY_OPTIONS = ["高", "中", "低", "自动"]
|
||||
RESOLUTION_OPTIONS = [
|
||||
"智能",
|
||||
"1024x1024(1K 正方形 1:1)",
|
||||
"1536x1024(1K 横版 3:2)",
|
||||
"1024x1536(1K 竖版 2:3)",
|
||||
"1360x1024(1K 横版 4:3)",
|
||||
"1024x1360(1K 竖版 3:4)",
|
||||
"1824x1024(1K 横版 16:9)",
|
||||
"1024x1824(1K 竖版 9:16)",
|
||||
"2048x2048(2K 正方形 1:1)",
|
||||
"3072x2048(2K 横版 3:2)",
|
||||
"2048x3072(2K 竖版 2:3)",
|
||||
"2736x2048(2K 横版 4:3)",
|
||||
"2048x2736(2K 竖版 3:4)",
|
||||
"3648x2048(2K 横版 16:9)",
|
||||
"2048x3648(2K 竖版 9:16)",
|
||||
"2880x2880(4K 正方形 1:1)",
|
||||
"3504x2336(4K 横版 3:2)",
|
||||
"2336x3504(4K 竖版 2:3)",
|
||||
"3264x2448(4K 横版 4:3)",
|
||||
"2448x3264(4K 竖版 3:4)",
|
||||
"3840x2160(4K 横版 16:9)",
|
||||
"2160x3840(4K 竖版 9:16)",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
optional_inputs = {}
|
||||
for image_index in range(1, 10):
|
||||
optional_inputs[f"参考图{image_index}"] = ("IMAGE", {
|
||||
"tooltip": "追加到每个批量任务末尾的固定参考图。",
|
||||
})
|
||||
|
||||
optional_inputs["遮罩"] = ("MASK", {
|
||||
"tooltip": "可选蒙版,会应用到每个任务的第一张参考图;请确保尺寸一致。",
|
||||
})
|
||||
optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, {
|
||||
"default": "不配对",
|
||||
"tooltip": "文件夹图片的组合方式;手动参考图只追加,不参与配对。",
|
||||
})
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "提示词;可用独占一行的 --- 分隔多条批量提示词。",
|
||||
}),
|
||||
"模型": (cls.MODEL_OPTIONS, {
|
||||
"default": "gpt-image-2-次卡",
|
||||
}),
|
||||
"网络": (NETWORK_ROUTE_OPTIONS, {
|
||||
"default": "全球加速",
|
||||
}),
|
||||
"分辨率": (cls.RESOLUTION_OPTIONS, {
|
||||
"default": "智能",
|
||||
}),
|
||||
"生图数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 8,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
}),
|
||||
"质量": (cls.QUALITY_OPTIONS, {
|
||||
"default": "自动",
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 2**31 - 1,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"control_after_generate": True,
|
||||
}),
|
||||
"图片格式": (cls.IMAGE_FORMATS, {
|
||||
"default": "原始",
|
||||
}),
|
||||
"文件夹1": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
}),
|
||||
"文件夹2": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
}),
|
||||
"文件夹3": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
}),
|
||||
"文件夹4": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
}),
|
||||
"文件夹5": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
}),
|
||||
"保存路径": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "为空时优先使用 ComfyUI 默认 output 目录。",
|
||||
}),
|
||||
},
|
||||
"optional": optional_inputs,
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("IMAGE",)
|
||||
FUNCTION = "process_batch"
|
||||
CATEGORY = "o1key/image"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def _load_folders(self, folders: List[str]) -> List[List[ImageInfo]]:
|
||||
image_lists = []
|
||||
for folder_index, folder in enumerate(folders, 1):
|
||||
if not folder or not folder.strip():
|
||||
continue
|
||||
try:
|
||||
loaded_images = load_images_from_folder(folder)
|
||||
if loaded_images:
|
||||
image_lists.append(loaded_images)
|
||||
except ValueError as error:
|
||||
print(f"[o1key GPT Image Batch] 文件夹{folder_index} 加载失败 - {error}")
|
||||
return image_lists
|
||||
|
||||
def _create_pairs(
|
||||
self,
|
||||
image_lists: List[List[ImageInfo]],
|
||||
pairing_mode: str,
|
||||
manual_images: Optional[List[ImageInfo]] = None,
|
||||
) -> List[Tuple[ImageInfo, ...]]:
|
||||
if pairing_mode == "不配对":
|
||||
if len(image_lists) > 1:
|
||||
raise ValueError("「不配对」模式只支持单个文件夹,请清空其他文件夹路径")
|
||||
|
||||
if image_lists and manual_images:
|
||||
return [
|
||||
(folder_image,) + tuple(manual_images)
|
||||
for folder_image in image_lists[0]
|
||||
]
|
||||
if image_lists:
|
||||
return [(folder_image,) for folder_image in image_lists[0]]
|
||||
return []
|
||||
|
||||
if not image_lists:
|
||||
return []
|
||||
|
||||
if len(image_lists) == 1:
|
||||
base_pairs = [(folder_image,) for folder_image in image_lists[0]]
|
||||
elif pairing_mode == "按相同图片命名":
|
||||
base_pairs = list(pair_images_by_name(*image_lists))
|
||||
else:
|
||||
base_pairs = list(pair_images_cartesian(*image_lists))
|
||||
|
||||
if manual_images:
|
||||
manual_tuple = tuple(manual_images)
|
||||
base_pairs = [pair + manual_tuple for pair in base_pairs]
|
||||
|
||||
return base_pairs
|
||||
|
||||
def _collect_manual_images(self, kwargs) -> List[ImageInfo]:
|
||||
manual_images = []
|
||||
for image_index in range(1, 10):
|
||||
key = f"参考图{image_index}"
|
||||
if key not in kwargs or kwargs[key] is None:
|
||||
continue
|
||||
for tensor_index, image in enumerate(tensor_to_pil(kwargs[key])):
|
||||
manual_images.append(ImageInfo(
|
||||
image=image,
|
||||
filename=f"manual_{image_index}_{tensor_index}",
|
||||
extension=".png",
|
||||
source_path="",
|
||||
))
|
||||
return manual_images
|
||||
|
||||
@staticmethod
|
||||
def _pair_to_tensors(pair: Tuple[ImageInfo, ...]) -> List:
|
||||
return [pil_to_tensor([image_info.image]) for image_info in pair]
|
||||
|
||||
@staticmethod
|
||||
def _resolve_size(分辨率: str) -> str:
|
||||
return _resolve_async_size(分辨率)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_model(模型: str) -> str:
|
||||
model_map = {
|
||||
"gpt-image-2-次卡": "gpt-image-2-c",
|
||||
"gpt-image-2-按量": "gpt-image-2",
|
||||
}
|
||||
return model_map.get(模型, 模型)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_quality(质量: str) -> str:
|
||||
quality_map = {"高": "high", "中": "medium", "低": "low", "自动": "auto"}
|
||||
return quality_map.get(质量, "auto")
|
||||
|
||||
@staticmethod
|
||||
def _resolve_output_format(图片格式: str) -> str:
|
||||
output_format_map = {
|
||||
"JPEG": "jpeg",
|
||||
"PNG": "png",
|
||||
"WebP": "webp",
|
||||
}
|
||||
return output_format_map.get(图片格式, "png")
|
||||
|
||||
@staticmethod
|
||||
def _ensure_output_folder(保存路径: str) -> str:
|
||||
output_folder = (保存路径 or "").strip()
|
||||
if not output_folder and _FOLDER_PATHS_AVAILABLE:
|
||||
output_folder = folder_paths.get_output_directory()
|
||||
print(f"[o1key GPT Image Batch] 未设置保存路径,使用 ComfyUI 默认 output 目录: {output_folder}")
|
||||
|
||||
if not output_folder:
|
||||
raise ValueError("未设置保存路径,且当前环境无法获取 ComfyUI 默认 output 目录")
|
||||
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
test_path = os.path.join(output_folder, ".write_test")
|
||||
with open(test_path, "w", encoding="utf-8") as test_file:
|
||||
test_file.write("test")
|
||||
os.remove(test_path)
|
||||
return output_folder
|
||||
|
||||
@staticmethod
|
||||
def _save_images(
|
||||
images: List[Image.Image],
|
||||
output_folder: str,
|
||||
image_format: str,
|
||||
base_filename: Optional[str] = None,
|
||||
) -> List[str]:
|
||||
format_ext_map = {"JPEG": ".jpg", "PNG": ".png", "WebP": ".webp"}
|
||||
save_ext = format_ext_map.get(image_format, ".png")
|
||||
saved_files = []
|
||||
|
||||
for image in images:
|
||||
if base_filename:
|
||||
counter = 0
|
||||
while True:
|
||||
suffix = "" if counter == 0 else f"+{counter}"
|
||||
filename = f"{base_filename}{suffix}{save_ext}"
|
||||
output_path = os.path.join(output_folder, filename)
|
||||
if not os.path.exists(output_path):
|
||||
break
|
||||
counter += 1
|
||||
else:
|
||||
output_path = generate_timestamp_filename(
|
||||
output_folder=output_folder,
|
||||
extension=save_ext,
|
||||
)
|
||||
|
||||
if image_format == "JPEG":
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
image.save(output_path, quality=100)
|
||||
elif image_format == "WebP":
|
||||
image.save(output_path, lossless=True)
|
||||
else:
|
||||
save_image(image, output_path)
|
||||
|
||||
saved_files.append(output_path)
|
||||
|
||||
return saved_files
|
||||
|
||||
def process_batch(
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str,
|
||||
网络: str,
|
||||
分辨率: str,
|
||||
生图数量: int,
|
||||
质量: str,
|
||||
seed: int,
|
||||
图片格式: str,
|
||||
文件夹1: str,
|
||||
文件夹2: str,
|
||||
文件夹3: str,
|
||||
文件夹4: str,
|
||||
文件夹5: str,
|
||||
保存路径: str = "",
|
||||
图片配对模式: str = "不配对",
|
||||
遮罩=None,
|
||||
**kwargs,
|
||||
):
|
||||
start_time = time.time()
|
||||
client = None
|
||||
|
||||
try:
|
||||
if not prompt or not prompt.strip():
|
||||
raise ValueError("提示词不能为空")
|
||||
|
||||
folders = [文件夹1, 文件夹2, 文件夹3, 文件夹4, 文件夹5]
|
||||
if not any(folder and folder.strip() for folder in folders):
|
||||
raise ValueError("请至少填写一个文件夹路径,该节点专为批量文件夹处理设计")
|
||||
|
||||
image_lists = self._load_folders(folders)
|
||||
total_folder_images = sum(len(image_list) for image_list in image_lists)
|
||||
if total_folder_images == 0:
|
||||
raise ValueError("文件夹中未找到任何图片,请检查文件夹路径是否正确")
|
||||
|
||||
manual_images = self._collect_manual_images(kwargs)
|
||||
pairs = self._create_pairs(
|
||||
image_lists=image_lists,
|
||||
pairing_mode=图片配对模式,
|
||||
manual_images=manual_images if manual_images else None,
|
||||
)
|
||||
if not pairs:
|
||||
raise ValueError("配对结果为空,请检查输入")
|
||||
|
||||
batch_prompts = parse_batch_prompts(prompt)
|
||||
prompts_per_task = None
|
||||
if batch_prompts:
|
||||
expanded_pairs = []
|
||||
expanded_prompts = []
|
||||
for pair in pairs:
|
||||
for batch_prompt in batch_prompts:
|
||||
expanded_pairs.append(pair)
|
||||
expanded_prompts.append(batch_prompt)
|
||||
pairs = expanded_pairs
|
||||
prompts_per_task = expanded_prompts
|
||||
|
||||
total_tasks = len(pairs)
|
||||
if batch_prompts:
|
||||
print(
|
||||
f"[o1key GPT Image Batch] 批量任务 | {图片配对模式} × "
|
||||
f"{len(batch_prompts)} 个提示词 | 共 {total_tasks} 任务"
|
||||
)
|
||||
else:
|
||||
print(f"[o1key GPT Image Batch] 批量任务 | {图片配对模式} | 共 {total_tasks} 任务")
|
||||
|
||||
output_folder = self._ensure_output_folder(保存路径)
|
||||
size = self._resolve_size(分辨率)
|
||||
model = self._resolve_model(模型)
|
||||
quality = self._resolve_quality(质量)
|
||||
output_format = self._resolve_output_format(图片格式)
|
||||
|
||||
client = GptImageClient()
|
||||
client.base_url = get_base_url_by_route(网络)
|
||||
|
||||
progress_bar = ProgressBar(total_tasks * 100) if _PROGRESS_BAR_AVAILABLE else None
|
||||
results = []
|
||||
all_saved_files = []
|
||||
|
||||
for task_index, pair in enumerate(pairs, 1):
|
||||
if _INTERRUPT_AVAILABLE and processing_interrupted():
|
||||
print("[o1key GPT Image Batch] 用户取消,已中断批量生成")
|
||||
raise InterruptProcessingException()
|
||||
|
||||
task_prompt = prompts_per_task[task_index - 1] if prompts_per_task else prompt
|
||||
base_filename = pair[0].filename if pair else None
|
||||
result = {
|
||||
"task_index": task_index,
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"saved_files": [],
|
||||
"error": None,
|
||||
}
|
||||
|
||||
try:
|
||||
pil_images = client.generate_image_async_sync(
|
||||
prompt=task_prompt,
|
||||
model=model,
|
||||
quality=quality,
|
||||
size=size,
|
||||
n=生图数量,
|
||||
seed=seed,
|
||||
image_tensor=self._pair_to_tensors(pair),
|
||||
mask_tensor=遮罩,
|
||||
output_format=output_format,
|
||||
progress_callback=_make_node_progress_callback(progress_bar, task_index, total_tasks),
|
||||
)
|
||||
saved_files = self._save_images(
|
||||
images=pil_images,
|
||||
output_folder=output_folder,
|
||||
image_format=图片格式,
|
||||
base_filename=base_filename,
|
||||
)
|
||||
result["success"] = bool(pil_images)
|
||||
result["generated_count"] = len(pil_images)
|
||||
result["saved_files"] = saved_files
|
||||
all_saved_files.extend(saved_files)
|
||||
print(f"[o1key GPT Image Batch] [{task_index}/{total_tasks}] ✓ {base_filename or 'task'}")
|
||||
except InterruptProcessingException:
|
||||
raise
|
||||
except Exception as error:
|
||||
error_msg = str(error).split("\n")[0]
|
||||
result["error"] = error_msg
|
||||
print(f"[o1key GPT Image Batch] [{task_index}/{total_tasks}] ❌ {base_filename or 'task'} → {error_msg}")
|
||||
|
||||
results.append(result)
|
||||
if progress_bar is not None:
|
||||
progress_bar.update_absolute(task_index * 100, total_tasks * 100)
|
||||
|
||||
success_count = sum(1 for result in results if result.get("success", False))
|
||||
total_generated = sum(result.get("generated_count", 0) for result in results)
|
||||
if success_count == 0:
|
||||
raise RuntimeError("所有批量任务均生成失败,无可用图像输出")
|
||||
|
||||
output_images = []
|
||||
for file_path in all_saved_files[-10:]:
|
||||
try:
|
||||
loaded_image = Image.open(file_path)
|
||||
loaded_image.load()
|
||||
output_images.append(loaded_image)
|
||||
except Exception as error:
|
||||
print(f"[o1key GPT Image Batch] 无法加载输出图片 {file_path} - {error}")
|
||||
|
||||
if not output_images:
|
||||
output_images = [Image.new("RGBA", (512, 512), (128, 128, 128, 255))]
|
||||
|
||||
output_tensor = GptImageClient._pil_list_to_tensor(output_images)
|
||||
elapsed = time.time() - start_time
|
||||
print("=" * 60)
|
||||
print(
|
||||
f"[o1key GPT Image Batch] 完成!耗时 {elapsed:.1f}s | "
|
||||
f"成功 {success_count}/{total_tasks} | 生成 {total_generated} 张"
|
||||
)
|
||||
print(f"[o1key GPT Image Batch] 保存路径: {output_folder}")
|
||||
if all_saved_files:
|
||||
print(f"[o1key GPT Image Batch] 最新保存文件: {all_saved_files[-1]}")
|
||||
|
||||
failed_results = [result for result in results if not result.get("success", False)]
|
||||
if failed_results:
|
||||
print(f"[o1key GPT Image Batch] 失败任务: {len(failed_results)} 个")
|
||||
for failed_result in failed_results[:3]:
|
||||
print(
|
||||
f" - #{failed_result.get('task_index')}: "
|
||||
f"{failed_result.get('error', '未知错误')}"
|
||||
)
|
||||
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as error:
|
||||
if str(error) == "未授权!":
|
||||
print("[o1key GPT Image Batch] 请联系作者授权后方可使用!")
|
||||
raise ValueError("未授权!") from None
|
||||
raise ValueError(str(error)) from None
|
||||
except RuntimeError as error:
|
||||
raise RuntimeError(str(error)) from None
|
||||
finally:
|
||||
if client is not None:
|
||||
try:
|
||||
balance_data = client.query_balance_sync()
|
||||
balance_info = client.format_balance_info(balance_data)
|
||||
print(f"[o1key GPT Image Batch] {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,398 @@
|
||||
"""
|
||||
Merged grid image splitter.
|
||||
|
||||
This node is designed for AI-generated contact sheets such as 3x3 or 2x3
|
||||
grids. Auto mode scores common layouts by looking for strong seams or flat
|
||||
separator bands near the expected grid lines, then crops each cell.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Sequence, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import pil_to_tensor, tensor_to_pil
|
||||
|
||||
|
||||
_AUTO_LAYOUTS: Sequence[Tuple[int, int]] = (
|
||||
(3, 3),
|
||||
(2, 3),
|
||||
(3, 2),
|
||||
(2, 2),
|
||||
(1, 2),
|
||||
(2, 1),
|
||||
(1, 3),
|
||||
(3, 1),
|
||||
(4, 4),
|
||||
(3, 4),
|
||||
(4, 3),
|
||||
)
|
||||
|
||||
_LAYOUTS = [
|
||||
"auto",
|
||||
"1x2",
|
||||
"2x1",
|
||||
"1x3",
|
||||
"3x1",
|
||||
"2x2",
|
||||
"2x3",
|
||||
"3x2",
|
||||
"3x3",
|
||||
"3x4",
|
||||
"4x3",
|
||||
"4x4",
|
||||
"custom",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _AxisCut:
|
||||
seam: int
|
||||
span_start: int
|
||||
span_end: int
|
||||
score: float
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _AxisPlan:
|
||||
intervals: List[Tuple[int, int]]
|
||||
cuts: List[_AxisCut]
|
||||
score: float
|
||||
|
||||
|
||||
def _to_float_array(image: Image.Image) -> np.ndarray:
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
return np.asarray(image).astype(np.float32) / 255.0
|
||||
|
||||
|
||||
def _axis_texture(arr: np.ndarray, axis: str) -> np.ndarray:
|
||||
if axis == "x":
|
||||
profile = arr.std(axis=(0, 2))
|
||||
else:
|
||||
profile = arr.std(axis=(1, 2))
|
||||
high = np.percentile(profile, 95) + 1e-6
|
||||
return np.clip(profile / high, 0.0, 1.0)
|
||||
|
||||
|
||||
def _axis_edge(arr: np.ndarray, axis: str) -> np.ndarray:
|
||||
if axis == "x":
|
||||
diff = np.abs(np.diff(arr, axis=1)).mean(axis=(0, 2))
|
||||
length = arr.shape[1]
|
||||
else:
|
||||
diff = np.abs(np.diff(arr, axis=0)).mean(axis=(1, 2))
|
||||
length = arr.shape[0]
|
||||
|
||||
padded = np.zeros(length, dtype=np.float32)
|
||||
if diff.size:
|
||||
padded[1:] = diff
|
||||
high = np.percentile(padded, 95) + 1e-6
|
||||
return np.clip(padded / high, 0.0, 1.5)
|
||||
|
||||
|
||||
def _smooth(profile: np.ndarray, radius: int = 2) -> np.ndarray:
|
||||
if radius <= 0 or profile.size < radius * 2 + 1:
|
||||
return profile
|
||||
kernel = np.ones(radius * 2 + 1, dtype=np.float32) / float(radius * 2 + 1)
|
||||
return np.convolve(profile, kernel, mode="same")
|
||||
|
||||
|
||||
def _separator_span(
|
||||
texture: np.ndarray,
|
||||
seam: int,
|
||||
search_px: int,
|
||||
min_separator_px: int,
|
||||
) -> Tuple[int, int]:
|
||||
length = texture.size
|
||||
if length <= 1:
|
||||
return 0, length
|
||||
|
||||
limit = max(1, min(search_px, length // 8))
|
||||
threshold = max(0.08, min(0.28, float(np.percentile(texture, 12)) * 1.8))
|
||||
|
||||
left = seam
|
||||
while left > 0 and seam - left < limit and texture[left - 1] <= threshold:
|
||||
left -= 1
|
||||
|
||||
right = seam
|
||||
while right < length and right - seam < limit and texture[right] <= threshold:
|
||||
right += 1
|
||||
|
||||
if right - left >= max(1, min_separator_px):
|
||||
return left, right
|
||||
|
||||
return seam, seam
|
||||
|
||||
|
||||
def _edge_trim(texture: np.ndarray, search_px: int, min_cell: int) -> Tuple[int, int]:
|
||||
length = texture.size
|
||||
if length <= 2:
|
||||
return 0, length
|
||||
|
||||
max_trim = max(0, min(search_px * 2, min_cell // 3, length // 6))
|
||||
if max_trim <= 0:
|
||||
return 0, length
|
||||
|
||||
threshold = max(0.08, min(0.24, float(np.percentile(texture, 12)) * 1.6))
|
||||
|
||||
start = 0
|
||||
while start < max_trim and texture[start] <= threshold:
|
||||
start += 1
|
||||
|
||||
end = length
|
||||
while length - end < max_trim and end > start + min_cell and texture[end - 1] <= threshold:
|
||||
end -= 1
|
||||
|
||||
return start, end
|
||||
|
||||
|
||||
def _axis_plan(
|
||||
arr: np.ndarray,
|
||||
cells: int,
|
||||
axis: str,
|
||||
search_px: int,
|
||||
crop_separators: bool,
|
||||
trim_outer: bool,
|
||||
min_separator_px: int,
|
||||
) -> _AxisPlan:
|
||||
length = arr.shape[1] if axis == "x" else arr.shape[0]
|
||||
if cells <= 1:
|
||||
return _AxisPlan(intervals=[(0, length)], cuts=[], score=0.0)
|
||||
|
||||
raw_texture = _axis_texture(arr, axis)
|
||||
raw_edge = _axis_edge(arr, axis)
|
||||
texture = _smooth(raw_texture, radius=2)
|
||||
edge = _smooth(raw_edge, radius=1)
|
||||
evidence = np.maximum(edge, (1.0 - texture) * 0.75)
|
||||
exact_evidence = np.maximum(raw_edge, (1.0 - raw_texture) * 0.75)
|
||||
|
||||
cuts: List[_AxisCut] = []
|
||||
scores: List[float] = []
|
||||
for idx in range(1, cells):
|
||||
expected = round(length * idx / cells)
|
||||
start = max(1, expected - search_px)
|
||||
end = min(length - 1, expected + search_px)
|
||||
if start >= end:
|
||||
seam = expected
|
||||
score = 0.0
|
||||
else:
|
||||
window = evidence[start:end + 1]
|
||||
offset = int(window.argmax())
|
||||
coarse = start + offset
|
||||
fine_start = max(start, coarse - 2)
|
||||
fine_end = min(end, coarse + 2)
|
||||
fine_window = exact_evidence[fine_start:fine_end + 1]
|
||||
seam = fine_start + int(fine_window.argmax())
|
||||
score = float(window[offset])
|
||||
|
||||
span_start, span_end = _separator_span(
|
||||
raw_texture,
|
||||
seam,
|
||||
search_px=search_px,
|
||||
min_separator_px=min_separator_px,
|
||||
)
|
||||
cuts.append(_AxisCut(seam=seam, span_start=span_start, span_end=span_end, score=score))
|
||||
scores.append(score)
|
||||
|
||||
min_cell = max(1, length // cells)
|
||||
outer_start, outer_end = _edge_trim(raw_texture, search_px, min_cell) if trim_outer else (0, length)
|
||||
|
||||
intervals: List[Tuple[int, int]] = []
|
||||
cursor = outer_start
|
||||
for cut in cuts:
|
||||
split_start = cut.span_start if crop_separators else cut.seam
|
||||
split_end = cut.span_end if crop_separators else cut.seam
|
||||
intervals.append((cursor, split_start))
|
||||
cursor = split_end
|
||||
intervals.append((cursor, outer_end))
|
||||
|
||||
cleaned: List[Tuple[int, int]] = []
|
||||
for start, end in intervals:
|
||||
start = max(0, min(length - 1, int(start)))
|
||||
end = max(start + 1, min(length, int(end)))
|
||||
cleaned.append((start, end))
|
||||
|
||||
return _AxisPlan(
|
||||
intervals=cleaned,
|
||||
cuts=cuts,
|
||||
score=float(np.mean(scores)) if scores else 0.0,
|
||||
)
|
||||
|
||||
|
||||
def _parse_layout(layout: str, custom_rows: int, custom_cols: int) -> Tuple[int, int]:
|
||||
if layout == "custom":
|
||||
return max(1, int(custom_rows)), max(1, int(custom_cols))
|
||||
rows_text, cols_text = layout.split("x", 1)
|
||||
return int(rows_text), int(cols_text)
|
||||
|
||||
|
||||
def _fallback_layout(width: int, height: int) -> Tuple[int, int]:
|
||||
aspect = width / max(1, height)
|
||||
if 0.82 <= aspect <= 1.22:
|
||||
return 3, 3
|
||||
if aspect > 1.22:
|
||||
return 2, 3
|
||||
return 3, 2
|
||||
|
||||
|
||||
def _choose_auto_layout(
|
||||
arr: np.ndarray,
|
||||
search_px: int,
|
||||
crop_separators: bool,
|
||||
trim_outer: bool,
|
||||
min_separator_px: int,
|
||||
) -> Tuple[int, int, _AxisPlan, _AxisPlan, float, bool]:
|
||||
height, width = arr.shape[:2]
|
||||
best = None
|
||||
|
||||
for rows, cols in _AUTO_LAYOUTS:
|
||||
x_plan = _axis_plan(arr, cols, "x", search_px, crop_separators, trim_outer, min_separator_px)
|
||||
y_plan = _axis_plan(arr, rows, "y", search_px, crop_separators, trim_outer, min_separator_px)
|
||||
score = (x_plan.score + y_plan.score) / 2.0
|
||||
|
||||
# Prefer common 3x3 / 2x3 / 3x2 layouts when the image gives weak signals.
|
||||
if (rows, cols) in ((3, 3), (2, 3), (3, 2)):
|
||||
score += 0.025
|
||||
|
||||
if best is None or score > best[0]:
|
||||
best = (score, rows, cols, x_plan, y_plan)
|
||||
|
||||
assert best is not None
|
||||
score, rows, cols, x_plan, y_plan = best
|
||||
confident = score >= 0.22
|
||||
|
||||
if confident:
|
||||
return rows, cols, x_plan, y_plan, score, True
|
||||
|
||||
rows, cols = _fallback_layout(width, height)
|
||||
x_plan = _axis_plan(arr, cols, "x", search_px, crop_separators, trim_outer, min_separator_px)
|
||||
y_plan = _axis_plan(arr, rows, "y", search_px, crop_separators, trim_outer, min_separator_px)
|
||||
return rows, cols, x_plan, y_plan, score, False
|
||||
|
||||
|
||||
def _normalize_sizes(crops: List[Image.Image]) -> List[Image.Image]:
|
||||
min_w = min(crop.width for crop in crops)
|
||||
min_h = min(crop.height for crop in crops)
|
||||
normalized = []
|
||||
for crop in crops:
|
||||
left = max(0, (crop.width - min_w) // 2)
|
||||
top = max(0, (crop.height - min_h) // 2)
|
||||
normalized.append(crop.crop((left, top, left + min_w, top + min_h)))
|
||||
return normalized
|
||||
|
||||
|
||||
def _split_one(
|
||||
image: Image.Image,
|
||||
layout: str,
|
||||
custom_rows: int,
|
||||
custom_cols: int,
|
||||
search_px: int,
|
||||
crop_separators: bool,
|
||||
trim_outer: bool,
|
||||
min_separator_px: int,
|
||||
) -> Tuple[List[Image.Image], str]:
|
||||
arr = _to_float_array(image)
|
||||
|
||||
if layout == "auto":
|
||||
rows, cols, x_plan, y_plan, confidence, confident = _choose_auto_layout(
|
||||
arr,
|
||||
search_px=search_px,
|
||||
crop_separators=crop_separators,
|
||||
trim_outer=trim_outer,
|
||||
min_separator_px=min_separator_px,
|
||||
)
|
||||
mode_note = "auto" if confident else "auto-low-confidence-fallback"
|
||||
else:
|
||||
rows, cols = _parse_layout(layout, custom_rows, custom_cols)
|
||||
x_plan = _axis_plan(arr, cols, "x", search_px, crop_separators, trim_outer, min_separator_px)
|
||||
y_plan = _axis_plan(arr, rows, "y", search_px, crop_separators, trim_outer, min_separator_px)
|
||||
confidence = (x_plan.score + y_plan.score) / 2.0
|
||||
mode_note = "manual"
|
||||
|
||||
crops: List[Image.Image] = []
|
||||
for y0, y1 in y_plan.intervals:
|
||||
for x0, x1 in x_plan.intervals:
|
||||
crops.append(image.crop((x0, y0, x1, y1)))
|
||||
|
||||
crops = _normalize_sizes(crops)
|
||||
info = (
|
||||
f"{mode_note}: {rows}x{cols}, cells={len(crops)}, "
|
||||
f"confidence={confidence:.3f}, "
|
||||
f"x={x_plan.intervals}, y={y_plan.intervals}"
|
||||
)
|
||||
return crops, info
|
||||
|
||||
|
||||
class O1keyGridSplitter:
|
||||
"""Split AI-generated grid/contact-sheet images into individual cells."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"图像": ("IMAGE",),
|
||||
"布局": (_LAYOUTS, {"default": "auto"}),
|
||||
"自定义行数": ("INT", {"default": 3, "min": 1, "max": 12, "step": 1}),
|
||||
"自定义列数": ("INT", {"default": 3, "min": 1, "max": 12, "step": 1}),
|
||||
"搜索范围px": ("INT", {"default": 32, "min": 0, "max": 256, "step": 1}),
|
||||
"裁掉分隔线": ("BOOLEAN", {"default": True}),
|
||||
"裁掉外边距": ("BOOLEAN", {"default": True}),
|
||||
"最小分隔线px": ("INT", {"default": 2, "min": 0, "max": 64, "step": 1}),
|
||||
"最大输出张数": ("INT", {"default": 16, "min": 1, "max": 144, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "STRING")
|
||||
RETURN_NAMES = ("切割图像", "检测信息")
|
||||
FUNCTION = "split_grid"
|
||||
CATEGORY = "o1key/image"
|
||||
DESCRIPTION = (
|
||||
"智能切割 AI 生成的九宫格、六宫格等合并图。"
|
||||
"自动模式会检测常见布局;没有明显分隔线时建议手动选择布局。"
|
||||
)
|
||||
|
||||
def split_grid(
|
||||
self,
|
||||
图像: torch.Tensor,
|
||||
布局: str = "auto",
|
||||
自定义行数: int = 3,
|
||||
自定义列数: int = 3,
|
||||
搜索范围px: int = 32,
|
||||
裁掉分隔线: bool = True,
|
||||
裁掉外边距: bool = True,
|
||||
最小分隔线px: int = 2,
|
||||
最大输出张数: int = 16,
|
||||
):
|
||||
source_images = tensor_to_pil(图像)
|
||||
all_crops: List[Image.Image] = []
|
||||
info_lines: List[str] = []
|
||||
|
||||
for batch_index, image in enumerate(source_images, start=1):
|
||||
crops, info = _split_one(
|
||||
image=image,
|
||||
layout=布局,
|
||||
custom_rows=自定义行数,
|
||||
custom_cols=自定义列数,
|
||||
search_px=搜索范围px,
|
||||
crop_separators=裁掉分隔线,
|
||||
trim_outer=裁掉外边距,
|
||||
min_separator_px=最小分隔线px,
|
||||
)
|
||||
if len(crops) > 最大输出张数:
|
||||
raise ValueError(
|
||||
f"合并图切割:检测到 {len(crops)} 张,超过最大输出张数 {最大输出张数}。"
|
||||
"请调大最大输出张数,或检查布局设置。"
|
||||
)
|
||||
all_crops.extend(crops)
|
||||
info_lines.append(f"batch {batch_index}: {info}")
|
||||
|
||||
if not all_crops:
|
||||
raise ValueError("合并图切割:没有生成任何切片。")
|
||||
|
||||
all_crops = _normalize_sizes(all_crops)
|
||||
print("[o1key 合并图切割] " + " | ".join(info_lines))
|
||||
return (pil_to_tensor(all_crops), "\n".join(info_lines))
|
||||
@@ -0,0 +1,165 @@
|
||||
"""
|
||||
o1key Grok Image 节点
|
||||
支持 Grok Image / Grok Image Pro 模型的文生图和图生图
|
||||
"""
|
||||
|
||||
import time
|
||||
from ..clients.grok_image_client import GrokImageClient
|
||||
from ..utils.image_utils import parse_batch_prompts
|
||||
from ..utils.config import NETWORK_ROUTE_OPTIONS
|
||||
|
||||
try:
|
||||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||||
_INTERRUPT_AVAILABLE = True
|
||||
except ImportError:
|
||||
_INTERRUPT_AVAILABLE = False
|
||||
processing_interrupted = lambda: False
|
||||
InterruptProcessingException = RuntimeError
|
||||
|
||||
_ASPECT_RATIOS = [
|
||||
"auto", "1:1", "16:9", "9:16", "4:3", "3:4",
|
||||
"3:2", "2:3", "2:1", "1:2", "19.5:9", "9:19.5", "20:9", "9:20",
|
||||
]
|
||||
|
||||
|
||||
class O1keyGrokImage:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
optional_inputs = {}
|
||||
for i in range(1, 4):
|
||||
optional_inputs[f"参考图{i}"] = ("IMAGE", {
|
||||
"tooltip": f"Optional reference image {i}",
|
||||
})
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "文本提示词,用 --- 独占一行分隔批量提示词",
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"模型": (["Grok Image", "Grok Image Pro"], {
|
||||
"default": "Grok Image Pro",
|
||||
}),
|
||||
"宽高比": (_ASPECT_RATIOS, {
|
||||
"default": "auto",
|
||||
}),
|
||||
"分辨率": (["1k", "2k"], {
|
||||
"default": "1k",
|
||||
}),
|
||||
"生图数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 4,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {
|
||||
"default": NETWORK_ROUTE_OPTIONS[0],
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 2**31 - 1,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"control_after_generate": True,
|
||||
}),
|
||||
**optional_inputs,
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("IMAGE",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "o1key/image"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str = "Grok Image Pro",
|
||||
宽高比: str = "auto",
|
||||
分辨率: str = "1k",
|
||||
生图数量: int = 1,
|
||||
网络线路: str = "全球加速",
|
||||
seed: int = 0,
|
||||
**kwargs,
|
||||
):
|
||||
start_time = time.time()
|
||||
|
||||
reference_tensors = []
|
||||
for i in range(1, 4):
|
||||
key = f"参考图{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
reference_tensors.append(kwargs[key])
|
||||
image_list = reference_tensors if reference_tensors else None
|
||||
|
||||
try:
|
||||
client = GrokImageClient(route=网络线路)
|
||||
except ValueError as e:
|
||||
if str(e) == "未授权!":
|
||||
print("[o1key Grok Image] 请联系作者授权后方可使用!")
|
||||
raise ValueError("未授权!") from None
|
||||
raise
|
||||
|
||||
try:
|
||||
batch_prompts = parse_batch_prompts(prompt)
|
||||
all_pil_images = []
|
||||
|
||||
if batch_prompts:
|
||||
total = len(batch_prompts)
|
||||
print(f"[o1key Grok Image] 批量模式 | {total} 条提示词 | 每条生成 {生图数量} 张")
|
||||
for idx, p in enumerate(batch_prompts, 1):
|
||||
if _INTERRUPT_AVAILABLE and processing_interrupted():
|
||||
print("[o1key Grok Image] 用户取消")
|
||||
raise InterruptProcessingException()
|
||||
try:
|
||||
pil_images = client.run_sync(
|
||||
prompt=p, model=模型, aspect_ratio=宽高比,
|
||||
resolution=分辨率, n=生图数量, image_list=image_list,
|
||||
)
|
||||
all_pil_images.extend(pil_images)
|
||||
snippet = p[:30] + ("..." if len(p) >= 30 else "")
|
||||
print(f"[o1key Grok Image] [{idx}/{total}] done: {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 Grok Image] [{idx}/{total}] fail: {snippet} → {error_msg}")
|
||||
else:
|
||||
if not prompt or not prompt.strip():
|
||||
raise ValueError("提示词不能为空")
|
||||
pil_images = client.run_sync(
|
||||
prompt=prompt, model=模型, aspect_ratio=宽高比,
|
||||
resolution=分辨率, n=生图数量, image_list=image_list,
|
||||
)
|
||||
all_pil_images.extend(pil_images)
|
||||
|
||||
if not all_pil_images:
|
||||
raise RuntimeError("所有提示词均生成失败,无可用图像输出")
|
||||
|
||||
output_tensor = GrokImageClient._pil_list_to_tensor(all_pil_images)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(
|
||||
f"[o1key Grok Image] 完成!耗时 {elapsed:.1f}s,"
|
||||
f"输出 {output_tensor.shape[0]} 张 "
|
||||
f"{output_tensor.shape[2]}x{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 Grok Image] {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,283 @@
|
||||
"""
|
||||
Grok Video node.
|
||||
|
||||
Submits a /v1/videos task, polls until completion, downloads the mp4,
|
||||
and returns ComfyUI's native VIDEO object.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
from ..clients.grok_video_client import GrokVideoClient
|
||||
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.image_utils import encode_images_for_request_body_limit, tensor_to_pil
|
||||
|
||||
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:
|
||||
ProgressBar = None
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from comfy_api.input_impl import VideoFromFile
|
||||
except Exception:
|
||||
try:
|
||||
from comfy_api.latest import InputImpl
|
||||
VideoFromFile = InputImpl.VideoFromFile
|
||||
except Exception:
|
||||
VideoFromFile = None
|
||||
|
||||
|
||||
MODEL_OPTIONS = ["grok-imagine-video-1.5-preview", "grok-imagine-1.0-video"]
|
||||
ASPECT_RATIO_OPTIONS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"]
|
||||
QUALITY_OPTIONS = ["720p"]
|
||||
QUALITY_VALUE_MAP = {
|
||||
"720p": "high",
|
||||
}
|
||||
MODEL_SECONDS_OPTIONS = {
|
||||
"grok-imagine-1.0-video": [6, 10, 12, 16, 20],
|
||||
}
|
||||
|
||||
MAX_REFERENCE_IMAGES = 3
|
||||
MAX_REQUEST_BODY_BYTES = 20 * 1024 * 1024
|
||||
|
||||
|
||||
def _get_output_dir() -> str:
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
base = folder_paths.get_output_directory()
|
||||
else:
|
||||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
comfy_root = os.path.dirname(os.path.dirname(plugin_dir))
|
||||
base = os.path.join(comfy_root, "output")
|
||||
|
||||
output_dir = os.path.join(base, "grok_video")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
return output_dir
|
||||
|
||||
|
||||
def _format_mb(size_bytes: int) -> str:
|
||||
return f"{size_bytes / 1024 / 1024:.2f}MB"
|
||||
|
||||
|
||||
def _image_tensor_to_first_pil(image_tensor):
|
||||
if image_tensor is None:
|
||||
return None
|
||||
|
||||
pil_images = tensor_to_pil(image_tensor)
|
||||
if not pil_images:
|
||||
return None
|
||||
|
||||
image = pil_images[0]
|
||||
if image.mode not in ("RGB", "L"):
|
||||
image = image.convert("RGB")
|
||||
return image
|
||||
|
||||
|
||||
def _collect_reference_images(**kwargs) -> List[object]:
|
||||
images = []
|
||||
for i in range(1, MAX_REFERENCE_IMAGES + 1):
|
||||
image = _image_tensor_to_first_pil(kwargs.get(f"参考图{i}"))
|
||||
if image is not None:
|
||||
images.append(image)
|
||||
return images
|
||||
|
||||
|
||||
def _to_data_urls(encoded_images) -> List[str]:
|
||||
return [f"data:{mime};base64,{b64}" for mime, b64 in encoded_images]
|
||||
|
||||
|
||||
def _encode_image_data_urls(
|
||||
images: List[object],
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
seconds: int,
|
||||
quality: str,
|
||||
) -> Optional[List[str]]:
|
||||
if not images:
|
||||
return None
|
||||
|
||||
def build_body(encoded_images):
|
||||
return GrokVideoClient.build_video_body(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seconds=seconds,
|
||||
quality=quality,
|
||||
images=_to_data_urls(encoded_images),
|
||||
)
|
||||
|
||||
encoded = encode_images_for_request_body_limit(
|
||||
images,
|
||||
build_body=build_body,
|
||||
max_body_bytes=MAX_REQUEST_BODY_BYTES,
|
||||
)
|
||||
data_urls = _to_data_urls(encoded)
|
||||
|
||||
return data_urls
|
||||
|
||||
|
||||
def _validate_request_body_size(body: dict) -> None:
|
||||
body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
|
||||
if body_size > MAX_REQUEST_BODY_BYTES:
|
||||
raise ValueError(
|
||||
f"Grok Video 请求体大小 {_format_mb(body_size)} 超过 "
|
||||
f"{_format_mb(MAX_REQUEST_BODY_BYTES)} 限制,请减少参考图片或降低图片尺寸。"
|
||||
)
|
||||
|
||||
|
||||
class O1keyGrokVideo:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"提示词": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
},
|
||||
),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": NETWORK_ROUTE_OPTIONS[0]}),
|
||||
"模型": (MODEL_OPTIONS, {"default": MODEL_OPTIONS[0]}),
|
||||
"宽高比": (ASPECT_RATIO_OPTIONS, {"default": "16:9"}),
|
||||
"秒数(按模型限制)": (
|
||||
"INT",
|
||||
{
|
||||
"default": 5,
|
||||
"min": 5,
|
||||
"max": 20,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
},
|
||||
),
|
||||
"画质": (QUALITY_OPTIONS, {"default": "720p"}),
|
||||
},
|
||||
"optional": {
|
||||
"参考图1": ("IMAGE",),
|
||||
"参考图2": ("IMAGE",),
|
||||
"参考图3": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("视频",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/Video"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Grok Video /v1/videos task node. Supports prompt plus up to "
|
||||
"three image references, multiple aspect ratios, model-specific seconds, 720p output."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
**kwargs,
|
||||
):
|
||||
if VideoFromFile is None:
|
||||
raise RuntimeError("当前 ComfyUI 版本不支持原生 VIDEO 输入实现 VideoFromFile。")
|
||||
|
||||
提示词 = kwargs.get("提示词", "")
|
||||
网络线路 = kwargs.get("网络线路", NETWORK_ROUTE_OPTIONS[0])
|
||||
模型 = kwargs.get("模型", MODEL_OPTIONS[0])
|
||||
宽高比 = kwargs.get("宽高比", "16:9")
|
||||
秒数 = kwargs.get("秒数(按模型限制)", kwargs.get("秒数(≤15s)", kwargs.get("秒数", 5)))
|
||||
画质 = kwargs.get("画质", "720p")
|
||||
|
||||
prompt = (提示词 or "").strip()
|
||||
if not prompt:
|
||||
raise ValueError("提示词不能为空。")
|
||||
if 模型 not in MODEL_OPTIONS:
|
||||
raise ValueError(f"模型仅支持: {', '.join(MODEL_OPTIONS)}")
|
||||
if 宽高比 not in ASPECT_RATIO_OPTIONS:
|
||||
raise ValueError(f"宽高比仅支持: {', '.join(ASPECT_RATIO_OPTIONS)}。")
|
||||
seconds = int(秒数)
|
||||
allowed_seconds = MODEL_SECONDS_OPTIONS.get(模型)
|
||||
if allowed_seconds is not None:
|
||||
if seconds not in allowed_seconds:
|
||||
raise ValueError(
|
||||
f"模型 {模型} 仅支持秒数: "
|
||||
f"{', '.join(str(s) for s in allowed_seconds)}。"
|
||||
"请修改为正确的秒数后再发起请求。"
|
||||
)
|
||||
elif seconds < 5 or seconds > 15:
|
||||
raise ValueError("秒数仅支持 5 到 15。")
|
||||
if 画质 not in QUALITY_OPTIONS:
|
||||
raise ValueError("画质仅支持 720p。")
|
||||
|
||||
quality = QUALITY_VALUE_MAP[画质]
|
||||
reference_images = _collect_reference_images(**kwargs)
|
||||
image_data_urls = _encode_image_data_urls(
|
||||
reference_images,
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
aspect_ratio=宽高比,
|
||||
seconds=seconds,
|
||||
quality=quality,
|
||||
)
|
||||
|
||||
request_body = GrokVideoClient.build_video_body(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
aspect_ratio=宽高比,
|
||||
seconds=seconds,
|
||||
quality=quality,
|
||||
images=image_data_urls,
|
||||
)
|
||||
_validate_request_body_size(request_body)
|
||||
|
||||
pbar = ProgressBar(100) if PROGRESS_BAR_AVAILABLE else None
|
||||
last_progress = [0]
|
||||
|
||||
def progress_callback(progress: int, status: str, elapsed: float):
|
||||
progress_value = max(0, min(100, int(progress or 0)))
|
||||
if pbar is not None and progress_value > last_progress[0]:
|
||||
pbar.update(progress_value - last_progress[0])
|
||||
last_progress[0] = progress_value
|
||||
|
||||
client = GrokVideoClient(base_url=get_base_url_by_route(网络线路))
|
||||
|
||||
try:
|
||||
result = client.generate_video_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
aspect_ratio=宽高比,
|
||||
seconds=seconds,
|
||||
quality=quality,
|
||||
output_dir=_get_output_dir(),
|
||||
images=image_data_urls,
|
||||
poll_interval=5,
|
||||
timeout=1200,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
if pbar is not None and last_progress[0] < 100:
|
||||
pbar.update(100 - last_progress[0])
|
||||
|
||||
video_path = result["video_path"]
|
||||
print(f"Grok Video:下载完成:{video_path}")
|
||||
return (VideoFromFile(video_path),)
|
||||
finally:
|
||||
try:
|
||||
balance_data = client.query_balance_sync()
|
||||
balance_info = client.format_balance_info(balance_data)
|
||||
print(f"Grok Video:{balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"O1keyGrokVideo": O1keyGrokVideo,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"O1keyGrokVideo": "Grok Video",
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
"""
|
||||
高级图像拼接节点
|
||||
支持最多 10 张图像按指定方向(上、下、左、右)依次拼接,
|
||||
支持调整图像大小匹配和添加间隔。
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor
|
||||
from ..utils.file_utils import load_images_from_folder
|
||||
|
||||
|
||||
# 间隔颜色映射
|
||||
SPACING_COLOR_MAP = {
|
||||
"white": (255, 255, 255),
|
||||
"black": (0, 0, 0),
|
||||
"red": (255, 0, 0),
|
||||
"green": (0, 255, 0),
|
||||
"blue": (0, 0, 255),
|
||||
}
|
||||
|
||||
|
||||
def _resize_to_match(img: Image.Image, ref: Image.Image, direction: str) -> Image.Image:
|
||||
"""
|
||||
按拼接方向将 img 缩放,使其与 ref 在垂直于拼接轴的尺寸上一致。
|
||||
|
||||
- 水平拼接 (right/left):统一高度
|
||||
- 垂直拼接 (down/up):统一宽度
|
||||
"""
|
||||
ref_w, ref_h = ref.size
|
||||
img_w, img_h = img.size
|
||||
|
||||
if direction in ("right", "left"):
|
||||
if img_h != ref_h:
|
||||
scale = ref_h / img_h
|
||||
new_w = max(1, int(img_w * scale))
|
||||
img = img.resize((new_w, ref_h), Image.LANCZOS)
|
||||
else:
|
||||
if img_w != ref_w:
|
||||
scale = ref_w / img_w
|
||||
new_h = max(1, int(img_h * scale))
|
||||
img = img.resize((ref_w, new_h), Image.LANCZOS)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def _make_spacer(ref: Image.Image, spacing_width: int,
|
||||
direction: str, color: Tuple[int, int, int]) -> Image.Image:
|
||||
"""创建间隔色块"""
|
||||
if direction in ("right", "left"):
|
||||
return Image.new("RGB", (spacing_width, ref.size[1]), color)
|
||||
else:
|
||||
return Image.new("RGB", (ref.size[0], spacing_width), color)
|
||||
|
||||
|
||||
def _stitch_two(img_a: Image.Image, img_b: Image.Image,
|
||||
direction: str, match_size: bool,
|
||||
spacing_width: int, spacing_color: Tuple[int, int, int]) -> Image.Image:
|
||||
"""
|
||||
将两张 PIL 图像按指定方向拼接。
|
||||
img_a 为基准图像,img_b 拼接在 img_a 的指定方向侧。
|
||||
direction="right" → img_b 在 img_a 右侧
|
||||
direction="left" → img_b 在 img_a 左侧
|
||||
direction="down" → img_b 在 img_a 下方
|
||||
direction="up" → img_b 在 img_a 上方
|
||||
"""
|
||||
if img_a.mode != "RGB":
|
||||
img_a = img_a.convert("RGB")
|
||||
if img_b.mode != "RGB":
|
||||
img_b = img_b.convert("RGB")
|
||||
|
||||
if match_size:
|
||||
img_b = _resize_to_match(img_b, img_a, direction)
|
||||
|
||||
if direction == "right":
|
||||
pieces = [img_a, img_b]
|
||||
elif direction == "left":
|
||||
pieces = [img_b, img_a]
|
||||
elif direction == "down":
|
||||
pieces = [img_a, img_b]
|
||||
else: # up
|
||||
pieces = [img_b, img_a]
|
||||
|
||||
if spacing_width > 0:
|
||||
interleaved: List[Image.Image] = []
|
||||
for idx, piece in enumerate(pieces):
|
||||
interleaved.append(piece)
|
||||
if idx < len(pieces) - 1:
|
||||
interleaved.append(_make_spacer(piece, spacing_width, direction, spacing_color))
|
||||
pieces = interleaved
|
||||
|
||||
if direction in ("right", "left"):
|
||||
total_w = sum(p.size[0] for p in pieces)
|
||||
max_h = max(p.size[1] for p in pieces)
|
||||
canvas = Image.new("RGB", (total_w, max_h), spacing_color)
|
||||
x = 0
|
||||
for piece in pieces:
|
||||
canvas.paste(piece, (x, 0))
|
||||
x += piece.size[0]
|
||||
else:
|
||||
max_w = max(p.size[0] for p in pieces)
|
||||
total_h = sum(p.size[1] for p in pieces)
|
||||
canvas = Image.new("RGB", (max_w, total_h), spacing_color)
|
||||
y = 0
|
||||
for piece in pieces:
|
||||
canvas.paste(piece, (0, y))
|
||||
y += piece.size[1]
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def _natural_sort_key(filename: str):
|
||||
"""按数字优先的文件名排序,使 1, 2, 3, 10 而非 1, 10, 2, 3"""
|
||||
try:
|
||||
return (0, int(filename))
|
||||
except ValueError:
|
||||
return (1, filename.lower())
|
||||
|
||||
|
||||
class ImageStitchPro:
|
||||
"""
|
||||
高级图像拼接节点
|
||||
|
||||
在 ComfyUI 原生拼接节点基础上扩展,支持同时输入最多 10 张图像,
|
||||
按指定方向依次拼接,并可在图像间添加任意颜色的间隔。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"方向": (["right", "down", "left", "up"], {"default": "down"}),
|
||||
"匹配图像尺寸": ("BOOLEAN", {"default": True}),
|
||||
"间距宽度": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 2}),
|
||||
"间距颜色": (["white", "black", "red", "green", "blue"], {"default": "white"}),
|
||||
},
|
||||
"optional": {
|
||||
"图1": ("IMAGE",),
|
||||
"图2": ("IMAGE",),
|
||||
"图3": ("IMAGE",),
|
||||
"图4": ("IMAGE",),
|
||||
"图5": ("IMAGE",),
|
||||
"图6": ("IMAGE",),
|
||||
"图7": ("IMAGE",),
|
||||
"图8": ("IMAGE",),
|
||||
"图9": ("IMAGE",),
|
||||
"图10": ("IMAGE",),
|
||||
"图11": ("IMAGE",),
|
||||
"图12": ("IMAGE",),
|
||||
"图片路径(可选)": ("STRING", {"default": "", "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("拼接图像",)
|
||||
FUNCTION = "stitch"
|
||||
CATEGORY = "image"
|
||||
|
||||
DESCRIPTION = (
|
||||
"高级图像拼接节点,支持最多 12 张图像按指定方向(右/下/左/上)依次拼接。\n"
|
||||
"可选择是否将后续图像缩放以匹配第一张图像的尺寸,并可在图像间添加彩色间隔。\n"
|
||||
"可选填「图片路径」:仅处理该文件夹内图片,按文件名顺序依次拼接;与输入端图片不可同时使用。"
|
||||
)
|
||||
|
||||
def stitch(
|
||||
self,
|
||||
方向: str = "down",
|
||||
匹配图像尺寸: bool = True,
|
||||
间距宽度: int = 0,
|
||||
间距颜色: str = "white",
|
||||
图1: Optional[torch.Tensor] = None,
|
||||
图2: Optional[torch.Tensor] = None,
|
||||
图3: Optional[torch.Tensor] = None,
|
||||
图4: Optional[torch.Tensor] = None,
|
||||
图5: Optional[torch.Tensor] = None,
|
||||
图6: Optional[torch.Tensor] = None,
|
||||
图7: Optional[torch.Tensor] = None,
|
||||
图8: Optional[torch.Tensor] = None,
|
||||
图9: Optional[torch.Tensor] = None,
|
||||
图10: Optional[torch.Tensor] = None,
|
||||
图11: Optional[torch.Tensor] = None,
|
||||
图12: Optional[torch.Tensor] = None,
|
||||
**kwargs: object,
|
||||
) -> Tuple[torch.Tensor]:
|
||||
|
||||
color = SPACING_COLOR_MAP.get(间距颜色, (255, 255, 255))
|
||||
raw_tensors = [图1, 图2, 图3, 图4, 图5, 图6, 图7, 图8, 图9, 图10, 图11, 图12]
|
||||
tensors = [t for t in raw_tensors if t is not None]
|
||||
has_input_images = len(tensors) > 0
|
||||
image_folder = (kwargs.get("图片路径(可选)") or "").strip()
|
||||
|
||||
if image_folder and has_input_images:
|
||||
raise ValueError("不可同时使用「图片路径(可选)」与输入端图片,请二选一。")
|
||||
|
||||
if image_folder:
|
||||
infos = load_images_from_folder(image_folder)
|
||||
if not infos:
|
||||
raise ValueError(f"文件夹中未找到可用的图片,或路径无效: {image_folder}")
|
||||
infos.sort(key=lambda x: _natural_sort_key(x.filename))
|
||||
pil_list = [info.image for info in infos]
|
||||
if len(pil_list) == 1:
|
||||
return (pil_to_tensor(pil_list),)
|
||||
base = pil_list[0]
|
||||
for next_img in pil_list[1:]:
|
||||
base = _stitch_two(
|
||||
base, next_img,
|
||||
direction=方向,
|
||||
match_size=匹配图像尺寸,
|
||||
spacing_width=间距宽度,
|
||||
spacing_color=color,
|
||||
)
|
||||
return (pil_to_tensor([base]),)
|
||||
else:
|
||||
if not has_input_images:
|
||||
raise ValueError("请至少接入一张图片,或填写「图片路径(可选)」中的文件夹路径。")
|
||||
|
||||
if len(tensors) == 1:
|
||||
return (tensors[0],)
|
||||
|
||||
pil_batches: List[List[Image.Image]] = [tensor_to_pil(t) for t in tensors]
|
||||
|
||||
batch_size = min(len(b) for b in pil_batches)
|
||||
result_images: List[Image.Image] = []
|
||||
|
||||
for i in range(batch_size):
|
||||
frames = [batch[i] for batch in pil_batches]
|
||||
base = frames[0]
|
||||
for next_img in frames[1:]:
|
||||
base = _stitch_two(
|
||||
base, next_img,
|
||||
direction=方向,
|
||||
match_size=匹配图像尺寸,
|
||||
spacing_width=间距宽度,
|
||||
spacing_color=color,
|
||||
)
|
||||
result_images.append(base)
|
||||
|
||||
return (pil_to_tensor(result_images),)
|
||||
@@ -0,0 +1,738 @@
|
||||
"""
|
||||
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 ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
|
||||
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": ""}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"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:
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
body["metadata"] = {"aspect_ratio": aspect_ratio}
|
||||
endpoint_type = "text2video"
|
||||
|
||||
# ── 保存路径(临时文件,避免与下游保存节点重复落盘)──────────────────
|
||||
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="kling_")
|
||||
|
||||
client = KlingClient()
|
||||
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
|
||||
|
||||
# ── 进度条 ────────────────────────────────────────────────────
|
||||
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": ""}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"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()
|
||||
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
|
||||
|
||||
# 进度条: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",),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
},
|
||||
"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()
|
||||
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
|
||||
|
||||
# ── 进度条 ────────────────────────────────────────────────────
|
||||
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": "图片宽高比预设",
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
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)
|
||||
@@ -0,0 +1,575 @@
|
||||
"""
|
||||
Nano Banana 节点 (V3)
|
||||
ComfyUI 自定义节点,用于调用异步生图模型
|
||||
使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动
|
||||
"""
|
||||
|
||||
import time
|
||||
import math
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||||
from ..utils.config import (
|
||||
NETWORK_ROUTE_OPTIONS,
|
||||
get_base_url_by_route,
|
||||
get_api_key_or_raise,
|
||||
)
|
||||
from ..utils.nano_banana_async import generate_nano_banana_async
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||||
INTERRUPT_AVAILABLE = True
|
||||
except ImportError:
|
||||
INTERRUPT_AVAILABLE = False
|
||||
InterruptProcessingException = RuntimeError
|
||||
processing_interrupted = lambda: False
|
||||
|
||||
REQUEST_LOG_ENABLED = False
|
||||
|
||||
_NODE = "Nano Banana"
|
||||
_REQUEST_TIMEOUT = 900
|
||||
_INTERRUPT_CHECK_INTERVAL = 0.2
|
||||
|
||||
_client_instance = None
|
||||
|
||||
|
||||
def _get_client():
|
||||
global _client_instance
|
||||
if _client_instance is None:
|
||||
_client_instance = GeminiAPIClient()
|
||||
return _client_instance
|
||||
|
||||
|
||||
async def _poll_interrupt():
|
||||
while True:
|
||||
await asyncio.sleep(_INTERRUPT_CHECK_INTERVAL)
|
||||
if INTERRUPT_AVAILABLE and processing_interrupted():
|
||||
return
|
||||
|
||||
|
||||
async def _run_with_interrupt(coro):
|
||||
if not INTERRUPT_AVAILABLE:
|
||||
return await coro
|
||||
|
||||
request_task = asyncio.ensure_future(coro)
|
||||
interrupt_task = asyncio.ensure_future(_poll_interrupt())
|
||||
|
||||
done, pending = await asyncio.wait(
|
||||
[request_task, interrupt_task],
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
if interrupt_task in done and request_task not in done:
|
||||
raise InterruptProcessingException()
|
||||
|
||||
return request_task.result()
|
||||
|
||||
|
||||
def _check_interrupt():
|
||||
if INTERRUPT_AVAILABLE and processing_interrupted():
|
||||
raise InterruptProcessingException()
|
||||
|
||||
|
||||
def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]:
|
||||
if pbar is None:
|
||||
return None
|
||||
|
||||
last_progress = [0.0]
|
||||
|
||||
def _on_progress(progress: float) -> None:
|
||||
try:
|
||||
progress = max(0.0, min(float(progress), 1.0))
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
if progress <= last_progress[0]:
|
||||
return
|
||||
pbar.update(progress - last_progress[0])
|
||||
last_progress[0] = progress
|
||||
|
||||
return _on_progress
|
||||
|
||||
|
||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.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]}x{img.size[1]}" for img in skipped)
|
||||
print(
|
||||
f"{node_label}: 丢弃 {len(skipped)} 张较小尺寸的图 ({sizes_str}),"
|
||||
f"仅输出最大尺寸 {base_size[0]}x{base_size[1]} 的 {len(matched)} 张"
|
||||
)
|
||||
|
||||
return pil_to_tensor(matched)
|
||||
|
||||
|
||||
MODEL_ID_MAP = {
|
||||
"Nano Banana Pro": "nano-banana-pro",
|
||||
"Nano Banana 2": "nano-banana-2",
|
||||
"Nano Banana": "nano-banana",
|
||||
}
|
||||
RESOLUTION_KEY_MAP = {
|
||||
"512px": "0.5k",
|
||||
"1K": "1k",
|
||||
"2K": "2k",
|
||||
"4K": "4k",
|
||||
}
|
||||
BILLING_SPECIAL_ONLY = {"nano-banana"}
|
||||
|
||||
|
||||
def _build_model_id(model_name: str, resolution: str, billing: str) -> str:
|
||||
base = MODEL_ID_MAP.get(model_name, "nano-banana-pro")
|
||||
|
||||
if base == "nano-banana":
|
||||
if billing == "官方":
|
||||
raise ValueError(f"模型 \"{model_name}\" 仅支持特价计费")
|
||||
return "nano-banana"
|
||||
|
||||
res_key = RESOLUTION_KEY_MAP.get(resolution, "2k")
|
||||
is_official = (billing == "官方")
|
||||
|
||||
if base == "nano-banana-pro" and res_key == "1k" and not is_official:
|
||||
return "nano-banana-pro"
|
||||
|
||||
if base == "nano-banana-2" and res_key == "0.5k":
|
||||
if is_official:
|
||||
raise ValueError("Nano Banana 2 的 512px 分辨率仅支持特价计费")
|
||||
return "nano-banana-2-0.5k"
|
||||
|
||||
model_id = f"{base}-{res_key}"
|
||||
if is_official:
|
||||
model_id += "-official"
|
||||
return model_id
|
||||
|
||||
|
||||
async def _generate_single(
|
||||
session: aiohttp.ClientSession,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
enable_grounding: bool = False,
|
||||
thinking_level: Optional[str] = None,
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
) -> List[Image.Image]:
|
||||
result_images, timing = await generate_nano_banana_async(
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
node_label="Nano Banana",
|
||||
request_log_enabled=REQUEST_LOG_ENABLED,
|
||||
check_interrupt=_check_interrupt,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
return result_images, timing["task_ms"], timing["parse_ms"]
|
||||
|
||||
|
||||
async def _generate_single_task(
|
||||
session: aiohttp.ClientSession,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: Optional[List[Image.Image]],
|
||||
global_task_index: int,
|
||||
enable_grounding: bool = False,
|
||||
thinking_level: Optional[str] = None,
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
) -> dict:
|
||||
result = {
|
||||
"global_task_index": global_task_index,
|
||||
"prompt": prompt,
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"output_images": [],
|
||||
"error": None,
|
||||
}
|
||||
try:
|
||||
gen_images, task_ms, parse_ms = await _generate_single(
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images if images else None,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
del task_ms, parse_ms
|
||||
result["output_images"] = gen_images
|
||||
result["success"] = True
|
||||
result["generated_count"] = len(gen_images)
|
||||
except InterruptProcessingException:
|
||||
raise
|
||||
except Exception as e:
|
||||
result["error"] = str(e)
|
||||
return result
|
||||
|
||||
|
||||
async def _process_batch_async(
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
prompts: List[str],
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images_per_prompt: int,
|
||||
input_images: Optional[List[Image.Image]],
|
||||
pbar=None,
|
||||
enable_grounding: bool = False,
|
||||
thinking_level: Optional[str] = None,
|
||||
) -> List[dict]:
|
||||
tasks_def = []
|
||||
for p_idx, prompt in enumerate(prompts):
|
||||
for sub_idx in range(images_per_prompt):
|
||||
tasks_def.append((p_idx, sub_idx, prompt))
|
||||
|
||||
total_tasks = len(tasks_def)
|
||||
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):
|
||||
_check_interrupt()
|
||||
start_idx = batch_idx * max_concurrent
|
||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||||
|
||||
tasks = []
|
||||
for i in range(start_idx, end_idx):
|
||||
_check_interrupt()
|
||||
_, _, prompt = tasks_def[i]
|
||||
task = asyncio.create_task(
|
||||
_generate_single_task(
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_images,
|
||||
global_task_index=i,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
progress_callback=_make_progress_callback(pbar),
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
batch_results = []
|
||||
for coro in asyncio.as_completed(tasks):
|
||||
_check_interrupt()
|
||||
result_data = None
|
||||
try:
|
||||
result = await coro
|
||||
if isinstance(result, Exception):
|
||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "output_images": [], "prompt": ""}
|
||||
else:
|
||||
result_data = result
|
||||
except InterruptProcessingException:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
raise
|
||||
except Exception as e:
|
||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "output_images": [], "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: [{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: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")
|
||||
|
||||
all_results.extend(batch_results)
|
||||
import gc; gc.collect()
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
return all_results
|
||||
|
||||
|
||||
class NanoBanana(io.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="NanoBanana",
|
||||
display_name="Nano Banana",
|
||||
category="image/generation",
|
||||
inputs=[
|
||||
io.String.Input(
|
||||
"prompt",
|
||||
default="一个中国女子的OOTD",
|
||||
multiline=True,
|
||||
),
|
||||
io.DynamicCombo.Input("模型", options=[
|
||||
io.DynamicCombo.Option("Nano Banana Pro", [
|
||||
io.Combo.Input("宽高比", options=[
|
||||
"智能", "1:1", "2:3", "3:2", "3:4", "4:3",
|
||||
"4:5", "5:4", "9:16", "16:9", "21:9",
|
||||
], default="智能"),
|
||||
io.Combo.Input("分辨率", options=["1K", "2K", "4K"], default="2K"),
|
||||
]),
|
||||
io.DynamicCombo.Option("Nano Banana 2", [
|
||||
io.Combo.Input("宽高比", options=[
|
||||
"智能", "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",
|
||||
], default="智能"),
|
||||
io.Combo.Input("分辨率", options=["512px", "1K", "2K", "4K"], default="2K"),
|
||||
io.Combo.Input("思考深度", options=["高", "低"], default="高"),
|
||||
]),
|
||||
io.DynamicCombo.Option("Nano Banana", [
|
||||
io.Combo.Input("宽高比", options=[
|
||||
"智能", "1:1", "2:3", "3:2", "3:4", "4:3",
|
||||
"4:5", "5:4", "9:16", "16:9", "21:9",
|
||||
], default="智能"),
|
||||
io.Combo.Input("分辨率", options=["1K"], default="1K"),
|
||||
]),
|
||||
]),
|
||||
io.Int.Input("生图数量", default=1, min=1, max=1000, step=1),
|
||||
io.Combo.Input("网络", options=NETWORK_ROUTE_OPTIONS, default="全球加速"),
|
||||
io.Combo.Input("计费", options=["特价", "官方"], default="特价"),
|
||||
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
|
||||
io.Int.Input("seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF),
|
||||
io.Image.Input("参考图1", optional=True),
|
||||
io.Image.Input("参考图2", optional=True),
|
||||
io.Image.Input("参考图3", optional=True),
|
||||
io.Image.Input("参考图4", optional=True),
|
||||
io.Image.Input("参考图5", optional=True),
|
||||
io.Image.Input("参考图6", optional=True),
|
||||
io.Image.Input("参考图7", optional=True),
|
||||
io.Image.Input("参考图8", optional=True),
|
||||
io.Image.Input("参考图9", optional=True),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="输出图像"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, prompt, 模型, 生图数量, 计费, 网络, 谷歌搜索, seed, **kwargs) -> io.NodeOutput:
|
||||
start_time = time.time()
|
||||
was_interrupted = False
|
||||
|
||||
model_name = 模型["模型"]
|
||||
宽高比 = 模型["宽高比"]
|
||||
分辨率 = 模型["分辨率"]
|
||||
思考深度 = 模型.get("思考深度")
|
||||
|
||||
enable_grounding = (谷歌搜索 == "打开")
|
||||
|
||||
thinking_level = None
|
||||
if model_name == "Nano Banana 2" and 思考深度:
|
||||
thinking_level = "High" if 思考深度 == "高" else "Low"
|
||||
|
||||
actual_model = _build_model_id(model_name, 分辨率, 计费)
|
||||
|
||||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
base_url = get_base_url_by_route(网络)
|
||||
|
||||
pbar = ProgressBar(生图数量) if PROGRESS_BAR_AVAILABLE else None
|
||||
|
||||
try:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed % (2**32))
|
||||
|
||||
input_images = []
|
||||
for i in range(1, 10):
|
||||
key = f"参考图{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
pil_imgs = tensor_to_pil(kwargs[key])
|
||||
input_images.extend(pil_imgs)
|
||||
|
||||
if len(input_images) > 14:
|
||||
raise ValueError(f"输入图像数量 {len(input_images)} 超过限制 14 张")
|
||||
|
||||
batch_prompts = parse_batch_prompts(prompt)
|
||||
|
||||
grounding_str = " | 谷歌搜索接地" if enable_grounding else ""
|
||||
thinking_str = f" | 思考:{thinking_level}" if thinking_level else ""
|
||||
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: {mode_str} | {分辨率} {宽高比} | 共{total_images}张{grounding_str}{thinking_str}")
|
||||
else:
|
||||
mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式"
|
||||
print(f"Nano Banana: {mode_str} | {分辨率} {宽高比} | {生图数量}张{grounding_str}{thinking_str}")
|
||||
|
||||
if batch_prompts or 生图数量 > 1:
|
||||
prompts = batch_prompts if batch_prompts else [prompt]
|
||||
images_per_prompt = 生图数量
|
||||
total_tasks = len(prompts) * images_per_prompt
|
||||
|
||||
if pbar is not None:
|
||||
pbar = ProgressBar(total_tasks)
|
||||
|
||||
def run_async_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
_run_with_interrupt(_process_batch_async(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
prompts=prompts,
|
||||
model=actual_model,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images_per_prompt=images_per_prompt,
|
||||
input_images=input_images,
|
||||
pbar=pbar,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
))
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async_in_thread)
|
||||
try:
|
||||
results = future.result(timeout=_REQUEST_TIMEOUT)
|
||||
except TimeoutError:
|
||||
raise RuntimeError(f"任务执行超时({_REQUEST_TIMEOUT}秒)")
|
||||
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
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_tasks} | 失败: {fail_count}")
|
||||
|
||||
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 io.NodeOutput(output_tensor)
|
||||
|
||||
else:
|
||||
def run_single():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
async def _do():
|
||||
connector = aiohttp.TCPConnector(ssl=False)
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
return await _generate_single(
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
model=actual_model,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images=input_images if input_images else None,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
progress_callback=_make_progress_callback(pbar),
|
||||
)
|
||||
return loop.run_until_complete(_run_with_interrupt(_do()))
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_single)
|
||||
generated_images, task_ms, parse_ms = future.result(timeout=_REQUEST_TIMEOUT)
|
||||
|
||||
output_tensor = _images_to_tensor_safe(generated_images, _NODE)
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||||
task_str = f"{task_ms/1000:.2f}s"
|
||||
parse_str = f"{parse_ms/1000:.2f}s"
|
||||
print(f"完成!总耗时 {time_str} | 异步任务 {task_str} | 解析 {parse_str} | 成功 {len(generated_images)}张")
|
||||
|
||||
import gc; gc.collect()
|
||||
return io.NodeOutput(output_tensor)
|
||||
|
||||
except InterruptProcessingException:
|
||||
was_interrupted = True
|
||||
print("Nano Banana: 用户取消")
|
||||
raise
|
||||
except ValueError as e:
|
||||
if str(e) == "未授权!":
|
||||
print("请联系作者授权后方可使用!")
|
||||
raise ValueError("未授权!") from None
|
||||
raise ValueError(str(e)) from None
|
||||
except RuntimeError as e:
|
||||
raise RuntimeError(str(e)) from None
|
||||
except Exception as e:
|
||||
raise RuntimeError(str(e)) from None
|
||||
finally:
|
||||
if not was_interrupted:
|
||||
try:
|
||||
client = _get_client()
|
||||
client.base_url = base_url
|
||||
balance_data = client.query_balance_sync()
|
||||
balance_info = client.format_balance_info(balance_data)
|
||||
print(f"Nano Banana: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
import gc; gc.collect()
|
||||
@@ -1,381 +0,0 @@
|
||||
"""
|
||||
Nano Banana Pro 节点
|
||||
ComfyUI 自定义节点,用于调用 Gemini 3 Pro 模型生成图像
|
||||
"""
|
||||
|
||||
import time
|
||||
import random
|
||||
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 ..clients.gemini_client import GeminiAPIClient
|
||||
from ..models_config import get_enabled_models, get_model_description
|
||||
|
||||
# 导入 ComfyUI 原生进度条
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ NanoBananaPro: comfy.utils.ProgressBar 不可用,将只使用终端进度显示")
|
||||
|
||||
|
||||
class NanoBananaPro:
|
||||
"""
|
||||
Nano Banana Pro 节点
|
||||
|
||||
功能:
|
||||
- 文生图:基于提示词生成图像
|
||||
- 图生图:基于输入图像和提示词生成新图像
|
||||
- 批量生成:支持并发生成多张图像
|
||||
|
||||
注意:
|
||||
- 支持的模型列表从 models_config.py 动态加载
|
||||
- 要添加/禁用模型,请编辑 models_config.py 文件
|
||||
"""
|
||||
|
||||
# 支持的模型列表(从配置文件动态加载)
|
||||
MODELS = None # 将在 INPUT_TYPES 中动态获取
|
||||
|
||||
# 支持的宽高比列表
|
||||
ASPECT_RATIOS = [
|
||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
||||
"2:3", "3:2", "4:5", "5:4", "21:9"
|
||||
]
|
||||
|
||||
# 支持的分辨率列表
|
||||
RESOLUTIONS = ["1K", "2K", "4K"]
|
||||
|
||||
def __init__(self):
|
||||
"""初始化节点"""
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
|
||||
ComfyUI 节点规范:
|
||||
- required: 必选参数
|
||||
- optional: 可选参数
|
||||
"""
|
||||
# 从配置文件动态获取启用的模型列表
|
||||
enabled_models = get_enabled_models()
|
||||
|
||||
# 如果没有启用的模型,使用空列表(会导致节点不可用,提示用户配置)
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
||||
|
||||
# 创建9个独立的图像输入
|
||||
optional_inputs = {}
|
||||
for i in range(1, 10): # 1-9
|
||||
optional_inputs[f"参考图{i}"] = ("IMAGE",)
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "一个中国女子的OOTD",
|
||||
"multiline": True
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0]
|
||||
}),
|
||||
"宽高比": (cls.ASPECT_RATIOS, {
|
||||
"default": "1:1"
|
||||
}),
|
||||
"分辨率": (cls.RESOLUTIONS, {
|
||||
"default": "2K"
|
||||
}),
|
||||
"生图数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 1000,
|
||||
"step": 1
|
||||
}),
|
||||
"像素缩放": ("BOOLEAN", {
|
||||
"default": False
|
||||
}),
|
||||
"分辨率像素": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.1,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"display": "number"
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff
|
||||
})
|
||||
},
|
||||
"optional": optional_inputs
|
||||
}
|
||||
|
||||
# 返回值类型
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("输出图像",)
|
||||
|
||||
# 执行函数名
|
||||
FUNCTION = "generate"
|
||||
|
||||
# 节点分类
|
||||
CATEGORY = "image/generation"
|
||||
|
||||
def resize_to_megapixels(
|
||||
self,
|
||||
image: Image.Image,
|
||||
target_megapixels: float
|
||||
) -> Image.Image:
|
||||
"""
|
||||
将图像缩放到指定的总像素数,保持纵横比
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
target_megapixels: 目标像素数(百万像素)
|
||||
|
||||
Returns:
|
||||
缩放后的 PIL Image
|
||||
|
||||
Example:
|
||||
>>> resized = self.resize_to_megapixels(img, 2.0) # 缩放到2百万像素
|
||||
"""
|
||||
# 计算当前像素数
|
||||
current_pixels = image.width * image.height
|
||||
target_pixels = int(target_megapixels * 1_000_000)
|
||||
|
||||
# 如果当前像素数已经接近目标,则不缩放
|
||||
if abs(current_pixels - target_pixels) / target_pixels < 0.05:
|
||||
return image
|
||||
|
||||
# 计算缩放比例
|
||||
scale = (target_pixels / current_pixels) ** 0.5
|
||||
|
||||
# 计算新尺寸
|
||||
new_width = int(image.width * scale)
|
||||
new_height = int(image.height * scale)
|
||||
|
||||
# 确保至少为1像素
|
||||
new_width = max(1, new_width)
|
||||
new_height = max(1, new_height)
|
||||
|
||||
# 使用 Lanczos 重采样
|
||||
resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
return resized_image
|
||||
|
||||
def validate_inputs(
|
||||
self,
|
||||
images: Optional[torch.Tensor],
|
||||
batch_size: int
|
||||
) -> None:
|
||||
"""
|
||||
验证输入参数
|
||||
|
||||
Args:
|
||||
images: 输入图像张量(可选)
|
||||
batch_size: 批次大小
|
||||
|
||||
Raises:
|
||||
ValueError: 如果输入参数不合法
|
||||
"""
|
||||
# 检查图像数量
|
||||
if images is not None:
|
||||
num_images = images.shape[0]
|
||||
if num_images > 14:
|
||||
raise ValueError(
|
||||
f"输入图像数量 {num_images} 超过限制 14 张,请减少输入图像数量"
|
||||
)
|
||||
|
||||
# 检查批次大小
|
||||
if batch_size < 1 or batch_size > 1000:
|
||||
raise ValueError(
|
||||
f"批次大小 {batch_size} 超出范围 [1, 1000]"
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str,
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
生图数量: int,
|
||||
像素缩放: bool,
|
||||
分辨率像素: float,
|
||||
seed: int,
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""
|
||||
生成图像
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
模型: 模型名称
|
||||
宽高比: 宽高比
|
||||
分辨率: 分辨率
|
||||
生图数量: 批次大小
|
||||
像素缩放: 是否启用像素缩放
|
||||
分辨率像素: 目标像素数(百万像素)
|
||||
seed: 随机种子
|
||||
**kwargs: 动态参考图输入 (参考图1-9)
|
||||
|
||||
Returns:
|
||||
生成的图像张量 (IMAGE,)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# 创建 ComfyUI 原生进度条
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(生图数量)
|
||||
|
||||
try:
|
||||
# 设置随机种子(用于本地随机操作)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed % (2**32))
|
||||
|
||||
# 初始化 API 客户端
|
||||
if self.client is None:
|
||||
try:
|
||||
self.client = GeminiAPIClient()
|
||||
except ValueError as e:
|
||||
raise ValueError(f"初始化失败: {str(e)}")
|
||||
|
||||
# 收集独立输入的参考图
|
||||
input_images = []
|
||||
for i in range(1, 10): # 1-9
|
||||
key = f"参考图{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
pil_imgs = tensor_to_pil(kwargs[key])
|
||||
input_images.extend(pil_imgs)
|
||||
|
||||
# 验证输入图像数量
|
||||
if input_images:
|
||||
if len(input_images) > 14:
|
||||
raise ValueError(
|
||||
f"输入图像数量 {len(input_images)} 超过限制 14 张,请减少输入图像数量"
|
||||
)
|
||||
|
||||
# 应用像素缩放(如果启用)
|
||||
if input_images and 像素缩放:
|
||||
scaled_images = []
|
||||
for img in input_images:
|
||||
scaled = self.resize_to_megapixels(img, 分辨率像素)
|
||||
scaled_images.append(scaled)
|
||||
input_images = scaled_images
|
||||
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)
|
||||
|
||||
# 统计变量
|
||||
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)
|
||||
|
||||
# 根据是否有批量提示词选择生成模式
|
||||
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)
|
||||
|
||||
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:
|
||||
# 单提示词模式
|
||||
print(f"Nano Banana Pro: {'图生图' if input_images else '文生图'}模式")
|
||||
print(f"Nano Banana Pro: 发送请求")
|
||||
print(f"Nano Banana Pro: 生图中...")
|
||||
|
||||
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 = pil_to_tensor(generated_images)
|
||||
|
||||
# 计算耗时
|
||||
elapsed = time.time() - start_time
|
||||
print(f"Nano Banana Pro: 完成生图 (耗时: {elapsed:.2f}s, 成功生成 {len(generated_images)} 张图像)")
|
||||
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as e:
|
||||
# 检测是否为授权错误
|
||||
if str(e) == "未授权!":
|
||||
print("请联系作者授权后方可使用!")
|
||||
else:
|
||||
# 用户输入错误
|
||||
print(f"Nano Banana Pro: 输入错误 - {str(e)}")
|
||||
raise
|
||||
|
||||
except RuntimeError as e:
|
||||
# API 或网络错误
|
||||
print(f"Nano Banana Pro: API 错误 - {str(e)}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
# 其他未知错误
|
||||
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 as e:
|
||||
print(f"Nano Banana Pro: ⚠️ 余额查询失败 - {str(e)}")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,254 @@
|
||||
"""
|
||||
Single-node new-api Veo 3.1 generator.
|
||||
|
||||
The node submits a /v1/videos task, waits for completion, downloads the mp4,
|
||||
and returns ComfyUI's native VIDEO object for the built-in Save Video node.
|
||||
"""
|
||||
|
||||
import os
|
||||
from io import BytesIO
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from ..clients.newapi_veo_client import NewAPIVeoClient
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
|
||||
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:
|
||||
ProgressBar = None
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from comfy_api.input_impl import VideoFromFile
|
||||
except Exception:
|
||||
VideoFromFile = None
|
||||
|
||||
|
||||
MODEL_OPTIONS = [
|
||||
"veo-3.1",
|
||||
]
|
||||
|
||||
DURATION_OPTIONS = ["4", "6", "8"]
|
||||
ASPECT_RATIO_OPTIONS = ["16:9", "9:16"]
|
||||
RESOLUTION_OPTIONS = ["720p", "1080p"]
|
||||
|
||||
TARGET_SIZE_MAP = {
|
||||
("720p", "16:9"): (1280, 720),
|
||||
("720p", "9:16"): (720, 1280),
|
||||
("1080p", "16:9"): (1920, 1080),
|
||||
("1080p", "9:16"): (1080, 1920),
|
||||
}
|
||||
|
||||
|
||||
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__)))
|
||||
comfy_root = os.path.dirname(os.path.dirname(plugin_dir))
|
||||
return os.path.join(comfy_root, "output")
|
||||
|
||||
|
||||
def _get_download_dir() -> str:
|
||||
output_dir = _get_output_dir()
|
||||
video_dir = os.path.join(output_dir, "newapi_veo")
|
||||
os.makedirs(video_dir, exist_ok=True)
|
||||
return video_dir
|
||||
|
||||
|
||||
def _fit_image_to_target(image, target_size: Tuple[int, int]):
|
||||
from PIL import Image as PILImage
|
||||
|
||||
target_w, target_h = target_size
|
||||
src_w, src_h = image.size
|
||||
src_ratio = src_w / src_h
|
||||
target_ratio = target_w / target_h
|
||||
|
||||
if src_w == target_w and src_h == target_h:
|
||||
return image
|
||||
|
||||
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)
|
||||
image = image.resize((new_w, target_h), resample=resample)
|
||||
left = max(0, (new_w - target_w) // 2)
|
||||
image = image.crop((left, 0, left + target_w, target_h))
|
||||
else:
|
||||
scale = target_w / src_w
|
||||
new_h = round(src_h * scale)
|
||||
image = image.resize((target_w, new_h), resample=resample)
|
||||
top = max(0, (new_h - target_h) // 2)
|
||||
image = image.crop((0, top, target_w, top + target_h))
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def _image_to_png_bytes(image_tensor, resolution: str, aspect_ratio: str) -> Optional[bytes]:
|
||||
if image_tensor is None:
|
||||
return None
|
||||
|
||||
pil_images = tensor_to_pil(image_tensor)
|
||||
if not pil_images:
|
||||
return None
|
||||
|
||||
image = pil_images[0]
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
|
||||
target_size = TARGET_SIZE_MAP.get((resolution, aspect_ratio))
|
||||
if target_size is not None:
|
||||
original_size = image.size
|
||||
image = _fit_image_to_target(image, target_size)
|
||||
if image.size != original_size:
|
||||
print(
|
||||
"NewAPI Veo: input image fitted "
|
||||
f"{original_size[0]}x{original_size[1]} -> {image.size[0]}x{image.size[1]}"
|
||||
)
|
||||
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
image_bytes = buffer.getvalue()
|
||||
print(
|
||||
"NewAPI Veo: input_reference PNG "
|
||||
f"{len(image_bytes) / 1024:.0f} KB ({image.size[0]}x{image.size[1]})"
|
||||
)
|
||||
return image_bytes
|
||||
|
||||
|
||||
class Google31Video:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"提示词": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "A cinematic shot of a small robot walking through a rainy neon street.",
|
||||
"multiline": True,
|
||||
},
|
||||
),
|
||||
"负向提示词": ("STRING", {"default": "", "multiline": True}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"模型": (MODEL_OPTIONS, {"default": MODEL_OPTIONS[0]}),
|
||||
"时长": (DURATION_OPTIONS, {"default": "8"}),
|
||||
"宽高比": (ASPECT_RATIO_OPTIONS, {"default": "16:9"}),
|
||||
"分辨率": (RESOLUTION_OPTIONS, {"default": "1080p"}),
|
||||
"生成音频": (["打开", "关闭"], {"default": "打开"}),
|
||||
"seed": (
|
||||
"INT",
|
||||
{
|
||||
"default": -1,
|
||||
"min": -1,
|
||||
"max": 0xFFFFFFFFFFFFFFFF,
|
||||
"step": 1,
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"参考图像": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("视频",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/Video"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Submit a new-api /v1/videos Veo 3.1 task, poll until complete, "
|
||||
"download the mp4, and output native VIDEO for ComfyUI Save Video."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
提示词: str,
|
||||
负向提示词: str,
|
||||
网络线路: str,
|
||||
模型: str,
|
||||
时长: str,
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
生成音频: str,
|
||||
seed: int,
|
||||
参考图像=None,
|
||||
):
|
||||
if VideoFromFile is None:
|
||||
raise RuntimeError("当前 ComfyUI 版本不支持原生 VIDEO 输入实现 VideoFromFile。")
|
||||
|
||||
prompt = (提示词 or "").strip()
|
||||
if not prompt:
|
||||
raise ValueError("提示词不能为空。")
|
||||
|
||||
duration_value = int(时长)
|
||||
if duration_value not in (4, 6, 8):
|
||||
raise ValueError("时长仅支持 4、6、8。")
|
||||
if 宽高比 not in ASPECT_RATIO_OPTIONS:
|
||||
raise ValueError("宽高比仅支持 16:9 或 9:16。")
|
||||
if 分辨率 not in RESOLUTION_OPTIONS:
|
||||
raise ValueError("分辨率仅支持 720p 或 1080p。")
|
||||
|
||||
output_dir = _get_download_dir()
|
||||
image_bytes = _image_to_png_bytes(参考图像, 分辨率, 宽高比)
|
||||
|
||||
pbar = ProgressBar(100) if PROGRESS_BAR_AVAILABLE else None
|
||||
last_progress = [0]
|
||||
last_status = [""]
|
||||
|
||||
def progress_callback(progress: int, status: str, elapsed: float):
|
||||
if status != last_status[0]:
|
||||
print(
|
||||
"NewAPI Veo: polling "
|
||||
f"status={status} | elapsed={elapsed:.0f}s"
|
||||
)
|
||||
last_status[0] = status
|
||||
|
||||
progress = max(0, min(100, int(progress or 0)))
|
||||
if pbar is not None and progress > last_progress[0]:
|
||||
pbar.update(progress - last_progress[0])
|
||||
last_progress[0] = progress
|
||||
|
||||
client = NewAPIVeoClient(base_url=get_base_url_by_route(网络线路))
|
||||
|
||||
result = client.generate_video_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
duration=duration_value,
|
||||
aspect_ratio=宽高比,
|
||||
resolution=分辨率,
|
||||
output_dir=output_dir,
|
||||
negative_prompt=负向提示词,
|
||||
generate_audio=(生成音频 == "打开"),
|
||||
image_bytes=image_bytes,
|
||||
poll_interval=10,
|
||||
timeout=900,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
video_path = result["video_path"]
|
||||
video = VideoFromFile(video_path)
|
||||
|
||||
print(
|
||||
"NewAPI Veo: completed "
|
||||
f"| task_id={result['task_id']} | video={video_path}"
|
||||
)
|
||||
|
||||
return (video,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Google31Video": Google31Video,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Google31Video": "Google 3.1 Video",
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
o1key 去背景节点
|
||||
基于 rembg 实现,支持 CPU 推理
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class O1keyRemoveBackground:
|
||||
"""
|
||||
移除图像背景,输出 RGBA 透明图层
|
||||
|
||||
基于 rembg (ISNet-General-Use) 模型,支持 CPU 推理。
|
||||
首次运行会自动下载模型(约 170MB)。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("RGBA图像",)
|
||||
FUNCTION = "remove_bg"
|
||||
CATEGORY = "o1key/image"
|
||||
|
||||
def remove_bg(self, image):
|
||||
from ..utils.rembg_utils import remove_background_tensor
|
||||
print("[o1key 去背景] 正在处理...")
|
||||
result = remove_background_tensor(image)
|
||||
print(f"[o1key 去背景] 完成,输出 {result.shape[0]} 张 RGBA")
|
||||
return (result,)
|
||||
@@ -0,0 +1,188 @@
|
||||
"""
|
||||
图像元数据去除节点
|
||||
|
||||
提供批量去除已有图片中元数据的功能
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
# 支持的图片格式
|
||||
SUPPORTED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.webp', '.bmp', '.tiff', '.tif'}
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 批量去除元数据
|
||||
# ============================================================================
|
||||
|
||||
class BatchCleanMetadata:
|
||||
"""
|
||||
批量去除文件夹中图片元数据的节点
|
||||
|
||||
功能:
|
||||
- 指定文件夹路径,批量处理其中所有图片
|
||||
- 去除 EXIF、PNG tEXt 块、ComfyUI 工作流等所有元数据
|
||||
- 支持保存到原目录(添加 _nometa 后缀)或覆盖原文件
|
||||
- 支持 PNG/JPG/JPEG/WEBP/BMP/TIFF 格式
|
||||
|
||||
使用场景:
|
||||
- 已经保存了一批含有 AI 元数据的图片,需要批量清理
|
||||
- 批量处理指定文件夹中的所有图片
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
|
||||
Returns:
|
||||
输入参数配置字典
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"文件夹路径": ("STRING", {"default": ""}),
|
||||
"覆盖原文件": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("处理结果",)
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "batch_clean"
|
||||
CATEGORY = "image"
|
||||
|
||||
DESCRIPTION = (
|
||||
"批量去除文件夹中图片的元数据。\n"
|
||||
"支持 PNG/JPG/JPEG/WEBP/BMP/TIFF 格式。\n"
|
||||
"默认在原文件名后添加 _nometa 后缀保存,也可选择覆盖原文件。"
|
||||
)
|
||||
|
||||
def batch_clean(
|
||||
self,
|
||||
文件夹路径: str,
|
||||
覆盖原文件: bool = False,
|
||||
) -> tuple:
|
||||
"""
|
||||
批量去除文件夹中图片的元数据
|
||||
|
||||
Args:
|
||||
文件夹路径: 待处理图片所在的文件夹路径
|
||||
覆盖原文件: 是否覆盖原文件(False 则添加 _nometa 后缀)
|
||||
|
||||
Returns:
|
||||
处理结果字符串
|
||||
|
||||
Raises:
|
||||
ValueError: 文件夹路径无效
|
||||
"""
|
||||
if not 文件夹路径 or not 文件夹路径.strip():
|
||||
raise ValueError("请输入文件夹路径")
|
||||
|
||||
folder = 文件夹路径.strip()
|
||||
|
||||
if not os.path.isdir(folder):
|
||||
raise ValueError(f"文件夹路径无效或不存在: {folder}")
|
||||
|
||||
# 扫描支持的图片文件
|
||||
files = []
|
||||
for f in sorted(os.listdir(folder)):
|
||||
ext = os.path.splitext(f)[1].lower()
|
||||
if ext in SUPPORTED_EXTENSIONS:
|
||||
files.append(f)
|
||||
|
||||
if not files:
|
||||
msg = f"文件夹中未找到支持的图片文件 ({', '.join(SUPPORTED_EXTENSIONS)})"
|
||||
print(f"批量去除元数据: {msg}")
|
||||
return (msg,)
|
||||
|
||||
print(f"批量去除元数据: 找到 {len(files)} 张图片,开始处理...")
|
||||
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
for f in files:
|
||||
try:
|
||||
src_path = os.path.join(folder, f)
|
||||
img = Image.open(src_path)
|
||||
|
||||
if 覆盖原文件:
|
||||
dst_path = src_path
|
||||
else:
|
||||
name, ext = os.path.splitext(f)
|
||||
dst_path = os.path.join(folder, f"{name}_nometa{ext}")
|
||||
|
||||
_save_image_clean(img, dst_path)
|
||||
success_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f"批量去除元数据: 处理 {f} 失败 - {str(e)}")
|
||||
fail_count += 1
|
||||
|
||||
# 构建结果消息
|
||||
if fail_count > 0:
|
||||
msg = f"处理完成: 成功 {success_count} 张, 失败 {fail_count} 张"
|
||||
else:
|
||||
msg = f"处理完成: 全部 {success_count} 张成功"
|
||||
|
||||
if not 覆盖原文件:
|
||||
msg += " (已添加 _nometa 后缀)"
|
||||
else:
|
||||
msg += " (已覆盖原文件)"
|
||||
|
||||
print(f"批量去除元数据: {msg}")
|
||||
|
||||
return (msg,)
|
||||
@@ -0,0 +1,94 @@
|
||||
"""保存图像节点 - 支持 PNG/JPEG/WebP 格式输出"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
import folder_paths
|
||||
from comfy.cli_args import args
|
||||
|
||||
|
||||
class SaveImageFormat:
|
||||
"""保存图像,支持 PNG / JPEG / WebP 三种格式"""
|
||||
|
||||
FORMATS = ["PNG", "JPEG", "WebP"]
|
||||
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
self.type = "output"
|
||||
self.compress_level = 4
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"图像": ("IMAGE",),
|
||||
"文件名前缀": ("STRING", {"default": "ComfyUI"}),
|
||||
"输出格式": (cls.FORMATS, {"default": "PNG"}),
|
||||
},
|
||||
"optional": {},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save_images"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "image"
|
||||
DESCRIPTION = "保存图像,支持 PNG / JPEG / WebP 格式输出。"
|
||||
|
||||
_EXT_MAP = {"PNG": ".png", "JPEG": ".jpg", "WebP": ".webp"}
|
||||
|
||||
def save_images(self, 图像=None, 文件名前缀="ComfyUI", 输出格式="PNG",
|
||||
prompt=None, extra_pnginfo=None):
|
||||
images = 图像
|
||||
filename_prefix = 文件名前缀
|
||||
format = 输出格式
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = \
|
||||
folder_paths.get_save_image_path(
|
||||
filename_prefix, self.output_dir,
|
||||
images[0].shape[1], images[0].shape[0]
|
||||
)
|
||||
|
||||
ext = self._EXT_MAP.get(format, ".png")
|
||||
results = []
|
||||
|
||||
for batch_number, image in enumerate(images):
|
||||
i = 255.0 * image.cpu().numpy()
|
||||
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
||||
|
||||
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
|
||||
file = f"{filename_with_batch_num}_{counter:05}_{ext}"
|
||||
|
||||
filepath = os.path.join(full_output_folder, file)
|
||||
|
||||
if format == "PNG":
|
||||
metadata = None
|
||||
if not args.disable_metadata:
|
||||
metadata = PngInfo()
|
||||
if prompt is not None:
|
||||
metadata.add_text("prompt", json.dumps(prompt))
|
||||
if extra_pnginfo is not None:
|
||||
for x in extra_pnginfo:
|
||||
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
||||
img.save(filepath, pnginfo=metadata,
|
||||
compress_level=self.compress_level)
|
||||
elif format == "JPEG":
|
||||
if img.mode == "RGBA":
|
||||
img = img.convert("RGB")
|
||||
img.save(filepath, quality=100, optimize=True)
|
||||
elif format == "WebP":
|
||||
img.save(filepath, lossless=True)
|
||||
|
||||
results.append({
|
||||
"filename": file,
|
||||
"subfolder": subfolder,
|
||||
"type": self.type,
|
||||
})
|
||||
counter += 1
|
||||
|
||||
return {"ui": {"images": results}}
|
||||
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
o1key SavePSD 节点
|
||||
将多个 IMAGE 图层合成为分层 PSD 文件
|
||||
手写 PSD 二进制格式,零外部依赖(仅 numpy + Pillow)
|
||||
"""
|
||||
|
||||
import os
|
||||
import struct
|
||||
import time
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
import folder_paths
|
||||
|
||||
|
||||
def _pad_even(data: bytes) -> bytes:
|
||||
if len(data) % 2:
|
||||
return data + b"\x00"
|
||||
return data
|
||||
|
||||
|
||||
def _pad4(data: bytes) -> bytes:
|
||||
return data + (b"\x00" * ((4 - (len(data) % 4)) % 4))
|
||||
|
||||
|
||||
def _pascal_name(name: str) -> bytes:
|
||||
raw = name.encode("macroman", errors="replace")[:255]
|
||||
data = bytes([len(raw)]) + raw
|
||||
return _pad4(data)
|
||||
|
||||
|
||||
def _unicode_name_block(name: str) -> bytes:
|
||||
payload = struct.pack(">I", len(name)) + name.encode("utf-16be")
|
||||
block = b"8BIM" + b"luni" + struct.pack(">I", len(payload)) + _pad_even(payload)
|
||||
return block
|
||||
|
||||
|
||||
def _layer_extra_data(name: str) -> bytes:
|
||||
data = b""
|
||||
data += struct.pack(">I", 0) # layer mask data length
|
||||
data += struct.pack(">I", 0) # layer blending ranges length
|
||||
data += _pascal_name(name)
|
||||
data += _unicode_name_block(name)
|
||||
return data
|
||||
|
||||
|
||||
def _alpha_bbox(rgba_arr: np.ndarray):
|
||||
"""找到 RGBA 数组中非透明区域的 bounding box。"""
|
||||
alpha = rgba_arr[:, :, 3]
|
||||
rows = np.any(alpha > 0, axis=1)
|
||||
cols = np.any(alpha > 0, axis=0)
|
||||
if not rows.any():
|
||||
return None
|
||||
top = int(np.argmax(rows))
|
||||
bottom = int(len(rows) - np.argmax(rows[::-1]))
|
||||
left = int(np.argmax(cols))
|
||||
right = int(len(cols) - np.argmax(cols[::-1]))
|
||||
return top, left, bottom, right
|
||||
|
||||
|
||||
def write_psd(filepath: str, layers: list, canvas_w: int, canvas_h: int):
|
||||
"""
|
||||
写入 PSD 文件。
|
||||
|
||||
layers: [(name, rgba_array), ...] 从底到顶排列
|
||||
rgba_array: numpy uint8 [H, W, 4]
|
||||
"""
|
||||
records = []
|
||||
channel_data_blocks = []
|
||||
layers_top_to_bottom = list(reversed(layers))
|
||||
|
||||
for name, rgba in layers_top_to_bottom:
|
||||
bbox = _alpha_bbox(rgba)
|
||||
if not bbox:
|
||||
continue
|
||||
top, left, bottom, right = bbox
|
||||
cropped = rgba[top:bottom, left:right]
|
||||
|
||||
# PLACEHOLDER_CHANNELS
|
||||
|
||||
channels = [
|
||||
(0, cropped[:, :, 0].tobytes(order="C")),
|
||||
(1, cropped[:, :, 1].tobytes(order="C")),
|
||||
(2, cropped[:, :, 2].tobytes(order="C")),
|
||||
(-1, cropped[:, :, 3].tobytes(order="C")),
|
||||
]
|
||||
channel_info = b""
|
||||
data_block = b""
|
||||
for channel_id, data in channels:
|
||||
channel_info += struct.pack(">hI", channel_id, 2 + len(data))
|
||||
data_block += struct.pack(">H", 0) + data # raw compression
|
||||
|
||||
extra = _layer_extra_data(name)
|
||||
record = b""
|
||||
record += struct.pack(">iiii", top, left, bottom, right)
|
||||
record += struct.pack(">H", len(channels))
|
||||
record += channel_info
|
||||
record += b"8BIM" + b"norm"
|
||||
record += bytes([255, 0, 0, 0]) # opacity=255, clipping, flags, filler
|
||||
record += struct.pack(">I", len(extra)) + extra
|
||||
records.append(record)
|
||||
channel_data_blocks.append(data_block)
|
||||
|
||||
if not records:
|
||||
raise ValueError("所有图层均为空(完全透明),无法生成 PSD")
|
||||
|
||||
# Layer and Mask Information
|
||||
layer_info = struct.pack(">h", len(records))
|
||||
layer_info += b"".join(records) + b"".join(channel_data_blocks)
|
||||
layer_info = _pad_even(layer_info)
|
||||
layer_info_block = struct.pack(">I", len(layer_info)) + layer_info
|
||||
global_mask = struct.pack(">I", 0)
|
||||
layer_mask_payload = layer_info_block + global_mask
|
||||
layer_and_mask = struct.pack(">I", len(layer_mask_payload)) + layer_mask_payload
|
||||
|
||||
# PLACEHOLDER_COMPOSITE
|
||||
|
||||
# Composite preview (flattened image for compatibility)
|
||||
comp = Image.new("RGBA", (canvas_w, canvas_h), (255, 255, 255, 255))
|
||||
for name, rgba in layers:
|
||||
layer_img = Image.fromarray(rgba, "RGBA")
|
||||
comp.alpha_composite(layer_img)
|
||||
comp_rgb = np.asarray(comp.convert("RGB"), dtype=np.uint8)
|
||||
composite_data = (
|
||||
struct.pack(">H", 0)
|
||||
+ comp_rgb[:, :, 0].tobytes(order="C")
|
||||
+ comp_rgb[:, :, 1].tobytes(order="C")
|
||||
+ comp_rgb[:, :, 2].tobytes(order="C")
|
||||
)
|
||||
|
||||
# Write PSD file
|
||||
with open(filepath, "wb") as f:
|
||||
# Header
|
||||
f.write(b"8BPS")
|
||||
f.write(struct.pack(">H", 1)) # version
|
||||
f.write(b"\x00" * 6) # reserved
|
||||
f.write(struct.pack(">HIIHH", 3, canvas_h, canvas_w, 8, 3))
|
||||
# Color Mode Data
|
||||
f.write(struct.pack(">I", 0))
|
||||
# Image Resources
|
||||
f.write(struct.pack(">I", 0))
|
||||
# Layer and Mask
|
||||
f.write(layer_and_mask)
|
||||
# Composite Image Data
|
||||
f.write(composite_data)
|
||||
|
||||
|
||||
# PLACEHOLDER_NODE
|
||||
|
||||
class O1keySavePSD:
|
||||
"""
|
||||
将多个 IMAGE 输入合成为分层 PSD 文件
|
||||
|
||||
每个输入作为独立图层,支持 RGBA 透明通道。
|
||||
图层从下到上排列(图层1在最底部)。
|
||||
使用 bbox 裁剪优化文件大小,包含合成预览层。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"批次图像": ("IMAGE", {
|
||||
"tooltip": "批次图像输入,每张图自动作为独立图层(支持RGBA透明)",
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"图层名称": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "每行一个图层名称,与图层顺序对应。留空则自动命名。",
|
||||
}),
|
||||
"文件名前缀": ("STRING", {
|
||||
"default": "o1key_layers",
|
||||
"tooltip": "输出 PSD 文件名前缀",
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("文件路径",)
|
||||
FUNCTION = "save_psd"
|
||||
CATEGORY = "o1key/image"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def save_psd(self, 批次图像, 图层名称: str = "", 文件名前缀: str = "o1key_layers", **kwargs):
|
||||
# 将批次 tensor [B, H, W, C] 拆为单张列表
|
||||
if 批次图像.dim() == 3:
|
||||
layer_tensors = [批次图像]
|
||||
else:
|
||||
layer_tensors = [批次图像[i] for i in range(批次图像.shape[0])]
|
||||
|
||||
names = [n.strip() for n in 图层名称.split("\n") if n.strip()]
|
||||
|
||||
# 确定画布尺寸
|
||||
max_h, max_w = 0, 0
|
||||
for t in layer_tensors:
|
||||
h, w = t.shape[0], t.shape[1]
|
||||
max_h = max(max_h, h)
|
||||
max_w = max(max_w, w)
|
||||
|
||||
# 转换为 [(name, rgba_array), ...] 格式
|
||||
layers = []
|
||||
for idx, tensor in enumerate(layer_tensors):
|
||||
arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||||
h, w = arr.shape[0], arr.shape[1]
|
||||
channels = arr.shape[2] if arr.ndim == 3 else 1
|
||||
|
||||
if channels == 3:
|
||||
rgba = np.zeros((max_h, max_w, 4), dtype=np.uint8)
|
||||
rgba[:h, :w, :3] = arr
|
||||
rgba[:h, :w, 3] = 255
|
||||
elif channels == 4:
|
||||
rgba = np.zeros((max_h, max_w, 4), dtype=np.uint8)
|
||||
rgba[:h, :w] = arr
|
||||
else:
|
||||
rgba = np.zeros((max_h, max_w, 4), dtype=np.uint8)
|
||||
rgba[:h, :w, 0] = rgba[:h, :w, 1] = rgba[:h, :w, 2] = arr[:, :, 0] if arr.ndim == 3 else arr
|
||||
rgba[:h, :w, 3] = 255
|
||||
|
||||
name = names[idx] if idx < len(names) else f"图层 {idx + 1}"
|
||||
layers.append((name, rgba))
|
||||
print(f"[o1key SavePSD] 图层 '{name}': {w}×{h}")
|
||||
|
||||
# 写入 PSD
|
||||
output_dir = folder_paths.get_output_directory()
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"{文件名前缀}_{timestamp}.psd"
|
||||
filepath = os.path.join(output_dir, filename)
|
||||
|
||||
write_psd(filepath, layers, max_w, max_h)
|
||||
|
||||
size_kb = os.path.getsize(filepath) / 1024
|
||||
print(f"[o1key SavePSD] 完成: {filepath} ({size_kb:.0f}KB, "
|
||||
f"{len(layers)} 层, {max_w}×{max_h})")
|
||||
return (filepath,)
|
||||
@@ -0,0 +1,490 @@
|
||||
"""
|
||||
Seedance 视频生成节点
|
||||
节点列表:
|
||||
- Seedance: 文生视频 / 图生视频 / 首尾帧生视频(根据图片输入自动切换模式)
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
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, pil_to_tensor
|
||||
from ..utils.r2_uploader import upload_video, upload_audio
|
||||
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
|
||||
from comfy_api.latest import InputImpl
|
||||
|
||||
|
||||
# ── 模型列表 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
_MODELS = [
|
||||
"doubao-seedance-2-0-260128",
|
||||
]
|
||||
|
||||
_RESOLUTIONS = ["720p", "1080p", "480p"]
|
||||
|
||||
_MAX_IMAGE_BYTES = 30 * 1024 * 1024
|
||||
_MAX_REQUEST_BODY_BYTES = 64 * 1024 * 1024
|
||||
|
||||
|
||||
# ── 模型能力判断 ──────────────────────────────────────────────────────────────
|
||||
|
||||
def _supports_camera_fixed(model: str) -> bool:
|
||||
"""2.0 系列不支持固定镜头"""
|
||||
return False # 当前仅 2.0 模型,均不支持
|
||||
|
||||
|
||||
# ── 工具函数 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _format_mb(size_bytes: int) -> str:
|
||||
return f"{size_bytes / 1024 / 1024:.2f}MB"
|
||||
|
||||
|
||||
def _tensor_to_base64_url(tensor, label: str = "图片") -> str:
|
||||
"""ComfyUI IMAGE tensor → data:image/png;base64,xxx"""
|
||||
pil_images = tensor_to_pil(tensor)
|
||||
image = pil_images[0]
|
||||
if image.mode == "RGBA":
|
||||
image = image.convert("RGB")
|
||||
|
||||
buffered = io.BytesIO()
|
||||
image.save(buffered, format="PNG")
|
||||
image_bytes = buffered.getvalue()
|
||||
image_size = len(image_bytes)
|
||||
|
||||
if image_size > _MAX_IMAGE_BYTES:
|
||||
raise ValueError(
|
||||
f"Seedance {label}大小 {_format_mb(image_size)} 超过单张图片 "
|
||||
f"{_format_mb(_MAX_IMAGE_BYTES)} 限制,请先压缩或缩小图片。"
|
||||
)
|
||||
|
||||
b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||||
return f"data:image/png;base64,{b64}"
|
||||
|
||||
|
||||
def _validate_request_body_size(body: dict, tag: str):
|
||||
body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
|
||||
if body_size > _MAX_REQUEST_BODY_BYTES:
|
||||
raise ValueError(
|
||||
f"{tag} 请求体大小 {_format_mb(body_size)} 超过 "
|
||||
f"{_format_mb(_MAX_REQUEST_BODY_BYTES)} 限制,请减少参考图片数量或降低图片尺寸。"
|
||||
)
|
||||
print(
|
||||
f"[{tag}] 请求体大小: {_format_mb(body_size)} "
|
||||
f"(限制 {_format_mb(_MAX_REQUEST_BODY_BYTES)})"
|
||||
)
|
||||
|
||||
|
||||
|
||||
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": ""}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"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,
|
||||
}
|
||||
|
||||
_validate_request_body_size(body, tag)
|
||||
|
||||
# 保存路径(临时文件,避免与下游保存节点重复落盘)
|
||||
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix=f"{file_prefix}_")
|
||||
|
||||
client = SeedanceClient()
|
||||
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
|
||||
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": ""}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"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)
|
||||
network_route = _first(kwargs.get("网络线路"), "全球加速")
|
||||
|
||||
# 参考图片: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 idx, img_tensor in enumerate(imgs, start=1):
|
||||
# 每个 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, f"参考图片{idx}")
|
||||
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
|
||||
|
||||
_validate_request_body_size(body, "Seedance多模态")
|
||||
|
||||
# ── 保存路径(临时文件,避免与下游保存节点重复落盘)──────────────────
|
||||
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="seedance_mm_")
|
||||
|
||||
client = SeedanceClient()
|
||||
client.base_url = get_base_url_by_route(network_route)
|
||||
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 多模态参考生视频",
|
||||
}
|
||||
@@ -0,0 +1,526 @@
|
||||
"""
|
||||
Sora 视频生成节点
|
||||
ComfyUI 自定义节点,调用 Sora API 生成视频
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from math import gcd
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..clients.sora_client import SoraClient
|
||||
from ..models_config import (
|
||||
get_enabled_sora_models,
|
||||
get_all_sora_seconds,
|
||||
get_all_sora_sizes,
|
||||
get_sora_supported_seconds,
|
||||
get_sora_supported_sizes,
|
||||
get_sora_seconds_with_labels,
|
||||
get_sora_sizes_with_labels,
|
||||
SORA_MODELS,
|
||||
)
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ SoraVideo: comfy.utils.ProgressBar 不可用,将只使用终端进度显示")
|
||||
|
||||
|
||||
def _size_to_display(size: str) -> str:
|
||||
"""
|
||||
将 'WxH' 格式的分辨率转换为友好显示名。
|
||||
|
||||
例如:
|
||||
"720x1280" → "720P 9:16"
|
||||
"1280x720" → "720P 16:9"
|
||||
"1024x1792" → "1K 4:7"
|
||||
"1792x1024" → "1K 7:4"
|
||||
|
||||
Args:
|
||||
size: 分辨率字符串,格式 "WxH"
|
||||
|
||||
Returns:
|
||||
友好显示名字符串
|
||||
"""
|
||||
parts = size.lower().split("x")
|
||||
w, h = int(parts[0]), int(parts[1])
|
||||
short_side = min(w, h)
|
||||
if short_side >= 3840:
|
||||
res = "4K"
|
||||
elif short_side >= 1920:
|
||||
res = "2K"
|
||||
elif short_side >= 1080:
|
||||
res = "1K"
|
||||
elif short_side >= 720:
|
||||
res = "720P"
|
||||
elif short_side >= 480:
|
||||
res = "480P"
|
||||
else:
|
||||
res = f"{short_side}P"
|
||||
g = gcd(w, h)
|
||||
ratio = f"{w // g}:{h // g}"
|
||||
return f"{res} {ratio} ({size})"
|
||||
|
||||
|
||||
def _build_size_display_map(sizes: list) -> dict:
|
||||
"""
|
||||
构建 显示名 → 实际值 映射字典。
|
||||
|
||||
Args:
|
||||
sizes: 实际分辨率列表,如 ["720x1280", "1280x720"]
|
||||
|
||||
Returns:
|
||||
字典,key 为显示名,value 为实际分辨率字符串
|
||||
"""
|
||||
mapping = {}
|
||||
for size in sizes:
|
||||
display = _size_to_display(size)
|
||||
if display in mapping:
|
||||
# 极少数情况下防止重名
|
||||
display = f"{display} ({size})"
|
||||
mapping[display] = size
|
||||
return mapping
|
||||
|
||||
|
||||
def _get_video_output_dir() -> str:
|
||||
"""获取视频输出目录: ComfyUI/output/video"""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
base = folder_paths.get_output_directory()
|
||||
else:
|
||||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
base = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output")
|
||||
video_dir = os.path.join(base, "video")
|
||||
os.makedirs(video_dir, exist_ok=True)
|
||||
return video_dir
|
||||
|
||||
|
||||
def _get_next_counter(directory: str, prefix: str) -> int:
|
||||
"""扫描目录,获取下一个可用的文件计数器"""
|
||||
if not os.path.exists(directory):
|
||||
return 1
|
||||
pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)")
|
||||
max_counter = 0
|
||||
for f in os.listdir(directory):
|
||||
m = pattern.match(f)
|
||||
if m:
|
||||
max_counter = max(max_counter, int(m.group(1)))
|
||||
return max_counter + 1
|
||||
|
||||
|
||||
def _fit_image_to_target(image, target_size: str):
|
||||
"""
|
||||
将参考图片按 "等比缩放覆盖 + 居中裁剪" 策略适配到目标分辨率。
|
||||
|
||||
策略 (Cover Crop):
|
||||
1. 比较图片宽高比和目标宽高比
|
||||
2. 等比缩放,使图片最短边刚好覆盖目标对应边(图片完全覆盖目标区域)
|
||||
3. 居中裁剪多余部分,得到精确目标尺寸
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
target_size: 目标分辨率字符串,格式 "WxH"(如 "720x1280")
|
||||
|
||||
Returns:
|
||||
适配后的 PIL Image 对象
|
||||
"""
|
||||
from PIL import Image as PILImage
|
||||
|
||||
# 解析目标尺寸
|
||||
parts = target_size.lower().split("x")
|
||||
target_w, target_h = int(parts[0]), int(parts[1])
|
||||
|
||||
src_w, src_h = image.size
|
||||
src_ratio = src_w / src_h
|
||||
target_ratio = target_w / target_h
|
||||
|
||||
# 宽高比一致且尺寸不超过目标,无需处理
|
||||
if abs(src_ratio - target_ratio) < 0.01 and src_w <= target_w and src_h <= target_h:
|
||||
return image
|
||||
|
||||
print(f"Sora: 参考图片 {src_w}x{src_h} (比例 {src_ratio:.2f}) → 目标 {target_w}x{target_h} (比例 {target_ratio:.2f})")
|
||||
|
||||
# 获取高质量重采样滤波器
|
||||
resample = PILImage.Resampling.LANCZOS if hasattr(PILImage, "Resampling") else PILImage.LANCZOS
|
||||
|
||||
# Cover Crop: 缩放使图片完全覆盖目标区域,然后居中裁剪
|
||||
if src_ratio > target_ratio:
|
||||
# 图片更宽:以高度为基准缩放,裁左右
|
||||
scale = target_h / src_h
|
||||
new_w = round(src_w * scale)
|
||||
new_h = target_h
|
||||
image = image.resize((new_w, new_h), resample=resample)
|
||||
# 居中裁剪宽度
|
||||
left = (new_w - target_w) // 2
|
||||
image = image.crop((left, 0, left + target_w, target_h))
|
||||
else:
|
||||
# 图片更高(或一样):以宽度为基准缩放,裁上下
|
||||
scale = target_w / src_w
|
||||
new_w = target_w
|
||||
new_h = round(src_h * scale)
|
||||
image = image.resize((new_w, new_h), resample=resample)
|
||||
# 居中裁剪高度
|
||||
top = (new_h - target_h) // 2
|
||||
image = image.crop((0, top, target_w, top + target_h))
|
||||
|
||||
print(f"Sora: 参考图片已适配为 {image.size[0]}x{image.size[1]}")
|
||||
return image
|
||||
|
||||
|
||||
def _compress_image_for_upload(
|
||||
image,
|
||||
target_size: Optional[str] = None,
|
||||
) -> bytes:
|
||||
"""
|
||||
将 PIL Image 适配目标分辨率并编码为 PNG 字节,用于上传。
|
||||
|
||||
============================================================
|
||||
⚠️ 已验证可用的标准做法,请勿随意修改以下编码逻辑!
|
||||
============================================================
|
||||
经过多轮调试(2026-02-28),以下参数组合为唯一验证成功的方案:
|
||||
|
||||
1. 图片格式:PNG(format="PNG")
|
||||
- 不可改为 JPEG —— API 会校验 Content-Type,抓包确认服务端使用 image/png
|
||||
- 不可使用 base64 字符串 —— 会报 "expected a file, got a string"
|
||||
- 不可使用 data URI —— 服务端不识别,返回 500
|
||||
|
||||
2. 图片尺寸:必须与视频分辨率完全一致(target_size)
|
||||
- 不可缩放降采样 —— 会报 "Inpaint image must match the requested width and height"
|
||||
- 尺寸由 _fit_image_to_target() 保证(等比缩放 + 居中裁剪)
|
||||
|
||||
3. 上传方式:由调用方(sora_client.py)以 multipart/form-data 文件字段上传
|
||||
- filename="reference.png", content_type="image/png"
|
||||
- 不可改回 application/json —— 服务端校验 input_reference 必须为 file 类型
|
||||
============================================================
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
target_size: 目标分辨率字符串 "WxH"(如 "720x1280")
|
||||
|
||||
Returns:
|
||||
PNG 格式的二进制字节
|
||||
"""
|
||||
from io import BytesIO
|
||||
|
||||
# 统一转换为 RGB(去除透明通道及其他模式)
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
|
||||
# 适配到目标分辨率(等比缩放 + 居中裁剪)
|
||||
# ⚠️ 必须保持此尺寸不变,API 强制要求参考图片与视频分辨率完全一致
|
||||
if target_size:
|
||||
image = _fit_image_to_target(image, target_size)
|
||||
|
||||
# ⚠️ 必须使用 PNG 格式,不可改为 JPEG 或其他格式
|
||||
buffered = BytesIO()
|
||||
image.save(buffered, format="PNG")
|
||||
size_kb = buffered.tell() / 1024
|
||||
print(f"Sora: 参考图片编码为 PNG,{size_kb:.0f} KB ({image.size[0]}x{image.size[1]})")
|
||||
return buffered.getvalue()
|
||||
|
||||
|
||||
class SoraVideo:
|
||||
"""
|
||||
Sora 视频生成节点
|
||||
|
||||
功能:
|
||||
- 文生视频:基于提示词生成视频
|
||||
- 图生视频:基于参考图片和提示词生成视频
|
||||
- 异步轮询:自动等待生成完成并下载
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
from ..models_config import SECONDS_DISPLAY_MAP, RESOLUTION_DISPLAY_MAP
|
||||
|
||||
enabled_models = get_enabled_sora_models()
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个 Sora 模型"]
|
||||
|
||||
# 构建秒数选项列表(按数字顺序排序)
|
||||
# 格式: ["4", "8", "10", "12", "15", "25(pro)"]
|
||||
all_seconds_display = []
|
||||
seen_seconds = set()
|
||||
for model_id in enabled_models:
|
||||
supported = get_sora_supported_seconds(model_id)
|
||||
for s in supported:
|
||||
if s not in seen_seconds:
|
||||
seen_seconds.add(s)
|
||||
display = SECONDS_DISPLAY_MAP.get(s, str(s))
|
||||
all_seconds_display.append((s, display))
|
||||
# 按秒数数值排序
|
||||
all_seconds_display = sorted(all_seconds_display, key=lambda x: x[0])
|
||||
seconds_options = [d for _, d in all_seconds_display] if all_seconds_display else ["4", "8", "12"]
|
||||
|
||||
# 构建分辨率选项列表(去重)
|
||||
# 格式: ["720P", "1080P"]
|
||||
seen_resolutions = set()
|
||||
for model_id in enabled_models:
|
||||
supported = get_sora_supported_sizes(model_id)
|
||||
for size in supported:
|
||||
if size in RESOLUTION_DISPLAY_MAP:
|
||||
res_name, _ = RESOLUTION_DISPLAY_MAP[size]
|
||||
seen_resolutions.add(res_name)
|
||||
resolution_options = sorted(list(seen_resolutions)) if seen_resolutions else ["720P"]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "A calico cat playing a piano on stage",
|
||||
"multiline": True,
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0],
|
||||
}),
|
||||
"分辨率": (resolution_options, {
|
||||
"default": resolution_options[0] if resolution_options else "720P",
|
||||
}),
|
||||
"宽高比": (["竖屏", "横屏"], {
|
||||
"default": "竖屏",
|
||||
}),
|
||||
"视频时长": (seconds_options, {
|
||||
"default": seconds_options[0] if seconds_options else "4",
|
||||
}),
|
||||
"生成数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"参考图片": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("预览视频",)
|
||||
FUNCTION = "generate_video"
|
||||
CATEGORY = "video/generation"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Sora 视频生成节点。\n"
|
||||
"支持文生视频和图生视频,自动轮询任务状态并下载视频。\n"
|
||||
"视频保存到 ComfyUI/output/video/ 目录。\n\n"
|
||||
"【模型说明】\n"
|
||||
"• sora-2:官方模型,支持 4/8/12秒、720P 分辨率\n"
|
||||
"• sora-2-pro:增强模型,支持全时长(含25秒)、1080P 分辨率\n\n"
|
||||
"【时长说明】\n"
|
||||
"• 25(pro):仅 sora-2-pro 支持的25秒时长\n\n"
|
||||
"【分辨率说明】\n"
|
||||
"• 720P:sora-2 和 sora-2-pro 均支持\n"
|
||||
"• 1080P:仅 sora-2-pro 支持的高清分辨率"
|
||||
)
|
||||
|
||||
def generate_video(
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str,
|
||||
**kwargs,
|
||||
) -> Tuple[str]:
|
||||
from ..models_config import SECONDS_DISPLAY_MAP, RESOLUTION_DISPLAY_MAP
|
||||
|
||||
视频时长_display = kwargs.pop("视频时长", "4")
|
||||
分辨率_display = kwargs.pop("分辨率", "720P")
|
||||
宽高比 = kwargs.pop("宽高比", "竖屏")
|
||||
生成数量 = kwargs.pop("生成数量", 1)
|
||||
seed = kwargs.pop("seed", 0)
|
||||
start_time = time.time()
|
||||
|
||||
# 解析秒数显示值(如 "25(pro)" → 25)
|
||||
seconds = 4 # 默认
|
||||
for actual, display in SECONDS_DISPLAY_MAP.items():
|
||||
if display == 视频时长_display:
|
||||
seconds = actual
|
||||
break
|
||||
# 如果找不到映射,尝试直接解析数字
|
||||
if seconds == 4 and 视频时长_display != "4":
|
||||
try:
|
||||
seconds = int(视频时长_display.replace("(pro)", ""))
|
||||
except ValueError:
|
||||
seconds = 4
|
||||
|
||||
# 根据分辨率和宽高比确定实际分辨率值
|
||||
分辨率 = "720x1280" # 默认
|
||||
for actual, (res_name, orientation) in RESOLUTION_DISPLAY_MAP.items():
|
||||
if res_name == 分辨率_display and orientation == 宽高比:
|
||||
分辨率 = actual
|
||||
break
|
||||
|
||||
# 检查参考图片
|
||||
ref_image = kwargs.get("参考图片")
|
||||
ref_image_bytes = None
|
||||
if ref_image is not None:
|
||||
pil_images = tensor_to_pil(ref_image)
|
||||
if pil_images:
|
||||
ref_image_bytes = _compress_image_for_upload(pil_images[0], target_size=分辨率)
|
||||
|
||||
mode_str = "图生视频 (含参考图)" if ref_image_bytes else "文生视频"
|
||||
# 获取用户友好的显示值用于日志
|
||||
seconds_display = SECONDS_DISPLAY_MAP.get(seconds, str(seconds))
|
||||
res_display = f"{分辨率_display} {宽高比}"
|
||||
if 生成数量 > 1:
|
||||
print(f"Sora: {mode_str} | 并发{生成数量}个 | {模型} | {seconds_display} | {res_display}")
|
||||
else:
|
||||
print(f"Sora: {mode_str} | {模型} | {seconds_display} | {res_display}")
|
||||
|
||||
# 校验参数兼容性
|
||||
supported_seconds = get_sora_supported_seconds(模型)
|
||||
if supported_seconds and seconds not in supported_seconds:
|
||||
# 构建带标签的支持时长列表
|
||||
supported_labels = []
|
||||
for s in supported_seconds:
|
||||
display = SECONDS_DISPLAY_MAP.get(s, str(s))
|
||||
supported_labels.append(display)
|
||||
raise ValueError(
|
||||
f"时长 {SECONDS_DISPLAY_MAP.get(seconds, str(seconds))} 与模型 \"{模型}\" 不兼容!\n"
|
||||
f"该模型支持的时长: {', '.join(supported_labels)}"
|
||||
)
|
||||
|
||||
supported_sizes = get_sora_supported_sizes(模型)
|
||||
if supported_sizes and 分辨率 not in supported_sizes:
|
||||
# 检查该分辨率是否为Pro独占
|
||||
pro_only_sizes = ["1024x1792", "1792x1024"]
|
||||
_, orientation = RESOLUTION_DISPLAY_MAP.get(分辨率, (分辨率, ""))
|
||||
extra_hint = f"\n提示:1080P {orientation} 为 sora-2-pro 独占,请切换模型或选择720P。" if 分辨率 in pro_only_sizes else ""
|
||||
raise ValueError(
|
||||
f"分辨率 \"{分辨率_display} {宽高比}\" 与模型 \"{模型}\" 不兼容!"
|
||||
f"支持的分辨率: {', '.join(supported_sizes)}" + extra_hint
|
||||
)
|
||||
|
||||
# 准备保存路径
|
||||
video_dir = _get_video_output_dir()
|
||||
counter = _get_next_counter(video_dir, "sora")
|
||||
|
||||
# ProgressBar
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(生成数量 if 生成数量 > 1 else 100)
|
||||
|
||||
try:
|
||||
if self.client is None:
|
||||
self.client = SoraClient()
|
||||
|
||||
if 生成数量 == 1:
|
||||
# ── 单个视频:保留详细进度(提交→轮询→下载)
|
||||
save_path = os.path.join(video_dir, f"sora_{counter:05d}.mp4")
|
||||
last_progress = [0]
|
||||
|
||||
def progress_callback(progress_pct: int):
|
||||
print(
|
||||
f"\rSora: 生成中... 进度: {progress_pct}%",
|
||||
end="", flush=True
|
||||
)
|
||||
if pbar is not None and progress_pct > last_progress[0]:
|
||||
pbar.update(progress_pct - last_progress[0])
|
||||
last_progress[0] = progress_pct
|
||||
|
||||
def on_stage(stage: str):
|
||||
if stage == "submitting":
|
||||
print("Sora: 正在提交视频生成任务...")
|
||||
elif stage.startswith("submitted:"):
|
||||
vid = stage.split(":", 1)[1]
|
||||
print(f"Sora: 视频任务已提交,ID: {vid}")
|
||||
elif stage == "polling":
|
||||
print("Sora: 等待视频生成...")
|
||||
elif stage == "downloading":
|
||||
print("") # 换行(结束 \r 行)
|
||||
print("Sora: 视频生成完成,正在下载...")
|
||||
|
||||
result_path = self.client.generate_video_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
seconds=seconds,
|
||||
size=分辨率,
|
||||
save_path=save_path,
|
||||
input_reference_bytes=ref_image_bytes,
|
||||
seed=seed,
|
||||
progress_callback=progress_callback,
|
||||
on_stage=on_stage,
|
||||
)
|
||||
result_paths = [result_path]
|
||||
|
||||
else:
|
||||
# ── 批量并发:同时提交多个任务
|
||||
save_paths = [
|
||||
os.path.join(video_dir, f"sora_{counter + i:05d}.mp4")
|
||||
for i in range(生成数量)
|
||||
]
|
||||
success_count = [0]
|
||||
fail_count = [0]
|
||||
|
||||
def batch_progress_callback(current: int, total: int, success: bool, error_msg):
|
||||
if success:
|
||||
success_count[0] += 1
|
||||
print(f"Sora: 第 {current}/{total} 个视频完成 ✓")
|
||||
else:
|
||||
fail_count[0] += 1
|
||||
print(f"Sora: 第 {current}/{total} 个视频失败 ✗")
|
||||
if error_msg:
|
||||
print(f"原始错误详情:\n{error_msg}")
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
print(f"Sora: 正在并发提交 {生成数量} 个视频任务,请耐心等待...")
|
||||
result_paths = self.client.generate_batch_videos_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
seconds=seconds,
|
||||
size=分辨率,
|
||||
save_paths=save_paths,
|
||||
input_reference_bytes=ref_image_bytes,
|
||||
seed=seed,
|
||||
progress_callback=batch_progress_callback,
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.2f}s" if elapsed >= 1 else f"{elapsed:.3f}s"
|
||||
print(f"Sora: 完成!总耗时 {time_str} | 已生成 {len(result_paths)} 个视频")
|
||||
for p in result_paths:
|
||||
print(f" → {p}")
|
||||
|
||||
output_path = "\n".join(result_paths)
|
||||
return (output_path,)
|
||||
|
||||
except ValueError as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nSora: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nSora: ❌ {error_msg}")
|
||||
raise RuntimeError(error_msg) from None
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nSora: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
finally:
|
||||
if self.client is not None:
|
||||
try:
|
||||
balance_data = self.client.query_balance_sync()
|
||||
balance_info = self.client.format_balance_info(balance_data)
|
||||
print(f"Sora: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,28 @@
|
||||
"""
|
||||
流式文本预览节点
|
||||
接收文本输入,支持 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": (文本,)}
|
||||
@@ -0,0 +1,513 @@
|
||||
"""
|
||||
全能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, NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.file_types import FileList
|
||||
|
||||
# ============================================================================
|
||||
# 模型配置
|
||||
# ============================================================================
|
||||
|
||||
SUPPORTED_MODELS = [
|
||||
"gpt-5.5",
|
||||
"gemini-3.1-pro-preview",
|
||||
"deepseek-v4-pro",
|
||||
"claude-opus-4-7",
|
||||
"claude-opus-4-6",
|
||||
"gemini-3.5-flash",
|
||||
"doubao-seed-2.0-pro",
|
||||
]
|
||||
|
||||
# 图片缩放最大尺寸
|
||||
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": {
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {
|
||||
"default": "全球加速"
|
||||
}),
|
||||
"模型": (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,
|
||||
网络线路: str = "全球加速",
|
||||
图片: Optional[torch.Tensor] = None,
|
||||
视频=None,
|
||||
文件: Optional[FileList] = None,
|
||||
令牌: str = "",
|
||||
node_id: str = "",
|
||||
) -> Tuple[str]:
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
self._ensure_config()
|
||||
self._base_url = get_base_url_by_route(网络线路)
|
||||
|
||||
# 如果用户传入了自定义令牌,则覆盖默认 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
|
||||
@@ -0,0 +1,422 @@
|
||||
"""
|
||||
Google Veo 视频生成节点
|
||||
ComfyUI 自定义节点,调用 Veo API 生成视频
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..clients.veo_client import VeoClient
|
||||
from ..models_config import (
|
||||
get_enabled_veo_models,
|
||||
VEO_MODELS,
|
||||
VEO_RESOLUTION_MAP,
|
||||
)
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ GoogleVeo: comfy.utils.ProgressBar 不可用,将只使用终端进度显示")
|
||||
|
||||
|
||||
def _get_video_output_dir() -> str:
|
||||
"""获取视频输出目录: ComfyUI/output/video"""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
base = folder_paths.get_output_directory()
|
||||
else:
|
||||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
base = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output")
|
||||
video_dir = os.path.join(base, "video")
|
||||
os.makedirs(video_dir, exist_ok=True)
|
||||
return video_dir
|
||||
|
||||
|
||||
def _get_next_counter(directory: str, prefix: str) -> int:
|
||||
"""扫描目录,获取下一个可用的文件计数器"""
|
||||
if not os.path.exists(directory):
|
||||
return 1
|
||||
pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)")
|
||||
max_counter = 0
|
||||
for f in os.listdir(directory):
|
||||
m = pattern.match(f)
|
||||
if m:
|
||||
max_counter = max(max_counter, int(m.group(1)))
|
||||
return max_counter + 1
|
||||
|
||||
|
||||
def _fit_image_to_target(image, target_size: str):
|
||||
"""
|
||||
将参考图片按 "等比缩放覆盖 + 居中裁剪" 策略适配到目标分辨率。
|
||||
"""
|
||||
from PIL import Image as PILImage
|
||||
|
||||
parts = target_size.lower().split("x")
|
||||
target_w, target_h = int(parts[0]), int(parts[1])
|
||||
|
||||
src_w, src_h = image.size
|
||||
src_ratio = src_w / src_h
|
||||
target_ratio = target_w / target_h
|
||||
|
||||
if abs(src_ratio - target_ratio) < 0.01 and src_w <= target_w and src_h <= target_h:
|
||||
return image
|
||||
|
||||
print(f"Veo: 参考图片 {src_w}x{src_h} (比例 {src_ratio:.2f}) → 目标 {target_w}x{target_h} (比例 {target_ratio:.2f})")
|
||||
|
||||
resample = PILImage.Resampling.LANCZOS if hasattr(PILImage, "Resampling") else PILImage.LANCZOS
|
||||
|
||||
if src_ratio > target_ratio:
|
||||
scale = target_h / src_h
|
||||
new_w = round(src_w * scale)
|
||||
new_h = target_h
|
||||
image = image.resize((new_w, new_h), resample=resample)
|
||||
left = (new_w - target_w) // 2
|
||||
image = image.crop((left, 0, left + target_w, target_h))
|
||||
else:
|
||||
scale = target_w / src_w
|
||||
new_w = target_w
|
||||
new_h = round(src_h * scale)
|
||||
image = image.resize((new_w, new_h), resample=resample)
|
||||
top = (new_h - target_h) // 2
|
||||
image = image.crop((0, top, target_w, top + target_h))
|
||||
|
||||
print(f"Veo: 参考图片已适配为 {image.size[0]}x{image.size[1]}")
|
||||
return image
|
||||
|
||||
|
||||
def _compress_image_to_bytes(image, target_size: Optional[str] = None) -> bytes:
|
||||
"""
|
||||
将 PIL Image 适配目标分辨率并编码为 PNG 字节
|
||||
"""
|
||||
from io import BytesIO
|
||||
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
|
||||
if target_size:
|
||||
image = _fit_image_to_target(image, target_size)
|
||||
|
||||
buffered = BytesIO()
|
||||
image.save(buffered, format="PNG")
|
||||
size_kb = buffered.tell() / 1024
|
||||
print(f"Veo: 参考图片编码为 PNG,{size_kb:.0f} KB ({image.size[0]}x{image.size[1]})")
|
||||
return buffered.getvalue()
|
||||
|
||||
|
||||
class GoogleVeo:
|
||||
"""
|
||||
Google Veo 视频生成节点
|
||||
|
||||
功能:
|
||||
- 文生视频:基于提示词生成视频
|
||||
- 图生视频:基于首帧/尾帧/参考图生成视频
|
||||
- 异步轮询:自动等待生成完成并下载
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
enabled_models = get_enabled_veo_models()
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用 Veo 模型"]
|
||||
|
||||
# 分辨率选项
|
||||
resolution_options = ["720p", "1080p", "4K"]
|
||||
|
||||
# 宽高比选项
|
||||
aspect_ratio_options = ["16:9", "9:16"]
|
||||
|
||||
# 视频秒数选项
|
||||
seconds_options = ["4", "6", "8"]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "A calico cat playing a piano on stage",
|
||||
"multiline": True,
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0] if enabled_models else "Veo3.1",
|
||||
}),
|
||||
"分辨率": (resolution_options, {
|
||||
"default": "720p",
|
||||
}),
|
||||
"宽高比": (aspect_ratio_options, {
|
||||
"default": "9:16",
|
||||
}),
|
||||
"视频时长": (seconds_options, {
|
||||
"default": "8",
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff,
|
||||
}),
|
||||
"生成数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"首帧": ("IMAGE",),
|
||||
"尾帧": ("IMAGE",),
|
||||
"参考图": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("预览视频",)
|
||||
FUNCTION = "generate_video"
|
||||
CATEGORY = "video/generation"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Google Veo 视频生成节点。\n"
|
||||
"支持文生视频和图生视频(图生视频支持首帧、尾帧、参考图)。\n"
|
||||
"视频保存到 ComfyUI/output/video/ 目录。\n\n"
|
||||
"【模型说明】\n"
|
||||
"• Veo3.1:Google 最新视频生成模型\n\n"
|
||||
"【分辨率说明】\n"
|
||||
"• 720p:标清\n"
|
||||
"• 1080p:高清\n"
|
||||
"• 4K:超高清\n\n"
|
||||
"【时长说明】\n"
|
||||
"• 4秒:短视频\n"
|
||||
"• 6秒:标准\n"
|
||||
"• 8秒:长视频(默认)\n\n"
|
||||
"【图生视频说明】\n"
|
||||
"• 首帧:视频开始的第一帧图像\n"
|
||||
"• 尾帧:视频结束时的最后一帧图像\n"
|
||||
"• 参考图:参考图像(与首帧/尾帧配合使用)\n"
|
||||
"• 至少需要提供首帧或参考图之一"
|
||||
)
|
||||
|
||||
def generate_video(
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str,
|
||||
**kwargs,
|
||||
) -> Tuple[str]:
|
||||
分辨率 = kwargs.pop("分辨率", "720p")
|
||||
宽高比 = kwargs.pop("宽高比", "9:16")
|
||||
视频时长 = kwargs.pop("视频时长", "8")
|
||||
seed = kwargs.pop("seed", 0)
|
||||
生成数量 = kwargs.pop("生成数量", 1)
|
||||
start_time = time.time()
|
||||
|
||||
# 解析视频时长
|
||||
seconds = int(视频时长)
|
||||
|
||||
# 解析分辨率和宽高比,映射到模型名称
|
||||
size_key = f"{分辨率}_{宽高比}"
|
||||
actual_size = VEO_RESOLUTION_MAP.get(size_key)
|
||||
if not actual_size:
|
||||
# 默认值
|
||||
actual_size = "720x1280" # 720p 9:16
|
||||
|
||||
# 检查是否有参考图输入
|
||||
首帧 = kwargs.get("首帧")
|
||||
尾帧 = kwargs.get("尾帧")
|
||||
参考图 = kwargs.get("参考图")
|
||||
|
||||
has_image = 首帧 is not None or 尾帧 is not None or 参考图 is not None
|
||||
|
||||
# 根据是否有图片选择模型前缀
|
||||
if has_image:
|
||||
model_prefix = "veo3.1"
|
||||
else:
|
||||
model_prefix = "veo3.1"
|
||||
|
||||
# 构建完整模型名称
|
||||
# 格式: veo3.1-portrait / veo3.1-landscape / veo3.1-portrait-fl / veo3.1-landscape-fl 等
|
||||
if 分辨率 == "720p":
|
||||
res_suffix = ""
|
||||
if 宽高比 == "9:16":
|
||||
orientation = "portrait"
|
||||
else:
|
||||
orientation = "landscape"
|
||||
elif 分辨率 == "1080p":
|
||||
res_suffix = "-hd"
|
||||
if 宽高比 == "9:16":
|
||||
orientation = "portrait"
|
||||
else:
|
||||
orientation = "landscape"
|
||||
else: # 4K
|
||||
res_suffix = "-4k"
|
||||
if 宽高比 == "9:16":
|
||||
orientation = "portrait"
|
||||
else:
|
||||
orientation = "landscape"
|
||||
|
||||
# 图生视频添加 -fl 后缀
|
||||
if has_image:
|
||||
model_suffix = f"-{orientation}-fl{res_suffix}"
|
||||
else:
|
||||
model_suffix = f"-{orientation}{res_suffix}"
|
||||
|
||||
model = f"{model_prefix}{model_suffix}"
|
||||
|
||||
# 准备图片字节
|
||||
first_frame_bytes = None
|
||||
last_frame_bytes = None
|
||||
reference_bytes = None
|
||||
|
||||
if 首帧 is not None:
|
||||
pil_images = tensor_to_pil(首帧)
|
||||
if pil_images:
|
||||
first_frame_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size)
|
||||
|
||||
if 尾帧 is not None:
|
||||
pil_images = tensor_to_pil(尾帧)
|
||||
if pil_images:
|
||||
last_frame_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size)
|
||||
|
||||
if 参考图 is not None:
|
||||
pil_images = tensor_to_pil(参考图)
|
||||
if pil_images:
|
||||
reference_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size)
|
||||
|
||||
mode_str = "图生视频" if has_image else "文生视频"
|
||||
print(f"Veo: {mode_str} | 并发{生成数量}个 | 模型: {model} | {seconds}秒 | {分辨率} {宽高比}")
|
||||
|
||||
# 准备保存路径
|
||||
video_dir = _get_video_output_dir()
|
||||
counter = _get_next_counter(video_dir, "veo")
|
||||
|
||||
# ProgressBar
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(生成数量 if 生成数量 > 1 else 100)
|
||||
|
||||
try:
|
||||
if self.client is None:
|
||||
self.client = VeoClient()
|
||||
|
||||
if 生成数量 == 1:
|
||||
save_path = os.path.join(video_dir, f"veo_{counter:05d}.mp4")
|
||||
last_progress = [0]
|
||||
|
||||
def progress_callback(progress_pct: int):
|
||||
print(
|
||||
f"\rVeo: 生成中... 进度: {progress_pct}%",
|
||||
end="", flush=True
|
||||
)
|
||||
if pbar is not None and progress_pct > last_progress[0]:
|
||||
pbar.update(progress_pct - last_progress[0])
|
||||
last_progress[0] = progress_pct
|
||||
|
||||
def on_stage(stage: str):
|
||||
if stage == "submitting":
|
||||
print("Veo: 正在提交视频生成任务...")
|
||||
elif stage.startswith("submitted:"):
|
||||
vid = stage.split(":", 1)[1]
|
||||
print(f"Veo: 视频任务已提交,ID: {vid}")
|
||||
elif stage == "polling":
|
||||
print("Veo: 等待视频生成...")
|
||||
elif stage == "downloading":
|
||||
print("")
|
||||
print("Veo: 视频生成完成,正在下载...")
|
||||
|
||||
result_path = self.client.generate_video_sync(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=actual_size,
|
||||
save_path=save_path,
|
||||
first_frame_bytes=first_frame_bytes,
|
||||
last_frame_bytes=last_frame_bytes,
|
||||
reference_bytes=reference_bytes,
|
||||
seed=seed,
|
||||
progress_callback=progress_callback,
|
||||
on_stage=on_stage,
|
||||
)
|
||||
result_paths = [result_path]
|
||||
|
||||
else:
|
||||
save_paths = [
|
||||
os.path.join(video_dir, f"veo_{counter + i:05d}.mp4")
|
||||
for i in range(生成数量)
|
||||
]
|
||||
success_count = [0]
|
||||
|
||||
def batch_progress_callback(current: int, total: int, success: bool, error_msg):
|
||||
if success:
|
||||
success_count[0] += 1
|
||||
print(f"Veo: 第 {current}/{total} 个视频完成 ✓")
|
||||
else:
|
||||
print(f"Veo: 第 {current}/{total} 个视频失败 ✗")
|
||||
if error_msg:
|
||||
print(f"原始错误详情:\n{error_msg}")
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
print(f"Veo: 正在并发提交 {生成数量} 个视频任务,请耐心等待...")
|
||||
result_paths = self.client.generate_batch_videos_sync(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=actual_size,
|
||||
save_paths=save_paths,
|
||||
first_frame_bytes=first_frame_bytes,
|
||||
last_frame_bytes=last_frame_bytes,
|
||||
reference_bytes=reference_bytes,
|
||||
seed=seed,
|
||||
progress_callback=batch_progress_callback,
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.2f}s" if elapsed >= 1 else f"{elapsed:.3f}s"
|
||||
print(f"Veo: 完成!总耗时 {time_str} | 已生成 {len(result_paths)} 个视频")
|
||||
for p in result_paths:
|
||||
print(f" → {p}")
|
||||
|
||||
output_path = "\n".join(result_paths)
|
||||
return (output_path,)
|
||||
|
||||
except ValueError as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nVeo: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nVeo: ❌ {error_msg}")
|
||||
raise RuntimeError(error_msg) from None
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nVeo: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
finally:
|
||||
if self.client is not None:
|
||||
try:
|
||||
balance_data = self.client.query_balance_sync()
|
||||
balance_info = self.client.format_balance_info(balance_data)
|
||||
print(f"Veo: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"GoogleVeo": GoogleVeo,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"GoogleVeo": "Google Veo - ab",
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
"""
|
||||
视频预览节点
|
||||
接收 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": "预览视频",
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
aiohttp>=3.9.0
|
||||
Pillow>=10.0.0
|
||||
requests>=2.31.0
|
||||
rembg[cpu]>=2.0.50
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Git integration checks for the sidebar updater."""
|
||||
|
||||
import importlib.util
|
||||
import subprocess
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
UPDATER_PATH = Path(__file__).resolve().parents[1] / "utils" / "updater.py"
|
||||
spec = importlib.util.spec_from_file_location("o1key_updater_under_test", UPDATER_PATH)
|
||||
updater = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(updater)
|
||||
|
||||
|
||||
class UpdaterTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.temp.cleanup)
|
||||
root = Path(self.temp.name)
|
||||
self.remote = root / "remote.git"
|
||||
self.author = root / "author"
|
||||
self.install = root / "install"
|
||||
self.git(root, "init", "--bare", str(self.remote))
|
||||
self.git(root, "clone", str(self.remote), str(self.author))
|
||||
self.git(self.author, "config", "user.email", "[email protected]")
|
||||
self.git(self.author, "config", "user.name", "Updater Test")
|
||||
self.git(self.author, "switch", "-c", "main")
|
||||
(self.author / "requirements.txt").write_text("requests>=2\n", encoding="utf-8")
|
||||
(self.author / "version.txt").write_text("1\n", encoding="utf-8")
|
||||
self.commit_and_push()
|
||||
self.git(root, "clone", "--branch", "main", str(self.remote), str(self.install))
|
||||
updater.PLUGIN_DIR = self.install
|
||||
|
||||
def git(self, cwd, *args):
|
||||
return subprocess.run(["git", *args], cwd=cwd, check=True, capture_output=True, text=True).stdout.strip()
|
||||
|
||||
def commit_and_push(self):
|
||||
self.git(self.author, "add", ".")
|
||||
self.git(self.author, "commit", "-m", "test update")
|
||||
self.git(self.author, "push", "origin", "main")
|
||||
|
||||
def test_fast_forward_and_requirements_change(self):
|
||||
self.assertFalse(updater.update_package()["updated"])
|
||||
(self.author / "version.txt").write_text("2\n", encoding="utf-8")
|
||||
(self.author / "requirements.txt").write_text("requests>=3\n", encoding="utf-8")
|
||||
self.commit_and_push()
|
||||
result = updater.update_package()
|
||||
self.assertTrue(result["updated"])
|
||||
self.assertTrue(result["requirements_changed"])
|
||||
self.assertEqual((self.install / "version.txt").read_text(encoding="utf-8"), "2\n")
|
||||
|
||||
def test_local_changes_are_preserved(self):
|
||||
(self.install / "version.txt").write_text("local\n", encoding="utf-8")
|
||||
with self.assertRaisesRegex(updater.UpdateError, "本地修改"):
|
||||
updater.update_package()
|
||||
self.assertEqual((self.install / "version.txt").read_text(encoding="utf-8"), "local\n")
|
||||
|
||||
def test_diverged_branch_is_rejected(self):
|
||||
self.git(self.install, "config", "user.email", "[email protected]")
|
||||
self.git(self.install, "config", "user.name", "Updater Test")
|
||||
(self.install / "version.txt").write_text("local commit\n", encoding="utf-8")
|
||||
self.git(self.install, "add", ".")
|
||||
self.git(self.install, "commit", "-m", "local")
|
||||
(self.author / "version.txt").write_text("remote commit\n", encoding="utf-8")
|
||||
self.commit_and_push()
|
||||
with self.assertRaisesRegex(updater.UpdateError, "已分叉"):
|
||||
updater.update_package()
|
||||
self.assertEqual((self.install / "version.txt").read_text(encoding="utf-8"), "local commit\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+62
-86
@@ -1,97 +1,73 @@
|
||||
@echo off
|
||||
chcp 65001 > nul
|
||||
echo ====================================
|
||||
echo Comfyui_o1key 插件更新工具
|
||||
echo ====================================
|
||||
echo.
|
||||
|
||||
:: 检查是否在 Git 仓库中
|
||||
if not exist ".git" (
|
||||
echo [错误] 当前目录不是 Git 仓库
|
||||
echo 请确保插件是通过 git clone 安装的
|
||||
pause
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
:: 保存当前版本
|
||||
if exist "version.txt" (
|
||||
set /p OLD_VERSION=<version.txt
|
||||
echo 当前版本: %OLD_VERSION%
|
||||
) else (
|
||||
set OLD_VERSION=未知
|
||||
echo 当前版本: 未知
|
||||
)
|
||||
setlocal EnableDelayedExpansion
|
||||
title comfyui_o1key Updater
|
||||
cd /d "%~dp0"
|
||||
|
||||
echo.
|
||||
echo [1/4] 检查远程更新...
|
||||
git fetch origin
|
||||
|
||||
:: 检查是否有更新
|
||||
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
|
||||
)
|
||||
|
||||
echo [ comfyui_o1key Updater ]
|
||||
echo.
|
||||
echo [2/4] 备份配置文件...
|
||||
if exist ".config" (
|
||||
copy /Y ".config" ".config.backup" > nul
|
||||
echo 已备份 .config 到 .config.backup
|
||||
)
|
||||
|
||||
:: Check git
|
||||
where git >nul 2>&1
|
||||
if errorlevel 1 ( set "ERR=Git not found in PATH." & goto :fail )
|
||||
|
||||
:: Check repo
|
||||
if not exist ".git" ( set "ERR=Not a git repo. Place this file in the plugin root." & goto :fail )
|
||||
|
||||
:: 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 [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 ^| SUCCESS ^|
|
||||
echo ^| %OLD% -> %NEW% ^|
|
||||
echo ^| Restart ComfyUI ^|
|
||||
echo +---------------------------+
|
||||
echo.
|
||||
echo [4/4] 更新依赖包...
|
||||
python -m pip install -r requirements.txt --upgrade --quiet
|
||||
if %errorlevel% neq 0 (
|
||||
echo [警告] 依赖包更新失败,请手动运行: pip install -r requirements.txt
|
||||
)
|
||||
pause & exit /b 0
|
||||
|
||||
:uptodate
|
||||
if exist ".config.bak" ( copy /y ".config.bak" ".config" >nul 2>&1 & del /f /q ".config.bak" >nul 2>&1 )
|
||||
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 ^| Already up to date ^|
|
||||
echo ^| %LOCAL:~0,7% (no change) ^|
|
||||
echo +---------------------------+
|
||||
echo.
|
||||
echo 请重启 ComfyUI 以使更改生效
|
||||
echo.
|
||||
pause
|
||||
goto :end
|
||||
pause & exit /b 0
|
||||
|
||||
:end
|
||||
exit /b 0
|
||||
: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.
|
||||
pause & exit /b 1
|
||||
|
||||
@@ -1,101 +0,0 @@
|
||||
#!/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 ""
|
||||
+2
-4
@@ -15,8 +15,7 @@ from .file_utils import (
|
||||
load_images_from_folder,
|
||||
pair_images_indexed,
|
||||
pair_images_cartesian,
|
||||
generate_output_filename,
|
||||
generate_batch_output_filenames,
|
||||
generate_timestamp_filename,
|
||||
save_image,
|
||||
get_folder_image_count
|
||||
)
|
||||
@@ -32,8 +31,7 @@ __all__ = [
|
||||
'load_images_from_folder',
|
||||
'pair_images_indexed',
|
||||
'pair_images_cartesian',
|
||||
'generate_output_filename',
|
||||
'generate_batch_output_filenames',
|
||||
'generate_timestamp_filename',
|
||||
'save_image',
|
||||
'get_folder_image_count'
|
||||
]
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
"""
|
||||
颜色去背景工具模块
|
||||
基于颜色距离计算实现精确可控的背景移除,不依赖 AI 模型。
|
||||
|
||||
支持模式:
|
||||
- white: 白色背景去除
|
||||
- white-preserve: 白色背景但保护浅色前景物体
|
||||
- corner: 自动采样四角颜色作为背景色
|
||||
- color: 指定任意颜色去除
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def background_to_alpha(
|
||||
image: Image.Image,
|
||||
bg_color: tuple = (255, 255, 255),
|
||||
tolerance: float = 8.0,
|
||||
feather: float = 45.0,
|
||||
strength: float = 1.0,
|
||||
min_alpha: int = 2,
|
||||
) -> Image.Image:
|
||||
"""
|
||||
将纯色背景转为透明。
|
||||
|
||||
对白色背景使用 white-to-alpha 恢复算法,保持彩色文字和抗锯齿边缘清晰。
|
||||
对其他颜色使用欧氏距离计算。
|
||||
"""
|
||||
rgba = np.asarray(image.convert("RGBA")).astype(np.float32)
|
||||
rgb = rgba[:, :, :3] / 255.0
|
||||
existing_alpha = rgba[:, :, 3] / 255.0
|
||||
bg = np.array(bg_color, dtype=np.float32) / 255.0
|
||||
|
||||
if max(bg_color) >= 245 and min(bg_color) >= 245:
|
||||
alpha = (1.0 - np.min(rgb, axis=2)) * float(strength)
|
||||
if tolerance > 0:
|
||||
dist = np.linalg.norm((1.0 - rgb) * 255.0, axis=2)
|
||||
gate = np.clip(
|
||||
(dist - float(tolerance)) / max(1.0, float(feather) * 0.25),
|
||||
0.0, 1.0,
|
||||
)
|
||||
alpha *= gate
|
||||
else:
|
||||
dist = np.linalg.norm((rgb - bg) * 255.0, axis=2)
|
||||
denom = max(1.0, float(feather))
|
||||
alpha = np.clip((dist - float(tolerance)) / denom, 0.0, 1.0)
|
||||
alpha *= float(strength)
|
||||
|
||||
alpha = np.clip(alpha, 0.0, 1.0) * existing_alpha
|
||||
alpha[alpha < (float(min_alpha) / 255.0)] = 0.0
|
||||
|
||||
# 从 alpha 混合中恢复前景色,避免白边
|
||||
out_rgb = rgb.copy()
|
||||
mask = alpha > 1e-6
|
||||
out_rgb[mask] = (rgb[mask] - bg * (1.0 - alpha[mask, None])) / alpha[mask, None]
|
||||
out_rgb = np.clip(out_rgb, 0.0, 1.0)
|
||||
|
||||
out = np.dstack([
|
||||
(out_rgb * 255.0).astype(np.uint8),
|
||||
(alpha * 255.0).astype(np.uint8),
|
||||
])
|
||||
return Image.fromarray(out, "RGBA")
|
||||
|
||||
|
||||
def corner_color(image: Image.Image, sample: int = 12) -> tuple:
|
||||
"""采样图片四角像素的中位数颜色,用于自动检测背景色。"""
|
||||
rgb = np.asarray(image.convert("RGB"))
|
||||
h, w = rgb.shape[:2]
|
||||
sample = max(1, min(sample, h, w))
|
||||
patches = [
|
||||
rgb[:sample, :sample],
|
||||
rgb[:sample, w - sample:],
|
||||
rgb[h - sample:, :sample],
|
||||
rgb[h - sample:, w - sample:],
|
||||
]
|
||||
merged = np.concatenate([p.reshape(-1, 3) for p in patches], axis=0)
|
||||
return tuple(np.median(merged, axis=0).astype(int))
|
||||
|
||||
|
||||
# PLACEHOLDER_PRESERVE
|
||||
|
||||
def preserve_light_foreground_to_alpha(
|
||||
image: Image.Image,
|
||||
tolerance: float = 10.0,
|
||||
preserve_opacity: float = 0.72,
|
||||
min_area_ratio: float = 0.00025,
|
||||
) -> Image.Image:
|
||||
"""
|
||||
白底去除 + 浅色前景保护。
|
||||
|
||||
适用于前景包含白色/浅色物体(白盘子、白帆、白色包装)的场景。
|
||||
使用 OpenCV 连通区域分析保护大面积浅色前景结构。
|
||||
如果 OpenCV 不可用,回退到普通 white-to-alpha。
|
||||
"""
|
||||
base = background_to_alpha(image, (255, 255, 255), tolerance=tolerance)
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
return base
|
||||
|
||||
rgb_u8 = np.asarray(image.convert("RGB"))
|
||||
h, w = rgb_u8.shape[:2]
|
||||
dist = np.sqrt(np.sum((255.0 - rgb_u8.astype(np.float32)) ** 2, axis=2))
|
||||
rough = (dist > float(tolerance)).astype(np.uint8) * 255
|
||||
|
||||
kernel_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
|
||||
kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (17, 17))
|
||||
rough = cv2.morphologyEx(rough, cv2.MORPH_OPEN, kernel_open, iterations=1)
|
||||
rough = cv2.morphologyEx(rough, cv2.MORPH_CLOSE, kernel_close, iterations=2)
|
||||
|
||||
count, labels, stats, _ = cv2.connectedComponentsWithStats(rough, 8)
|
||||
keep = np.zeros_like(rough)
|
||||
min_area = max(24, int(w * h * float(min_area_ratio)))
|
||||
for idx in range(1, count):
|
||||
if stats[idx, cv2.CC_STAT_AREA] >= min_area:
|
||||
keep[labels == idx] = 255
|
||||
|
||||
# PLACEHOLDER_FLOOD
|
||||
|
||||
flood = keep.copy()
|
||||
ff_mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
|
||||
cv2.floodFill(flood, ff_mask, (0, 0), 255)
|
||||
filled = cv2.bitwise_or(keep, cv2.bitwise_not(flood))
|
||||
soft = cv2.GaussianBlur(filled, (0, 0), 5).astype(np.float32) / 255.0
|
||||
|
||||
near_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (29, 29))
|
||||
near = cv2.dilate(
|
||||
(dist > (float(tolerance) * 0.65)).astype(np.uint8) * 255,
|
||||
near_kernel, iterations=1,
|
||||
)
|
||||
near = cv2.GaussianBlur(near, (0, 0), 8).astype(np.float32) / 255.0
|
||||
lift = np.minimum(soft, near) * float(preserve_opacity)
|
||||
|
||||
arr = np.asarray(base.convert("RGBA")).copy()
|
||||
alpha = arr[:, :, 3].astype(np.float32) / 255.0
|
||||
alpha = np.maximum(alpha, lift)
|
||||
alpha[alpha < (2.0 / 255.0)] = 0.0
|
||||
|
||||
original = np.asarray(image.convert("RGB"))
|
||||
very_light = (np.mean(original, axis=2) > 224) & (lift > 0.12)
|
||||
arr[:, :, :3][very_light] = original[very_light]
|
||||
arr[:, :, 3] = np.clip(alpha * 255.0, 0, 255).astype(np.uint8)
|
||||
return Image.fromarray(arr, "RGBA")
|
||||
|
||||
|
||||
def remove_background(
|
||||
image: Image.Image,
|
||||
mode: str = "white",
|
||||
bg_color: tuple = (255, 255, 255),
|
||||
tolerance: float = 8.0,
|
||||
feather: float = 45.0,
|
||||
strength: float = 1.0,
|
||||
) -> Image.Image:
|
||||
"""
|
||||
统一入口:根据模式移除背景。
|
||||
|
||||
mode:
|
||||
- white: 白色背景去除
|
||||
- white-preserve: 白底 + 保护浅色前景
|
||||
- corner: 自动采样四角颜色
|
||||
- color: 使用指定 bg_color
|
||||
"""
|
||||
if mode == "white":
|
||||
return background_to_alpha(image, (255, 255, 255), tolerance, feather, strength)
|
||||
elif mode == "white-preserve":
|
||||
return preserve_light_foreground_to_alpha(image, tolerance)
|
||||
elif mode == "corner":
|
||||
bg = corner_color(image)
|
||||
return background_to_alpha(image, bg, tolerance, feather, strength)
|
||||
elif mode == "color":
|
||||
return background_to_alpha(image, bg_color, tolerance, feather, strength)
|
||||
else:
|
||||
return image.convert("RGBA")
|
||||
+63
-24
@@ -12,6 +12,24 @@ 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"
|
||||
|
||||
# ============ 网络线路配置 ============
|
||||
NETWORK_ROUTES = {
|
||||
"全球加速": "https://api.o1key.cn",
|
||||
"CF加速": "https://cf-api.o1key.com",
|
||||
"美国直连": "https://api.o1key.com",
|
||||
}
|
||||
NETWORK_ROUTE_OPTIONS = ["全球加速", "CF加速", "美国直连"]
|
||||
|
||||
|
||||
def load_config(config_path: Optional[str] = None) -> Dict[str, str]:
|
||||
"""
|
||||
从配置文件加载所有配置项
|
||||
@@ -61,36 +79,16 @@ 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()
|
||||
api_key = config.get(key_name)
|
||||
|
||||
if api_key:
|
||||
return api_key
|
||||
|
||||
return None
|
||||
return config.get(key_name)
|
||||
|
||||
|
||||
def get_api_key_or_raise(key_name: str = "O1KEY_API_KEY") -> str:
|
||||
@@ -112,3 +110,44 @@ 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
|
||||
|
||||
|
||||
def get_base_url_by_route(route: str) -> str:
|
||||
"""根据网络线路选项返回对应域名,未匹配则走 config 垫底"""
|
||||
if isinstance(route, (list, tuple)):
|
||||
route = route[0] if route else None
|
||||
return NETWORK_ROUTES.get(route, get_api_base_url())
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
文件数据类型定义
|
||||
用于在 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,
|
||||
}
|
||||
+137
-90
@@ -4,8 +4,11 @@
|
||||
"""
|
||||
|
||||
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
|
||||
@@ -13,6 +16,45 @@ 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'}
|
||||
|
||||
@@ -133,6 +175,77 @@ 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, ...]]:
|
||||
@@ -166,100 +279,36 @@ def pair_images_cartesian(
|
||||
return list(product(*non_empty_lists))
|
||||
|
||||
|
||||
def generate_output_filename(
|
||||
source_images: List[ImageInfo],
|
||||
batch_index: int,
|
||||
output_folder: str,
|
||||
extension: str = ".png",
|
||||
task_id: Optional[str] = None
|
||||
) -> str:
|
||||
def generate_timestamp_filename(output_folder: str, prefix: str = "", extension: str = ".png", port_suffix: str = "") -> str:
|
||||
"""
|
||||
生成智能输出文件名
|
||||
|
||||
基于源图片文件名生成输出文件名,使用任务ID和时间戳确保并发安全。
|
||||
|
||||
Args:
|
||||
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 和时间戳
|
||||
"""
|
||||
# 构建基础文件名
|
||||
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)
|
||||
生成基于时间戳的文件名,确保按文件名排序 = 按生成时间排序。
|
||||
|
||||
格式:{prefix}{HHMMSS_YYYYMMDD_mmm}{port_suffix}{extension}
|
||||
例如:161700_20260322_001.png 或 去除ai_161700_20260322_001.png
|
||||
|
||||
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: 任务唯一标识符(用于并发场景)
|
||||
|
||||
output_folder: 输出目录
|
||||
prefix: 文件名前缀(如 "去除ai_")
|
||||
extension: 文件扩展名(如 ".png")
|
||||
port_suffix: 端口后缀(如 "_8189"),为空时自动获取
|
||||
|
||||
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
|
||||
Path(output_folder).mkdir(parents=True, exist_ok=True)
|
||||
if not port_suffix:
|
||||
port_suffix = _get_port_suffix()
|
||||
|
||||
date_part = datetime.now().strftime("%Y%m%d")
|
||||
time_part = datetime.now().strftime("%H%M%S")
|
||||
ms = random.randint(0, 999)
|
||||
|
||||
while True:
|
||||
filename = f"{prefix}{time_part}_{date_part}_{ms:03d}{port_suffix}{extension}"
|
||||
full_path = Path(output_folder) / filename
|
||||
if not full_path.exists():
|
||||
return str(full_path)
|
||||
ms = (ms + 1) % 1000
|
||||
|
||||
|
||||
def save_image(
|
||||
@@ -290,8 +339,6 @@ 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:
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
"""
|
||||
统一 HTTP 错误处理 & 退避重试模块
|
||||
|
||||
使用方式:
|
||||
1. 对于 aiohttp 请求,用 async_request_with_retry() 包裹 POST/GET 调用
|
||||
2. 对于已拿到 status code 的场景,调用 raise_for_status() 抛出友好错误
|
||||
|
||||
新增生图/视频节点时,请统一使用本模块处理 HTTP 错误。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
from typing import Any, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# 状态码 → 用户友好文案
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
HTTP_ERROR_MESSAGES = {
|
||||
429: "模型速率超限或额度不足!",
|
||||
502: "网关超时。请重试或将网络切换为美国直连",
|
||||
503: "模型超载。请稍后重试!",
|
||||
504: "网关超时。请稍后重试。",
|
||||
}
|
||||
|
||||
# 错误内容关键词 → 用户友好文案(优先于状态码匹配)
|
||||
ERROR_CONTENT_MESSAGES = {
|
||||
"Your request was rejected by the safety system": "请求被安全系统拦截:请调整提示词,避免敏感、违规、血腥、色情、仇恨、未成年人或真实人物等高风险内容。",
|
||||
"safety system": "请求被安全系统拦截:请调整提示词,避免敏感、违规、血腥、色情、仇恨、未成年人或真实人物等高风险内容。",
|
||||
"unexpected end of JSON input": "通常重试能解决;反复出现就降低分辨率、数量或换网络线路。",
|
||||
"The current model has a high load": "模型过载,请稍后重试!",
|
||||
"system error": "系统错误,请稍后重试。",
|
||||
}
|
||||
|
||||
# 可退避重试的状态码
|
||||
RETRYABLE_STATUS_CODES = {429, 502, 503, 504, 524}
|
||||
|
||||
# 退避重试默认参数
|
||||
DEFAULT_MAX_RETRIES = 3
|
||||
DEFAULT_BASE_DELAY = 2.0 # 首次重试等待秒数
|
||||
DEFAULT_MAX_DELAY = 30.0 # 最大等待秒数
|
||||
DEFAULT_BACKOFF_FACTOR = 2.0 # 指数退避因子
|
||||
|
||||
|
||||
def _extract_message_from_payload(payload: Any) -> str:
|
||||
if isinstance(payload, str):
|
||||
text = payload.strip()
|
||||
if not text:
|
||||
return ""
|
||||
if text.startswith("{") or text.startswith("["):
|
||||
try:
|
||||
return _extract_message_from_payload(json.loads(text))
|
||||
except Exception:
|
||||
return text
|
||||
return text
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
return ""
|
||||
|
||||
error = payload.get("error")
|
||||
if isinstance(error, dict):
|
||||
for key in ("message", "msg", "detail", "reason"):
|
||||
value = error.get(key)
|
||||
if value:
|
||||
return _extract_message_from_payload(value)
|
||||
elif error:
|
||||
return _extract_message_from_payload(error)
|
||||
|
||||
for key in ("message", "msg", "detail", "reason", "error_message"):
|
||||
value = payload.get(key)
|
||||
if value:
|
||||
return _extract_message_from_payload(value)
|
||||
|
||||
for key in ("data", "result", "response", "output"):
|
||||
value = payload.get(key)
|
||||
nested = _extract_message_from_payload(value)
|
||||
if nested:
|
||||
return nested
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def extract_structured_error_message(raw_message: str) -> str:
|
||||
if not isinstance(raw_message, str):
|
||||
return ""
|
||||
text = raw_message.strip()
|
||||
if not (text.startswith("{") or text.startswith("[")):
|
||||
return ""
|
||||
return _extract_message_from_payload(text)
|
||||
|
||||
|
||||
def get_friendly_message(status_code: int, raw_message: str = "") -> str:
|
||||
"""根据状态码/错误内容返回友好文案,未匹配则返回原始信息"""
|
||||
if status_code == 524:
|
||||
return "Gateway timed out while waiting for upstream image generation. Please retry, lower resolution/count, or switch network route."
|
||||
if raw_message:
|
||||
structured_message = extract_structured_error_message(raw_message)
|
||||
message_for_matching = structured_message or raw_message
|
||||
raw_message_lower = message_for_matching.lower()
|
||||
for keyword, friendly_msg in ERROR_CONTENT_MESSAGES.items():
|
||||
if keyword.lower() in raw_message_lower:
|
||||
return friendly_msg
|
||||
if structured_message:
|
||||
return structured_message
|
||||
if status_code == 500:
|
||||
return "服务器返回 500:上游生成失败或服务端临时异常。请稍后重试;如果多次出现,请降低分辨率/数量,或调整提示词。"
|
||||
friendly = HTTP_ERROR_MESSAGES.get(status_code)
|
||||
if friendly:
|
||||
return friendly
|
||||
return raw_message or f"请求失败 ({status_code})"
|
||||
|
||||
|
||||
def raise_for_status(status_code: int, raw_message: str = "", prefix: str = ""):
|
||||
"""根据状态码抛出带友好文案的 RuntimeError"""
|
||||
friendly = get_friendly_message(status_code, raw_message)
|
||||
full_msg = f"{prefix}{friendly}" if prefix else friendly
|
||||
raise RuntimeError(full_msg)
|
||||
|
||||
|
||||
def is_retryable(status_code: int) -> bool:
|
||||
return status_code in RETRYABLE_STATUS_CODES
|
||||
|
||||
|
||||
def _compute_delay(attempt: int, base_delay: float, max_delay: float, backoff_factor: float) -> float:
|
||||
"""计算第 attempt 次重试的等待时间(含 jitter)"""
|
||||
delay = base_delay * (backoff_factor ** attempt)
|
||||
delay = min(delay, max_delay)
|
||||
jitter = random.uniform(0, delay * 0.3)
|
||||
return delay + jitter
|
||||
|
||||
|
||||
async def async_request_with_retry(
|
||||
session: aiohttp.ClientSession,
|
||||
method: str,
|
||||
url: str,
|
||||
*,
|
||||
max_retries: int = DEFAULT_MAX_RETRIES,
|
||||
base_delay: float = DEFAULT_BASE_DELAY,
|
||||
max_delay: float = DEFAULT_MAX_DELAY,
|
||||
backoff_factor: float = DEFAULT_BACKOFF_FACTOR,
|
||||
prefix: str = "",
|
||||
**request_kwargs,
|
||||
) -> aiohttp.ClientResponse:
|
||||
"""
|
||||
带退避重试的 aiohttp 请求。
|
||||
|
||||
仅对 RETRYABLE_STATUS_CODES (429/502/503/504/524) 进行重试。
|
||||
超过最大重试次数后抛出友好 RuntimeError。
|
||||
成功时返回 response 对象(调用者需在 async with 外自行处理 body)。
|
||||
|
||||
用法示例:
|
||||
resp = await async_request_with_retry(session, "POST", url, json=body, headers=headers)
|
||||
data = await resp.json()
|
||||
"""
|
||||
last_status: Optional[int] = None
|
||||
last_message = ""
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
resp = await session.request(method, url, **request_kwargs)
|
||||
except (aiohttp.ClientError, asyncio.TimeoutError) as e:
|
||||
if attempt < max_retries:
|
||||
delay = _compute_delay(attempt, base_delay, max_delay, backoff_factor)
|
||||
print(f"{prefix}网络错误,{delay:.1f}s 后重试 ({attempt+1}/{max_retries})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
raise RuntimeError(f"{prefix}网络错误: {e}") from None
|
||||
|
||||
if resp.status == 200:
|
||||
return resp
|
||||
|
||||
last_status = resp.status
|
||||
try:
|
||||
last_message = await resp.text()
|
||||
except Exception:
|
||||
last_message = ""
|
||||
|
||||
if is_retryable(resp.status) and attempt < max_retries:
|
||||
delay = _compute_delay(attempt, base_delay, max_delay, backoff_factor)
|
||||
friendly = get_friendly_message(resp.status)
|
||||
print(f"{prefix}{friendly} {delay:.1f}s 后重试 ({attempt+1}/{max_retries})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
|
||||
break
|
||||
|
||||
if last_status and last_status in HTTP_ERROR_MESSAGES:
|
||||
raise_for_status(last_status, raw_message=last_message, prefix=prefix)
|
||||
|
||||
friendly = get_friendly_message(last_status or 0, last_message)
|
||||
raise RuntimeError(f"{prefix}{friendly}")
|
||||
+186
-11
@@ -5,7 +5,8 @@
|
||||
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from typing import List
|
||||
import json
|
||||
from typing import Callable, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -49,36 +50,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)
|
||||
|
||||
@@ -110,6 +111,180 @@ def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
|
||||
return base64.b64encode(img_bytes).decode('utf-8')
|
||||
|
||||
|
||||
_MAX_REQUEST_BODY_BYTES = 50 * 1024 * 1024 # 50MB 请求体上限
|
||||
|
||||
|
||||
def _encode_image_to_base64_with_quality(image: Image.Image, quality: int) -> str:
|
||||
buffered = BytesIO()
|
||||
working = image
|
||||
if working.mode != 'RGB':
|
||||
working = working.convert('RGB')
|
||||
|
||||
working.save(
|
||||
buffered,
|
||||
format="JPEG",
|
||||
quality=quality,
|
||||
optimize=True,
|
||||
subsampling=2,
|
||||
)
|
||||
return base64.b64encode(buffered.getvalue()).decode('utf-8')
|
||||
|
||||
|
||||
def encode_images_for_request_body_limit(
|
||||
images: List[Image.Image],
|
||||
build_body: Callable[[List[Tuple[str, str]]], dict],
|
||||
max_body_bytes: int = _MAX_REQUEST_BODY_BYTES,
|
||||
) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
为请求体编码图片,并保证完整 JSON 请求体不超过 max_body_bytes。
|
||||
|
||||
策略:
|
||||
- 先按原始 PNG 编码估算完整请求体;
|
||||
- 若超过限制,改用 JPEG 质量压缩,逐步降低 quality;
|
||||
- 全程不缩放图片尺寸。
|
||||
|
||||
Returns:
|
||||
[(mime_type, base64), ...]
|
||||
"""
|
||||
encoded = [("image/png", encode_image_to_base64(img, format="PNG")) for img in images]
|
||||
body_size = len(json.dumps(build_body(encoded)).encode("utf-8"))
|
||||
if body_size <= max_body_bytes:
|
||||
return encoded
|
||||
|
||||
for quality in [95, 90, 85, 80, 75, 70, 65, 60, 55, 50, 45, 40, 35, 30, 25, 20, 15, 10, 5, 1]:
|
||||
encoded = [
|
||||
("image/jpeg", _encode_image_to_base64_with_quality(img, quality))
|
||||
for img in images
|
||||
]
|
||||
body_size = len(json.dumps(build_body(encoded)).encode("utf-8"))
|
||||
if body_size <= max_body_bytes:
|
||||
print(
|
||||
f"输入图片已通过 JPEG 质量压缩控制请求体积: "
|
||||
f"quality={quality}, 请求体积={body_size / 1024 / 1024:.2f}MB "
|
||||
f"(限制 {max_body_bytes / 1024 / 1024:.0f}MB)"
|
||||
)
|
||||
return encoded
|
||||
|
||||
raise ValueError(
|
||||
f"请求体超过 {max_body_bytes / 1024 / 1024:.0f}MB,"
|
||||
"即使压缩到最低图片质量仍无法满足限制;请减少参考图数量或输入图片内容复杂度"
|
||||
)
|
||||
|
||||
|
||||
_MAX_IMAGE_BYTES = 10 * 1024 * 1024 # 10MB 单张图片上限
|
||||
|
||||
|
||||
def _encode_image_to_bytes(image: Image.Image, format: str = "PNG", quality: int = None) -> bytes:
|
||||
buffered = BytesIO()
|
||||
working = image
|
||||
|
||||
if format.upper() == "JPEG" and working.mode != 'RGB':
|
||||
working = working.convert('RGB')
|
||||
elif working.mode == 'RGBA':
|
||||
working = working.convert('RGB')
|
||||
|
||||
save_kwargs = {"format": format}
|
||||
if quality is not None:
|
||||
save_kwargs.update({
|
||||
"quality": quality,
|
||||
"optimize": True,
|
||||
"subsampling": 2,
|
||||
})
|
||||
|
||||
working.save(buffered, **save_kwargs)
|
||||
return buffered.getvalue()
|
||||
|
||||
|
||||
def encode_images_for_image_size_limit(
|
||||
images: List[Image.Image],
|
||||
max_image_bytes: int = _MAX_IMAGE_BYTES,
|
||||
) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
将图片编码为 base64,并保证每张编码前的图片文件体积不超过 max_image_bytes。
|
||||
|
||||
策略:
|
||||
- 先尝试 PNG 原图尺寸编码;
|
||||
- 单张超过限制时,改用 JPEG 质量压缩;
|
||||
- 全程不缩放图片尺寸。
|
||||
|
||||
Returns:
|
||||
[(mime_type, base64), ...]
|
||||
"""
|
||||
encoded = []
|
||||
|
||||
for idx, img in enumerate(images, start=1):
|
||||
png_bytes = _encode_image_to_bytes(img, format="PNG")
|
||||
if len(png_bytes) <= max_image_bytes:
|
||||
encoded.append(("image/png", base64.b64encode(png_bytes).decode('utf-8')))
|
||||
continue
|
||||
|
||||
for quality in [95, 90, 85, 80, 75, 70, 65, 60, 55, 50, 45, 40, 35, 30, 25, 20, 15, 10, 5, 1]:
|
||||
jpg_bytes = _encode_image_to_bytes(img, format="JPEG", quality=quality)
|
||||
if len(jpg_bytes) <= max_image_bytes:
|
||||
print(
|
||||
f"输入图片 {idx} 已通过 JPEG 质量压缩控制单图体积: "
|
||||
f"quality={quality}, 图片体积={len(jpg_bytes) / 1024 / 1024:.2f}MB "
|
||||
f"(限制 {max_image_bytes / 1024 / 1024:.0f}MB),尺寸保持 {img.width}x{img.height}"
|
||||
)
|
||||
encoded.append(("image/jpeg", base64.b64encode(jpg_bytes).decode('utf-8')))
|
||||
break
|
||||
else:
|
||||
raise ValueError(
|
||||
f"输入图片 {idx} 超过 {max_image_bytes / 1024 / 1024:.0f}MB,"
|
||||
"即使压缩到最低图片质量仍无法满足限制;请减少图片内容复杂度或手动处理图片"
|
||||
)
|
||||
|
||||
return encoded
|
||||
|
||||
|
||||
def encode_image_to_base64_limited(
|
||||
image: Image.Image,
|
||||
format: str = "PNG",
|
||||
max_bytes: int = _MAX_IMAGE_BYTES,
|
||||
) -> str:
|
||||
"""
|
||||
将 PIL Image 编码为 base64,若超过 max_bytes 则自动缩放直到满足限制。
|
||||
|
||||
策略:等比缩放,每轮缩小到上一轮的 80%,最多 10 轮。
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
format: 图像格式,默认 PNG
|
||||
max_bytes: base64 字符串最大字节数,默认 10MB
|
||||
|
||||
Returns:
|
||||
base64 编码的字符串(保证 <= max_bytes)
|
||||
"""
|
||||
working = image
|
||||
if working.mode == 'RGBA':
|
||||
working = working.convert('RGB')
|
||||
|
||||
for attempt in range(10):
|
||||
buffered = BytesIO()
|
||||
working.save(buffered, format=format)
|
||||
b64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
|
||||
|
||||
if len(b64) <= max_bytes:
|
||||
if attempt > 0:
|
||||
print(
|
||||
f"图片已自动缩放: {image.width}x{image.height} → "
|
||||
f"{working.width}x{working.height} "
|
||||
f"({len(b64) / 1024 / 1024:.2f}MB)"
|
||||
)
|
||||
return b64
|
||||
|
||||
# 缩放到 80%
|
||||
scale = 0.8
|
||||
new_w = max(1, int(working.width * scale))
|
||||
new_h = max(1, int(working.height * scale))
|
||||
working = working.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
||||
|
||||
# 兜底:返回最后一次编码结果
|
||||
buffered = BytesIO()
|
||||
working.save(buffered, format=format)
|
||||
return base64.b64encode(buffered.getvalue()).decode('utf-8')
|
||||
|
||||
|
||||
def decode_base64_to_pil(base64_string: str) -> Image.Image:
|
||||
"""
|
||||
将 base64 字符串解码为 PIL Image
|
||||
@@ -190,4 +365,4 @@ def parse_batch_prompts(prompt: str) -> List[str]:
|
||||
if not filtered_prompts:
|
||||
raise ValueError("批量提示词模式下,所有提示词都为空,请至少提供一个有效的提示词")
|
||||
|
||||
return filtered_prompts
|
||||
return filtered_prompts
|
||||
|
||||
@@ -0,0 +1,675 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from io import BytesIO
|
||||
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple
|
||||
|
||||
import aiohttp
|
||||
from PIL import Image
|
||||
|
||||
from .http_error import (
|
||||
DEFAULT_BACKOFF_FACTOR,
|
||||
DEFAULT_BASE_DELAY,
|
||||
DEFAULT_MAX_DELAY,
|
||||
DEFAULT_MAX_RETRIES,
|
||||
RETRYABLE_STATUS_CODES,
|
||||
_compute_delay,
|
||||
extract_structured_error_message,
|
||||
get_friendly_message,
|
||||
)
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
|
||||
|
||||
_MAX_BODY_BYTES = 20_000_000
|
||||
_BODY_TARGET_BYTES = int(_MAX_BODY_BYTES * 0.8)
|
||||
_SUBMIT_ENDPOINT = "/async/v1/generateImage"
|
||||
_TASK_ENDPOINT = "/async/v1/tasks/{task_id}"
|
||||
_POLL_SCHEDULE = [5.0, 20.0]
|
||||
_POLL_INTERVAL = 3.0
|
||||
_MAX_WAIT_SECONDS = 900.0
|
||||
_INTERRUPT_STEP = 0.2
|
||||
_RUNNING_PROGRESS_MAX = 0.99
|
||||
_POLL_LOG_ENABLED = False
|
||||
|
||||
_SUCCESS_STATUSES = {"success", "succeed", "succeeded", "completed", "done", "finished"}
|
||||
_FAILURE_STATUSES = {
|
||||
"failure",
|
||||
"fail",
|
||||
"failed",
|
||||
"error",
|
||||
"expired",
|
||||
"timeout",
|
||||
"timed_out",
|
||||
"cancel",
|
||||
"canceled",
|
||||
"cancelled",
|
||||
"rejected",
|
||||
}
|
||||
_RUNNING_STATUSES = {
|
||||
"submitted",
|
||||
"queued",
|
||||
"pending",
|
||||
"running",
|
||||
"processing",
|
||||
"in_progress",
|
||||
"in-progress",
|
||||
"created",
|
||||
}
|
||||
|
||||
|
||||
def _headers(api_key: str) -> Dict[str, str]:
|
||||
return {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
|
||||
def _json_dumps(body: Dict[str, Any]) -> str:
|
||||
return json.dumps(body, ensure_ascii=False, separators=(",", ":"))
|
||||
|
||||
|
||||
def _json_size(body: Dict[str, Any]) -> int:
|
||||
return len(_json_dumps(body).encode("utf-8"))
|
||||
|
||||
|
||||
def _scale_images(images: List[Image.Image], scale: float) -> List[Image.Image]:
|
||||
if scale >= 1.0:
|
||||
return images
|
||||
scaled = []
|
||||
for img in images:
|
||||
new_w = max(1, int(img.width * scale))
|
||||
new_h = max(1, int(img.height * scale))
|
||||
scaled.append(img.resize((new_w, new_h), Image.Resampling.LANCZOS))
|
||||
return scaled
|
||||
|
||||
|
||||
def _encode_image_data_url(
|
||||
image: Image.Image,
|
||||
image_format: str,
|
||||
quality: Optional[int] = None,
|
||||
) -> str:
|
||||
buffered = BytesIO()
|
||||
working = image
|
||||
fmt = image_format.upper()
|
||||
save_kwargs = {"format": fmt}
|
||||
|
||||
if fmt == "JPEG":
|
||||
if working.mode != "RGB":
|
||||
working = working.convert("RGB")
|
||||
save_kwargs.update({"quality": quality or 90, "optimize": True, "subsampling": 2})
|
||||
mime_type = "image/jpeg"
|
||||
else:
|
||||
if working.mode == "RGBA":
|
||||
working = working.convert("RGB")
|
||||
mime_type = "image/png"
|
||||
|
||||
working.save(buffered, **save_kwargs)
|
||||
encoded = base64.b64encode(buffered.getvalue()).decode("ascii")
|
||||
return f"data:{mime_type};base64,{encoded}"
|
||||
|
||||
|
||||
def _encode_image_data_urls(
|
||||
images: Sequence[Image.Image],
|
||||
image_format: str,
|
||||
quality: Optional[int] = None,
|
||||
) -> List[str]:
|
||||
return [_encode_image_data_url(img, image_format, quality) for img in images]
|
||||
|
||||
|
||||
def _fit_image_data_urls_to_body_limit(
|
||||
images: Sequence[Image.Image],
|
||||
build_body: Callable[[List[str]], Dict[str, Any]],
|
||||
) -> Tuple[List[str], List[Image.Image], str, int]:
|
||||
working_images = list(images)
|
||||
image_urls = _encode_image_data_urls(working_images, "PNG")
|
||||
body_size = _json_size(build_body(image_urls))
|
||||
if body_size <= _BODY_TARGET_BYTES:
|
||||
return image_urls, working_images, "PNG", body_size
|
||||
|
||||
for _ in range(10):
|
||||
if body_size <= _BODY_TARGET_BYTES:
|
||||
break
|
||||
ratio = _BODY_TARGET_BYTES / max(body_size, 1)
|
||||
scale = min(0.98, ratio ** 0.5)
|
||||
working_images = _scale_images(working_images, scale)
|
||||
image_urls = _encode_image_data_urls(working_images, "PNG")
|
||||
body_size = _json_size(build_body(image_urls))
|
||||
|
||||
return image_urls, working_images, "PNG", body_size
|
||||
|
||||
|
||||
def _shorten_base64_for_log(value: Any, max_len: int = 160) -> Any:
|
||||
if isinstance(value, dict):
|
||||
result = {}
|
||||
for key, item in value.items():
|
||||
if key in ("data", "b64_json", "base64", "image_base64") and isinstance(item, str) and len(item) > max_len:
|
||||
result[key] = f"<base64 data, {len(item)} chars>"
|
||||
else:
|
||||
result[key] = _shorten_base64_for_log(item, max_len)
|
||||
return result
|
||||
if isinstance(value, list):
|
||||
return [_shorten_base64_for_log(item, max_len) for item in value]
|
||||
if isinstance(value, str) and value.startswith("data:image") and len(value) > max_len:
|
||||
return f"<data image url, {len(value)} chars>"
|
||||
return value
|
||||
|
||||
|
||||
def _log_body(label: str, text_or_body: Any) -> None:
|
||||
if isinstance(text_or_body, str):
|
||||
try:
|
||||
text_or_body = json.loads(text_or_body)
|
||||
except Exception:
|
||||
print(f"{label}\n{text_or_body}")
|
||||
return
|
||||
print(
|
||||
f"{label}\n"
|
||||
f"{json.dumps(_shorten_base64_for_log(text_or_body), ensure_ascii=False, indent=2)}"
|
||||
)
|
||||
|
||||
|
||||
def build_nano_banana_submit_body(
|
||||
model: str,
|
||||
prompt: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
enable_grounding: bool = False,
|
||||
thinking_level: Optional[str] = None,
|
||||
request_log_enabled: bool = False,
|
||||
node_label: str = "Nano Banana",
|
||||
) -> Dict[str, Any]:
|
||||
def _make_body(image_urls: List[str]) -> Dict[str, Any]:
|
||||
body: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"size": resolution,
|
||||
}
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
body["aspect_ratio"] = aspect_ratio
|
||||
if image_urls:
|
||||
body["images"] = image_urls
|
||||
if enable_grounding:
|
||||
body["google_search"] = True
|
||||
if thinking_level:
|
||||
body["thinking_level"] = thinking_level
|
||||
return body
|
||||
|
||||
working_images = list(images or [])
|
||||
image_urls: List[str] = []
|
||||
|
||||
if working_images:
|
||||
image_urls, working_images, _, _ = _fit_image_data_urls_to_body_limit(
|
||||
working_images,
|
||||
_make_body,
|
||||
)
|
||||
|
||||
body = _make_body(image_urls)
|
||||
body_size = _json_size(body)
|
||||
|
||||
if working_images and body_size > _MAX_BODY_BYTES:
|
||||
raise ValueError(
|
||||
f"Request body exceeds the 20MB limit after compression "
|
||||
f"({body_size / 1_000_000:.2f}MB). Reduce reference image count, "
|
||||
"image complexity, or prompt length."
|
||||
)
|
||||
if not working_images and body_size > _MAX_BODY_BYTES:
|
||||
raise ValueError(
|
||||
f"Request body exceeds the 20MB limit ({body_size / 1_000_000:.2f}MB). "
|
||||
"Shorten the prompt or system instructions."
|
||||
)
|
||||
|
||||
if request_log_enabled:
|
||||
print(
|
||||
f"[{node_label} 异步请求体] {body_size / 1024:.1f}KB\n"
|
||||
f"{json.dumps(_shorten_base64_for_log(body), ensure_ascii=False, indent=2)}"
|
||||
)
|
||||
|
||||
return body
|
||||
|
||||
|
||||
async def _interruptible_sleep(
|
||||
seconds: float,
|
||||
check_interrupt: Optional[Callable[[], None]] = None,
|
||||
) -> None:
|
||||
elapsed = 0.0
|
||||
while elapsed < seconds:
|
||||
if check_interrupt:
|
||||
check_interrupt()
|
||||
delay = min(_INTERRUPT_STEP, seconds - elapsed)
|
||||
await asyncio.sleep(delay)
|
||||
elapsed += delay
|
||||
if check_interrupt:
|
||||
check_interrupt()
|
||||
|
||||
|
||||
def _payload_sources(payload: Dict[str, Any]) -> Iterable[Dict[str, Any]]:
|
||||
queue = [payload]
|
||||
seen = set()
|
||||
while queue:
|
||||
current = queue.pop(0)
|
||||
if not isinstance(current, dict):
|
||||
continue
|
||||
obj_id = id(current)
|
||||
if obj_id in seen:
|
||||
continue
|
||||
seen.add(obj_id)
|
||||
yield current
|
||||
for key in ("data", "result", "response", "output", "task_result", "content"):
|
||||
value = current.get(key)
|
||||
if isinstance(value, dict):
|
||||
queue.append(value)
|
||||
|
||||
|
||||
def _extract_task_id(payload: Dict[str, Any]) -> str:
|
||||
for source in _payload_sources(payload):
|
||||
for key in ("task_id", "taskId", "id"):
|
||||
value = source.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
raise RuntimeError(f"提交响应中未找到 task_id: {payload}")
|
||||
|
||||
|
||||
def _extract_status(payload: Dict[str, Any]) -> str:
|
||||
statuses = []
|
||||
for source in _payload_sources(payload):
|
||||
for key in ("status", "task_status", "state", "task_state"):
|
||||
value = source.get(key)
|
||||
if value is not None and str(value).strip():
|
||||
statuses.append(str(value).strip())
|
||||
|
||||
for status in statuses:
|
||||
normalized = status.lower()
|
||||
if normalized in _FAILURE_STATUSES or any(
|
||||
token in normalized for token in ("fail", "error", "reject", "timeout", "cancel")
|
||||
):
|
||||
return status
|
||||
for status in statuses:
|
||||
if status.lower() in _RUNNING_STATUSES:
|
||||
return status
|
||||
for status in statuses:
|
||||
if status.lower() in _SUCCESS_STATUSES:
|
||||
return status
|
||||
return statuses[0] if statuses else ""
|
||||
|
||||
|
||||
def _coerce_progress_fraction(value: Any) -> Optional[float]:
|
||||
if value is None or isinstance(value, bool):
|
||||
return None
|
||||
|
||||
if isinstance(value, (int, float)):
|
||||
progress = float(value)
|
||||
elif isinstance(value, str):
|
||||
text = value.strip()
|
||||
if not text:
|
||||
return None
|
||||
has_percent_suffix = text.endswith("%")
|
||||
if has_percent_suffix:
|
||||
text = text[:-1].strip()
|
||||
try:
|
||||
progress = float(text)
|
||||
except ValueError:
|
||||
return None
|
||||
if has_percent_suffix:
|
||||
progress /= 100.0
|
||||
else:
|
||||
return None
|
||||
|
||||
if progress > 1.0:
|
||||
progress /= 100.0
|
||||
return max(0.0, min(progress, 1.0))
|
||||
|
||||
|
||||
def _extract_progress(payload: Dict[str, Any]) -> Optional[float]:
|
||||
for source in _payload_sources(payload):
|
||||
for key in ("progress", "percentage", "percent"):
|
||||
progress = _coerce_progress_fraction(source.get(key))
|
||||
if progress is not None:
|
||||
return progress
|
||||
|
||||
for key in ("progressInfo", "progress_info"):
|
||||
info = source.get(key)
|
||||
if not isinstance(info, dict):
|
||||
continue
|
||||
for field in ("progress", "percentage", "percent"):
|
||||
progress = _coerce_progress_fraction(info.get(field))
|
||||
if progress is not None:
|
||||
return progress
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _is_failure_status(normalized_status: str) -> bool:
|
||||
return normalized_status in _FAILURE_STATUSES or any(
|
||||
token in normalized_status for token in ("fail", "error", "reject", "timeout", "cancel")
|
||||
)
|
||||
|
||||
|
||||
def _extract_error_message(payload: Dict[str, Any]) -> str:
|
||||
for source in _payload_sources(payload):
|
||||
error = source.get("error")
|
||||
if isinstance(error, dict):
|
||||
for key in ("message", "msg", "detail", "reason", "code"):
|
||||
value = error.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
elif error:
|
||||
message = extract_structured_error_message(str(error))
|
||||
return message or str(error)
|
||||
|
||||
for key in (
|
||||
"fail_reason",
|
||||
"failure_reason",
|
||||
"task_status_msg",
|
||||
"status_msg",
|
||||
"error_message",
|
||||
"message",
|
||||
"msg",
|
||||
"reason",
|
||||
"detail",
|
||||
):
|
||||
value = source.get(key)
|
||||
if value:
|
||||
message = extract_structured_error_message(str(value))
|
||||
return message or str(value)
|
||||
return "未知错误"
|
||||
|
||||
|
||||
async def _submit_task(
|
||||
session: aiohttp.ClientSession,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
body: Dict[str, Any],
|
||||
node_label: str,
|
||||
log_body_enabled: bool = False,
|
||||
) -> str:
|
||||
url = f"{base_url}{_SUBMIT_ENDPOINT}"
|
||||
timeout = aiohttp.ClientTimeout(total=120, connect=30, sock_read=120)
|
||||
last_status = None
|
||||
last_text = ""
|
||||
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
try:
|
||||
async with session.post(
|
||||
url,
|
||||
headers=_headers(api_key),
|
||||
data=_json_dumps(body).encode("utf-8"),
|
||||
timeout=timeout,
|
||||
) as resp:
|
||||
text = await resp.text()
|
||||
if log_body_enabled:
|
||||
_log_body(f"[{node_label} 异步提交响应] HTTP {resp.status}", text)
|
||||
if resp.status not in (200, 201, 202):
|
||||
last_status = resp.status
|
||||
last_text = text
|
||||
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
|
||||
friendly = get_friendly_message(resp.status)
|
||||
delay = _compute_delay(
|
||||
attempt,
|
||||
DEFAULT_BASE_DELAY,
|
||||
DEFAULT_MAX_DELAY,
|
||||
DEFAULT_BACKOFF_FACTOR,
|
||||
)
|
||||
print(f"{node_label}: {friendly} {delay:.1f}s 后重试提交 ({attempt + 1}/{DEFAULT_MAX_RETRIES})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
raise RuntimeError(get_friendly_message(resp.status, text))
|
||||
|
||||
try:
|
||||
data = json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError(f"提交响应 JSON 解析失败: {text[:500]}") from None
|
||||
task_id = _extract_task_id(data)
|
||||
status = _extract_status(data) or "SUBMITTED"
|
||||
print(f"{node_label}: 异步任务已提交 | task_id={task_id} | status={status}")
|
||||
return task_id
|
||||
except (
|
||||
aiohttp.ClientConnectorError,
|
||||
aiohttp.ClientOSError,
|
||||
aiohttp.ServerDisconnectedError,
|
||||
asyncio.TimeoutError,
|
||||
) as e:
|
||||
last_text = str(e)
|
||||
if attempt < DEFAULT_MAX_RETRIES:
|
||||
delay = _compute_delay(
|
||||
attempt,
|
||||
DEFAULT_BASE_DELAY,
|
||||
DEFAULT_MAX_DELAY,
|
||||
DEFAULT_BACKOFF_FACTOR,
|
||||
)
|
||||
print(f"{node_label}: 网络连接失败,{delay:.1f}s 后重试提交 ({attempt + 1}/{DEFAULT_MAX_RETRIES})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
raise RuntimeError(
|
||||
f"网络连接失败,无法连接 {url}: {str(e)}。请切换节点里的网络线路,或检查 VPN/代理/防火墙。"
|
||||
) from None
|
||||
|
||||
raise RuntimeError(get_friendly_message(last_status or 0, last_text))
|
||||
|
||||
|
||||
async def _poll_task(
|
||||
session: aiohttp.ClientSession,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
task_id: str,
|
||||
node_label: str,
|
||||
check_interrupt: Optional[Callable[[], None]] = None,
|
||||
log_body_enabled: bool = False,
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
url = f"{base_url}{_TASK_ENDPOINT.format(task_id=task_id)}"
|
||||
start_time = time.time()
|
||||
last_poll_at = start_time
|
||||
poll_count = 0
|
||||
|
||||
while True:
|
||||
if check_interrupt:
|
||||
check_interrupt()
|
||||
|
||||
if poll_count < len(_POLL_SCHEDULE):
|
||||
next_poll_at = start_time + _POLL_SCHEDULE[poll_count]
|
||||
else:
|
||||
next_poll_at = last_poll_at + _POLL_INTERVAL
|
||||
|
||||
sleep_time = next_poll_at - time.time()
|
||||
if sleep_time > 0:
|
||||
await _interruptible_sleep(sleep_time, check_interrupt=check_interrupt)
|
||||
|
||||
last_poll_at = time.time()
|
||||
elapsed = last_poll_at - start_time
|
||||
if elapsed > _MAX_WAIT_SECONDS:
|
||||
raise RuntimeError(f"任务 {task_id} 超时(>{int(_MAX_WAIT_SECONDS)}秒),请稍后用 task_id 查询结果")
|
||||
|
||||
poll_count += 1
|
||||
async with session.get(url, headers=_headers(api_key)) as resp:
|
||||
text = await resp.text()
|
||||
if log_body_enabled:
|
||||
_log_body(f"[{node_label} 任务查询响应 #{poll_count}] HTTP {resp.status}", text)
|
||||
if resp.status != 200:
|
||||
raise RuntimeError(get_friendly_message(resp.status, text))
|
||||
try:
|
||||
payload = json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError(f"任务查询响应 JSON 解析失败: {text[:500]}") from None
|
||||
|
||||
status = _extract_status(payload) or "UNKNOWN"
|
||||
normalized = status.lower()
|
||||
progress = _extract_progress(payload)
|
||||
is_failure = _is_failure_status(normalized)
|
||||
if _POLL_LOG_ENABLED:
|
||||
progress_text = ""
|
||||
if progress is not None and not is_failure:
|
||||
displayed_progress = 1.0 if normalized in _SUCCESS_STATUSES else min(progress, _RUNNING_PROGRESS_MAX)
|
||||
progress_text = f" | progress={displayed_progress * 100:.0f}%"
|
||||
print(f"{node_label}: 查询任务 #{poll_count} | task_id={task_id} | status={status}{progress_text}")
|
||||
|
||||
if normalized in _SUCCESS_STATUSES:
|
||||
if progress_callback:
|
||||
progress_callback(1.0)
|
||||
return payload
|
||||
if is_failure:
|
||||
raise RuntimeError(f"任务失败: {_extract_error_message(payload)}")
|
||||
if normalized not in _RUNNING_STATUSES:
|
||||
raise RuntimeError(f"未知任务状态 {status}: {payload}")
|
||||
if progress_callback and progress is not None:
|
||||
progress_callback(min(progress, _RUNNING_PROGRESS_MAX))
|
||||
|
||||
|
||||
async def _image_from_url_or_data(
|
||||
value: str,
|
||||
session: aiohttp.ClientSession,
|
||||
) -> Optional[Image.Image]:
|
||||
if not value:
|
||||
return None
|
||||
if value.startswith("data:image"):
|
||||
try:
|
||||
_, b64_data = value.split(",", 1)
|
||||
return Image.open(BytesIO(base64.b64decode(b64_data))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"data URL 图片解码失败: {e}") from None
|
||||
if value.startswith("http"):
|
||||
async with session.get(value, allow_redirects=True) as resp:
|
||||
if resp.status != 200:
|
||||
raise RuntimeError(f"图片下载失败 ({resp.status}): {value}")
|
||||
img_bytes = await resp.read()
|
||||
return Image.open(BytesIO(img_bytes)).convert("RGB")
|
||||
return None
|
||||
|
||||
|
||||
async def _parse_direct_images(
|
||||
payload: Dict[str, Any],
|
||||
session: aiohttp.ClientSession,
|
||||
) -> List[Image.Image]:
|
||||
images: List[Image.Image] = []
|
||||
|
||||
async def _try_item(item: Any) -> None:
|
||||
if isinstance(item, str):
|
||||
img = await _image_from_url_or_data(item, session)
|
||||
if img:
|
||||
images.append(img)
|
||||
return
|
||||
if not isinstance(item, dict):
|
||||
return
|
||||
|
||||
for key in ("url", "image_url", "result_url", "download_url"):
|
||||
img = await _image_from_url_or_data(str(item.get(key) or ""), session)
|
||||
if img:
|
||||
images.append(img)
|
||||
return
|
||||
|
||||
b64_data = item.get("b64_json") or item.get("base64") or item.get("image_base64")
|
||||
if b64_data:
|
||||
images.append(Image.open(BytesIO(base64.b64decode(str(b64_data)))).convert("RGB"))
|
||||
return
|
||||
|
||||
for inline_key in ("inline_data", "inlineData"):
|
||||
inline = item.get(inline_key)
|
||||
if isinstance(inline, dict) and inline.get("data"):
|
||||
images.append(Image.open(BytesIO(base64.b64decode(str(inline["data"])))).convert("RGB"))
|
||||
return
|
||||
|
||||
for source in _payload_sources(payload):
|
||||
for key in ("image_url", "result_url", "url", "download_url"):
|
||||
img = await _image_from_url_or_data(str(source.get(key) or ""), session)
|
||||
if img:
|
||||
images.append(img)
|
||||
|
||||
for key in ("images", "output_images", "outputs"):
|
||||
value = source.get(key)
|
||||
if isinstance(value, list):
|
||||
for item in value:
|
||||
await _try_item(item)
|
||||
elif value:
|
||||
await _try_item(value)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
async def _parse_task_images(
|
||||
task_payload: Dict[str, Any],
|
||||
session: aiohttp.ClientSession,
|
||||
api_key: str,
|
||||
) -> List[Image.Image]:
|
||||
direct_images = await _parse_direct_images(task_payload, session)
|
||||
if direct_images:
|
||||
return direct_images
|
||||
|
||||
client = GeminiAPIClient(api_key=api_key)
|
||||
last_error = None
|
||||
for source in _payload_sources(task_payload):
|
||||
if "candidates" not in source:
|
||||
continue
|
||||
try:
|
||||
images, _ = await client.parse_response_async(source, session=session)
|
||||
if images:
|
||||
return [img.convert("RGB") for img in images]
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
|
||||
if last_error is not None:
|
||||
raise RuntimeError(str(last_error)) from None
|
||||
raise RuntimeError(f"任务成功但未找到图片结果: {task_payload}")
|
||||
|
||||
|
||||
async def generate_nano_banana_async(
|
||||
session: aiohttp.ClientSession,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
enable_grounding: bool = False,
|
||||
thinking_level: Optional[str] = None,
|
||||
node_label: str = "Nano Banana",
|
||||
request_log_enabled: bool = False,
|
||||
check_interrupt: Optional[Callable[[], None]] = None,
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
) -> Tuple[List[Image.Image], Dict[str, Any]]:
|
||||
if check_interrupt:
|
||||
check_interrupt()
|
||||
|
||||
body = build_nano_banana_submit_body(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
request_log_enabled=request_log_enabled,
|
||||
node_label=node_label,
|
||||
)
|
||||
|
||||
task_start = time.time()
|
||||
task_id = await _submit_task(
|
||||
session,
|
||||
base_url,
|
||||
api_key,
|
||||
body,
|
||||
node_label,
|
||||
log_body_enabled=request_log_enabled,
|
||||
)
|
||||
task_payload = await _poll_task(
|
||||
session,
|
||||
base_url,
|
||||
api_key,
|
||||
task_id,
|
||||
node_label,
|
||||
check_interrupt=check_interrupt,
|
||||
log_body_enabled=request_log_enabled,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
task_done = time.time()
|
||||
|
||||
parse_start = time.time()
|
||||
images_list = await _parse_task_images(task_payload, session, api_key)
|
||||
parse_done = time.time()
|
||||
|
||||
return images_list, {
|
||||
"task_id": task_id,
|
||||
"task_ms": (task_done - task_start) * 1000,
|
||||
"parse_ms": (parse_done - parse_start) * 1000,
|
||||
"request_bytes": _json_size(body),
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,73 @@
|
||||
"""
|
||||
背景移除工具模块
|
||||
基于 rembg 库实现,支持 CPU 推理
|
||||
"""
|
||||
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
_session = None
|
||||
|
||||
|
||||
def _get_session():
|
||||
"""懒加载 rembg session,避免启动时加载模型"""
|
||||
global _session
|
||||
if _session is None:
|
||||
try:
|
||||
import folder_paths
|
||||
models_dir = os.path.join(folder_paths.models_dir, "rembg")
|
||||
os.makedirs(models_dir, exist_ok=True)
|
||||
os.environ["U2NET_HOME"] = models_dir
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
from rembg import new_session
|
||||
_session = new_session("isnet-general-use")
|
||||
print("[o1key] rembg 模型加载完成 (isnet-general-use)")
|
||||
except ImportError:
|
||||
raise RuntimeError(
|
||||
"未安装 rembg,请执行: pip install rembg[cpu]>=2.0.50"
|
||||
)
|
||||
return _session
|
||||
|
||||
|
||||
def remove_background_pil(image: Image.Image) -> Image.Image:
|
||||
"""
|
||||
移除 PIL Image 背景,返回 RGBA 图像(背景透明)
|
||||
"""
|
||||
from rembg import remove
|
||||
session = _get_session()
|
||||
result = remove(image, session=session)
|
||||
return result.convert("RGBA")
|
||||
|
||||
|
||||
def remove_background_tensor(tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
移除 ComfyUI IMAGE tensor 的背景
|
||||
输入: [B, H, W, C] (3或4通道)
|
||||
输出: [B, H, W, 4] RGBA tensor
|
||||
"""
|
||||
from rembg import remove
|
||||
session = _get_session()
|
||||
|
||||
results = []
|
||||
batch_size = tensor.shape[0]
|
||||
|
||||
for i in range(batch_size):
|
||||
frame = tensor[i] # [H, W, C]
|
||||
arr = (frame.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||||
|
||||
if arr.shape[2] == 4:
|
||||
pil_img = Image.fromarray(arr, mode="RGBA")
|
||||
else:
|
||||
pil_img = Image.fromarray(arr, mode="RGB")
|
||||
|
||||
result = remove(pil_img, session=session)
|
||||
result_rgba = result.convert("RGBA")
|
||||
|
||||
result_arr = np.array(result_rgba).astype(np.float32) / 255.0
|
||||
results.append(torch.from_numpy(result_arr))
|
||||
|
||||
return torch.stack(results, dim=0)
|
||||
+79
-14
@@ -39,12 +39,15 @@ def check_for_updates() -> bool:
|
||||
if not os.path.exists(git_dir):
|
||||
return False
|
||||
|
||||
# 执行 git fetch
|
||||
# 执行 git fetch(禁止弹出认证弹框,失败时静默处理)
|
||||
env = os.environ.copy()
|
||||
env['GIT_TERMINAL_PROMPT'] = '0'
|
||||
subprocess.run(
|
||||
['git', 'fetch', 'origin'],
|
||||
cwd=plugin_dir,
|
||||
capture_output=True,
|
||||
timeout=10
|
||||
timeout=10,
|
||||
env=env
|
||||
)
|
||||
|
||||
# 检查本地和远程版本
|
||||
@@ -68,16 +71,78 @@ 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():
|
||||
"""通知用户有更新可用"""
|
||||
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")
|
||||
"""通知用户有更新可用(前端弹窗)"""
|
||||
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
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Safely fast-forward a Git installation of this node package."""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
PLUGIN_DIR = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
class UpdateError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _git(*args, timeout=60, check=True):
|
||||
env = os.environ.copy()
|
||||
env["GIT_TERMINAL_PROMPT"] = "0"
|
||||
env["GCM_INTERACTIVE"] = "Never"
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", *args],
|
||||
cwd=PLUGIN_DIR,
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
timeout=timeout,
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
raise UpdateError("未找到 Git,请先安装 Git。") from exc
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
raise UpdateError("Git 操作超时,请检查网络后重试。") from exc
|
||||
if check and result.returncode:
|
||||
detail = (result.stderr or result.stdout).strip().splitlines()
|
||||
raise UpdateError(detail[-1] if detail else "Git 操作失败。")
|
||||
return result
|
||||
|
||||
|
||||
def update_package():
|
||||
"""Update origin/main without discarding local changes or switching branches."""
|
||||
if not (PLUGIN_DIR / ".git").exists():
|
||||
raise UpdateError("当前节点包不是 Git 安装。请通过 Git 安装后再使用界面更新。")
|
||||
|
||||
branch = _git("symbolic-ref", "--quiet", "--short", "HEAD", check=False)
|
||||
if branch.returncode or branch.stdout.strip() != "main":
|
||||
raise UpdateError("当前不在 main 分支,请手动检查分支后更新。")
|
||||
|
||||
if _git("status", "--porcelain", "--untracked-files=no").stdout.strip():
|
||||
raise UpdateError("节点包有本地修改,请先保存或处理修改后再更新。")
|
||||
|
||||
old_commit = _git("rev-parse", "HEAD").stdout.strip()
|
||||
old_requirements = _git("show", "HEAD:requirements.txt", check=False).stdout
|
||||
_git("fetch", "origin", "main")
|
||||
new_commit = _git("rev-parse", "FETCH_HEAD").stdout.strip()
|
||||
if old_commit == new_commit:
|
||||
return {"updated": False, "version": old_commit[:7], "requirements_changed": False}
|
||||
|
||||
if _git("merge-base", "--is-ancestor", "HEAD", "FETCH_HEAD", check=False).returncode:
|
||||
raise UpdateError("本地与 origin/main 已分叉,无法安全快进。请手动处理。")
|
||||
|
||||
_git("merge", "--ff-only", "FETCH_HEAD")
|
||||
requirements_changed = old_requirements != (PLUGIN_DIR / "requirements.txt").read_text(encoding="utf-8")
|
||||
return {
|
||||
"updated": True,
|
||||
"version": new_commit[:7],
|
||||
"requirements_changed": requirements_changed,
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
import asyncio
|
||||
from typing import Any, Dict
|
||||
|
||||
|
||||
try:
|
||||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||||
INTERRUPT_AVAILABLE = True
|
||||
except Exception:
|
||||
INTERRUPT_AVAILABLE = False
|
||||
processing_interrupted = lambda: False
|
||||
|
||||
class InterruptProcessingException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
SUCCESS_STATUSES = {"succeed", "succeeded", "success", "completed", "done", "finished"}
|
||||
FAILURE_STATUSES = {
|
||||
"fail",
|
||||
"failed",
|
||||
"failure",
|
||||
"error",
|
||||
"expired",
|
||||
"timeout",
|
||||
"timed_out",
|
||||
"cancel",
|
||||
"canceled",
|
||||
"cancelled",
|
||||
"rejected",
|
||||
}
|
||||
|
||||
|
||||
def check_interrupt() -> None:
|
||||
if INTERRUPT_AVAILABLE and processing_interrupted():
|
||||
raise InterruptProcessingException()
|
||||
|
||||
|
||||
async def interruptible_sleep(seconds: float, step: float = 0.2) -> None:
|
||||
elapsed = 0.0
|
||||
while elapsed < seconds:
|
||||
check_interrupt()
|
||||
delay = min(step, seconds - elapsed)
|
||||
await asyncio.sleep(delay)
|
||||
elapsed += delay
|
||||
check_interrupt()
|
||||
|
||||
|
||||
async def run_with_interrupt(coro, step: float = 0.2):
|
||||
task = asyncio.ensure_future(coro)
|
||||
try:
|
||||
while not task.done():
|
||||
check_interrupt()
|
||||
await asyncio.wait({task}, timeout=step)
|
||||
check_interrupt()
|
||||
return await task
|
||||
except InterruptProcessingException:
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except BaseException:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> Dict[str, Any]:
|
||||
return value if isinstance(value, dict) else {}
|
||||
|
||||
|
||||
def _nested_payloads(payload: Dict[str, Any]):
|
||||
root = _as_dict(payload)
|
||||
data = _as_dict(root.get("data"))
|
||||
inner = _as_dict(data.get("data"))
|
||||
return root, data, inner
|
||||
|
||||
|
||||
def extract_status(payload: Dict[str, Any]) -> str:
|
||||
root, data, inner = _nested_payloads(payload)
|
||||
keys = ("status", "task_status", "state", "task_state")
|
||||
statuses = []
|
||||
for source in (data, inner, root):
|
||||
for key in keys:
|
||||
value = source.get(key)
|
||||
if value is not None and str(value).strip():
|
||||
statuses.append(str(value).strip().lower())
|
||||
for status in statuses:
|
||||
if status in FAILURE_STATUSES or any(
|
||||
token in status for token in ("fail", "error", "reject", "timeout", "cancel")
|
||||
):
|
||||
return status
|
||||
for status in statuses:
|
||||
if status in SUCCESS_STATUSES:
|
||||
return status
|
||||
return statuses[0] if statuses else ""
|
||||
|
||||
|
||||
def extract_progress(payload: Dict[str, Any]) -> int:
|
||||
root, data, inner = _nested_payloads(payload)
|
||||
for source in (data, inner, root):
|
||||
value = source.get("progress")
|
||||
if value is None:
|
||||
continue
|
||||
try:
|
||||
return max(0, min(100, int(float(str(value).strip().rstrip("%")))))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
return 0
|
||||
|
||||
|
||||
def extract_error_message(payload: Dict[str, Any], default: str = "未知错误") -> str:
|
||||
root, data, inner = _nested_payloads(payload)
|
||||
keys = (
|
||||
"fail_reason",
|
||||
"failure_reason",
|
||||
"task_status_msg",
|
||||
"status_msg",
|
||||
"error_message",
|
||||
"message",
|
||||
"msg",
|
||||
"reason",
|
||||
"detail",
|
||||
"details",
|
||||
)
|
||||
for source in (data, inner, root):
|
||||
error = source.get("error")
|
||||
if isinstance(error, dict):
|
||||
for key in ("message", "msg", "detail", "reason", "code"):
|
||||
value = error.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
elif error:
|
||||
return str(error)
|
||||
|
||||
for key in keys:
|
||||
value = source.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
return default
|
||||
|
||||
|
||||
def extract_video_url(payload: Dict[str, Any]) -> str | None:
|
||||
root, data, inner = _nested_payloads(payload)
|
||||
for source in (data, inner, root):
|
||||
for key in ("video_url", "result_url", "url", "download_url"):
|
||||
value = source.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
|
||||
result = _as_dict(source.get("result"))
|
||||
for key in ("video_url", "result_url", "url", "download_url"):
|
||||
value = result.get(key)
|
||||
if value:
|
||||
return str(value)
|
||||
|
||||
content = _as_dict(source.get("content"))
|
||||
value = content.get("video_url") or content.get("url")
|
||||
if value:
|
||||
return str(value)
|
||||
|
||||
task_result = _as_dict(source.get("task_result"))
|
||||
videos = task_result.get("videos")
|
||||
if isinstance(videos, list) and videos:
|
||||
first = _as_dict(videos[0])
|
||||
value = first.get("url") or first.get("video_url")
|
||||
if value:
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
|
||||
def is_success_status(status: str) -> bool:
|
||||
return status in SUCCESS_STATUSES
|
||||
|
||||
|
||||
def is_failure_status(status: str, payload: Dict[str, Any] | None = None) -> bool:
|
||||
if status in FAILURE_STATUSES:
|
||||
return True
|
||||
if any(token in status for token in ("fail", "error", "reject", "timeout", "cancel")):
|
||||
return True
|
||||
if payload is None:
|
||||
return False
|
||||
root, data, inner = _nested_payloads(payload)
|
||||
failure_keys = ("error", "fail_reason", "failure_reason", "task_status_msg", "error_message")
|
||||
return any(any(source.get(key) for key in failure_keys) for source in (data, inner, root))
|
||||
+1
-1
@@ -1 +1 @@
|
||||
v1.10.0
|
||||
v1.10.3
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { api } from "../../../scripts/api.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "o1key.assetToggle",
|
||||
settings: [
|
||||
{
|
||||
id: "o1key.AssetSave",
|
||||
name: "资产保存",
|
||||
tooltip: "持久性保存生图记录",
|
||||
type: "boolean",
|
||||
defaultValue: true,
|
||||
},
|
||||
],
|
||||
async init() {
|
||||
const enabled = () => {
|
||||
try {
|
||||
return app.ui.settings.getSettingValue("o1key.AssetSave", true);
|
||||
} catch {
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
// --- Patch getHistory: 合并真实 jobs 与持久化历史 ---
|
||||
const _origGetHistory = api.getHistory.bind(api);
|
||||
api.getHistory = async function (maxItems = 200, opts = {}) {
|
||||
const real = await _origGetHistory(maxItems, opts);
|
||||
if (!enabled()) return real;
|
||||
try {
|
||||
const offset = opts?.offset || 0;
|
||||
const resp = await fetch(
|
||||
`/o1key/output_history?limit=${maxItems}&offset=${offset}`
|
||||
);
|
||||
if (!resp.ok) return real;
|
||||
const data = await resp.json();
|
||||
const persisted = data.jobs || [];
|
||||
if (!persisted.length) return real;
|
||||
if (!real || !Array.isArray(real) || !real.length) {
|
||||
// 仅持久化数据时也补充 priority
|
||||
const t = data.pagination?.total || persisted.length;
|
||||
return persisted.map((j, i) => ({
|
||||
...j,
|
||||
priority: j.priority ?? t - i,
|
||||
}));
|
||||
}
|
||||
// 合并去重:以 id 为 key,真实 jobs 优先
|
||||
const seen = new Set(real.map((j) => j.id));
|
||||
const merged = [...real];
|
||||
for (const job of persisted) {
|
||||
if (!seen.has(job.id)) {
|
||||
merged.push(job);
|
||||
}
|
||||
}
|
||||
// 按时间倒序
|
||||
merged.sort(
|
||||
(a, b) => (b.create_time || 0) - (a.create_time || 0)
|
||||
);
|
||||
const result = merged.slice(0, maxItems);
|
||||
// 补充 priority 字段(队列面板依赖此字段排序)
|
||||
const total = data.pagination?.total || result.length;
|
||||
for (let idx = 0; idx < result.length; idx++) {
|
||||
if (result[idx].priority == null) {
|
||||
result[idx] = {
|
||||
...result[idx],
|
||||
priority: total - idx,
|
||||
};
|
||||
}
|
||||
}
|
||||
return result;
|
||||
} catch (e) {
|
||||
return real;
|
||||
}
|
||||
};
|
||||
|
||||
// --- Patch deleteItem: 同时删除 o1key 持久化记录和文件 ---
|
||||
const _origDeleteItem = api.deleteItem.bind(api);
|
||||
api.deleteItem = async function (type, id) {
|
||||
const result = await _origDeleteItem(type, id);
|
||||
if (type === "history" && enabled()) {
|
||||
try {
|
||||
await fetch("/o1key/delete_history", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ delete: [id] }),
|
||||
});
|
||||
} catch (e) {}
|
||||
}
|
||||
return result;
|
||||
};
|
||||
|
||||
// --- Patch getJobDetail: 真实 API 失败时回退到本地路由 ---
|
||||
const _origGetJobDetail = api.getJobDetail.bind(api);
|
||||
api.getJobDetail = async function (jobId) {
|
||||
const real = await _origGetJobDetail(jobId);
|
||||
if (real) return real;
|
||||
if (!enabled()) return undefined;
|
||||
try {
|
||||
const resp = await fetch(
|
||||
`/o1key/job_detail/${encodeURIComponent(jobId)}`
|
||||
);
|
||||
if (!resp.ok) return undefined;
|
||||
return await resp.json();
|
||||
} catch (e) {
|
||||
return undefined;
|
||||
}
|
||||
};
|
||||
},
|
||||
});
|
||||
+1201
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,41 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "o1key.dotGrid",
|
||||
async setup() {
|
||||
function createBlackTile() {
|
||||
const size = 64;
|
||||
const c = document.createElement("canvas");
|
||||
c.width = size;
|
||||
c.height = size;
|
||||
const ctx = c.getContext("2d");
|
||||
ctx.fillStyle = "#1a1a1a";
|
||||
ctx.fillRect(0, 0, size, size);
|
||||
return c;
|
||||
}
|
||||
|
||||
// Hook immediately so the first draw already uses our tile
|
||||
const orig = LGraphCanvas.prototype.drawBackCanvas;
|
||||
LGraphCanvas.prototype.drawBackCanvas = function () {
|
||||
if (!this._pattern || !this._pattern_img) {
|
||||
const ctx = this.bgcanvas?.getContext("2d");
|
||||
if (ctx) {
|
||||
const t = createBlackTile();
|
||||
this._pattern = ctx.createPattern(t, "repeat");
|
||||
this._pattern_img = t;
|
||||
}
|
||||
}
|
||||
return orig.apply(this, arguments);
|
||||
};
|
||||
|
||||
// Also apply to current canvas instance if already exists
|
||||
const canvas = app.canvas;
|
||||
if (canvas?.bgcanvas) {
|
||||
const bgCtx = canvas.bgcanvas.getContext("2d");
|
||||
const tile = createBlackTile();
|
||||
canvas._pattern = bgCtx.createPattern(tile, "repeat");
|
||||
canvas._pattern_img = tile;
|
||||
canvas.draw(true, true);
|
||||
}
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,122 @@
|
||||
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() {},
|
||||
});
|
||||
}
|
||||
};
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,221 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "o1key.hideSidebarItems",
|
||||
async setup() {
|
||||
const hide = () => {
|
||||
// 隐藏侧边栏的"说明"、"应用"、"模型"、"节点"、"模板"按钮
|
||||
const hiddenLabels = ["说明", "帮助", "help", "应用", "apps", "模型", "models", "节点", "nodes", "模板", "templates", "template"];
|
||||
document.querySelectorAll(".side-bar-button, [class*='sidebar'] button, .p-togglebutton").forEach(btn => {
|
||||
const label = [
|
||||
btn.getAttribute("aria-label"),
|
||||
btn.getAttribute("title"),
|
||||
btn.getAttribute("data-title"),
|
||||
btn.getAttribute("data-label"),
|
||||
btn.textContent,
|
||||
].filter(Boolean).join(" ").toLowerCase();
|
||||
if (hiddenLabels.some(k => label.includes(k))) {
|
||||
btn.style.display = "none";
|
||||
}
|
||||
});
|
||||
// 隐藏左上角下拉菜单中的"帮助"项
|
||||
document.querySelectorAll(".p-menuitem, .p-menu-item, [class*='menu'] li, [class*='Menu'] li").forEach(item => {
|
||||
const text = item.textContent || "";
|
||||
if (text.trim() === "帮助" || text.trim() === "Help") {
|
||||
item.style.display = "none";
|
||||
}
|
||||
});
|
||||
// 隐藏登录/注册弹框(Google/Github 登录对话框)
|
||||
document.querySelectorAll("[class*='dialog'], [class*='Dialog'], [class*='modal'], [class*='Modal']").forEach(dialog => {
|
||||
const text = dialog.textContent || "";
|
||||
if ((text.includes("Google") || text.includes("Github")) &&
|
||||
(text.includes("登录") || text.includes("注册"))) {
|
||||
dialog.style.display = "none";
|
||||
const mask = dialog.previousElementSibling;
|
||||
if (mask && mask.className && mask.className.includes("mask")) {
|
||||
mask.style.display = "none";
|
||||
}
|
||||
}
|
||||
});
|
||||
// 在右侧内容区隐藏"登录/注册"按钮并注入 API Key
|
||||
injectApiKeyPanel();
|
||||
};
|
||||
|
||||
async function injectApiKeyPanel() {
|
||||
// 找到右侧内容区中包含"我的用户设置"的区域
|
||||
let contentArea = null;
|
||||
document.querySelectorAll("h1, h2, h3, h4, span, div").forEach(el => {
|
||||
const t = (el.textContent || "").trim();
|
||||
if (t === "我的用户设置" || t === "My User Settings") {
|
||||
contentArea = el.closest("div");
|
||||
}
|
||||
});
|
||||
if (!contentArea) return;
|
||||
|
||||
// 隐藏"登录/注册"按钮和"登录您的账户"文字
|
||||
contentArea.querySelectorAll("button, a, span, p, div").forEach(el => {
|
||||
const t = (el.textContent || "").trim();
|
||||
if (t.includes("登录") || t.includes("注册") || t === "Sign In" || t === "Sign Up" || t.includes("登录您的账户") || t.includes("Log in")) {
|
||||
if (el.tagName === "BUTTON" || el.tagName === "A" || t.includes("登录您的账户")) {
|
||||
el.style.display = "none";
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// 检查是否已注入(DOM 中已存在则跳过)
|
||||
if (document.querySelector("#o1key-apikey-box")) return;
|
||||
|
||||
// 创建 API Key 输入区域
|
||||
const box = document.createElement("div");
|
||||
box.id = "o1key-apikey-box";
|
||||
box.style.cssText = "margin-top:24px;padding:20px;border:1px solid #444;border-radius:8px;background:#1e1e1e;";
|
||||
box.innerHTML = `
|
||||
<div style="font-weight:bold;font-size:15px;margin-bottom:6px;color:#eee;">O1Key API 密钥</div>
|
||||
<div style="font-size:12px;color:#999;margin-bottom:14px;">输入您的 API 密钥,测试通过后方可保存</div>
|
||||
<div style="display:flex;gap:8px;align-items:center;">
|
||||
<input id="o1key-apikey-input" type="text" placeholder="请输入 API 密钥"
|
||||
autocomplete="off" autocorrect="off" autocapitalize="off" spellcheck="false"
|
||||
data-form-type="other" data-lpignore="true" name="o1key-key-field"
|
||||
style="flex:1;padding:8px 12px;border:1px solid #555;border-radius:4px;background:#111;color:#eee;font-size:13px;" />
|
||||
<button id="o1key-apikey-test"
|
||||
style="padding:8px 14px;border:none;border-radius:4px;background:#47a;color:#fff;cursor:pointer;font-size:13px;white-space:nowrap;">测试令牌</button>
|
||||
<button id="o1key-apikey-save" disabled
|
||||
style="padding:8px 14px;border:none;border-radius:4px;background:#555;color:#999;cursor:not-allowed;font-size:13px;white-space:nowrap;">保存</button>
|
||||
<button id="o1key-apikey-clear"
|
||||
style="padding:8px 14px;border:none;border-radius:4px;background:#a44;color:#fff;cursor:pointer;font-size:13px;white-space:nowrap;">清空密钥</button>
|
||||
</div>
|
||||
<div id="o1key-apikey-status" style="margin-top:10px;font-size:12px;color:#999;"></div>
|
||||
`;
|
||||
contentArea.appendChild(box);
|
||||
|
||||
// 加载当前状态
|
||||
try {
|
||||
const resp = await fetch("/o1key/api_key");
|
||||
const data = await resp.json();
|
||||
const status = box.querySelector("#o1key-apikey-status");
|
||||
if (data.has_key) {
|
||||
status.textContent = "当前密钥: " + data.masked;
|
||||
status.style.color = "#3b8";
|
||||
} else {
|
||||
status.textContent = "尚未配置 API 密钥";
|
||||
status.style.color = "#a84";
|
||||
}
|
||||
} catch (e) {}
|
||||
|
||||
const saveBtn = box.querySelector("#o1key-apikey-save");
|
||||
const testBtn = box.querySelector("#o1key-apikey-test");
|
||||
let testPassed = false;
|
||||
|
||||
// 输入变化时重置测试状态
|
||||
box.querySelector("#o1key-apikey-input").addEventListener("input", () => {
|
||||
testPassed = false;
|
||||
saveBtn.disabled = true;
|
||||
saveBtn.style.background = "#555";
|
||||
saveBtn.style.color = "#999";
|
||||
saveBtn.style.cursor = "not-allowed";
|
||||
});
|
||||
|
||||
// 测试令牌按钮
|
||||
testBtn.addEventListener("click", async () => {
|
||||
const input = box.querySelector("#o1key-apikey-input");
|
||||
const status = box.querySelector("#o1key-apikey-status");
|
||||
const key = input.value.trim();
|
||||
if (!key) { status.textContent = "请输入密钥"; status.style.color = "#a44"; return; }
|
||||
testBtn.disabled = true;
|
||||
testBtn.textContent = "验证中...";
|
||||
status.textContent = "正在验证密钥...";
|
||||
status.style.color = "#999";
|
||||
try {
|
||||
const resp = await fetch("/o1key/test_key", {
|
||||
method: "POST",
|
||||
headers: {"Content-Type": "application/json"},
|
||||
body: JSON.stringify({api_key: key})
|
||||
});
|
||||
const data = await resp.json();
|
||||
if (data.valid) {
|
||||
testPassed = true;
|
||||
status.textContent = "验证通过,可以保存";
|
||||
status.style.color = "#3b8";
|
||||
saveBtn.disabled = false;
|
||||
saveBtn.style.background = "#3b8";
|
||||
saveBtn.style.color = "#fff";
|
||||
saveBtn.style.cursor = "pointer";
|
||||
} else {
|
||||
testPassed = false;
|
||||
status.textContent = data.error || "验证失败";
|
||||
status.style.color = "#a44";
|
||||
}
|
||||
} catch (e) {
|
||||
status.textContent = "网络错误";
|
||||
status.style.color = "#a44";
|
||||
}
|
||||
testBtn.disabled = false;
|
||||
testBtn.textContent = "测试令牌";
|
||||
});
|
||||
|
||||
// 保存按钮(仅测试通过后可用)
|
||||
saveBtn.addEventListener("click", async () => {
|
||||
if (!testPassed) return;
|
||||
const input = box.querySelector("#o1key-apikey-input");
|
||||
const status = box.querySelector("#o1key-apikey-status");
|
||||
const key = input.value.trim();
|
||||
try {
|
||||
const resp = await fetch("/o1key/api_key", {
|
||||
method: "POST",
|
||||
headers: {"Content-Type": "application/json"},
|
||||
body: JSON.stringify({api_key: key})
|
||||
});
|
||||
const data = await resp.json();
|
||||
if (data.success) {
|
||||
status.textContent = "密钥已保存";
|
||||
status.style.color = "#3b8";
|
||||
input.value = "";
|
||||
testPassed = false;
|
||||
saveBtn.disabled = true;
|
||||
saveBtn.style.background = "#555";
|
||||
saveBtn.style.color = "#999";
|
||||
saveBtn.style.cursor = "not-allowed";
|
||||
} else {
|
||||
status.textContent = data.error || "保存失败";
|
||||
status.style.color = "#a44";
|
||||
}
|
||||
} catch (e) {
|
||||
status.textContent = "网络错误";
|
||||
status.style.color = "#a44";
|
||||
}
|
||||
});
|
||||
|
||||
// 清空密钥按钮
|
||||
const clearBtn = box.querySelector("#o1key-apikey-clear");
|
||||
clearBtn.addEventListener("click", async () => {
|
||||
if (!confirm("确定要清空 API 密钥吗?")) return;
|
||||
const status = box.querySelector("#o1key-apikey-status");
|
||||
try {
|
||||
const resp = await fetch("/o1key/api_key", { method: "DELETE" });
|
||||
const data = await resp.json();
|
||||
if (data.success) {
|
||||
status.textContent = "API 密钥已清空";
|
||||
status.style.color = "#a84";
|
||||
box.querySelector("#o1key-apikey-input").value = "";
|
||||
testPassed = false;
|
||||
saveBtn.disabled = true;
|
||||
saveBtn.style.background = "#555";
|
||||
saveBtn.style.color = "#999";
|
||||
saveBtn.style.cursor = "not-allowed";
|
||||
} else {
|
||||
status.textContent = data.error || "清空失败";
|
||||
status.style.color = "#a44";
|
||||
}
|
||||
} catch (e) {
|
||||
status.textContent = "网络错误";
|
||||
status.style.color = "#a44";
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
const observer = new MutationObserver(hide);
|
||||
observer.observe(document.body, { childList: true, subtree: true });
|
||||
setTimeout(hide, 1000);
|
||||
setTimeout(hide, 3000);
|
||||
},
|
||||
});
|
||||
Vendored
+1
File diff suppressed because one or more lines are too long
@@ -0,0 +1,767 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { api } from "../../../scripts/api.js";
|
||||
|
||||
const STORAGE_KEY = "o1key-notes";
|
||||
const SEEDED_KEY = "o1key-notes-seeded-v2";
|
||||
const NOTES_API = "/o1key/notes";
|
||||
const STYLE_ID = "o1key-notes-styles";
|
||||
|
||||
let notes = [];
|
||||
let searchText = "";
|
||||
let activeFilter = "all";
|
||||
let noteContainer = null;
|
||||
let editingNoteId = null;
|
||||
let draftNote = null;
|
||||
let pendingDelete = false;
|
||||
|
||||
const SAMPLE_NOTES = [
|
||||
{
|
||||
title: "产品主图:高级玻璃质感",
|
||||
tags: ["产品图", "玻璃", "灯光"],
|
||||
content: `Clean studio lighting, translucent glass material, subtle caustics, soft shadow, premium product photography, 85mm lens, minimal background, high detail.
|
||||
|
||||
使用方式:
|
||||
1. 把产品图作为参考图输入
|
||||
2. 保留主体轮廓,只调整材质和灯光
|
||||
3. 如果玻璃过亮,降低 "caustics" 权重`
|
||||
},
|
||||
{
|
||||
title: "Nano Banana 参考图经验",
|
||||
tags: ["Nano Banana", "参考图"],
|
||||
content: `参考图越多越容易跑偏,主体一致性优先用 1-3 张图。
|
||||
|
||||
复杂场景建议分两步:
|
||||
1. 先生成主体和构图
|
||||
2. 再用局部或参考图做材质、背景、文字等精修
|
||||
|
||||
如果提示词和参考图冲突,模型通常会优先参考图。`
|
||||
},
|
||||
{
|
||||
title: "电商模特换装模板",
|
||||
tags: ["电商", "模特", "换装"],
|
||||
content: `Keep face identity and original pose. Replace the outfit with: [服装描述].
|
||||
Realistic fabric texture, natural folds, accurate seams, studio e-commerce photography, clean background, consistent lighting.
|
||||
|
||||
Avoid changing body shape, face, hairstyle, camera angle, or hand position.`
|
||||
},
|
||||
{
|
||||
title: "常用负向词",
|
||||
tags: ["负向词", "通用"],
|
||||
content: "low quality, blurry, deformed hands, extra fingers, bad anatomy, distorted text, watermark, logo, oversaturated, plastic skin, broken geometry"
|
||||
},
|
||||
{
|
||||
title: "图片批量命名经验",
|
||||
tags: ["批量", "工作流"],
|
||||
content: `批量生图前先确认保存路径和命名规则。
|
||||
|
||||
推荐流程:
|
||||
1. 小批量跑 2-3 张确认风格
|
||||
2. 固定提示词和参考图
|
||||
3. 再放大批量数量
|
||||
|
||||
这样失败成本最低,也更容易定位是哪一环导致跑偏。`
|
||||
}
|
||||
];
|
||||
|
||||
const CSS = `
|
||||
#o1key-notes-root{position:absolute;inset:0;display:flex;flex-direction:column;min-height:0;color:#ddd;background:var(--comfy-menu-bg,#202020);overflow:hidden;font-family:inherit;--o1n-blue:#4f8cff;--o1n-blue-2:#7eb8f7;--o1n-blue-soft:rgba(79,140,255,.14)}
|
||||
#o1key-notes-header{padding:12px 14px 10px;border-bottom:1px solid rgba(255,255,255,.06);flex-shrink:0}
|
||||
.o1n-title-row{display:flex;align-items:center;justify-content:space-between;margin-bottom:10px}
|
||||
.o1n-title{font-size:14px;font-weight:700;color:#eee;letter-spacing:.2px}
|
||||
.o1n-head-actions{display:flex;gap:5px}
|
||||
.o1n-icon-btn{width:28px;height:28px;border:1px solid rgba(255,255,255,.09);border-radius:7px;background:rgba(255,255,255,.035);color:#888;display:flex;align-items:center;justify-content:center;cursor:pointer;transition:all .15s;flex-shrink:0}
|
||||
.o1n-icon-btn:hover{color:#ddd;background:rgba(255,255,255,.08);border-color:rgba(255,255,255,.16)}
|
||||
.o1n-icon-btn.primary{background:var(--o1n-blue);color:#fff;border-color:transparent}
|
||||
.o1n-icon-btn.primary:hover{background:#6ca0ff;color:#fff}
|
||||
#o1n-search-box{height:34px;border:1px solid rgba(255,255,255,.1);border-radius:8px;background:rgba(255,255,255,.045);display:flex;align-items:center;gap:8px;padding:0 10px;transition:border-color .15s}
|
||||
#o1n-search-box:focus-within{border-color:rgba(79,140,255,.42)}
|
||||
#o1n-search{flex:1;background:transparent;border:0;outline:0;color:#ddd;font-size:12px;min-width:0}
|
||||
#o1n-search::placeholder{color:#666}
|
||||
#o1n-panel-status{height:16px;margin-top:7px;color:#777;font-size:11px;line-height:16px;white-space:nowrap;overflow:hidden;text-overflow:ellipsis}
|
||||
.o1n-tag-strip{display:flex;gap:6px;margin-top:10px;overflow-x:auto;padding-bottom:1px}
|
||||
.o1n-tag-strip::-webkit-scrollbar{display:none}
|
||||
.o1n-tag-filter{height:26px;padding:0 9px;border-radius:7px;border:1px solid rgba(255,255,255,.09);background:transparent;color:#929292;font-size:12px;display:flex;align-items:center;gap:5px;cursor:pointer;white-space:nowrap;transition:all .15s}
|
||||
.o1n-tag-filter:hover{color:#cfcfcf;border-color:rgba(255,255,255,.16)}
|
||||
.o1n-tag-filter.active{color:var(--o1n-blue-2);border-color:rgba(79,140,255,.36);background:var(--o1n-blue-soft)}
|
||||
#o1key-notes-list{flex:1;min-height:0;overflow:auto;padding:8px 10px 12px;display:flex;flex-direction:column;gap:6px}
|
||||
#o1key-notes-list::-webkit-scrollbar,#o1n-content::-webkit-scrollbar{width:4px}
|
||||
#o1key-notes-list::-webkit-scrollbar-thumb,#o1n-content::-webkit-scrollbar-thumb{background:rgba(255,255,255,.1);border-radius:2px}
|
||||
.o1n-item{border:1px solid transparent;border-radius:8px;padding:10px;background:transparent;cursor:pointer;transition:all .12s}
|
||||
.o1n-item:hover{background:rgba(255,255,255,.04);border-color:rgba(255,255,255,.07)}
|
||||
.o1n-item.active{background:rgba(255,255,255,.065);border-color:rgba(255,255,255,.12);box-shadow:inset 2px 0 0 var(--o1n-blue)}
|
||||
.o1n-item-top{display:flex;align-items:center;gap:8px;margin-bottom:6px}
|
||||
.o1n-item-title{font-size:13px;color:#e0e0e0;white-space:nowrap;overflow:hidden;text-overflow:ellipsis;flex:1;min-width:0}
|
||||
.o1n-item-edit{font-size:11px;color:#777;flex-shrink:0}
|
||||
.o1n-item-text{font-size:12px;line-height:1.45;color:#8b8b8b;display:-webkit-box;-webkit-line-clamp:3;-webkit-box-orient:vertical;overflow:hidden;word-break:break-word}
|
||||
.o1n-item-meta{margin-top:8px;display:flex;align-items:center;justify-content:space-between;gap:8px;color:#666;font-size:10px}
|
||||
.o1n-item-tags{display:flex;gap:4px;overflow:hidden;min-width:0}
|
||||
.o1n-tag{color:var(--o1n-blue-2);background:rgba(79,140,255,.1);border:1px solid rgba(79,140,255,.2);border-radius:5px;padding:2px 5px;white-space:nowrap;max-width:120px;overflow:hidden;text-overflow:ellipsis}
|
||||
.o1n-empty{height:100%;display:flex;align-items:center;justify-content:center;text-align:center;color:#666;font-size:13px;line-height:1.6;padding:20px}
|
||||
.o1n-editor-backdrop{position:absolute;inset:0;background:rgba(0,0,0,.34);display:flex;align-items:stretch;justify-content:flex-end;z-index:30;animation:o1n-fade .12s ease}
|
||||
.o1n-editor{width:100%;height:100%;background:var(--comfy-menu-bg,#202020);border-left:1px solid rgba(79,140,255,.28);box-shadow:0 22px 70px rgba(0,0,0,.46);display:flex;flex-direction:column;min-height:0}
|
||||
.o1n-editor-top{height:48px;padding:0 14px;border-bottom:1px solid rgba(255,255,255,.08);display:flex;align-items:center;justify-content:space-between;flex-shrink:0}
|
||||
.o1n-editor-title{font-size:13px;font-weight:700;color:#eee}
|
||||
.o1n-editor-body{padding:12px;display:flex;flex-direction:column;gap:9px;flex:1;min-height:0}
|
||||
#o1n-title-input,#o1n-content{width:100%;border:1px solid rgba(255,255,255,.1);border-radius:8px;background:rgba(255,255,255,.045);color:#ddd;outline:0;font-family:inherit;transition:border-color .15s}
|
||||
#o1n-title-input{height:34px;padding:0 10px;font-size:13px;font-weight:700;flex-shrink:0}
|
||||
#o1n-content{flex:1;min-height:160px;resize:none;padding:11px 12px;font-size:13px;line-height:1.58}
|
||||
#o1n-title-input:focus,#o1n-content:focus{border-color:rgba(79,140,255,.42)}
|
||||
.o1n-tag-editor{min-height:74px;border:1px solid rgba(255,255,255,.1);border-radius:8px;background:rgba(255,255,255,.035);padding:7px;flex-shrink:0}
|
||||
.o1n-tag-row{display:flex;align-items:center;gap:6px;flex-wrap:wrap}
|
||||
.o1n-tag-token{height:24px;padding:0 7px;border-radius:6px;border:1px solid rgba(79,140,255,.25);background:rgba(79,140,255,.1);color:var(--o1n-blue-2);display:inline-flex;align-items:center;gap:6px;font-size:12px}
|
||||
.o1n-tag-remove{color:#8baee8;font-size:13px;line-height:1;cursor:pointer}
|
||||
.o1n-tag-input{height:24px;min-width:104px;flex:1;border:1px dashed rgba(255,255,255,.14);border-radius:6px;background:transparent;color:#ddd;padding:0 7px;font-size:12px;outline:0}
|
||||
.o1n-tag-input::placeholder{color:#777}
|
||||
.o1n-tag-hint{margin-top:7px;color:#666;font-size:11px;line-height:1.35}
|
||||
.o1n-editor-actions{display:grid;grid-template-columns:1fr 1fr 1fr;gap:7px;padding:12px;border-top:1px solid rgba(255,255,255,.08);flex-shrink:0}
|
||||
.o1n-action{height:32px;border:1px solid rgba(255,255,255,.1);border-radius:7px;background:rgba(255,255,255,.035);color:#aaa;font-size:12px;display:flex;align-items:center;justify-content:center;gap:6px;cursor:pointer;transition:all .15s}
|
||||
.o1n-action:hover{color:#e5e5e5;background:rgba(255,255,255,.075);border-color:rgba(255,255,255,.16)}
|
||||
.o1n-action.primary{background:var(--o1n-blue);color:#fff;border-color:transparent;font-weight:700}
|
||||
.o1n-action.primary:hover{background:#6ca0ff;color:#fff}
|
||||
.o1n-action.accent{background:#d8b45b;color:#171717;border-color:transparent;font-weight:700}
|
||||
.o1n-action.accent:hover{background:#e3c474;color:#111}
|
||||
.o1n-action.danger:hover{background:rgba(215,101,101,.18);border-color:rgba(215,101,101,.32);color:#f0a0a0}
|
||||
.o1n-delete-confirm{grid-column:1 / -1;display:none;align-items:center;gap:7px;padding:7px 8px;border:1px solid rgba(215,101,101,.24);border-radius:7px;background:rgba(215,101,101,.08);color:#ccc;font-size:12px}
|
||||
.o1n-delete-confirm.show{display:flex}
|
||||
.o1n-delete-confirm span{flex:1;min-width:0}
|
||||
.o1n-mini-btn{height:24px;border:1px solid rgba(255,255,255,.1);border-radius:6px;background:rgba(255,255,255,.06);color:#bbb;padding:0 8px;font-size:11px;cursor:pointer}
|
||||
.o1n-mini-btn.danger{background:rgba(215,101,101,.5);border-color:rgba(215,101,101,.3);color:#fff}
|
||||
#o1n-status{grid-column:1 / -1;height:16px;color:#777;font-size:11px;line-height:16px;white-space:nowrap;overflow:hidden;text-overflow:ellipsis}
|
||||
#o1n-import-file{display:none}
|
||||
@keyframes o1n-fade{from{opacity:0}to{opacity:1}}
|
||||
`;
|
||||
|
||||
function genId() {
|
||||
return Date.now().toString(36) + Math.random().toString(36).slice(2, 8);
|
||||
}
|
||||
|
||||
function now() {
|
||||
return Date.now();
|
||||
}
|
||||
|
||||
function parseTags(value) {
|
||||
if (Array.isArray(value)) return value.map(String).map(s => s.trim()).filter(Boolean);
|
||||
return String(value || "")
|
||||
.split(/[,,、\s]+/)
|
||||
.map(s => s.trim())
|
||||
.filter(Boolean);
|
||||
}
|
||||
|
||||
function makeNote(data = {}) {
|
||||
const ts = now();
|
||||
return {
|
||||
id: data.id || genId(),
|
||||
title: String(data.title || "未命名笔记"),
|
||||
tags: parseTags(data.tags || data.category || ""),
|
||||
content: String(data.content || ""),
|
||||
createdAt: data.createdAt || ts,
|
||||
updatedAt: data.updatedAt || ts,
|
||||
};
|
||||
}
|
||||
|
||||
function cloneNote(note) {
|
||||
return {
|
||||
...note,
|
||||
tags: [...(note.tags || [])],
|
||||
};
|
||||
}
|
||||
|
||||
function formatDate(ts) {
|
||||
const d = new Date(ts || now());
|
||||
const today = new Date();
|
||||
const pad = n => String(n).padStart(2, "0");
|
||||
if (d.toDateString() === today.toDateString()) {
|
||||
return `今天 ${pad(d.getHours())}:${pad(d.getMinutes())}`;
|
||||
}
|
||||
return `${pad(d.getMonth() + 1)}/${pad(d.getDate())}`;
|
||||
}
|
||||
|
||||
function escapeHtml(value) {
|
||||
return String(value ?? "")
|
||||
.replace(/&/g, "&")
|
||||
.replace(/</g, "<")
|
||||
.replace(/>/g, ">")
|
||||
.replace(/"/g, """)
|
||||
.replace(/'/g, "'");
|
||||
}
|
||||
|
||||
function icon(name, size = 14) {
|
||||
const icons = {
|
||||
plus: `<path d="M12 5v14M5 12h14"/>`,
|
||||
search: `<circle cx="11" cy="11" r="7"/><path d="M20 20l-3.5-3.5"/>`,
|
||||
copy: `<rect x="9" y="9" width="13" height="13" rx="2"/><path d="M5 15H4a2 2 0 01-2-2V4a2 2 0 012-2h9a2 2 0 012 2v1"/>`,
|
||||
insert: `<path d="M12 5v14"/><path d="M19 12H5"/>`,
|
||||
save: `<path d="M19 21H5a2 2 0 01-2-2V5a2 2 0 012-2h11l5 5v11a2 2 0 01-2 2z"/><path d="M17 21v-8H7v8"/><path d="M7 3v5h8"/>`,
|
||||
trash: `<path d="M3 6h18"/><path d="M8 6V4h8v2"/><path d="M19 6l-1 14H6L5 6"/>`,
|
||||
download: `<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4"/><path d="M7 10l5 5 5-5"/><path d="M12 15V3"/>`,
|
||||
upload: `<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4"/><path d="M17 8l-5-5-5 5"/><path d="M12 3v12"/>`,
|
||||
node: `<rect x="3" y="4" width="7" height="7" rx="1"/><rect x="14" y="13" width="7" height="7" rx="1"/><path d="M10 7.5h3a4 4 0 014 4V13"/>`,
|
||||
close: `<path d="M18 6L6 18M6 6l12 12"/>`,
|
||||
};
|
||||
return `<svg width="${size}" height="${size}" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">${icons[name] || icons.plus}</svg>`;
|
||||
}
|
||||
|
||||
function injectStyles() {
|
||||
if (document.getElementById(STYLE_ID)) return;
|
||||
const el = document.createElement("style");
|
||||
el.id = STYLE_ID;
|
||||
el.textContent = CSS;
|
||||
document.head.appendChild(el);
|
||||
}
|
||||
|
||||
function readCachedNotes() {
|
||||
try {
|
||||
const raw = localStorage.getItem(STORAGE_KEY);
|
||||
if (raw === null) return { found: false, notes: [] };
|
||||
const parsed = JSON.parse(raw);
|
||||
return { found: true, notes: Array.isArray(parsed) ? parsed.map(makeNote) : [] };
|
||||
} catch {
|
||||
return { found: false, notes: [] };
|
||||
}
|
||||
}
|
||||
|
||||
function writeCachedNotes() {
|
||||
try {
|
||||
localStorage.setItem(STORAGE_KEY, JSON.stringify(notes));
|
||||
} catch {}
|
||||
}
|
||||
|
||||
function hasSeededNotes() {
|
||||
try {
|
||||
return !!localStorage.getItem(SEEDED_KEY);
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
function markSeededNotes() {
|
||||
try {
|
||||
localStorage.setItem(SEEDED_KEY, "1");
|
||||
} catch {}
|
||||
}
|
||||
|
||||
function createSeedNotes() {
|
||||
return SAMPLE_NOTES.map(makeNote);
|
||||
}
|
||||
|
||||
async function loadNotes() {
|
||||
const cached = readCachedNotes();
|
||||
|
||||
try {
|
||||
const resp = await api.fetchApi(NOTES_API);
|
||||
if (!resp.ok) throw new Error(`HTTP ${resp.status}`);
|
||||
const data = await resp.json();
|
||||
|
||||
if (data.exists) {
|
||||
notes = Array.isArray(data.notes) ? data.notes.map(makeNote) : [];
|
||||
writeCachedNotes();
|
||||
return;
|
||||
}
|
||||
|
||||
notes = cached.found ? cached.notes : createSeedNotes();
|
||||
markSeededNotes();
|
||||
writeCachedNotes();
|
||||
await persistNotesToFile();
|
||||
return;
|
||||
} catch (e) {
|
||||
console.warn("[o1key notes] file storage unavailable, using localStorage", e);
|
||||
}
|
||||
|
||||
notes = cached.found ? cached.notes : [];
|
||||
if (!cached.found && !hasSeededNotes()) {
|
||||
notes = createSeedNotes();
|
||||
markSeededNotes();
|
||||
writeCachedNotes();
|
||||
}
|
||||
}
|
||||
|
||||
function saveNotes() {
|
||||
writeCachedNotes();
|
||||
void persistNotesToFile();
|
||||
}
|
||||
|
||||
async function persistNotesToFile() {
|
||||
try {
|
||||
const resp = await api.fetchApi(NOTES_API, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ notes }),
|
||||
});
|
||||
if (!resp.ok) throw new Error(`HTTP ${resp.status}`);
|
||||
return true;
|
||||
} catch (e) {
|
||||
console.warn("[o1key notes] failed to save notes file", e);
|
||||
setPanelStatus("笔记文件保存失败,已保存在浏览器缓存");
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
function allTags() {
|
||||
const counts = new Map();
|
||||
for (const note of notes) {
|
||||
for (const tag of note.tags || []) counts.set(tag, (counts.get(tag) || 0) + 1);
|
||||
}
|
||||
return [...counts.entries()].sort((a, b) => b[1] - a[1]).slice(0, 12);
|
||||
}
|
||||
|
||||
function filteredNotes() {
|
||||
const q = searchText.trim().toLowerCase();
|
||||
return notes.filter(note => {
|
||||
if (activeFilter.startsWith("tag:") && !(note.tags || []).includes(activeFilter.slice(4))) return false;
|
||||
if (!q) return true;
|
||||
return [note.title, note.content, ...(note.tags || [])].join("\n").toLowerCase().includes(q);
|
||||
}).sort((a, b) => (b.updatedAt || 0) - (a.updatedAt || 0));
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "o1key.notePanel",
|
||||
async setup() {
|
||||
await loadNotes();
|
||||
app.extensionManager.registerSidebarTab({
|
||||
id: "o1key-notes",
|
||||
title: "笔记",
|
||||
icon: "pi pi-pencil",
|
||||
type: "custom",
|
||||
render: (container) => {
|
||||
injectStyles();
|
||||
renderNotePanel(container);
|
||||
},
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
function renderNotePanel(container) {
|
||||
container.innerHTML = "";
|
||||
container.style.position = "relative";
|
||||
container.style.height = "100%";
|
||||
container.style.overflow = "hidden";
|
||||
|
||||
const root = document.createElement("div");
|
||||
root.id = "o1key-notes-root";
|
||||
root.innerHTML = `
|
||||
<div id="o1key-notes-header">
|
||||
<div class="o1n-title-row">
|
||||
<div class="o1n-title">笔记</div>
|
||||
<div class="o1n-head-actions">
|
||||
<button class="o1n-icon-btn" id="o1n-import" title="导入 JSON">${icon("upload")}</button>
|
||||
<button class="o1n-icon-btn" id="o1n-export" title="导出 JSON">${icon("download")}</button>
|
||||
<button class="o1n-icon-btn primary" id="o1n-new" title="新建笔记">${icon("plus")}</button>
|
||||
</div>
|
||||
</div>
|
||||
<div id="o1n-search-box">${icon("search", 13)}<input id="o1n-search" placeholder="搜索笔记内容或标签" value="${escapeHtml(searchText)}"></div>
|
||||
<div class="o1n-tag-strip" id="o1n-tag-strip"></div>
|
||||
<div id="o1n-panel-status"></div>
|
||||
</div>
|
||||
<div id="o1key-notes-list"></div>
|
||||
<input id="o1n-import-file" type="file" accept=".json,application/json">
|
||||
`;
|
||||
container.appendChild(root);
|
||||
noteContainer = root;
|
||||
bindPanelEvents(root);
|
||||
renderTagFilters();
|
||||
renderList();
|
||||
renderEditor();
|
||||
}
|
||||
|
||||
function bindPanelEvents(root) {
|
||||
root.querySelector("#o1n-new").addEventListener("click", createNote);
|
||||
root.querySelector("#o1n-export").addEventListener("click", exportNotes);
|
||||
root.querySelector("#o1n-import").addEventListener("click", () => root.querySelector("#o1n-import-file").click());
|
||||
root.querySelector("#o1n-import-file").addEventListener("change", importNotes);
|
||||
root.querySelector("#o1n-search").addEventListener("input", (e) => {
|
||||
searchText = e.target.value;
|
||||
renderList();
|
||||
});
|
||||
}
|
||||
|
||||
function renderTagFilters() {
|
||||
const row = noteContainer?.querySelector("#o1n-tag-strip");
|
||||
if (!row) return;
|
||||
const tags = allTags();
|
||||
row.innerHTML = `<button class="o1n-tag-filter${activeFilter === "all" ? " active" : ""}" data-filter="all">全部 <span>${notes.length}</span></button>` +
|
||||
tags.map(([tag, count]) => {
|
||||
const filter = `tag:${tag}`;
|
||||
return `<button class="o1n-tag-filter${activeFilter === filter ? " active" : ""}" data-filter="${escapeHtml(filter)}">#${escapeHtml(tag)} <span>${count}</span></button>`;
|
||||
}).join("");
|
||||
|
||||
row.querySelectorAll("[data-filter]").forEach(btn => {
|
||||
btn.addEventListener("click", () => {
|
||||
activeFilter = btn.dataset.filter;
|
||||
renderTagFilters();
|
||||
renderList();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
function renderList() {
|
||||
const list = noteContainer.querySelector("#o1key-notes-list");
|
||||
const visible = filteredNotes();
|
||||
|
||||
if (!visible.length) {
|
||||
list.innerHTML = `<div class="o1n-empty">没有匹配的笔记</div>`;
|
||||
return;
|
||||
}
|
||||
|
||||
list.innerHTML = visible.map(note => {
|
||||
const tags = (note.tags || []).slice(0, 3).map(tag => `<span class="o1n-tag">#${escapeHtml(tag)}</span>`).join("");
|
||||
return `<div class="o1n-item${note.id === editingNoteId ? " active" : ""}" data-id="${note.id}">
|
||||
<div class="o1n-item-top">
|
||||
<div class="o1n-item-title">${escapeHtml(note.title || "未命名笔记")}</div>
|
||||
<div class="o1n-item-edit">编辑</div>
|
||||
</div>
|
||||
<div class="o1n-item-text">${escapeHtml(note.content || "空笔记")}</div>
|
||||
<div class="o1n-item-meta"><div class="o1n-item-tags">${tags || `<span class="o1n-tag">#未分类</span>`}</div><span>${formatDate(note.updatedAt)}</span></div>
|
||||
</div>`;
|
||||
}).join("");
|
||||
|
||||
list.querySelectorAll(".o1n-item").forEach(item => {
|
||||
item.addEventListener("click", () => openEditor(item.dataset.id));
|
||||
});
|
||||
}
|
||||
|
||||
function renderEditor() {
|
||||
noteContainer.querySelector(".o1n-editor-backdrop")?.remove();
|
||||
if (!draftNote) return;
|
||||
pendingDelete = false;
|
||||
|
||||
const overlay = document.createElement("div");
|
||||
overlay.className = "o1n-editor-backdrop";
|
||||
overlay.innerHTML = `
|
||||
<div class="o1n-editor">
|
||||
<div class="o1n-editor-top">
|
||||
<div class="o1n-editor-title">${editingNoteId ? "编辑笔记" : "新建笔记"}</div>
|
||||
<button class="o1n-icon-btn" id="o1n-close-editor" title="关闭">${icon("close", 13)}</button>
|
||||
</div>
|
||||
<div class="o1n-editor-body">
|
||||
<input id="o1n-title-input" value="${escapeHtml(draftNote.title)}" placeholder="笔记标题">
|
||||
<div class="o1n-tag-editor">
|
||||
<div class="o1n-tag-row">
|
||||
${(draftNote.tags || []).map(tag => `<span class="o1n-tag-token">${escapeHtml(tag)} <span class="o1n-tag-remove" data-tag="${escapeHtml(tag)}">×</span></span>`).join("")}
|
||||
<input class="o1n-tag-input" id="o1n-tag-input" placeholder="+ 添加标签">
|
||||
</div>
|
||||
<div class="o1n-tag-hint">输入标签后按 Enter,或用逗号分隔多个标签。</div>
|
||||
</div>
|
||||
<textarea id="o1n-content" placeholder="记录提示词、参数经验、踩坑结论...">${escapeHtml(draftNote.content)}</textarea>
|
||||
</div>
|
||||
<div class="o1n-editor-actions">
|
||||
<button class="o1n-action" id="o1n-cancel">取消</button>
|
||||
<button class="o1n-action" id="o1n-copy">${icon("copy", 13)}复制</button>
|
||||
<button class="o1n-action accent" id="o1n-insert">${icon("insert", 13)}插入</button>
|
||||
<button class="o1n-action" id="o1n-save-from-node">${icon("node", 13)}从节点保存</button>
|
||||
<button class="o1n-action danger" id="o1n-delete">${icon("trash", 13)}删除</button>
|
||||
<button class="o1n-action primary" id="o1n-save">${icon("save", 13)}保存</button>
|
||||
<div class="o1n-delete-confirm" id="o1n-delete-confirm">
|
||||
<span>确定删除这条笔记?</span>
|
||||
<button class="o1n-mini-btn danger" id="o1n-delete-yes">删除</button>
|
||||
<button class="o1n-mini-btn" id="o1n-delete-no">取消</button>
|
||||
</div>
|
||||
<div id="o1n-status"></div>
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
noteContainer.appendChild(overlay);
|
||||
bindEditorEvents(overlay);
|
||||
}
|
||||
|
||||
function bindEditorEvents(overlay) {
|
||||
overlay.querySelector("#o1n-close-editor").addEventListener("click", closeEditor);
|
||||
overlay.querySelector("#o1n-cancel").addEventListener("click", closeEditor);
|
||||
overlay.querySelector("#o1n-title-input").addEventListener("input", updateDraftFromEditor);
|
||||
overlay.querySelector("#o1n-content").addEventListener("input", updateDraftFromEditor);
|
||||
overlay.querySelectorAll(".o1n-tag-remove").forEach(btn => {
|
||||
btn.addEventListener("click", () => removeDraftTag(btn.dataset.tag));
|
||||
});
|
||||
overlay.querySelector("#o1n-tag-input").addEventListener("keydown", handleTagInputKeydown);
|
||||
overlay.querySelector("#o1n-tag-input").addEventListener("blur", commitTagInput);
|
||||
overlay.querySelector("#o1n-copy").addEventListener("click", copyDraftContent);
|
||||
overlay.querySelector("#o1n-insert").addEventListener("click", insertDraftContent);
|
||||
overlay.querySelector("#o1n-save-from-node").addEventListener("click", saveFromCurrentNode);
|
||||
overlay.querySelector("#o1n-delete").addEventListener("click", requestDeleteEditingNote);
|
||||
overlay.querySelector("#o1n-delete-yes").addEventListener("click", deleteEditingNote);
|
||||
overlay.querySelector("#o1n-delete-no").addEventListener("click", hideDeleteConfirm);
|
||||
overlay.querySelector("#o1n-save").addEventListener("click", saveDraft);
|
||||
}
|
||||
|
||||
function openEditor(id) {
|
||||
const note = notes.find(n => n.id === id);
|
||||
if (!note) return;
|
||||
editingNoteId = id;
|
||||
draftNote = cloneNote(note);
|
||||
renderList();
|
||||
renderEditor();
|
||||
}
|
||||
|
||||
function closeEditor() {
|
||||
editingNoteId = null;
|
||||
draftNote = null;
|
||||
pendingDelete = false;
|
||||
renderList();
|
||||
renderEditor();
|
||||
}
|
||||
|
||||
function createNote() {
|
||||
editingNoteId = null;
|
||||
draftNote = makeNote({ title: "新笔记", tags: ["未分类"], content: "" });
|
||||
renderList();
|
||||
renderEditor();
|
||||
noteContainer.querySelector("#o1n-title-input")?.focus();
|
||||
}
|
||||
|
||||
function updateDraftFromEditor() {
|
||||
if (!draftNote) return;
|
||||
draftNote.title = noteContainer.querySelector("#o1n-title-input")?.value.trim() || "未命名笔记";
|
||||
draftNote.content = noteContainer.querySelector("#o1n-content")?.value || "";
|
||||
draftNote.updatedAt = now();
|
||||
}
|
||||
|
||||
function addTagsFromText(value) {
|
||||
if (!draftNote) return false;
|
||||
const incoming = parseTags(value);
|
||||
if (!incoming.length) return false;
|
||||
const existing = new Set(draftNote.tags || []);
|
||||
for (const tag of incoming) existing.add(tag);
|
||||
draftNote.tags = [...existing];
|
||||
draftNote.updatedAt = now();
|
||||
renderDraftTags();
|
||||
return true;
|
||||
}
|
||||
|
||||
function commitTagInput() {
|
||||
const input = noteContainer?.querySelector("#o1n-tag-input");
|
||||
if (!input) return;
|
||||
if (addTagsFromText(input.value)) input.value = "";
|
||||
}
|
||||
|
||||
function handleTagInputKeydown(e) {
|
||||
if (e.key !== "Enter" && e.key !== "," && e.key !== ",") return;
|
||||
e.preventDefault();
|
||||
commitTagInput();
|
||||
}
|
||||
|
||||
function removeDraftTag(tag) {
|
||||
if (!draftNote) return;
|
||||
draftNote.tags = (draftNote.tags || []).filter(t => t !== tag);
|
||||
draftNote.updatedAt = now();
|
||||
renderDraftTags();
|
||||
}
|
||||
|
||||
function renderDraftTags() {
|
||||
const row = noteContainer?.querySelector(".o1n-tag-row");
|
||||
if (!row || !draftNote) return;
|
||||
row.innerHTML = `
|
||||
${(draftNote.tags || []).map(tag => `<span class="o1n-tag-token">${escapeHtml(tag)} <span class="o1n-tag-remove" data-tag="${escapeHtml(tag)}">×</span></span>`).join("")}
|
||||
<input class="o1n-tag-input" id="o1n-tag-input" placeholder="+ 添加标签">
|
||||
`;
|
||||
row.querySelectorAll(".o1n-tag-remove").forEach(btn => {
|
||||
btn.addEventListener("click", () => removeDraftTag(btn.dataset.tag));
|
||||
});
|
||||
const input = row.querySelector("#o1n-tag-input");
|
||||
input.addEventListener("keydown", handleTagInputKeydown);
|
||||
input.addEventListener("blur", commitTagInput);
|
||||
input.focus();
|
||||
}
|
||||
|
||||
function saveDraft() {
|
||||
if (!draftNote) return;
|
||||
updateDraftFromEditor();
|
||||
draftNote.tags = draftNote.tags?.length ? draftNote.tags : ["未分类"];
|
||||
if (editingNoteId) {
|
||||
const idx = notes.findIndex(n => n.id === editingNoteId);
|
||||
if (idx >= 0) notes[idx] = cloneNote(draftNote);
|
||||
} else {
|
||||
draftNote.id = genId();
|
||||
draftNote.createdAt = now();
|
||||
draftNote.updatedAt = now();
|
||||
notes.unshift(cloneNote(draftNote));
|
||||
editingNoteId = draftNote.id;
|
||||
}
|
||||
saveNotes();
|
||||
renderTagFilters();
|
||||
closeEditor();
|
||||
}
|
||||
|
||||
function requestDeleteEditingNote() {
|
||||
if (!editingNoteId) {
|
||||
closeEditor();
|
||||
return;
|
||||
}
|
||||
pendingDelete = true;
|
||||
const box = noteContainer?.querySelector("#o1n-delete-confirm");
|
||||
box?.classList.add("show");
|
||||
setStatus("");
|
||||
}
|
||||
|
||||
function hideDeleteConfirm() {
|
||||
pendingDelete = false;
|
||||
noteContainer?.querySelector("#o1n-delete-confirm")?.classList.remove("show");
|
||||
}
|
||||
|
||||
function deleteEditingNote() {
|
||||
if (!editingNoteId) {
|
||||
closeEditor();
|
||||
return;
|
||||
}
|
||||
const note = notes.find(n => n.id === editingNoteId);
|
||||
if (!note) return;
|
||||
notes = notes.filter(n => n.id !== editingNoteId);
|
||||
saveNotes();
|
||||
renderTagFilters();
|
||||
closeEditor();
|
||||
}
|
||||
|
||||
async function copyDraftContent() {
|
||||
updateDraftFromEditor();
|
||||
try {
|
||||
await navigator.clipboard.writeText(draftNote?.content || "");
|
||||
setStatus("已复制");
|
||||
} catch {
|
||||
setStatus("复制失败,请手动选择内容");
|
||||
}
|
||||
}
|
||||
|
||||
function insertDraftContent() {
|
||||
updateDraftFromEditor();
|
||||
const text = draftNote?.content || "";
|
||||
if (!text.trim()) {
|
||||
setStatus("当前笔记内容为空");
|
||||
return;
|
||||
}
|
||||
if (insertIntoFocusedInput(text) || insertIntoSelectedNode(text)) {
|
||||
setStatus("已插入");
|
||||
return;
|
||||
}
|
||||
navigator.clipboard?.writeText(text).catch(() => {});
|
||||
setStatus("未找到可插入位置,已复制内容");
|
||||
}
|
||||
|
||||
function setStatus(message) {
|
||||
const el = noteContainer?.querySelector("#o1n-status");
|
||||
if (!el) return;
|
||||
el.textContent = message;
|
||||
clearTimeout(setStatus._timer);
|
||||
setStatus._timer = setTimeout(() => {
|
||||
if (el.textContent === message) el.textContent = "";
|
||||
}, 2200);
|
||||
}
|
||||
|
||||
function setPanelStatus(message) {
|
||||
const el = noteContainer?.querySelector("#o1n-panel-status");
|
||||
if (!el) return;
|
||||
el.textContent = message;
|
||||
clearTimeout(setPanelStatus._timer);
|
||||
setPanelStatus._timer = setTimeout(() => {
|
||||
if (el.textContent === message) el.textContent = "";
|
||||
}, 2600);
|
||||
}
|
||||
|
||||
function insertIntoFocusedInput(text) {
|
||||
const el = document.activeElement;
|
||||
if (!el || noteContainer.contains(el)) return false;
|
||||
if (!(el instanceof HTMLTextAreaElement || el instanceof HTMLInputElement)) return false;
|
||||
const start = el.selectionStart ?? el.value.length;
|
||||
const end = el.selectionEnd ?? el.value.length;
|
||||
const before = el.value.slice(0, start);
|
||||
const after = el.value.slice(end);
|
||||
const spacer = before && !before.endsWith("\n") ? "\n" : "";
|
||||
el.value = before + spacer + text + after;
|
||||
const pos = (before + spacer + text).length;
|
||||
el.setSelectionRange(pos, pos);
|
||||
el.dispatchEvent(new Event("input", { bubbles: true }));
|
||||
el.dispatchEvent(new Event("change", { bubbles: true }));
|
||||
return true;
|
||||
}
|
||||
|
||||
function findSelectedNode() {
|
||||
const selectedNodes = app.canvas?.selected_nodes;
|
||||
if (selectedNodes) {
|
||||
const values = Array.isArray(selectedNodes) ? selectedNodes : Object.values(selectedNodes);
|
||||
if (values.length) return values[0];
|
||||
}
|
||||
return app.canvas?.selected_node || null;
|
||||
}
|
||||
|
||||
function isPromptWidget(widget) {
|
||||
const name = String(widget?.name || "").toLowerCase();
|
||||
const value = widget?.value;
|
||||
if (typeof value !== "string") return false;
|
||||
return (
|
||||
name.includes("prompt") ||
|
||||
name.includes("提示词") ||
|
||||
name.includes("正向") ||
|
||||
name === "text" ||
|
||||
name === "文本"
|
||||
);
|
||||
}
|
||||
|
||||
function insertIntoSelectedNode(text) {
|
||||
const node = findSelectedNode();
|
||||
if (!node?.widgets?.length) return false;
|
||||
const widget = node.widgets.find(isPromptWidget) || node.widgets.find(w => typeof w.value === "string");
|
||||
if (!widget) return false;
|
||||
const current = widget.value || "";
|
||||
const spacer = current && !current.endsWith("\n") ? "\n" : "";
|
||||
widget.value = current + spacer + text;
|
||||
widget.callback?.(widget.value, app.canvas, node, widget);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
return true;
|
||||
}
|
||||
|
||||
function getPromptWidgetsFromSelectedNode() {
|
||||
const node = findSelectedNode();
|
||||
if (!node?.widgets?.length) return [];
|
||||
return node.widgets.filter(w => typeof w.value === "string" && String(w.value).trim());
|
||||
}
|
||||
|
||||
function saveFromCurrentNode() {
|
||||
const widgets = getPromptWidgetsFromSelectedNode();
|
||||
if (!widgets.length) {
|
||||
setStatus("当前没有选中包含文本的节点");
|
||||
return;
|
||||
}
|
||||
|
||||
const preferred = widgets.find(isPromptWidget) || widgets[0];
|
||||
const node = findSelectedNode();
|
||||
draftNote = makeNote({
|
||||
title: `${node?.title || node?.type || "节点"}:${preferred.name || "提示词"}`,
|
||||
tags: [node?.title || node?.type || "节点"],
|
||||
content: preferred.value,
|
||||
});
|
||||
editingNoteId = null;
|
||||
renderEditor();
|
||||
setStatus("已读取当前节点内容,保存后写入笔记");
|
||||
}
|
||||
|
||||
function exportNotes() {
|
||||
const blob = new Blob([JSON.stringify(notes, null, 2)], { type: "application/json" });
|
||||
const url = URL.createObjectURL(blob);
|
||||
const a = document.createElement("a");
|
||||
a.href = url;
|
||||
a.download = `o1key-notes-${new Date().toISOString().slice(0, 10)}.json`;
|
||||
document.body.appendChild(a);
|
||||
a.click();
|
||||
a.remove();
|
||||
URL.revokeObjectURL(url);
|
||||
}
|
||||
|
||||
async function importNotes(e) {
|
||||
const file = e.target.files?.[0];
|
||||
e.target.value = "";
|
||||
if (!file) return;
|
||||
try {
|
||||
const imported = JSON.parse(await file.text());
|
||||
const list = Array.isArray(imported) ? imported : imported.notes;
|
||||
if (!Array.isArray(list)) throw new Error("Invalid notes file");
|
||||
const existing = new Set(notes.map(n => n.id));
|
||||
const normalized = list.map(makeNote).map(n => {
|
||||
if (existing.has(n.id)) n.id = genId();
|
||||
return n;
|
||||
});
|
||||
notes = [...normalized, ...notes];
|
||||
saveNotes();
|
||||
renderTagFilters();
|
||||
renderList();
|
||||
setPanelStatus(`已导入 ${normalized.length} 条笔记`);
|
||||
} catch {
|
||||
setPanelStatus("导入失败,请确认 JSON 格式");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,797 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { api } from "../../../scripts/api.js";
|
||||
|
||||
// Fabric.js local loader
|
||||
let fabricLoaded = false;
|
||||
function loadFabric() {
|
||||
if (fabricLoaded) return Promise.resolve();
|
||||
return new Promise((resolve, reject) => {
|
||||
const script = document.createElement("script");
|
||||
// Load from local extension directory (allowed by CSP 'self')
|
||||
script.src = new URL("./lib/fabric.min.js", import.meta.url).href;
|
||||
script.onload = () => { fabricLoaded = true; resolve(); };
|
||||
script.onerror = () => reject(new Error("Failed to load Fabric.js"));
|
||||
document.head.appendChild(script);
|
||||
});
|
||||
}
|
||||
|
||||
// --- Styles ---
|
||||
const STYLES = `
|
||||
.pb-overlay { position:fixed; inset:0; z-index:99999; background:rgba(0,0,0,0.85); display:flex; flex-direction:column; align-items:center; justify-content:center; }
|
||||
.pb-toolbar { display:flex; gap:6px; padding:10px 16px; background:#1e1e1e; border-radius:8px; margin-bottom:10px; align-items:center; flex-wrap:wrap; }
|
||||
.pb-toolbar button { background:#333; color:#eee; border:1px solid #555; border-radius:4px; padding:6px 12px; cursor:pointer; font-size:13px; transition:all .15s; }
|
||||
.pb-toolbar button:hover { background:#444; }
|
||||
.pb-toolbar button.active { background:#0066ff; border-color:#0066ff; color:#fff; }
|
||||
.pb-toolbar .pb-sep { width:1px; height:24px; background:#555; margin:0 4px; }
|
||||
.pb-toolbar input[type=color] { width:32px; height:28px; border:none; padding:0; cursor:pointer; border-radius:4px; }
|
||||
.pb-toolbar input[type=range] { width:80px; accent-color:#0066ff; }
|
||||
.pb-toolbar input.pb-mosaic-size { width:90px; }
|
||||
.pb-toolbar .pb-label { color:#aaa; font-size:12px; }
|
||||
.pb-canvas-wrap { border:2px solid #444; border-radius:4px; overflow:hidden; }
|
||||
.pb-actions { display:flex; gap:10px; margin-top:10px; }
|
||||
.pb-actions button { padding:8px 24px; border-radius:6px; font-size:14px; cursor:pointer; border:none; }
|
||||
.pb-actions .pb-cancel { background:#555; color:#eee; }
|
||||
.pb-actions .pb-confirm { background:#0066ff; color:#fff; }
|
||||
`;
|
||||
|
||||
function injectStyles() {
|
||||
if (document.getElementById("pb-styles")) return;
|
||||
const el = document.createElement("style");
|
||||
el.id = "pb-styles";
|
||||
el.textContent = STYLES;
|
||||
document.head.appendChild(el);
|
||||
}
|
||||
|
||||
function clamp(value, min, max) {
|
||||
return Math.max(min, Math.min(max, value));
|
||||
}
|
||||
|
||||
const PB_EXPORT_PROPS = ["pbTool"];
|
||||
|
||||
function clampPointer(canvas, pointer) {
|
||||
return {
|
||||
x: clamp(pointer.x, 0, canvas.width),
|
||||
y: clamp(pointer.y, 0, canvas.height),
|
||||
};
|
||||
}
|
||||
|
||||
function makeCircleCursor(size) {
|
||||
const cursorSize = clamp(Math.round(size), 18, 80);
|
||||
const center = cursorSize / 2;
|
||||
const hotspot = Math.round(center);
|
||||
const radius = Math.max(3, center - 2);
|
||||
const svg = `
|
||||
<svg xmlns="http://www.w3.org/2000/svg" width="${cursorSize}" height="${cursorSize}" viewBox="0 0 ${cursorSize} ${cursorSize}">
|
||||
<circle cx="${center}" cy="${center}" r="${radius}" fill="none" stroke="black" stroke-width="3"/>
|
||||
<circle cx="${center}" cy="${center}" r="${radius}" fill="none" stroke="white" stroke-width="1.5"/>
|
||||
</svg>
|
||||
`.trim();
|
||||
return `url("data:image/svg+xml;charset=UTF-8,${encodeURIComponent(svg)}") ${hotspot} ${hotspot}, crosshair`;
|
||||
}
|
||||
|
||||
function applyMosaicCursor(canvas, state) {
|
||||
const cursor = makeCircleCursor(state.mosaicSize);
|
||||
canvas.defaultCursor = cursor;
|
||||
canvas.hoverCursor = cursor;
|
||||
}
|
||||
|
||||
function configureMosaicObject(obj) {
|
||||
obj.set({
|
||||
selectable: false,
|
||||
hasControls: false,
|
||||
hasBorders: false,
|
||||
lockMovementX: true,
|
||||
lockMovementY: true,
|
||||
lockScalingX: true,
|
||||
lockScalingY: true,
|
||||
lockRotation: true,
|
||||
perPixelTargetFind: true,
|
||||
objectCaching: false,
|
||||
});
|
||||
obj.pbTool = "mosaic";
|
||||
return obj;
|
||||
}
|
||||
|
||||
function normalizeMosaicObjects(canvas) {
|
||||
canvas.getObjects().forEach((obj) => {
|
||||
if (obj.pbTool === "mosaic" || obj.type === "image") {
|
||||
configureMosaicObject(obj);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// --- Get current image URL from node ---
|
||||
function getImageUrl(node) {
|
||||
if (node.imgs && node.imgs.length > 0) {
|
||||
return node.imgs[node.imageIndex ?? 0].src;
|
||||
}
|
||||
const widget = node.widgets?.find(w => w.name === "image");
|
||||
if (!widget?.value) return null;
|
||||
const val = String(widget.value);
|
||||
const match = val.match(/^(.+?)(?:\s*\[(\w+)\])?$/);
|
||||
if (!match) return null;
|
||||
const filename = match[1];
|
||||
const type = match[2] || "input";
|
||||
const parts = filename.split("/");
|
||||
const name = parts.pop();
|
||||
const subfolder = parts.join("/");
|
||||
return `/view?filename=${encodeURIComponent(name)}&type=${type}&subfolder=${encodeURIComponent(subfolder)}`;
|
||||
}
|
||||
|
||||
// --- History Manager ---
|
||||
class HistoryManager {
|
||||
constructor(canvas) {
|
||||
this.canvas = canvas;
|
||||
this.stack = [];
|
||||
this.index = -1;
|
||||
this.locked = false;
|
||||
}
|
||||
save() {
|
||||
if (this.locked) return;
|
||||
this.index++;
|
||||
this.stack.length = this.index;
|
||||
this.stack.push(this.canvas.toJSON(PB_EXPORT_PROPS));
|
||||
}
|
||||
undo() {
|
||||
if (this.index <= 0) return;
|
||||
this.index--;
|
||||
this._restore();
|
||||
}
|
||||
redo() {
|
||||
if (this.index >= this.stack.length - 1) return;
|
||||
this.index++;
|
||||
this._restore();
|
||||
}
|
||||
_restore() {
|
||||
this.locked = true;
|
||||
this.canvas.loadFromJSON(this.stack[this.index], () => {
|
||||
normalizeMosaicObjects(this.canvas);
|
||||
this.canvas.renderAll();
|
||||
this.locked = false;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// --- Shape drawing handler ---
|
||||
function setupShapeDrawing(canvas, state) {
|
||||
let startX, startY, shape;
|
||||
|
||||
canvas.on("mouse:down", (opt) => {
|
||||
if (state.tool === "select" || state.tool === "brush" || state.tool === "eraser" || state.tool === "mosaic") return;
|
||||
const ptr = canvas.getPointer(opt.e);
|
||||
startX = ptr.x;
|
||||
startY = ptr.y;
|
||||
state.drawing = true;
|
||||
|
||||
const opts = { left: startX, top: startY, fill: "transparent", stroke: state.color, strokeWidth: state.width, selectable: true };
|
||||
|
||||
if (state.tool === "rect") {
|
||||
shape = new fabric.Rect({ ...opts, width: 0, height: 0 });
|
||||
} else if (state.tool === "circle") {
|
||||
shape = new fabric.Ellipse({ ...opts, rx: 0, ry: 0 });
|
||||
} else if (state.tool === "line") {
|
||||
shape = new fabric.Line([startX, startY, startX, startY], { stroke: state.color, strokeWidth: state.width, selectable: true });
|
||||
}
|
||||
if (shape) canvas.add(shape);
|
||||
});
|
||||
|
||||
canvas.on("mouse:move", (opt) => {
|
||||
if (!state.drawing || !shape) return;
|
||||
const ptr = canvas.getPointer(opt.e);
|
||||
const dx = ptr.x - startX;
|
||||
const dy = ptr.y - startY;
|
||||
|
||||
if (state.tool === "rect") {
|
||||
shape.set({ left: dx > 0 ? startX : ptr.x, top: dy > 0 ? startY : ptr.y, width: Math.abs(dx), height: Math.abs(dy) });
|
||||
} else if (state.tool === "circle") {
|
||||
shape.set({ left: dx > 0 ? startX : ptr.x, top: dy > 0 ? startY : ptr.y, rx: Math.abs(dx) / 2, ry: Math.abs(dy) / 2 });
|
||||
} else if (state.tool === "line") {
|
||||
shape.set({ x2: ptr.x, y2: ptr.y });
|
||||
}
|
||||
canvas.renderAll();
|
||||
});
|
||||
|
||||
canvas.on("mouse:up", () => {
|
||||
if (!state.drawing || !shape) return;
|
||||
state.drawing = false;
|
||||
shape = null;
|
||||
});
|
||||
}
|
||||
|
||||
// --- Mosaic brush handler ---
|
||||
function createMosaicSource(sourceImg, canvas, blockSize) {
|
||||
const scaleX = canvas.width / sourceImg.width;
|
||||
const scaleY = canvas.height / sourceImg.height;
|
||||
const sourceBlockSize = Math.max(1, Math.round(blockSize / Math.min(scaleX, scaleY)));
|
||||
const smallW = Math.max(1, Math.ceil(sourceImg.width / sourceBlockSize));
|
||||
const smallH = Math.max(1, Math.ceil(sourceImg.height / sourceBlockSize));
|
||||
|
||||
const smallCanvas = document.createElement("canvas");
|
||||
smallCanvas.width = smallW;
|
||||
smallCanvas.height = smallH;
|
||||
const smallCtx = smallCanvas.getContext("2d");
|
||||
smallCtx.imageSmoothingEnabled = true;
|
||||
smallCtx.drawImage(sourceImg, 0, 0, sourceImg.width, sourceImg.height, 0, 0, smallW, smallH);
|
||||
|
||||
const pixelCanvas = document.createElement("canvas");
|
||||
pixelCanvas.width = sourceImg.width;
|
||||
pixelCanvas.height = sourceImg.height;
|
||||
const pixelCtx = pixelCanvas.getContext("2d");
|
||||
pixelCtx.imageSmoothingEnabled = false;
|
||||
pixelCtx.drawImage(smallCanvas, 0, 0, smallW, smallH, 0, 0, sourceImg.width, sourceImg.height);
|
||||
|
||||
return pixelCanvas;
|
||||
}
|
||||
|
||||
function getMosaicBrushBounds(sourceImg, canvas, centerX, centerY, size) {
|
||||
const scaleX = canvas.width / sourceImg.width;
|
||||
const scaleY = canvas.height / sourceImg.height;
|
||||
const sourceX = centerX / scaleX;
|
||||
const sourceY = centerY / scaleY;
|
||||
const radiusX = Math.max(1, (size / 2) / scaleX);
|
||||
const radiusY = Math.max(1, (size / 2) / scaleY);
|
||||
|
||||
return {
|
||||
centerX: sourceX,
|
||||
centerY: sourceY,
|
||||
radius: Math.max(radiusX, radiusY),
|
||||
left: clamp(Math.floor(sourceX - radiusX), 0, sourceImg.width),
|
||||
top: clamp(Math.floor(sourceY - radiusY), 0, sourceImg.height),
|
||||
right: clamp(Math.ceil(sourceX + radiusX), 0, sourceImg.width),
|
||||
bottom: clamp(Math.ceil(sourceY + radiusY), 0, sourceImg.height),
|
||||
};
|
||||
}
|
||||
|
||||
function expandMosaicBounds(bounds, brushBounds) {
|
||||
if (!bounds.left && bounds.left !== 0) {
|
||||
bounds.left = brushBounds.left;
|
||||
bounds.top = brushBounds.top;
|
||||
bounds.right = brushBounds.right;
|
||||
bounds.bottom = brushBounds.bottom;
|
||||
return;
|
||||
}
|
||||
bounds.left = Math.min(bounds.left, brushBounds.left);
|
||||
bounds.top = Math.min(bounds.top, brushBounds.top);
|
||||
bounds.right = Math.max(bounds.right, brushBounds.right);
|
||||
bounds.bottom = Math.max(bounds.bottom, brushBounds.bottom);
|
||||
}
|
||||
|
||||
function paintMosaicStamp(sourceImg, canvas, mosaicSource, strokeCtx, bounds, point) {
|
||||
const brushBounds = getMosaicBrushBounds(sourceImg, canvas, point.x, point.y, point.size);
|
||||
if (brushBounds.right <= brushBounds.left || brushBounds.bottom <= brushBounds.top) return false;
|
||||
|
||||
strokeCtx.save();
|
||||
strokeCtx.beginPath();
|
||||
strokeCtx.arc(brushBounds.centerX, brushBounds.centerY, brushBounds.radius, 0, Math.PI * 2);
|
||||
strokeCtx.clip();
|
||||
strokeCtx.drawImage(
|
||||
mosaicSource,
|
||||
brushBounds.left,
|
||||
brushBounds.top,
|
||||
brushBounds.right - brushBounds.left,
|
||||
brushBounds.bottom - brushBounds.top,
|
||||
brushBounds.left,
|
||||
brushBounds.top,
|
||||
brushBounds.right - brushBounds.left,
|
||||
brushBounds.bottom - brushBounds.top
|
||||
);
|
||||
strokeCtx.restore();
|
||||
|
||||
expandMosaicBounds(bounds, brushBounds);
|
||||
return true;
|
||||
}
|
||||
|
||||
function paintMosaicLine(sourceImg, canvas, mosaicSource, strokeCtx, bounds, from, to, size) {
|
||||
const dx = to.x - from.x;
|
||||
const dy = to.y - from.y;
|
||||
const distance = Math.hypot(dx, dy);
|
||||
const step = Math.max(2, size * 0.25);
|
||||
const steps = Math.max(1, Math.ceil(distance / step));
|
||||
let painted = false;
|
||||
|
||||
for (let i = 1; i <= steps; i++) {
|
||||
const t = i / steps;
|
||||
painted = paintMosaicStamp(sourceImg, canvas, mosaicSource, strokeCtx, bounds, {
|
||||
x: from.x + dx * t,
|
||||
y: from.y + dy * t,
|
||||
size,
|
||||
}) || painted;
|
||||
}
|
||||
|
||||
return painted;
|
||||
}
|
||||
|
||||
function createMosaicStrokeObject(strokeCanvas, bounds, canvas, sourceImg) {
|
||||
const width = bounds.right - bounds.left;
|
||||
const height = bounds.bottom - bounds.top;
|
||||
if (width < 1 || height < 1) return Promise.resolve(null);
|
||||
|
||||
const cropCanvas = document.createElement("canvas");
|
||||
cropCanvas.width = width;
|
||||
cropCanvas.height = height;
|
||||
cropCanvas.getContext("2d").drawImage(
|
||||
strokeCanvas,
|
||||
bounds.left,
|
||||
bounds.top,
|
||||
width,
|
||||
height,
|
||||
0,
|
||||
0,
|
||||
width,
|
||||
height
|
||||
);
|
||||
|
||||
const scaleX = canvas.width / sourceImg.width;
|
||||
const scaleY = canvas.height / sourceImg.height;
|
||||
const dataUrl = cropCanvas.toDataURL("image/png");
|
||||
return new Promise(resolve => {
|
||||
fabric.Image.fromURL(dataUrl, (imgObj) => {
|
||||
configureMosaicObject(imgObj).set({
|
||||
left: bounds.left * scaleX,
|
||||
top: bounds.top * scaleY,
|
||||
scaleX,
|
||||
scaleY,
|
||||
});
|
||||
resolve(imgObj);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
function setupMosaicDrawing(canvas, state, sourceImg, history) {
|
||||
let lastPoint, mosaicSource, strokeCanvas, strokeCtx, strokePreview, strokeBounds, strokePainted;
|
||||
|
||||
const refreshPreview = () => {
|
||||
if (!strokePreview) return;
|
||||
strokePreview.dirty = true;
|
||||
canvas.requestRenderAll();
|
||||
};
|
||||
|
||||
canvas.on("mouse:down", (opt) => {
|
||||
if (state.tool !== "mosaic") return;
|
||||
const ptr = clampPointer(canvas, canvas.getPointer(opt.e));
|
||||
lastPoint = ptr;
|
||||
state.drawing = true;
|
||||
strokePainted = false;
|
||||
strokeBounds = {};
|
||||
mosaicSource = createMosaicSource(sourceImg, canvas, state.mosaicSize);
|
||||
|
||||
strokeCanvas = document.createElement("canvas");
|
||||
strokeCanvas.width = sourceImg.width;
|
||||
strokeCanvas.height = sourceImg.height;
|
||||
strokeCtx = strokeCanvas.getContext("2d");
|
||||
strokeCtx.imageSmoothingEnabled = false;
|
||||
|
||||
strokePreview = new fabric.Image(strokeCanvas, {
|
||||
left: 0,
|
||||
top: 0,
|
||||
scaleX: canvas.width / sourceImg.width,
|
||||
scaleY: canvas.height / sourceImg.height,
|
||||
selectable: false,
|
||||
evented: false,
|
||||
excludeFromExport: true,
|
||||
objectCaching: false,
|
||||
});
|
||||
|
||||
history.locked = true;
|
||||
canvas.add(strokePreview);
|
||||
history.locked = false;
|
||||
|
||||
strokePainted = paintMosaicStamp(sourceImg, canvas, mosaicSource, strokeCtx, strokeBounds, {
|
||||
x: ptr.x,
|
||||
y: ptr.y,
|
||||
size: state.mosaicSize,
|
||||
}) || strokePainted;
|
||||
refreshPreview();
|
||||
});
|
||||
|
||||
canvas.on("mouse:move", (opt) => {
|
||||
if (state.tool !== "mosaic" || !state.drawing || !strokeCtx || !lastPoint) return;
|
||||
const ptr = clampPointer(canvas, canvas.getPointer(opt.e));
|
||||
strokePainted = paintMosaicLine(
|
||||
sourceImg,
|
||||
canvas,
|
||||
mosaicSource,
|
||||
strokeCtx,
|
||||
strokeBounds,
|
||||
lastPoint,
|
||||
ptr,
|
||||
state.mosaicSize
|
||||
) || strokePainted;
|
||||
lastPoint = ptr;
|
||||
refreshPreview();
|
||||
});
|
||||
|
||||
canvas.on("mouse:up", async () => {
|
||||
if (!state.drawing || !strokePreview) return;
|
||||
|
||||
history.locked = true;
|
||||
canvas.remove(strokePreview);
|
||||
history.locked = false;
|
||||
state.drawing = false;
|
||||
lastPoint = null;
|
||||
|
||||
if (!strokePainted) {
|
||||
strokeCanvas = null;
|
||||
strokeCtx = null;
|
||||
strokePreview = null;
|
||||
mosaicSource = null;
|
||||
canvas.renderAll();
|
||||
return;
|
||||
}
|
||||
|
||||
const mosaicObj = await createMosaicStrokeObject(strokeCanvas, strokeBounds, canvas, sourceImg);
|
||||
mosaicSource = null;
|
||||
strokeCanvas = null;
|
||||
strokeCtx = null;
|
||||
strokePreview = null;
|
||||
|
||||
if (mosaicObj) {
|
||||
canvas.add(mosaicObj);
|
||||
canvas.discardActiveObject();
|
||||
}
|
||||
canvas.renderAll();
|
||||
});
|
||||
}
|
||||
|
||||
// --- Open Paint Modal ---
|
||||
async function openPaintModal(node) {
|
||||
injectStyles();
|
||||
await loadFabric();
|
||||
|
||||
if (!node.properties) node.properties = {};
|
||||
|
||||
// Detect if user switched to a different image — reset saved state
|
||||
const widget = node.widgets?.find(w => w.name === "image");
|
||||
const currentVal = widget?.value ? String(widget.value) : "";
|
||||
const originalImg = node.properties.paintBrushOriginal;
|
||||
if (originalImg && currentVal !== originalImg && currentVal !== "painted_" + originalImg) {
|
||||
// Image changed, clear old paint state
|
||||
delete node.properties.paintBrushCanvas;
|
||||
delete node.properties.paintBrushOriginal;
|
||||
}
|
||||
|
||||
const savedState = node.properties.paintBrushCanvas;
|
||||
const storedOriginal = node.properties.paintBrushOriginal;
|
||||
|
||||
// If re-editing, use the original image as background; otherwise use current
|
||||
let bgUrl;
|
||||
if (savedState && storedOriginal) {
|
||||
bgUrl = `/view?filename=${encodeURIComponent(storedOriginal)}&type=input&subfolder=`;
|
||||
} else {
|
||||
bgUrl = getImageUrl(node);
|
||||
}
|
||||
if (!bgUrl) { alert("请先加载一张图片"); return; }
|
||||
|
||||
// Save original image name on first paint
|
||||
if (!storedOriginal) {
|
||||
node.properties.paintBrushOriginal = currentVal;
|
||||
}
|
||||
|
||||
// Create overlay
|
||||
const overlay = document.createElement("div");
|
||||
overlay.className = "pb-overlay";
|
||||
|
||||
const maxW = window.innerWidth * 0.8;
|
||||
const maxH = window.innerHeight * 0.75;
|
||||
|
||||
// Load image to get dimensions
|
||||
const img = await new Promise((resolve, reject) => {
|
||||
const i = new Image();
|
||||
i.crossOrigin = "anonymous";
|
||||
i.onload = () => resolve(i);
|
||||
i.onerror = () => reject(new Error("图片加载失败"));
|
||||
i.src = bgUrl;
|
||||
});
|
||||
|
||||
const scale = Math.min(maxW / img.width, maxH / img.height, 1);
|
||||
const cw = Math.round(img.width * scale);
|
||||
const ch = Math.round(img.height * scale);
|
||||
|
||||
const state = { tool: "brush", color: "#ff0000", width: 4, mosaicSize: 14, drawing: false };
|
||||
|
||||
const toolbar = buildToolbar(state);
|
||||
overlay.appendChild(toolbar);
|
||||
|
||||
const wrap = document.createElement("div");
|
||||
wrap.className = "pb-canvas-wrap";
|
||||
const canvasEl = document.createElement("canvas");
|
||||
canvasEl.width = cw;
|
||||
canvasEl.height = ch;
|
||||
wrap.appendChild(canvasEl);
|
||||
overlay.appendChild(wrap);
|
||||
document.body.appendChild(overlay);
|
||||
|
||||
// Init Fabric canvas
|
||||
const canvas = new fabric.Canvas(canvasEl, { width: cw, height: ch, isDrawingMode: true });
|
||||
canvas.freeDrawingBrush.color = state.color;
|
||||
canvas.freeDrawingBrush.width = state.width;
|
||||
|
||||
// Set background image (always the original)
|
||||
await new Promise(resolve => {
|
||||
canvas.setBackgroundImage(bgUrl, () => { canvas.renderAll(); resolve(); }, {
|
||||
scaleX: cw / img.width, scaleY: ch / img.height, crossOrigin: "anonymous"
|
||||
});
|
||||
});
|
||||
|
||||
// Restore previous drawing objects if re-editing
|
||||
if (savedState) {
|
||||
await new Promise(resolve => {
|
||||
canvas.loadFromJSON(savedState, () => {
|
||||
normalizeMosaicObjects(canvas);
|
||||
// Re-apply background since loadFromJSON may clear it
|
||||
canvas.setBackgroundImage(bgUrl, () => { canvas.renderAll(); resolve(); }, {
|
||||
scaleX: cw / img.width, scaleY: ch / img.height, crossOrigin: "anonymous"
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
// History
|
||||
const history = new HistoryManager(canvas);
|
||||
setTimeout(() => history.save(), 300);
|
||||
canvas.on("object:added", () => history.save());
|
||||
canvas.on("object:modified", () => history.save());
|
||||
|
||||
// Shape drawing
|
||||
setupShapeDrawing(canvas, state);
|
||||
setupMosaicDrawing(canvas, state, img, history);
|
||||
wireToolbar(toolbar, canvas, state, history);
|
||||
|
||||
// Actions buttons
|
||||
const actions = document.createElement("div");
|
||||
actions.className = "pb-actions";
|
||||
actions.innerHTML = `<button class="pb-cancel">取消</button><button class="pb-confirm">确认</button>`;
|
||||
overlay.appendChild(actions);
|
||||
|
||||
const close = () => { canvas.dispose(); overlay.remove(); };
|
||||
actions.querySelector(".pb-cancel").onclick = close;
|
||||
overlay.addEventListener("keydown", (e) => { if (e.key === "Escape") close(); });
|
||||
overlay.tabIndex = 0;
|
||||
overlay.focus();
|
||||
|
||||
// Keyboard shortcuts
|
||||
overlay.addEventListener("keydown", (e) => {
|
||||
if (e.ctrlKey && e.key === "z" && !e.shiftKey) { e.preventDefault(); history.undo(); }
|
||||
if (e.ctrlKey && (e.key === "Z" || (e.key === "z" && e.shiftKey))) { e.preventDefault(); history.redo(); }
|
||||
if (e.key === "Delete" || e.key === "Backspace") {
|
||||
const active = canvas.getActiveObject();
|
||||
if (active) { canvas.remove(active); canvas.discardActiveObject(); history.save(); }
|
||||
}
|
||||
});
|
||||
|
||||
// Confirm - save painted image and store canvas state for re-editing
|
||||
actions.querySelector(".pb-confirm").onclick = async () => {
|
||||
// Save canvas objects (without background) for future re-editing
|
||||
node.properties.paintBrushCanvas = canvas.toJSON(PB_EXPORT_PROPS);
|
||||
node.graph?.change?.();
|
||||
await savePaintedImage(canvas, node, img.width, img.height);
|
||||
close();
|
||||
};
|
||||
}
|
||||
|
||||
// --- Save painted image ---
|
||||
async function savePaintedImage(canvas, node, origW, origH) {
|
||||
// Export at original resolution
|
||||
const exportCanvas = document.createElement("canvas");
|
||||
exportCanvas.width = origW;
|
||||
exportCanvas.height = origH;
|
||||
const ctx = exportCanvas.getContext("2d");
|
||||
|
||||
const dataUrl = canvas.toDataURL({ format: "png", multiplier: origW / canvas.width });
|
||||
const exportImg = await new Promise((resolve) => {
|
||||
const i = new Image();
|
||||
i.onload = () => resolve(i);
|
||||
i.src = dataUrl;
|
||||
});
|
||||
ctx.drawImage(exportImg, 0, 0, origW, origH);
|
||||
|
||||
// Get original filename for naming (use stored original, not current painted name)
|
||||
const origName = (node.properties.paintBrushOriginal || "").split("/").pop() || "image.png";
|
||||
const paintedName = "painted_" + origName;
|
||||
|
||||
const blob = await new Promise(r => exportCanvas.toBlob(r, "image/png"));
|
||||
const formData = new FormData();
|
||||
formData.append("image", blob, paintedName);
|
||||
formData.append("type", "input");
|
||||
formData.append("overwrite", "true");
|
||||
|
||||
const resp = await api.fetchApi("/upload/image", { method: "POST", body: formData });
|
||||
const data = await resp.json();
|
||||
|
||||
// Add to combo options so the value persists across refresh
|
||||
const widget = node.widgets.find(w => w.name === "image");
|
||||
if (widget) {
|
||||
if (Array.isArray(widget.options?.values) && !widget.options.values.includes(data.name)) {
|
||||
widget.options.values.push(data.name);
|
||||
}
|
||||
widget.value = data.name;
|
||||
if (widget.callback) {
|
||||
widget.callback(data.name);
|
||||
}
|
||||
}
|
||||
// Mark graph as changed to trigger workflow auto-save
|
||||
node.graph?.change?.();
|
||||
app.graph.setDirtyCanvas(true, true);
|
||||
}
|
||||
|
||||
// --- Build Toolbar ---
|
||||
function buildToolbar(state) {
|
||||
const toolbar = document.createElement("div");
|
||||
toolbar.className = "pb-toolbar";
|
||||
toolbar.innerHTML = `
|
||||
<button data-tool="brush" class="active">画笔</button>
|
||||
<button data-tool="rect">矩形</button>
|
||||
<button data-tool="circle">圆形</button>
|
||||
<button data-tool="line">直线</button>
|
||||
<button data-tool="mosaic">马赛克</button>
|
||||
<button data-tool="eraser">橡皮擦</button>
|
||||
<span class="pb-sep"></span>
|
||||
<span class="pb-label">颜色</span>
|
||||
<input type="color" class="pb-color" value="${state.color}">
|
||||
<span class="pb-label">线宽</span>
|
||||
<input type="range" class="pb-width" min="1" max="40" value="${state.width}">
|
||||
<span class="pb-label">块大小</span>
|
||||
<input type="range" class="pb-mosaic-size" min="4" max="80" value="${state.mosaicSize}">
|
||||
<span class="pb-sep"></span>
|
||||
<button data-action="undo">撤回</button>
|
||||
<button data-action="clear">清空</button>
|
||||
`;
|
||||
return toolbar;
|
||||
}
|
||||
|
||||
// --- Wire Toolbar ---
|
||||
function wireToolbar(toolbar, canvas, state, history) {
|
||||
// Tool buttons
|
||||
toolbar.querySelectorAll("[data-tool]").forEach(btn => {
|
||||
btn.onclick = () => {
|
||||
toolbar.querySelectorAll("[data-tool]").forEach(b => b.classList.remove("active"));
|
||||
btn.classList.add("active");
|
||||
state.tool = btn.dataset.tool;
|
||||
|
||||
if (state.tool === "brush") {
|
||||
canvas.isDrawingMode = true;
|
||||
canvas.selection = false;
|
||||
canvas.defaultCursor = "default";
|
||||
canvas.hoverCursor = "move";
|
||||
canvas.freeDrawingBrush.color = state.color;
|
||||
canvas.freeDrawingBrush.width = state.width;
|
||||
} else if (state.tool === "eraser") {
|
||||
// Eraser: click on object to delete it
|
||||
canvas.isDrawingMode = false;
|
||||
canvas.selection = true;
|
||||
canvas.defaultCursor = "crosshair";
|
||||
canvas.hoverCursor = "pointer";
|
||||
} else if (state.tool === "mosaic") {
|
||||
canvas.isDrawingMode = false;
|
||||
canvas.selection = false;
|
||||
canvas.discardActiveObject();
|
||||
applyMosaicCursor(canvas, state);
|
||||
} else {
|
||||
canvas.isDrawingMode = false;
|
||||
canvas.selection = false;
|
||||
canvas.defaultCursor = "default";
|
||||
canvas.hoverCursor = "move";
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
// Eraser: remove object on click
|
||||
canvas.on("mouse:down", (opt) => {
|
||||
if (state.tool !== "eraser") return;
|
||||
const target = canvas.findTarget(opt.e);
|
||||
if (target) {
|
||||
canvas.remove(target);
|
||||
canvas.discardActiveObject();
|
||||
canvas.renderAll();
|
||||
history.save();
|
||||
}
|
||||
});
|
||||
|
||||
// Color picker
|
||||
toolbar.querySelector(".pb-color").oninput = (e) => {
|
||||
state.color = e.target.value;
|
||||
if (canvas.isDrawingMode) {
|
||||
canvas.freeDrawingBrush.color = state.color;
|
||||
}
|
||||
};
|
||||
|
||||
// Width slider
|
||||
toolbar.querySelector(".pb-width").oninput = (e) => {
|
||||
state.width = parseInt(e.target.value);
|
||||
if (canvas.isDrawingMode) {
|
||||
canvas.freeDrawingBrush.width = state.width;
|
||||
}
|
||||
};
|
||||
|
||||
toolbar.querySelector(".pb-mosaic-size").oninput = (e) => {
|
||||
state.mosaicSize = parseInt(e.target.value);
|
||||
if (state.tool === "mosaic") {
|
||||
applyMosaicCursor(canvas, state);
|
||||
}
|
||||
};
|
||||
|
||||
// Action buttons
|
||||
toolbar.querySelector("[data-action=undo]").onclick = () => history.undo();
|
||||
toolbar.querySelector("[data-action=clear]").onclick = () => {
|
||||
canvas.getObjects().forEach(obj => canvas.remove(obj));
|
||||
canvas.renderAll();
|
||||
history.save();
|
||||
};
|
||||
}
|
||||
|
||||
// --- Extension Registration ---
|
||||
app.registerExtension({
|
||||
name: "o1key.paintBrush",
|
||||
|
||||
// Register command for toolbar button
|
||||
commands: [
|
||||
{
|
||||
id: "o1key.PaintBrush",
|
||||
icon: "pi pi-pencil",
|
||||
label: "画笔",
|
||||
tooltip: "画笔",
|
||||
function: () => {
|
||||
const selectedNodes = app.canvas.selected_nodes;
|
||||
const node = selectedNodes ? Object.values(selectedNodes)[0] : null;
|
||||
if (node) openPaintModal(node);
|
||||
}
|
||||
}
|
||||
],
|
||||
|
||||
// Add title attribute to our toolbar button for native tooltip
|
||||
// Hide "节点信息" button and reorder paint brush next to mask editor
|
||||
setup() {
|
||||
const observer = new MutationObserver(() => {
|
||||
const toolbox = document.querySelector('[class*="selection-toolbox"], [class*="SelectionToolbox"]');
|
||||
if (!toolbox) return;
|
||||
|
||||
// Hide "节点信息" button (matches by aria-label or title)
|
||||
toolbox.querySelectorAll("button").forEach(btn => {
|
||||
const label = btn.title || btn.getAttribute("aria-label") || btn.textContent || "";
|
||||
if (label.includes("节点信息") || label.includes("Node Info") || label.includes("Info")) {
|
||||
btn.style.display = "none";
|
||||
}
|
||||
});
|
||||
|
||||
// Find our paint brush button and move it next to mask editor
|
||||
const pencilIcon = toolbox.querySelector('[class*="pi-pencil"]');
|
||||
if (pencilIcon) {
|
||||
const paintBtn = pencilIcon.closest("button");
|
||||
if (paintBtn && !paintBtn.title) paintBtn.title = "画笔";
|
||||
|
||||
// Find mask editor button (has mask/pen-tool icon)
|
||||
const allBtns = Array.from(toolbox.querySelectorAll("button"));
|
||||
const maskBtn = allBtns.find(b => {
|
||||
const cls = b.innerHTML || "";
|
||||
return cls.includes("mask") || cls.includes("pen-tool") || cls.includes("Mask");
|
||||
}) || allBtns[0];
|
||||
|
||||
if (maskBtn && paintBtn && paintBtn.previousElementSibling !== maskBtn) {
|
||||
maskBtn.after(paintBtn);
|
||||
}
|
||||
}
|
||||
});
|
||||
observer.observe(document.body, { childList: true, subtree: true });
|
||||
},
|
||||
|
||||
// Show button in toolbar when LoadImage node is selected
|
||||
getSelectionToolboxCommands(item) {
|
||||
if (item?.comfyClass === "LoadImage" || item?.type === "LoadImage") {
|
||||
return ["o1key.PaintBrush"];
|
||||
}
|
||||
return [];
|
||||
},
|
||||
|
||||
beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData.name !== "LoadImage") return;
|
||||
|
||||
const origMenu = nodeType.prototype.getExtraMenuOptions;
|
||||
nodeType.prototype.getExtraMenuOptions = function (canvasRef, options) {
|
||||
origMenu?.call(this, canvasRef, options);
|
||||
options.unshift({
|
||||
content: "画笔 (Paint)",
|
||||
callback: () => openPaintModal(this)
|
||||
});
|
||||
};
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,29 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "o1key.renameConsole",
|
||||
async setup() {
|
||||
const rename = () => {
|
||||
document.querySelectorAll(".side-bar-button, [class*='sidebar'] button, .p-togglebutton, button").forEach(btn => {
|
||||
const label = btn.getAttribute("aria-label") || "";
|
||||
const text = btn.textContent || "";
|
||||
if (label === "控制台" || label === "Console" || text.trim() === "控制台" || text.trim() === "Console") {
|
||||
if (label === "控制台" || label === "Console") {
|
||||
btn.setAttribute("aria-label", "日志");
|
||||
}
|
||||
const span = btn.querySelector("span");
|
||||
if (span && (span.textContent.trim() === "控制台" || span.textContent.trim() === "Console")) {
|
||||
span.textContent = "日志";
|
||||
} else if (!span && (btn.textContent.trim() === "控制台" || btn.textContent.trim() === "Console")) {
|
||||
btn.textContent = "日志";
|
||||
}
|
||||
}
|
||||
});
|
||||
};
|
||||
|
||||
const observer = new MutationObserver(rename);
|
||||
observer.observe(document.body, { childList: true, subtree: true });
|
||||
setTimeout(rename, 1000);
|
||||
setTimeout(rename, 3000);
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,18 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "o1key.renameTab",
|
||||
async setup() {
|
||||
const rename = () => {
|
||||
if (document.title.includes("ComfyUI")) {
|
||||
document.title = document.title.replace("ComfyUI", "o1key");
|
||||
}
|
||||
};
|
||||
|
||||
rename();
|
||||
new MutationObserver(rename).observe(
|
||||
document.querySelector("title") || document.head,
|
||||
{ childList: true, subtree: true, characterData: true }
|
||||
);
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,175 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "o1key.restartButton",
|
||||
async setup() {
|
||||
function inject() {
|
||||
if (document.querySelector("#o1k-restart-btn") && document.querySelector("#o1k-update-btn")) return;
|
||||
|
||||
const allBtns = document.querySelectorAll("button, .p-togglebutton, .side-bar-button");
|
||||
let logBtn = null;
|
||||
for (const btn of allBtns) {
|
||||
const label = (btn.getAttribute("aria-label") || "") + (btn.textContent || "");
|
||||
if (label.includes("日志") || label.includes("Console") || label.includes("控制台") || label.includes("Logs")) {
|
||||
logBtn = btn;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!logBtn || !logBtn.parentNode) return;
|
||||
|
||||
function makeButton(id, label, title, icon) {
|
||||
const btn = logBtn.cloneNode(false);
|
||||
btn.id = id;
|
||||
btn.setAttribute("aria-label", label);
|
||||
btn.title = title;
|
||||
const logStyle = window.getComputedStyle(logBtn);
|
||||
btn.style.display = "flex";
|
||||
btn.style.flexDirection = "column";
|
||||
btn.style.alignItems = "center";
|
||||
btn.style.justifyContent = "center";
|
||||
btn.style.gap = logStyle.gap || "4px";
|
||||
const iconSpan = document.createElement("span");
|
||||
iconSpan.innerHTML = icon;
|
||||
const textSpan = document.createElement("span");
|
||||
textSpan.textContent = label;
|
||||
btn.append(iconSpan, textSpan);
|
||||
return btn;
|
||||
}
|
||||
|
||||
let restartBtn = document.querySelector("#o1k-restart-btn");
|
||||
if (!restartBtn) {
|
||||
restartBtn = makeButton("o1k-restart-btn", "重启", "重启 ComfyUI",
|
||||
`<svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 2v6h-6"/><path d="M3 12a9 9 0 0 1 15-6.7L21 8"/><path d="M3 22v-6h6"/><path d="M21 12a9 9 0 0 1-15 6.7L3 16"/></svg>`);
|
||||
restartBtn.addEventListener("click", async () => {
|
||||
if (!confirm("确定要重启 ComfyUI 吗?")) return;
|
||||
restartBtn.style.opacity = "0.5";
|
||||
restartBtn.style.pointerEvents = "none";
|
||||
await disableExperimentalAssetApi();
|
||||
try { await fetch("/o1key/restart", { method: "POST" }); } catch {}
|
||||
pollUntilReady();
|
||||
});
|
||||
logBtn.parentNode.insertBefore(restartBtn, logBtn);
|
||||
}
|
||||
|
||||
if (!document.querySelector("#o1k-update-btn")) {
|
||||
const updateBtn = makeButton("o1k-update-btn", "更新", "更新 comfyui_o1key 节点包",
|
||||
`<svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M12 3v12"/><path d="m7 10 5 5 5-5"/><path d="M4 18v3h16v-3"/></svg>`);
|
||||
let updating = false;
|
||||
updateBtn.addEventListener("click", async () => {
|
||||
if (updating) return;
|
||||
if (!confirm("从 origin/main 拉取 comfyui_o1key 最新版本?")) return;
|
||||
updating = true;
|
||||
updateBtn.disabled = true;
|
||||
updateBtn.style.opacity = "0.5";
|
||||
updateBtn.title = "正在更新...";
|
||||
try {
|
||||
const response = await fetch("/o1key/update", {
|
||||
method: "POST",
|
||||
headers: { "X-O1Key-Update": "1" },
|
||||
});
|
||||
const result = await response.json();
|
||||
if (!response.ok) throw new Error(result.error || "更新失败");
|
||||
if (!result.updated) {
|
||||
alert(`已是最新版本(${result.version})。`);
|
||||
} else {
|
||||
const dependencies = result.requirements_changed
|
||||
? "\n依赖列表已变化,请先在 ComfyUI 的 Python 环境中执行 pip install -r requirements.txt。"
|
||||
: "";
|
||||
alert(`更新完成(${result.version})。${dependencies}\n请点击“重启”使新版本生效。`);
|
||||
}
|
||||
} catch (error) {
|
||||
alert(`更新失败:${error.message}`);
|
||||
} finally {
|
||||
updating = false;
|
||||
updateBtn.disabled = false;
|
||||
updateBtn.style.opacity = "";
|
||||
updateBtn.title = "更新 comfyui_o1key 节点包";
|
||||
}
|
||||
});
|
||||
restartBtn.after(updateBtn);
|
||||
}
|
||||
}
|
||||
|
||||
async function disableExperimentalAssetApi() {
|
||||
if (!(await shouldDisableExperimentalAssetApi())) return;
|
||||
try {
|
||||
await fetch("/api/settings/Comfy.Assets.UseAssetAPI", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(false),
|
||||
signal: AbortSignal.timeout(2000),
|
||||
});
|
||||
} catch {}
|
||||
}
|
||||
|
||||
async function shouldDisableExperimentalAssetApi() {
|
||||
try {
|
||||
const r = await fetch("/api/settings/Comfy.Assets.UseAssetAPI", {
|
||||
cache: "no-store",
|
||||
signal: AbortSignal.timeout(2000),
|
||||
});
|
||||
if (!r.ok || !(await r.json())) return false;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
return !(await fetchOk("/api/assets/seed/status", 2000));
|
||||
}
|
||||
|
||||
async function fetchOk(url, timeout = 2500) {
|
||||
try {
|
||||
const r = await fetch(url, {
|
||||
cache: "no-store",
|
||||
signal: AbortSignal.timeout(timeout),
|
||||
});
|
||||
return r.ok;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
async function comfyReady() {
|
||||
const [statsOk, modelFoldersOk] = await Promise.all([
|
||||
fetchOk("/api/system_stats"),
|
||||
fetchOk("/api/experiment/models"),
|
||||
]);
|
||||
return statsOk && modelFoldersOk;
|
||||
}
|
||||
|
||||
function pollUntilReady() {
|
||||
let attempts = 0;
|
||||
const maxAttempts = 80;
|
||||
const minRestartWaitMs = 5000;
|
||||
const startedAt = Date.now();
|
||||
let sawUnavailable = false;
|
||||
const interval = setInterval(async () => {
|
||||
attempts++;
|
||||
if (attempts > maxAttempts) { clearInterval(interval); forceReload(); return; }
|
||||
const ready = await comfyReady();
|
||||
if (!ready) {
|
||||
sawUnavailable = true;
|
||||
return;
|
||||
}
|
||||
if (!sawUnavailable && Date.now() - startedAt < minRestartWaitMs) return;
|
||||
|
||||
clearInterval(interval);
|
||||
await disableExperimentalAssetApi();
|
||||
setTimeout(forceReload, 800);
|
||||
}, 1500);
|
||||
}
|
||||
|
||||
function forceReload() {
|
||||
window.onbeforeunload = null;
|
||||
Object.defineProperty(BeforeUnloadEvent.prototype, "returnValue", {
|
||||
get() { return ""; },
|
||||
set() {}
|
||||
});
|
||||
location.reload();
|
||||
}
|
||||
|
||||
const observer = new MutationObserver(inject);
|
||||
observer.observe(document.body, { childList: true, subtree: true });
|
||||
setTimeout(inject, 2000);
|
||||
setTimeout(inject, 4000);
|
||||
setTimeout(inject, 8000);
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,210 @@
|
||||
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);
|
||||
}
|
||||
};
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,90 @@
|
||||
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");
|
||||
});
|
||||
@@ -0,0 +1,146 @@
|
||||
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();
|
||||
}
|
||||
};
|
||||
},
|
||||
});
|
||||
@@ -1,134 +0,0 @@
|
||||
# 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 并附上错误信息
|
||||
|
||||
---
|
||||
|
||||
**享受自动更新带来的便利!** 🎉
|
||||
Reference in New Issue
Block a user