Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5fccb3e8eb | ||
|
|
30e9603f77 | ||
|
|
5d9aff9ca7 | ||
|
|
3f0f4099fb | ||
|
|
c974df1b5e | ||
|
|
815c2c598e | ||
|
|
228c4c5141 | ||
|
|
8eae3da298 | ||
|
|
69279c654d | ||
|
|
c491731c99 |
@@ -132,45 +132,28 @@ O1KEY_API_KEY=你的API密钥
|
||||
|
||||
## 🔄 更新插件
|
||||
|
||||
自动更新脚本**已改为从国内镜像(Gitee)拉取**,国内用户无需科学上网即可更新。
|
||||
### 界面更新
|
||||
|
||||
### 方法一:自动更新(推荐)⭐
|
||||
在 ComfyUI 左侧功能栏点击「更新」(位于「重启」下方)。按钮会从当前 Git 仓库的 `origin/main` 拉取最新版本。完成后点击「重启」使新版本生效。
|
||||
|
||||
**Windows 用户:**
|
||||
1. 进入插件目录:`ComfyUI\custom_nodes\comfyui_o1key`
|
||||
2. 双击运行 `自动更新插件(win).bat`
|
||||
3. 等待更新完成
|
||||
4. 重启 ComfyUI
|
||||
界面更新需要通过 Git 安装、处于 `main` 分支,且节点包文件没有本地修改。更新仅允许快进,不会覆盖本地修改或删除配置。ZIP 安装、分支分叉或网络连接失败时,界面会显示原因,需要手动处理。
|
||||
|
||||
如果提示依赖列表已变化,请在 ComfyUI 使用的 Python 环境中执行:
|
||||
|
||||
**Linux/Mac 用户:**
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes/comfyui_o1key
|
||||
chmod +x "自动更新插件(mac).sh" # 首次运行需要添加执行权限
|
||||
./"自动更新插件(mac).sh"
|
||||
python -m pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 方法二:手动更新
|
||||
### 手动更新
|
||||
|
||||
从 Gitee 镜像拉取(国内推荐):
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes/comfyui_o1key
|
||||
git remote get-url gitee &>/dev/null || git remote add gitee https://gitee.com/resonLzy/comfyui_o1key.git
|
||||
git pull gitee main
|
||||
pip install -r requirements.txt --upgrade
|
||||
git pull --ff-only origin main
|
||||
python -m pip install -r requirements.txt
|
||||
```
|
||||
|
||||
从 GitHub 拉取:
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes/comfyui_o1key
|
||||
git pull origin main
|
||||
pip install -r requirements.txt --upgrade
|
||||
```
|
||||
|
||||
**💡 提示:**
|
||||
- 自动更新脚本会自动备份和恢复你的 `.config` 配置文件
|
||||
- 更新会保留环境变量中配置的 API 密钥
|
||||
- 更新检查在每次启动 ComfyUI 时自动进行(不会影响性能)
|
||||
- 如果发现新版本,终端会显示更新提示
|
||||
更新保留环境变量中配置的 API 密钥。启动时仍会检查是否有新版本。
|
||||
|
||||
---
|
||||
|
||||
|
||||
+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()
|
||||
"""
|
||||
+627
-8
@@ -11,9 +11,72 @@ Comfyui_o1key - ComfyUI 自定义节点集合
|
||||
|
||||
|
||||
import ssl
|
||||
import logging
|
||||
import asyncio
|
||||
|
||||
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, SaveCleanImage, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, MultiResPreview, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, KVideoFirstLast, KVideoImage2Video
|
||||
from .nodes import K3Video, K3VideoFirstLast, K3MotionControl, K3MotionVideoCheck, NanoBananaV2, NanoBananaV2Batch
|
||||
# 屏蔽 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:
|
||||
_install_asyncio_win10054_filter(asyncio.get_event_loop())
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
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 请求超时,请稍后重试或检查网络。"
|
||||
@@ -52,28 +115,32 @@ _wrap_generate_for_error_display(NanoBananaV2Batch)
|
||||
|
||||
# ComfyUI 节点注册
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NanoBananaPro": NanoBananaPro,
|
||||
"NanoBanana": NanoBananaPro,
|
||||
"BatchNanoBananaPro": BatchNanoBananaPro,
|
||||
"GoogleGemini": GoogleGemini,
|
||||
"LoadFile": LoadFile,
|
||||
"ImageStitchPro": ImageStitchPro,
|
||||
"SaveCleanImage": SaveCleanImage,
|
||||
|
||||
"BatchCleanMetadata": BatchCleanMetadata,
|
||||
"VideoPreview": VideoPreview,
|
||||
"GoogleVeo": GoogleVeo,
|
||||
"Google31Video": Google31Video,
|
||||
"FluxImageEdit": FluxImageEdit,
|
||||
"UniversalLLMChat": UniversalLLMChat,
|
||||
"KlingVideo": KlingVideo,
|
||||
"KlingFirstLastFrame": KlingFirstLastFrame,
|
||||
"KlingMotionControlTest": KlingMotionControlTest,
|
||||
"AspectRatioPreset": AspectRatioPreset,
|
||||
"MultiResPreview": MultiResPreview,
|
||||
|
||||
"BatchImagesO1key": BatchImagesO1key,
|
||||
"Seedance": Seedance,
|
||||
"SeedanceMultiModal": SeedanceMultiModal,
|
||||
"StreamPreview": StreamPreview,
|
||||
"DoubaoImage": DoubaoImage,
|
||||
"O1keyGPTImage": O1keyGPTImage,
|
||||
"O1keyGPTImageBatch": O1keyGPTImageBatch,
|
||||
"O1keyGrokImage": O1keyGrokImage,
|
||||
"O1keyGrokVideo": O1keyGrokVideo,
|
||||
"KVideoFirstLast": KVideoFirstLast,
|
||||
"KVideoImage2Video": KVideoImage2Video,
|
||||
"K3Video": K3Video,
|
||||
@@ -82,31 +149,40 @@ NODE_CLASS_MAPPINGS = {
|
||||
"K3MotionVideoCheck": K3MotionVideoCheck,
|
||||
"NanoBananaV2": NanoBananaV2,
|
||||
"NanoBananaV2Batch": NanoBananaV2Batch,
|
||||
"SaveImageFormat": SaveImageFormat,
|
||||
"O1keySavePSD": O1keySavePSD,
|
||||
"O1keyRemoveBackground": O1keyRemoveBackground,
|
||||
"O1keyColorRemoveBG": O1keyColorRemoveBG,
|
||||
"O1keyGridSplitter": O1keyGridSplitter,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NanoBananaPro": "Nano Banana",
|
||||
"NanoBanana": "Nano Banana",
|
||||
"BatchNanoBananaPro": "批量 Nano Banana",
|
||||
"GoogleGemini": "Google Gemini",
|
||||
"LoadFile": "加载文件",
|
||||
"ImageStitchPro": "图像拼接 Pro",
|
||||
"SaveCleanImage": "保存图像(防AI识别)",
|
||||
|
||||
"BatchCleanMetadata": "批量任务(防AI识别)",
|
||||
"VideoPreview": "预览视频",
|
||||
"GoogleVeo": "Google Veo - ab",
|
||||
"Google31Video": "Google 3.1 Video",
|
||||
"FluxImageEdit": "Flux2 图像编辑",
|
||||
"UniversalLLMChat": "全能LLM对话助手",
|
||||
"KlingVideo": "文/图生视频 自研模型",
|
||||
"KlingFirstLastFrame": "首尾帧生视频 自研模型",
|
||||
"KlingMotionControlTest": "动作控制 自研模型",
|
||||
"AspectRatioPreset": "图片宽高比预设",
|
||||
"MultiResPreview": "预览图像(v2)",
|
||||
|
||||
"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 图生视频 自研",
|
||||
@@ -115,6 +191,11 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"K3MotionVideoCheck": "视频时长检测 K3",
|
||||
"NanoBananaV2": "Nano Banana V2",
|
||||
"NanoBananaV2Batch": "Nano Banana V2(批量)",
|
||||
"SaveImageFormat": "保存图像(格式转换)",
|
||||
"O1keySavePSD": "保存 PSD(分层)",
|
||||
"O1keyRemoveBackground": "去背景(rembg)",
|
||||
"O1keyColorRemoveBG": "颜色去背景",
|
||||
"O1keyGridSplitter": "合并图智能切割",
|
||||
}
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
@@ -126,11 +207,549 @@ 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
|
||||
|
||||
+3
-1
@@ -9,6 +9,8 @@ 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', 'SoraClient', 'KlingClient', 'VeoClient', 'OpenAIAPIClient']
|
||||
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'NewAPIVeoClient', 'GrokVideoClient', 'OpenAIAPIClient']
|
||||
|
||||
+39
-46
@@ -14,6 +14,8 @@ 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):
|
||||
@@ -38,7 +40,8 @@ class BaseAPIClient(ABC):
|
||||
Args:
|
||||
base_url: API 基础 URL
|
||||
api_key: API 密钥
|
||||
max_request_size: 最大请求体大小(字节),默认 100MB
|
||||
max_request_size: 兼容参数;基类不再用它限制 JSON 请求体,
|
||||
部分子类仍用它作为上传文件大小限制
|
||||
"""
|
||||
self.base_url = base_url
|
||||
self.api_key = api_key
|
||||
@@ -114,24 +117,6 @@ class BaseAPIClient(ABC):
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
def check_request_size(self, request_body: Dict[str, Any]) -> None:
|
||||
"""
|
||||
检查请求体大小是否超过限制
|
||||
|
||||
Args:
|
||||
request_body: 请求体字典
|
||||
|
||||
Raises:
|
||||
ValueError: 如果请求体超过限制
|
||||
"""
|
||||
request_json = json.dumps(request_body)
|
||||
request_size = len(request_json.encode('utf-8'))
|
||||
|
||||
if request_size > self.max_request_size:
|
||||
raise ValueError(
|
||||
"请求体积超过100MB限制,请调整分辨率或减少图片数量"
|
||||
)
|
||||
|
||||
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
|
||||
"""
|
||||
子类可重写:为指定 HTTP 状态码返回自定义错误文案。
|
||||
@@ -183,9 +168,6 @@ class BaseAPIClient(ABC):
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token)
|
||||
|
||||
# 检查请求大小
|
||||
self.check_request_size(request_body)
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = self._make_session()
|
||||
@@ -206,7 +188,8 @@ class BaseAPIClient(ABC):
|
||||
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
raise RuntimeError(error_text)
|
||||
# 返回状态码和错误文本,由外层处理重试
|
||||
return {"_error": True, "_status": response.status, "_text": error_text}
|
||||
|
||||
wait_start = time.time()
|
||||
response_data = await response.json()
|
||||
@@ -231,6 +214,10 @@ class BaseAPIClient(ABC):
|
||||
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())
|
||||
@@ -240,7 +227,6 @@ class BaseAPIClient(ABC):
|
||||
return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
|
||||
# 取消未完成的任务
|
||||
for t in pending:
|
||||
t.cancel()
|
||||
try:
|
||||
@@ -248,14 +234,35 @@ class BaseAPIClient(ABC):
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
# 判断是哪个先完成
|
||||
if interrupt_task in done and request_task not in done:
|
||||
raise InterruptProcessingException()
|
||||
|
||||
# 请求完成,取出结果(可能含异常)
|
||||
return request_task.result()
|
||||
result = request_task.result()
|
||||
else:
|
||||
return await _do_request()
|
||||
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
|
||||
@@ -352,30 +359,16 @@ class BaseAPIClient(ABC):
|
||||
custom = self.get_http_error_message(429, error_message)
|
||||
if custom is not None:
|
||||
raise RuntimeError(custom)
|
||||
raise RuntimeError(
|
||||
f"请求频率超限 (429 Too Many Requests)\n"
|
||||
f"API 返回错误:{error_message}\n"
|
||||
f"建议:等待一段时间后重试"
|
||||
)
|
||||
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(
|
||||
f"服务暂时不可用 (503 Service Unavailable)\n"
|
||||
f"API 返回错误:{error_message}\n"
|
||||
f"建议:稍后重试"
|
||||
)
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[503])
|
||||
elif response.status == 504:
|
||||
raise RuntimeError(
|
||||
f"API 请求超时 (504 Gateway Timeout)\n"
|
||||
f"API 返回错误:{error_message}\n"
|
||||
f"建议:稍后重试"
|
||||
)
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[504])
|
||||
elif response.status == 502:
|
||||
raise RuntimeError(
|
||||
"糟糕!请求到上游时遇到超时或过载!别担心,过会儿再次点击运行即可!"
|
||||
)
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[502])
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"API 请求失败 (状态码: {response.status})\n"
|
||||
|
||||
@@ -21,6 +21,7 @@ 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
|
||||
|
||||
|
||||
# ── 固定端点 ──────────────────────────────────────────────────────────────────
|
||||
@@ -245,7 +246,9 @@ class DoubaoImageClient:
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
||||
|
||||
# 3. 发送 POST 请求
|
||||
# 3. 发送 POST 请求(带退避重试)
|
||||
last_status = None
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
t0 = time.time()
|
||||
async with session.post(
|
||||
url,
|
||||
@@ -256,7 +259,15 @@ class DoubaoImageClient:
|
||||
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", {})
|
||||
@@ -279,6 +290,12 @@ class DoubaoImageClient:
|
||||
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 复用)
|
||||
|
||||
@@ -14,6 +14,7 @@ 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
|
||||
|
||||
|
||||
# 显示名 → 实际请求值的映射
|
||||
@@ -112,6 +113,8 @@ class FluxEditClient:
|
||||
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]}"
|
||||
|
||||
@@ -44,7 +44,7 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||||
|
||||
@property
|
||||
def api_base_url(self) -> str:
|
||||
return get_async_api_base_url()
|
||||
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(
|
||||
@@ -72,7 +72,9 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||||
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),
|
||||
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:
|
||||
@@ -85,6 +87,23 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||||
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:
|
||||
@@ -124,19 +143,44 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||||
|
||||
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"):
|
||||
val = response.get(field)
|
||||
if val is not None and isinstance(val, (int, float)):
|
||||
return val / 100.0 if val > 1 else float(val)
|
||||
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"):
|
||||
val = progress_info.get(field)
|
||||
if val is not None and isinstance(val, (int, float)):
|
||||
return val / 100.0 if val > 1 else float(val)
|
||||
for field in ("progress", "percentage", "percent"):
|
||||
progress = _coerce(progress_info.get(field))
|
||||
if progress is not None:
|
||||
return progress
|
||||
|
||||
return None
|
||||
|
||||
|
||||
+119
-81
@@ -34,8 +34,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
|
||||
super().__init__(
|
||||
base_url=get_api_base_url(),
|
||||
api_key=api_key,
|
||||
max_request_size=100 * 1024 * 1024
|
||||
api_key=api_key
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@@ -187,6 +186,8 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
image_compression: str = None,
|
||||
thinking_level: str = None,
|
||||
request_log_enabled: bool = True,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
@@ -203,21 +204,90 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
Returns:
|
||||
请求体字典
|
||||
"""
|
||||
import json
|
||||
|
||||
_MAX_BODY_BYTES = 20 * 1024 * 1024 # 20 MB
|
||||
_BODY_TARGET_BYTES = int(_MAX_BODY_BYTES * 0.9)
|
||||
|
||||
parts = []
|
||||
|
||||
# 添加文本部分
|
||||
parts.append({"text": prompt})
|
||||
|
||||
image_config = {"imageSize": resolution}
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
image_config["aspectRatio"] = aspect_ratio
|
||||
|
||||
def _build_request_body(body_parts: List[dict]) -> Dict[str, Any]:
|
||||
body = {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": body_parts
|
||||
}
|
||||
],
|
||||
"generationConfig": {
|
||||
"responseModalities": ["IMAGE"],
|
||||
"imageConfig": image_config
|
||||
}
|
||||
}
|
||||
|
||||
# 添加思考深度配置
|
||||
if thinking_level:
|
||||
body["generationConfig"]["thinkingConfig"] = {
|
||||
"thinkingLevel": thinking_level,
|
||||
"includeThoughts": True
|
||||
}
|
||||
|
||||
# 添加图片压缩参数
|
||||
if image_compression:
|
||||
body["image_compression"] = image_compression
|
||||
|
||||
# 添加 Google Search Grounding(如果启用)
|
||||
# 新异步接口要求直接放在请求体顶层:{"google_search": true}
|
||||
if enable_grounding or enable_image_search:
|
||||
body["google_search"] = True
|
||||
|
||||
return body
|
||||
|
||||
def _request_size(body: Dict[str, Any]) -> int:
|
||||
return len(json.dumps(body).encode("utf-8"))
|
||||
|
||||
def _format_size(size: int) -> str:
|
||||
if size < 1024 * 1024:
|
||||
return f"{size / 1024:.2f}KB"
|
||||
return f"{size / 1024 / 1024:.2f}MB"
|
||||
|
||||
def _shorten_base64_for_log(obj, max_len: int = 200):
|
||||
if isinstance(obj, dict):
|
||||
result = {}
|
||||
for key, value in obj.items():
|
||||
if key == "data" and isinstance(value, str) and len(value) > max_len:
|
||||
result[key] = f"<base64 data, {len(value)} chars>"
|
||||
else:
|
||||
result[key] = _shorten_base64_for_log(value, max_len)
|
||||
return result
|
||||
if isinstance(obj, list):
|
||||
return [_shorten_base64_for_log(item, max_len) for item in obj]
|
||||
return obj
|
||||
|
||||
def _log_original_request_body(body: Dict[str, Any]) -> None:
|
||||
body_size = _request_size(body)
|
||||
print(
|
||||
f"\n{'=' * 60}\n"
|
||||
f"[原始请求体日志] 请求体积: {_format_size(body_size)} "
|
||||
f"(inline_data.data 已折叠显示 base64 长度)\n"
|
||||
f"{json.dumps(_shorten_base64_for_log(body), ensure_ascii=False, indent=2)}\n"
|
||||
f"{'=' * 60}\n"
|
||||
)
|
||||
|
||||
original_request_logged = False
|
||||
|
||||
# 添加图像部分(如果有)
|
||||
if images:
|
||||
working_images = list(images)
|
||||
|
||||
# 编码一次,估算大小,超限则迭代缩放
|
||||
for _attempt in range(10):
|
||||
def _build_image_parts(src_images: List[Image.Image]) -> List[dict]:
|
||||
img_parts = []
|
||||
for img in working_images:
|
||||
for img in src_images:
|
||||
img_base64 = encode_image_to_base64(img)
|
||||
img_parts.append({
|
||||
"inline_data": {
|
||||
@@ -225,86 +295,49 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
"data": img_base64
|
||||
}
|
||||
})
|
||||
return img_parts
|
||||
|
||||
# 估算完整 body 大小(不含工具字段,工具字段很小可忽略)
|
||||
estimated = self._estimate_body_size(
|
||||
parts + img_parts,
|
||||
{
|
||||
"generationConfig": {
|
||||
"responseModalities": ["IMAGE"],
|
||||
"imageConfig": {
|
||||
"aspectRatio": aspect_ratio,
|
||||
"imageSize": resolution
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
def _estimate_with_images(img_parts: List[dict]) -> int:
|
||||
return _request_size(_build_request_body(parts + img_parts))
|
||||
|
||||
if estimated <= _MAX_BODY_BYTES:
|
||||
parts.extend(img_parts)
|
||||
if _attempt > 0:
|
||||
orig_sizes = ", ".join(
|
||||
f"{img.width}×{img.height}" for img in images
|
||||
)
|
||||
new_sizes = ", ".join(
|
||||
f"{img.width}×{img.height}" for img in working_images
|
||||
)
|
||||
size_mb = estimated / (1024 * 1024)
|
||||
print(
|
||||
f"Nano Banana Pro: 输入图片已自动缩放以控制请求体积\n"
|
||||
f" 原始尺寸: {orig_sizes}\n"
|
||||
f" 缩放后: {new_sizes}\n"
|
||||
f" 请求体积: {size_mb:.2f}MB(限制 20MB)"
|
||||
)
|
||||
break
|
||||
else:
|
||||
# 按像素面积比推算需要的线性缩放系数,留 5% 余量
|
||||
ratio = (_MAX_BODY_BYTES * 0.95) / estimated
|
||||
working_images = list(images)
|
||||
img_parts = _build_image_parts(working_images)
|
||||
original_request_body = _build_request_body(parts + img_parts)
|
||||
if request_log_enabled:
|
||||
_log_original_request_body(original_request_body)
|
||||
original_request_logged = True
|
||||
estimated = _estimate_with_images(img_parts)
|
||||
|
||||
if estimated > _MAX_BODY_BYTES:
|
||||
ratio = _BODY_TARGET_BYTES / estimated
|
||||
scale = ratio ** 0.5 # 面积比 → 线性比
|
||||
working_images = self._scale_images_to_fit(working_images, scale)
|
||||
else:
|
||||
# 10 轮后仍超限,使用最后一次结果(极端情况兜底)
|
||||
img_parts = _build_image_parts(working_images)
|
||||
estimated = _estimate_with_images(img_parts)
|
||||
|
||||
orig_sizes = ", ".join(f"{img.width}×{img.height}" for img in images)
|
||||
new_sizes = ", ".join(f"{img.width}×{img.height}" for img in working_images)
|
||||
size_mb = estimated / (1024 * 1024)
|
||||
target_mb = _BODY_TARGET_BYTES / (1024 * 1024)
|
||||
print(
|
||||
f"Nano Banana Pro: 输入图片已按请求体目标大小自动缩放\n"
|
||||
f" 原始尺寸: {orig_sizes}\n"
|
||||
f" 缩放后: {new_sizes}\n"
|
||||
f" 请求体积: {size_mb:.2f}MB(目标 {target_mb:.2f}MB,限制 20MB)"
|
||||
)
|
||||
|
||||
parts.extend(img_parts)
|
||||
|
||||
# 构建请求体
|
||||
request_body = {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": parts
|
||||
}
|
||||
],
|
||||
"generationConfig": {
|
||||
"responseModalities": ["IMAGE"],
|
||||
"imageConfig": {
|
||||
"aspectRatio": aspect_ratio,
|
||||
"imageSize": resolution
|
||||
}
|
||||
}
|
||||
}
|
||||
request_body = _build_request_body(parts)
|
||||
if request_log_enabled and not original_request_logged:
|
||||
_log_original_request_body(request_body)
|
||||
|
||||
# 添加图片压缩参数
|
||||
if image_compression:
|
||||
request_body["image_compression"] = image_compression
|
||||
|
||||
# 添加 Google Search Grounding 工具(如果启用)
|
||||
# 注意:enable_image_search=True 时会自动隐含 enable_grounding
|
||||
if enable_grounding or enable_image_search:
|
||||
if enable_image_search:
|
||||
# 同时启用网页搜索和图片搜索(仅 nano-banana-2 / gemini-3.1-flash-image-preview 支持)
|
||||
request_body["tools"] = [
|
||||
{
|
||||
"google_search": {
|
||||
"searchTypes": {
|
||||
"webSearch": {},
|
||||
"imageSearch": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
else:
|
||||
# 仅启用网页搜索(通用)
|
||||
request_body["tools"] = [{"google_search": {}}]
|
||||
request_size = _request_size(request_body)
|
||||
if request_size > _MAX_BODY_BYTES:
|
||||
raise ValueError(
|
||||
f"请求体超过 20MB 限制(当前 {request_size / 1024 / 1024:.2f}MB),"
|
||||
"已停止提交;请减少参考图数量、降低图片复杂度或缩短提示词"
|
||||
)
|
||||
|
||||
return request_body
|
||||
|
||||
@@ -564,6 +597,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
image_format: str = "base64",
|
||||
thinking_level: str = None,
|
||||
) -> tuple[List[Image.Image], Dict[str, Any]]:
|
||||
"""
|
||||
单次异步生成请求(极简单行日志)
|
||||
@@ -602,6 +636,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
resolution=resolution,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
build_time = time.time() - build_start
|
||||
|
||||
@@ -737,6 +772,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
image_format: str = "base64",
|
||||
thinking_level: str = None,
|
||||
) -> List[Image.Image]:
|
||||
"""
|
||||
批量全并发生成 - 改进版:支持分批处理和内存管理
|
||||
@@ -801,6 +837,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
image_format=image_format,
|
||||
thinking_level=thinking_level,
|
||||
),
|
||||
name=f"task_{task_index}"
|
||||
)
|
||||
@@ -873,6 +910,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
image_format: str = "base64",
|
||||
thinking_level: str = None,
|
||||
) -> List[Image.Image]:
|
||||
"""
|
||||
同步生成接口(用于 ComfyUI)
|
||||
@@ -906,6 +944,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
image_format=image_format,
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
|
||||
return self.run_async_in_thread(coro)
|
||||
@@ -1091,4 +1130,3 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
)
|
||||
|
||||
return self.run_async_in_thread(coro)
|
||||
|
||||
+970
-58
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())
|
||||
+40
-45
@@ -10,6 +10,17 @@ 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:
|
||||
@@ -49,10 +60,12 @@ class KlingClient:
|
||||
) -> Dict[str, Any]:
|
||||
url = f"{self.base_url}{self.ENDPOINTS[endpoint_type]}"
|
||||
|
||||
async with session.post(url, json=body, headers=self._headers()) as resp:
|
||||
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()
|
||||
if resp.status != 200:
|
||||
raise RuntimeError(f"提交失败 ({resp.status}): {text}")
|
||||
return json.loads(text)
|
||||
|
||||
# ── 轮询状态 ──────────────────────────────────────────────────────
|
||||
@@ -68,6 +81,7 @@ class KlingClient:
|
||||
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:
|
||||
@@ -76,36 +90,22 @@ class KlingClient:
|
||||
|
||||
data = result.get("data", {})
|
||||
inner_data = data.get("data", {}) if isinstance(data, dict) else {}
|
||||
status = (
|
||||
data.get("status") or
|
||||
inner_data.get("task_status") or
|
||||
result.get("status") or
|
||||
""
|
||||
)
|
||||
status = status.lower() if status else ""
|
||||
status = extract_status(result)
|
||||
|
||||
progress_str = data.get("progress", "0%")
|
||||
try:
|
||||
progress_pct = int(str(progress_str).replace("%", "").strip())
|
||||
except (ValueError, AttributeError):
|
||||
progress_pct = 0
|
||||
progress_pct = extract_progress(result)
|
||||
|
||||
print(f"[视频生成] 生成中 {progress_pct}%")
|
||||
|
||||
if on_progress:
|
||||
on_progress(progress_pct)
|
||||
|
||||
if status in ("success", "completed", "done", "finished", "succeed"):
|
||||
if is_success_status(status):
|
||||
return result
|
||||
elif status in ("failed", "fail"):
|
||||
error_info = result.get("error", {})
|
||||
if isinstance(error_info, dict):
|
||||
error_msg = error_info.get("message", "未知错误")
|
||||
else:
|
||||
error_msg = str(error_info)
|
||||
elif is_failure_status(status, result):
|
||||
error_msg = extract_error_message(result)
|
||||
raise RuntimeError(f"生成失败:{error_msg}")
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
await interruptible_sleep(interval)
|
||||
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
|
||||
|
||||
# ── 下载视频 ──────────────────────────────────────────────────────
|
||||
@@ -117,12 +117,14 @@ class KlingClient:
|
||||
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
|
||||
|
||||
@@ -200,19 +202,15 @@ class KlingClient:
|
||||
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}"
|
||||
async with session.post(create_url, json=body, headers=headers) as resp:
|
||||
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()
|
||||
if resp.status != 200:
|
||||
# 尝试提取友好错误信息
|
||||
try:
|
||||
err = json.loads(text)
|
||||
msg = err.get("error", {}).get("message") or err.get("message") or text
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"动作控制提交失败 ({resp.status}): {msg}")
|
||||
create_resp = json.loads(text)
|
||||
|
||||
video_id = create_resp.get("id")
|
||||
@@ -224,6 +222,7 @@ class KlingClient:
|
||||
# 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:
|
||||
@@ -235,29 +234,24 @@ class KlingClient:
|
||||
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
|
||||
status_resp = json.loads(text)
|
||||
|
||||
status = status_resp.get("status", "").lower()
|
||||
progress_raw = status_resp.get("progress", 0)
|
||||
try:
|
||||
progress_pct = int(str(progress_raw).rstrip("%").strip())
|
||||
except (ValueError, AttributeError):
|
||||
progress_pct = 0
|
||||
status = extract_status(status_resp)
|
||||
progress_pct = extract_progress(status_resp)
|
||||
|
||||
print(f"[动作控制] 生成中 {progress_pct}%")
|
||||
if on_progress:
|
||||
on_progress(progress_pct)
|
||||
|
||||
if status == "completed":
|
||||
if is_success_status(status):
|
||||
break
|
||||
if status == "failed":
|
||||
error_info = status_resp.get("error", {})
|
||||
error_msg = (error_info.get("message", "未知错误")
|
||||
if isinstance(error_info, dict) else str(error_info))
|
||||
if is_failure_status(status, status_resp):
|
||||
error_msg = extract_error_message(status_resp)
|
||||
raise RuntimeError(f"动作控制生成失败:{error_msg}")
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
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)}"
|
||||
@@ -277,14 +271,15 @@ class KlingClient:
|
||||
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())
|
||||
@@ -204,13 +204,15 @@ class OpenAIAPIClient(BaseAPIClient):
|
||||
"extra_body": {
|
||||
"google": {
|
||||
"image_config": {
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"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]:
|
||||
|
||||
+26
-24
@@ -11,6 +11,17 @@ 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:
|
||||
@@ -47,17 +58,12 @@ class SeedanceClient:
|
||||
) -> str:
|
||||
"""提交视频生成任务,返回 task_id"""
|
||||
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
|
||||
async with session.post(url, json=body, headers=self._headers()) as resp:
|
||||
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()
|
||||
if resp.status != 200:
|
||||
try:
|
||||
err = json.loads(text)
|
||||
msg = (err.get("error", {}).get("message")
|
||||
or err.get("message")
|
||||
or text)
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"提交失败 ({resp.status}): {msg}")
|
||||
data = json.loads(text)
|
||||
|
||||
# new-api 返回字段:id / task_id
|
||||
@@ -79,6 +85,7 @@ class SeedanceClient:
|
||||
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:
|
||||
@@ -95,20 +102,16 @@ class SeedanceClient:
|
||||
# new-api 包装格式:真实数据在 result["data"] 里
|
||||
inner = result.get("data") or result
|
||||
|
||||
status = (inner.get("status") or "").lower()
|
||||
status = extract_status(result)
|
||||
|
||||
# 解析进度
|
||||
progress_raw = inner.get("progress", "0")
|
||||
try:
|
||||
progress_pct = int(str(progress_raw).rstrip("%").strip())
|
||||
except (ValueError, AttributeError):
|
||||
progress_pct = 0
|
||||
progress_pct = extract_progress(result)
|
||||
|
||||
print(f"[Seedance] 生成中 {progress_pct}%")
|
||||
if on_progress:
|
||||
on_progress(progress_pct)
|
||||
|
||||
if status in self.SUCCESS_STATUSES:
|
||||
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 {}
|
||||
@@ -129,15 +132,11 @@ class SeedanceClient:
|
||||
)
|
||||
return video_url, last_frame_url
|
||||
|
||||
if status in self.FAILURE_STATUSES:
|
||||
reason = (
|
||||
inner.get("fail_reason")
|
||||
or (inner.get("error") or {}).get("message")
|
||||
or "未知错误"
|
||||
)
|
||||
if is_failure_status(status, result):
|
||||
reason = extract_error_message(result)
|
||||
raise RuntimeError(f"视频生成失败:{reason}")
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
await interruptible_sleep(interval)
|
||||
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
|
||||
|
||||
# ── 3. 下载视频 ────────────────────────────────────────────────────
|
||||
@@ -150,12 +149,14 @@ class SeedanceClient:
|
||||
) -> 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
|
||||
|
||||
@@ -173,6 +174,7 @@ class SeedanceClient:
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
|
||||
# 提交
|
||||
check_interrupt()
|
||||
if on_stage:
|
||||
on_stage("submitting")
|
||||
task_id = await self.submit_async(body, session)
|
||||
|
||||
@@ -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>
|
||||
+3
-16
@@ -114,11 +114,11 @@ GEMINI_MODELS = [
|
||||
|
||||
GEMINI_FLASH_MODELS = [
|
||||
{
|
||||
"id": "gemini-3-flash-preview",
|
||||
"description": "Gemini 3 Flash,快速多模态文本生成,通过 thinkingConfig 控制思考等级",
|
||||
"id": "gemini-3.5-flash",
|
||||
"description": "Gemini 3.5 Flash,快速多模态文本生成,通过 thinkingConfig 控制思考等级",
|
||||
"enabled": True,
|
||||
"endpoint_type": "standard",
|
||||
"endpoint": "/v1beta/models/gemini-3-flash-preview:generateContent",
|
||||
"endpoint": "/v1beta/models/gemini-3.5-flash:generateContent",
|
||||
"thinking_config": {
|
||||
"低": "low",
|
||||
"中": "medium",
|
||||
@@ -137,19 +137,6 @@ GEMINI_FLASH_MODELS = [
|
||||
"中": "high"
|
||||
}
|
||||
},
|
||||
|
||||
{
|
||||
"id": "gemini-3.1-flash-lite-preview",
|
||||
"description": "Gemini 3.1 Flash Lite,轻量级多模态文本生成,通过 thinkingConfig 控制思考等级",
|
||||
"enabled": True,
|
||||
"endpoint_type": "standard",
|
||||
"endpoint": "/v1beta/models/gemini-3.1-flash-lite-preview:generateContent",
|
||||
"thinking_config": {
|
||||
"低": "low",
|
||||
"中": "medium",
|
||||
"高": "high"
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
|
||||
+34
-34
@@ -13,9 +13,21 @@ import tempfile
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_async_api_base_url
|
||||
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
|
||||
@@ -106,6 +118,7 @@ class K3MotionControl:
|
||||
"参考图片": ("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}),
|
||||
@@ -123,9 +136,9 @@ class K3MotionControl:
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/KVideo"
|
||||
|
||||
async def generate(self, 参考图片, 参考视频, 提示词, 保留原声, 角色朝向, 模式, 模型, 时长, seed, **kwargs):
|
||||
async def generate(self, 参考图片, 参考视频, 提示词, 保留原声, 角色朝向, 模式, 模型, 时长, 网络线路, seed, **kwargs):
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_async_api_base_url()
|
||||
base_url = get_base_url_by_route(网络线路)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
@@ -181,8 +194,10 @@ class K3MotionControl:
|
||||
|
||||
# ── 图片 & 视频上传 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(参考视频)
|
||||
|
||||
# ── 构建请求体 ────────────────────────────────────────────────
|
||||
@@ -205,21 +220,16 @@ class K3MotionControl:
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
|
||||
# 1. 提交任务
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
async with session.post(
|
||||
create_url,
|
||||
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,
|
||||
) as resp:
|
||||
headers=headers, prefix="K3 动作控制提交: "
|
||||
))
|
||||
check_interrupt()
|
||||
text = await resp.text()
|
||||
if resp.status != 200:
|
||||
try:
|
||||
err = json.loads(text)
|
||||
msg = err.get("message") or text
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"K3 动作控制提交失败 ({resp.status}): {msg}")
|
||||
create_resp = json.loads(text)
|
||||
|
||||
# task_id 兼容扁平结构和 data 嵌套结构
|
||||
@@ -238,7 +248,8 @@ class K3MotionControl:
|
||||
video_result_url = None
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(interval)
|
||||
await interruptible_sleep(interval)
|
||||
check_interrupt()
|
||||
async with session.get(status_url, headers=headers) as resp:
|
||||
text = await resp.text()
|
||||
if resp.status != 200:
|
||||
@@ -252,30 +263,17 @@ class K3MotionControl:
|
||||
|
||||
# 兼容扁平结构和 data 嵌套结构
|
||||
data = sr.get("data", sr)
|
||||
status = (data.get("status") or sr.get("status") or "").lower()
|
||||
status = extract_status(sr)
|
||||
|
||||
pct_raw = data.get("progress", 0)
|
||||
try:
|
||||
pct = int(str(pct_raw).rstrip("%").strip())
|
||||
except (ValueError, AttributeError):
|
||||
pct = 0
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K3 动作控制] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if status in ("success", "completed", "done", "finished", "succeed"):
|
||||
video_result_url = (
|
||||
data.get("video_url")
|
||||
or data.get("result_url")
|
||||
or data.get("url")
|
||||
or (data.get("result", {}) or {}).get("url")
|
||||
or sr.get("video_url")
|
||||
or sr.get("url")
|
||||
)
|
||||
if is_success_status(status):
|
||||
video_result_url = extract_video_url(sr)
|
||||
break
|
||||
elif status in ("failed", "fail"):
|
||||
err_info = data.get("error") or sr.get("error") or {}
|
||||
err_msg = (err_info.get("message", "未知错误")
|
||||
if isinstance(err_info, dict) else str(err_info))
|
||||
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)
|
||||
@@ -284,6 +282,7 @@ class K3MotionControl:
|
||||
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:
|
||||
@@ -291,6 +290,7 @@ class K3MotionControl:
|
||||
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")
|
||||
|
||||
+31
-31
@@ -11,8 +11,20 @@ import tempfile
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_async_api_base_url
|
||||
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
|
||||
@@ -113,6 +125,7 @@ class K3Video:
|
||||
"时长": ([5, 10, 15], {"default": 5}),
|
||||
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
|
||||
"模式": (_MODES, {"default": "720p"}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"seed": ("INT", {
|
||||
"default": 0, "min": 0, "max": 2147483647,
|
||||
"tooltip": "seed 仅控制节点是否重新运行,结果本身不可复现。",
|
||||
@@ -137,9 +150,9 @@ class K3Video:
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/KVideo"
|
||||
|
||||
async def generate(self, 多镜头, 起始帧, 提示词, 负向提示词, 时长, 生成音频, 模式, seed, **kwargs):
|
||||
async def generate(self, 多镜头, 起始帧, 提示词, 负向提示词, 时长, 生成音频, 模式, 网络线路, seed, **kwargs):
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_async_api_base_url()
|
||||
base_url = get_base_url_by_route(网络线路)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
@@ -244,17 +257,14 @@ class K3Video:
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
|
||||
# 1. 提交
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
async with session.post(create_url, json=body, headers=headers) as resp:
|
||||
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()
|
||||
if resp.status != 200:
|
||||
try:
|
||||
err = json.loads(text)
|
||||
msg = err.get("error", {}).get("message") or err.get("message") or text
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"K3 提交失败 ({resp.status}): {msg}")
|
||||
create_resp = json.loads(text)
|
||||
|
||||
task_id = (
|
||||
@@ -272,6 +282,7 @@ class K3Video:
|
||||
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:
|
||||
@@ -284,39 +295,27 @@ class K3Video:
|
||||
sr = json.loads(text)
|
||||
|
||||
data = sr.get("data", sr)
|
||||
status = (data.get("status") or sr.get("status") or "").lower()
|
||||
status = extract_status(sr)
|
||||
|
||||
pct_raw = data.get("progress", 0)
|
||||
try:
|
||||
pct = int(str(pct_raw).rstrip("%").strip())
|
||||
except (ValueError, AttributeError):
|
||||
pct = 0
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K3 {tag}] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if status in ("success", "completed", "done", "finished", "succeed"):
|
||||
video_url = (
|
||||
data.get("video_url")
|
||||
or data.get("result_url")
|
||||
or data.get("url")
|
||||
or (data.get("result", {}) or {}).get("url")
|
||||
or sr.get("video_url")
|
||||
or sr.get("url")
|
||||
)
|
||||
if is_success_status(status):
|
||||
video_url = extract_video_url(sr)
|
||||
break
|
||||
if status in ("failed", "fail"):
|
||||
err_info = data.get("error") or sr.get("error") or {}
|
||||
err_msg = (err_info.get("message", "未知错误")
|
||||
if isinstance(err_info, dict) else str(err_info))
|
||||
if is_failure_status(status, sr):
|
||||
err_msg = extract_error_message(sr)
|
||||
raise RuntimeError(f"K3 生成失败:{err_msg}")
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
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:
|
||||
@@ -324,6 +323,7 @@ class K3Video:
|
||||
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")
|
||||
|
||||
+31
-31
@@ -10,8 +10,20 @@ import tempfile
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_async_api_base_url
|
||||
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
|
||||
@@ -94,6 +106,7 @@ class K3VideoFirstLast:
|
||||
"时长": ([5, 10, 15], {"default": 5}),
|
||||
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
|
||||
"模式": (_MODES, {"default": "720p"}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"seed": ("INT", {
|
||||
"default": 0, "min": 0, "max": 2147483647,
|
||||
"tooltip": "seed 仅控制节点是否重新运行,结果本身不可复现。",
|
||||
@@ -109,9 +122,9 @@ class K3VideoFirstLast:
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/KVideo"
|
||||
|
||||
async def generate(self, 起始帧, 提示词, 负向提示词, 时长, 生成音频, 模式, seed, 尾帧=None):
|
||||
async def generate(self, 起始帧, 提示词, 负向提示词, 时长, 生成音频, 模式, 网络线路, seed, 尾帧=None):
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_async_api_base_url()
|
||||
base_url = get_base_url_by_route(网络线路)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
@@ -183,17 +196,14 @@ class K3VideoFirstLast:
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
|
||||
# 1. 提交
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
async with session.post(create_url, json=body, headers=headers) as resp:
|
||||
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()
|
||||
if resp.status != 200:
|
||||
try:
|
||||
err = json.loads(text)
|
||||
msg = err.get("error", {}).get("message") or err.get("message") or text
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"K3 首尾帧提交失败 ({resp.status}): {msg}")
|
||||
create_resp = json.loads(text)
|
||||
|
||||
task_id = (
|
||||
@@ -211,6 +221,7 @@ class K3VideoFirstLast:
|
||||
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:
|
||||
@@ -223,39 +234,27 @@ class K3VideoFirstLast:
|
||||
sr = json.loads(text)
|
||||
|
||||
data = sr.get("data", sr)
|
||||
status = (data.get("status") or sr.get("status") or "").lower()
|
||||
status = extract_status(sr)
|
||||
|
||||
pct_raw = data.get("progress", 0)
|
||||
try:
|
||||
pct = int(str(pct_raw).rstrip("%").strip())
|
||||
except (ValueError, AttributeError):
|
||||
pct = 0
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K3 首尾帧] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if status in ("success", "completed", "done", "finished", "succeed"):
|
||||
video_url = (
|
||||
data.get("video_url")
|
||||
or data.get("result_url")
|
||||
or data.get("url")
|
||||
or (data.get("result", {}) or {}).get("url")
|
||||
or sr.get("video_url")
|
||||
or sr.get("url")
|
||||
)
|
||||
if is_success_status(status):
|
||||
video_url = extract_video_url(sr)
|
||||
break
|
||||
if status in ("failed", "fail"):
|
||||
err_info = data.get("error") or sr.get("error") or {}
|
||||
err_msg = (err_info.get("message", "未知错误")
|
||||
if isinstance(err_info, dict) else str(err_info))
|
||||
if is_failure_status(status, sr):
|
||||
err_msg = extract_error_message(sr)
|
||||
raise RuntimeError(f"K3 首尾帧生成失败:{err_msg}")
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
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:
|
||||
@@ -263,6 +262,7 @@ class K3VideoFirstLast:
|
||||
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")
|
||||
|
||||
+37
-34
@@ -10,8 +10,20 @@ import tempfile
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
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
|
||||
@@ -55,6 +67,7 @@ class KVideoFirstLast:
|
||||
"模式": (["1080p"],),
|
||||
"时长": ([5, 10],),
|
||||
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"seed": ("INT", {
|
||||
"default": 0, "min": 0, "max": 2147483647,
|
||||
"tooltip": "seed 仅控制节点是否重新运行,结果本身不可复现。",
|
||||
@@ -70,9 +83,9 @@ class KVideoFirstLast:
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/KVideo"
|
||||
|
||||
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 尾帧=None, seed=0):
|
||||
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 网络线路="全球加速", 尾帧=None, seed=0):
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_api_base_url()
|
||||
base_url = get_base_url_by_route(网络线路)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
@@ -97,10 +110,13 @@ class KVideoFirstLast:
|
||||
"mode": mode_api,
|
||||
"duration": 时长,
|
||||
}
|
||||
if 生成音频 == "打开":
|
||||
body["generate_audio"] = True
|
||||
metadata = {}
|
||||
if 尾帧 is not None:
|
||||
body["metadata"] = {"image_tail": _image_to_base64(尾帧, scale)}
|
||||
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"))
|
||||
@@ -152,17 +168,14 @@ class KVideoFirstLast:
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
|
||||
# 1. 提交
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
async with session.post(create_url, json=body, headers=headers) as resp:
|
||||
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()
|
||||
if resp.status != 200:
|
||||
try:
|
||||
err = json.loads(text)
|
||||
msg = err.get("error", {}).get("message") or err.get("message") or text
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"K26 提交失败 ({resp.status}): {msg}")
|
||||
create_resp = json.loads(text)
|
||||
|
||||
task_id = (
|
||||
@@ -180,6 +193,7 @@ class KVideoFirstLast:
|
||||
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:
|
||||
@@ -192,40 +206,28 @@ class KVideoFirstLast:
|
||||
sr = json.loads(text)
|
||||
|
||||
data = sr.get("data", sr)
|
||||
status = (data.get("status") or sr.get("status") or "").lower()
|
||||
status = extract_status(sr)
|
||||
|
||||
pct_raw = data.get("progress", 0)
|
||||
try:
|
||||
pct = int(str(pct_raw).rstrip("%").strip())
|
||||
except (ValueError, AttributeError):
|
||||
pct = 0
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K26 图生视频] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if status in ("success", "completed", "done", "finished", "succeed"):
|
||||
if is_success_status(status):
|
||||
# 提取视频 URL
|
||||
video_url = (
|
||||
data.get("video_url")
|
||||
or data.get("result_url")
|
||||
or data.get("url")
|
||||
or (data.get("result", {}) or {}).get("url")
|
||||
or sr.get("video_url")
|
||||
or sr.get("url")
|
||||
)
|
||||
video_url = extract_video_url(sr)
|
||||
break
|
||||
if status in ("failed", "fail"):
|
||||
err_info = data.get("error") or sr.get("error") or {}
|
||||
err_msg = (err_info.get("message", "未知错误")
|
||||
if isinstance(err_info, dict) else str(err_info))
|
||||
if is_failure_status(status, sr):
|
||||
err_msg = extract_error_message(sr)
|
||||
raise RuntimeError(f"K26 生成失败:{err_msg}")
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
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:
|
||||
@@ -233,6 +235,7 @@ class KVideoFirstLast:
|
||||
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")
|
||||
|
||||
@@ -11,8 +11,20 @@ import tempfile
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
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
|
||||
@@ -56,6 +68,7 @@ class KVideoImage2Video:
|
||||
"模式": (["720p", "1080p"], {"default": "720p"}),
|
||||
"时长": ([5, 10], {"default": 5}),
|
||||
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"seed": ("INT", {
|
||||
"default": 0, "min": 0, "max": 2147483647,
|
||||
"tooltip": "seed 仅控制节点是否重新运行,结果本身不可复现。",
|
||||
@@ -68,9 +81,12 @@ class KVideoImage2Video:
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/KVideo"
|
||||
|
||||
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", seed=0):
|
||||
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 网络线路="全球加速", seed=0):
|
||||
if 模式 == "720p" and 生成音频 == "打开":
|
||||
raise RuntimeError("K26 仅1080p支持音频,请将模式切换为1080p或关闭生成音频。")
|
||||
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_api_base_url()
|
||||
base_url = get_base_url_by_route(网络线路)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
@@ -96,7 +112,7 @@ class KVideoImage2Video:
|
||||
"duration": 时长,
|
||||
}
|
||||
if 生成音频 == "打开":
|
||||
body["generate_audio"] = True
|
||||
body["metadata"] = {"sound": "on"}
|
||||
|
||||
body_str = json.dumps(body, ensure_ascii=False)
|
||||
body_size = len(body_str.encode("utf-8"))
|
||||
@@ -148,12 +164,13 @@ class KVideoImage2Video:
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
|
||||
# 1. 提交
|
||||
check_interrupt()
|
||||
_stage("submitting")
|
||||
create_url = f"{base_url}{_ENDPOINT_CREATE}"
|
||||
async with session.post(create_url, json=body, headers=headers) as resp:
|
||||
if resp.status != 200:
|
||||
err_text = await resp.text()
|
||||
raise RuntimeError(f"提交失败 ({resp.status}): {err_text}")
|
||||
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")
|
||||
@@ -168,8 +185,9 @@ class KVideoImage2Video:
|
||||
video_url = None
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(interval)
|
||||
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()
|
||||
@@ -177,40 +195,27 @@ class KVideoImage2Video:
|
||||
sr = await resp.json()
|
||||
|
||||
data = sr.get("data", {}) or {}
|
||||
status = (sr.get("status") or data.get("status") or "").lower()
|
||||
status = extract_status(sr)
|
||||
|
||||
pct_raw = str(data.get("progress", 0)).strip().rstrip('%')
|
||||
try:
|
||||
pct = max(0, min(100, int(float(pct_raw))))
|
||||
except (ValueError, TypeError):
|
||||
pct = 0
|
||||
pct = extract_progress(sr)
|
||||
print(f"[K26 图生视频] 生成中 {pct}%")
|
||||
_progress(pct)
|
||||
|
||||
if status in ("success", "completed", "done", "finished", "succeed"):
|
||||
if is_success_status(status):
|
||||
# 提取视频 URL
|
||||
video_url = (
|
||||
data.get("video_url")
|
||||
or data.get("result_url")
|
||||
or data.get("url")
|
||||
or (data.get("result", {}) or {}).get("url")
|
||||
or sr.get("video_url")
|
||||
or sr.get("url")
|
||||
)
|
||||
video_url = extract_video_url(sr)
|
||||
break
|
||||
if status in ("failed", "fail"):
|
||||
err_info = data.get("error") or sr.get("error") or {}
|
||||
err_msg = (err_info.get("message", "未知错误")
|
||||
if isinstance(err_info, dict) else str(err_info))
|
||||
if is_failure_status(status, sr):
|
||||
err_msg = extract_error_message(sr)
|
||||
raise RuntimeError(f"K26 生成失败:{err_msg}")
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
interval = min(interval * 1.5, _POLL_MAX)
|
||||
|
||||
if not video_url:
|
||||
raise RuntimeError(f"API 未返回视频 URL,响应:{sr}")
|
||||
|
||||
# 3. 下载
|
||||
check_interrupt()
|
||||
_stage("downloading")
|
||||
async with session.get(video_url, allow_redirects=True) as resp:
|
||||
if resp.status != 200:
|
||||
@@ -218,6 +223,7 @@ class KVideoImage2Video:
|
||||
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")
|
||||
|
||||
+14
-6
@@ -4,27 +4,35 @@
|
||||
"""
|
||||
|
||||
from .stream_preview import StreamPreview
|
||||
from .nano_banana_pro import NanoBananaPro
|
||||
from .batch_nano_banana_pro import BatchNanoBananaPro
|
||||
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 SaveCleanImage, BatchCleanMetadata
|
||||
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 .multi_res_preview import MultiResPreview
|
||||
from .batch_images_o1key import BatchImagesO1key
|
||||
from .seedance_video import Seedance, SeedanceMultiModal
|
||||
from .nano_banana_v2 import NanoBananaV2, NanoBananaV2Batch, AsyncImageGenerator, BatchAsyncImageGenerator
|
||||
from .doubao_image import DoubaoImage
|
||||
from .gpt_image import O1keyGPTImage
|
||||
from .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__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'MultiResPreview', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator']
|
||||
__all__ = ['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']
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
"""
|
||||
批量 Nano Banana Pro 节点
|
||||
批量 Nano Banana 节点
|
||||
ComfyUI 自定义节点,用于批量处理图像生成任务
|
||||
支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存
|
||||
"""
|
||||
@@ -10,12 +10,14 @@ import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Optional, Tuple, List
|
||||
from typing import Callable, Optional, Tuple, List
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||||
from ..utils.file_utils import (
|
||||
ImageInfo,
|
||||
@@ -25,9 +27,10 @@ from ..utils.file_utils import (
|
||||
generate_timestamp_filename,
|
||||
save_image,
|
||||
)
|
||||
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
|
||||
from ..models_config import (
|
||||
get_enabled_models,
|
||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
||||
get_model_supported_resolutions, get_all_supported_resolutions
|
||||
)
|
||||
@@ -67,7 +70,57 @@ DEBUG_LOG_ENABLED = False
|
||||
REQUEST_LOG_ENABLED = False
|
||||
# ============================================================================
|
||||
|
||||
_NODE = "Nano Banana Pro"
|
||||
_NODE = "Nano Banana"
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
async def _generate_single_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,
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
thinking_level: Optional[str] = None,
|
||||
) -> List[Image.Image]:
|
||||
result_images, _ = 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="BatchNanoBananaPro",
|
||||
request_log_enabled=REQUEST_LOG_ENABLED,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
return result_images
|
||||
|
||||
|
||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
||||
@@ -96,9 +149,9 @@ def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.
|
||||
return pil_to_tensor(matched)
|
||||
|
||||
|
||||
class BatchNanoBananaPro:
|
||||
class BatchNanoBananaPro(io.ComfyNode):
|
||||
"""
|
||||
批量 Nano Banana Pro 节点
|
||||
批量 Nano Banana 节点
|
||||
|
||||
功能:
|
||||
- 从多个文件夹加载图片
|
||||
@@ -116,26 +169,142 @@ class BatchNanoBananaPro:
|
||||
- 要添加/禁用模型,请编辑 models_config.py 文件
|
||||
"""
|
||||
|
||||
# 支持的模型列表(从配置文件动态加载)
|
||||
MODELS = None # 将在 INPUT_TYPES 中动态获取
|
||||
|
||||
# 支持的宽高比列表(全量:所有启用模型的并集,动态加载)
|
||||
# 实际渲染时通过 get_all_supported_aspect_ratios() 获取
|
||||
ASPECT_RATIOS = [
|
||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
||||
"1:4", "4:1", "1:8", "8:1"
|
||||
]
|
||||
|
||||
# 支持的分辨率列表(全量兜底,实际由 get_all_supported_resolutions() 动态生成)
|
||||
RESOLUTIONS = ["512px", "1K", "2K", "4K"]
|
||||
# 模型展示名到基础 ID 的映射
|
||||
MODEL_DISPLAY_NAMES = ["Nano Banana Pro", "Nano Banana 2", "Nano Banana"]
|
||||
MODEL_ID_MAP = {
|
||||
"Nano Banana Pro": "nano-banana-pro",
|
||||
"Nano Banana 2": "nano-banana-2",
|
||||
"Nano Banana": "nano-banana",
|
||||
}
|
||||
# 计费后缀映射
|
||||
BILLING_SUFFIX = {
|
||||
"特价": "-次卡",
|
||||
"官方": "-官方计费",
|
||||
}
|
||||
RESOLUTION_KEY_MAP = {
|
||||
"512px": "0.5k",
|
||||
"1K": "1k",
|
||||
"2K": "2k",
|
||||
"4K": "4k",
|
||||
}
|
||||
# 仅支持特价的模型
|
||||
BILLING_SPECIAL_ONLY = {"nano-banana"}
|
||||
|
||||
# 配对模式
|
||||
PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"]
|
||||
|
||||
def __init__(self):
|
||||
"""初始化节点"""
|
||||
self.client = None
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
normal_aspect_ratios = [
|
||||
"智能", "1:1", "2:3", "3:2", "3:4", "4:3",
|
||||
"4:5", "5:4", "9:16", "16:9", "21:9",
|
||||
]
|
||||
nano_banana_2_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",
|
||||
]
|
||||
|
||||
return io.Schema(
|
||||
node_id="BatchNanoBananaPro",
|
||||
display_name="批量 Nano Banana",
|
||||
category="image/batch",
|
||||
inputs=[
|
||||
io.String.Input(
|
||||
"prompt",
|
||||
default="一个中国女子的OOTD",
|
||||
multiline=True,
|
||||
),
|
||||
io.DynamicCombo.Input("模型", options=[
|
||||
io.DynamicCombo.Option("Nano Banana Pro", [
|
||||
io.Combo.Input("宽高比", options=normal_aspect_ratios, default="智能"),
|
||||
io.Combo.Input("分辨率", options=["1K", "2K", "4K"], default="2K"),
|
||||
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
|
||||
]),
|
||||
io.DynamicCombo.Option("Nano Banana 2", [
|
||||
io.Combo.Input("宽高比", options=nano_banana_2_aspect_ratios, default="智能"),
|
||||
io.Combo.Input("分辨率", options=["512px", "1K", "2K", "4K"], default="2K"),
|
||||
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
|
||||
io.Combo.Input("思考深度", options=["高", "低"], default="高"),
|
||||
]),
|
||||
io.DynamicCombo.Option("Nano Banana", [
|
||||
io.Combo.Input("宽高比", options=normal_aspect_ratios, default="智能"),
|
||||
io.Combo.Input("分辨率", options=["1K"], default="1K"),
|
||||
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
|
||||
]),
|
||||
]),
|
||||
io.Combo.Input("图片格式", options=["原始", "JPEG", "PNG", "WebP"], default="原始"),
|
||||
io.Combo.Input("计费", options=["特价", "官方"], default="特价"),
|
||||
io.Combo.Input("网络", options=NETWORK_ROUTE_OPTIONS, default="全球加速"),
|
||||
io.Int.Input("seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF),
|
||||
io.String.Input("文件夹1", default="", multiline=False),
|
||||
io.String.Input("文件夹2", default="", multiline=False),
|
||||
io.String.Input("文件夹3", default="", multiline=False),
|
||||
io.String.Input("文件夹4", default="", multiline=False),
|
||||
io.String.Input("文件夹5", default="", multiline=False),
|
||||
io.String.Input("保存路径", default="", multiline=False),
|
||||
io.Combo.Input("图片配对模式", options=cls.PAIRING_MODES, default="不配对"),
|
||||
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),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="输出图像"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
prompt,
|
||||
模型,
|
||||
图片格式,
|
||||
计费,
|
||||
网络,
|
||||
seed,
|
||||
文件夹1,
|
||||
文件夹2,
|
||||
文件夹3,
|
||||
文件夹4,
|
||||
文件夹5,
|
||||
保存路径,
|
||||
图片配对模式,
|
||||
**kwargs,
|
||||
) -> io.NodeOutput:
|
||||
model_name = 模型["模型"]
|
||||
宽高比 = 模型.get("宽高比", "智能")
|
||||
分辨率 = 模型.get("分辨率", "2K")
|
||||
思考深度 = 模型.get("思考深度")
|
||||
谷歌搜索 = 模型.get("谷歌搜索", "关闭")
|
||||
if 思考深度:
|
||||
kwargs["思考深度"] = 思考深度
|
||||
kwargs["谷歌搜索"] = 谷歌搜索
|
||||
|
||||
node = cls()
|
||||
output_tensor, = node.process_batch(
|
||||
prompt=prompt,
|
||||
文件夹1=文件夹1,
|
||||
文件夹2=文件夹2,
|
||||
文件夹3=文件夹3,
|
||||
文件夹4=文件夹4,
|
||||
文件夹5=文件夹5,
|
||||
seed=seed,
|
||||
图片配对模式=图片配对模式,
|
||||
模型=model_name,
|
||||
计费=计费,
|
||||
宽高比=宽高比,
|
||||
分辨率=分辨率,
|
||||
图片格式=图片格式,
|
||||
网络=网络,
|
||||
保存路径=保存路径,
|
||||
**kwargs,
|
||||
)
|
||||
return io.NodeOutput(output_tensor)
|
||||
|
||||
def resize_to_megapixels(
|
||||
self,
|
||||
@@ -181,69 +350,46 @@ class BatchNanoBananaPro:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
|
||||
ComfyUI 节点规范:
|
||||
- required: 必选参数
|
||||
- optional: 可选参数
|
||||
"""
|
||||
# 从配置文件动态获取启用的模型列表
|
||||
enabled_models = get_enabled_models()
|
||||
|
||||
# 如果没有启用的模型,使用空列表(会导致节点不可用,提示用户配置)
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
||||
|
||||
# 动态获取所有启用模型支持的宽高比(去重合并)
|
||||
all_aspect_ratios = get_all_supported_aspect_ratios()
|
||||
if not all_aspect_ratios:
|
||||
all_aspect_ratios = cls.ASPECT_RATIOS
|
||||
all_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"]
|
||||
|
||||
# 动态获取所有启用模型支持的分辨率(去重合并)
|
||||
all_resolutions = get_all_supported_resolutions()
|
||||
if not all_resolutions:
|
||||
all_resolutions = cls.RESOLUTIONS
|
||||
all_resolutions = ["512px", "1K", "2K", "4K"]
|
||||
|
||||
# 创建9个独立的图像输入
|
||||
# 创建5个独立的图像输入
|
||||
optional_inputs = {}
|
||||
for i in range(1, 10): # 1-9
|
||||
for i in range(1, 6): # 1-5
|
||||
optional_inputs[f"参考图{i}"] = ("IMAGE",)
|
||||
|
||||
# 图片配对模式移到可选参数
|
||||
optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, {
|
||||
"default": "不配对"
|
||||
})
|
||||
|
||||
optional_inputs["代理端口(如7897)"] = ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"placeholder": "本地代理端口,如 7897(Clash Verge)或 10808(v2rayN),留空不使用"
|
||||
})
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "一个中国女子的OOTD",
|
||||
"multiline": True
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0]
|
||||
"模型": (cls.MODEL_DISPLAY_NAMES, {
|
||||
"default": cls.MODEL_DISPLAY_NAMES[0]
|
||||
}),
|
||||
"宽高比": (all_aspect_ratios, {
|
||||
"default": "1:1"
|
||||
"宽高比": (["智能"] + all_aspect_ratios, {
|
||||
"default": "智能"
|
||||
}),
|
||||
"分辨率": (all_resolutions, {
|
||||
"default": "2K"
|
||||
}),
|
||||
"谷歌搜索(联网)": (["关闭", "打开"], {
|
||||
"default": "关闭"
|
||||
"图片格式": (["原始", "JPEG", "PNG", "WebP"], {
|
||||
"default": "原始"
|
||||
}),
|
||||
"图片搜索(联网)": (["关闭", "打开"], {
|
||||
"default": "关闭"
|
||||
"计费": (["特价", "官方"], {
|
||||
"default": "特价"
|
||||
}),
|
||||
"返回格式": (["url", "base64"], {
|
||||
"default": "url"
|
||||
"网络": (NETWORK_ROUTE_OPTIONS, {
|
||||
"default": "全球加速"
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
@@ -270,22 +416,6 @@ class BatchNanoBananaPro:
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹6": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹7": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹8": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹9": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"保存路径": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
@@ -403,8 +533,9 @@ class BatchNanoBananaPro:
|
||||
|
||||
async def _generate_single_task(
|
||||
self,
|
||||
client: GeminiAPIClient,
|
||||
session: aiohttp.ClientSession,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
@@ -412,27 +543,14 @@ class BatchNanoBananaPro:
|
||||
images: List[ImageInfo],
|
||||
output_folder: str,
|
||||
task_index: int,
|
||||
enable_grounding: bool = True,
|
||||
enable_image_search: bool = False,
|
||||
enable_grounding: bool = False,
|
||||
base_filename: str = None,
|
||||
image_format: str = "url",
|
||||
image_format: str = "原始",
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
thinking_level: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
执行单个生成任务
|
||||
|
||||
Args:
|
||||
client: API 客户端
|
||||
session: aiohttp 会话
|
||||
prompt: 提示词
|
||||
model: 模型名称
|
||||
resolution: 分辨率
|
||||
aspect_ratio: 宽高比
|
||||
images: 输入图片列表
|
||||
output_folder: 输出文件夹
|
||||
task_index: 任务索引
|
||||
|
||||
Returns:
|
||||
包含结果信息的字典
|
||||
执行单个生成任务(异步生图接口)
|
||||
"""
|
||||
result = {
|
||||
"task_index": task_index,
|
||||
@@ -440,7 +558,7 @@ class BatchNanoBananaPro:
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"saved_files": [],
|
||||
"output_images": [], # 无保存路径时存储内存图片
|
||||
"output_images": [],
|
||||
"error": None
|
||||
}
|
||||
|
||||
@@ -448,26 +566,23 @@ class BatchNanoBananaPro:
|
||||
# 准备输入图片
|
||||
input_pil_images = [info.image for info in images]
|
||||
|
||||
# 调用 API 生成图片(固定生成1次)
|
||||
# 调用异步生图接口生成图片
|
||||
generated_images = []
|
||||
try:
|
||||
gen_result = await client.generate_single_async(
|
||||
gen_images = await _generate_single_async(
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_pil_images,
|
||||
session=session,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
images=input_pil_images if input_pil_images else None,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
image_format=image_format,
|
||||
progress_callback=progress_callback,
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
if gen_result:
|
||||
# 正确解包元组:第一个元素是图像列表,第二个是计时信息
|
||||
images_list, timing_info = gen_result
|
||||
generated_images.extend(images_list)
|
||||
generated_images.extend(gen_images)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_msg = str(e)
|
||||
@@ -490,29 +605,54 @@ class BatchNanoBananaPro:
|
||||
|
||||
# 保存生成的图片到磁盘(始终保存)
|
||||
import os
|
||||
|
||||
# 确定保存扩展名
|
||||
_FORMAT_EXT_MAP = {"JPEG": ".jpg", "PNG": ".png", "WebP": ".webp"}
|
||||
save_ext = _FORMAT_EXT_MAP.get(image_format, ".png")
|
||||
|
||||
for i, gen_img in enumerate(generated_images):
|
||||
# 格式转换:非"原始"时检测并转换
|
||||
if image_format != "原始":
|
||||
src_format = (gen_img.format or "").upper()
|
||||
target_upper = image_format.upper()
|
||||
# JPEG 格式名在 PIL 中为 "JPEG"
|
||||
if src_format == "JPG":
|
||||
src_format = "JPEG"
|
||||
need_convert = (src_format != target_upper)
|
||||
if need_convert:
|
||||
if target_upper in ("JPEG", "WEBP") and gen_img.mode in ("RGBA", "LA", "P"):
|
||||
gen_img = gen_img.convert("RGB")
|
||||
|
||||
# 使用文件夹1图片的名称,如果重名则+1
|
||||
if base_filename:
|
||||
base_name = base_filename
|
||||
counter = 0
|
||||
while True:
|
||||
if counter == 0:
|
||||
filename = f"{base_name}.png"
|
||||
filename = f"{base_name}{save_ext}"
|
||||
else:
|
||||
filename = f"{base_name}+{counter}.png"
|
||||
filename = f"{base_name}+{counter}{save_ext}"
|
||||
output_path = os.path.join(output_folder, filename)
|
||||
if not os.path.exists(output_path):
|
||||
break
|
||||
counter += 1
|
||||
else:
|
||||
# 如果没有base_filename,使用时间戳
|
||||
output_path = generate_timestamp_filename(
|
||||
output_folder=output_folder,
|
||||
extension=".png"
|
||||
extension=save_ext
|
||||
)
|
||||
|
||||
# 保存时不压缩
|
||||
if image_format == "JPEG":
|
||||
if gen_img.mode != "RGB":
|
||||
gen_img = gen_img.convert("RGB")
|
||||
gen_img.save(output_path, quality=100)
|
||||
elif image_format == "WebP":
|
||||
gen_img.save(output_path, lossless=True)
|
||||
else:
|
||||
save_image(gen_img, output_path)
|
||||
|
||||
result["saved_files"].append(output_path)
|
||||
# 立即释放内存
|
||||
gen_img = None
|
||||
|
||||
# 只有生成了图片才标记为成功
|
||||
@@ -533,38 +673,21 @@ class BatchNanoBananaPro:
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
output_folder: str,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
pbar=None,
|
||||
prompts_per_task: Optional[List[str]] = None,
|
||||
enable_grounding: bool = True,
|
||||
enable_image_search: bool = False,
|
||||
image_format: str = "url",
|
||||
enable_grounding: bool = False,
|
||||
image_format: str = "原始",
|
||||
thinking_level: Optional[str] = None,
|
||||
) -> List[dict]:
|
||||
"""
|
||||
异步批量处理所有任务 - 改进版:支持分批保存
|
||||
|
||||
Args:
|
||||
pairs: 配对后的图片组合
|
||||
prompt: 提示词(单提示词模式时使用)
|
||||
model: 模型名称
|
||||
resolution: 分辨率
|
||||
aspect_ratio: 宽高比
|
||||
output_folder: 输出文件夹
|
||||
pbar: ComfyUI 进度条
|
||||
prompts_per_task: 每个任务对应的提示词列表(批量提示词模式时使用)
|
||||
|
||||
Returns:
|
||||
所有任务的结果列表
|
||||
异步批量处理所有任务(异步生图接口)
|
||||
"""
|
||||
if self.client is None:
|
||||
self.client = GeminiAPIClient()
|
||||
|
||||
total_tasks = len(pairs)
|
||||
|
||||
max_concurrent = 50
|
||||
|
||||
# 分批保存的批次大小(与并发数一致)
|
||||
save_batch_size = 10
|
||||
|
||||
print(f"BatchNanoBananaPro: 检测到 {total_tasks} 个任务")
|
||||
|
||||
all_results = []
|
||||
@@ -617,8 +740,9 @@ class BatchNanoBananaPro:
|
||||
|
||||
task = asyncio.create_task(
|
||||
self._generate_single_task(
|
||||
client=self.client,
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
prompt=task_prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
@@ -627,9 +751,10 @@ class BatchNanoBananaPro:
|
||||
output_folder=output_folder,
|
||||
task_index=start_idx + i,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
base_filename=base_filename,
|
||||
image_format=image_format,
|
||||
progress_callback=_make_progress_callback(pbar),
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
@@ -669,13 +794,12 @@ class BatchNanoBananaPro:
|
||||
else:
|
||||
fail_count += 1
|
||||
|
||||
# 更新 ComfyUI 原生进度条
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
# 大任务额外显示百分比里程碑
|
||||
if show_milestone and milestone_index < len(milestones):
|
||||
progress = completed / total_tasks
|
||||
if pbar is not None and getattr(pbar, "total", 0):
|
||||
progress = pbar.current / pbar.total
|
||||
else:
|
||||
progress = success_count / total_tasks
|
||||
if progress >= milestones[milestone_index]:
|
||||
percentage = int(milestones[milestone_index] * 100)
|
||||
print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<")
|
||||
@@ -721,15 +845,14 @@ class BatchNanoBananaPro:
|
||||
文件夹3: str,
|
||||
文件夹4: str,
|
||||
文件夹5: str,
|
||||
文件夹6: str,
|
||||
文件夹7: str,
|
||||
文件夹8: str,
|
||||
文件夹9: str,
|
||||
seed: int,
|
||||
图片配对模式: str,
|
||||
模型: str,
|
||||
计费: str,
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
图片格式: str,
|
||||
网络: str,
|
||||
保存路径: str = "",
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
@@ -738,14 +861,14 @@ class BatchNanoBananaPro:
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
文件夹1-9: 图片文件夹路径
|
||||
文件夹1-5: 图片文件夹路径
|
||||
seed: 随机种子
|
||||
保存路径: 输出保存路径
|
||||
图片配对模式: 1:1 或 1*N
|
||||
模型: 模型名称
|
||||
宽高比: 输出宽高比
|
||||
分辨率: 输出分辨率
|
||||
**kwargs: 动态参考图输入 (参考图1-9)
|
||||
**kwargs: 动态参考图输入 (参考图1-5)
|
||||
|
||||
Returns:
|
||||
输出图像张量
|
||||
@@ -753,10 +876,32 @@ class BatchNanoBananaPro:
|
||||
start_time = time.time()
|
||||
|
||||
# 从 kwargs 提取搜索参数(界面显示为「关闭/打开」,转为 bool 供调用)
|
||||
enable_grounding: bool = (kwargs.pop("谷歌搜索(联网)", "关闭") == "打开")
|
||||
enable_image_search: bool = (kwargs.pop("图片搜索(联网)", "关闭") == "打开")
|
||||
proxy_port: str = kwargs.pop("代理端口(如7897)", "")
|
||||
image_format: str = kwargs.pop("返回格式", "url")
|
||||
enable_grounding: bool = kwargs.get("谷歌搜索", "关闭") == "打开"
|
||||
|
||||
# 拼接实际模型 ID
|
||||
base_model_id = self.MODEL_ID_MAP.get(模型, "nano-banana-pro")
|
||||
思考深度 = kwargs.get("思考深度", "高")
|
||||
thinking_level = None
|
||||
if base_model_id == "nano-banana-2":
|
||||
thinking_level = "High" if 思考深度 == "高" else "Low"
|
||||
|
||||
if base_model_id == "nano-banana":
|
||||
if 计费 == "官方":
|
||||
raise ValueError(f"模型 \"{模型}\" 仅支持特价计费")
|
||||
模型 = "nano-banana"
|
||||
else:
|
||||
res_key = self.RESOLUTION_KEY_MAP.get(分辨率, "2k")
|
||||
is_official = (计费 == "官方")
|
||||
if base_model_id == "nano-banana-pro" and res_key == "1k" and not is_official:
|
||||
模型 = "nano-banana-pro"
|
||||
elif base_model_id == "nano-banana-2" and res_key == "0.5k":
|
||||
if is_official:
|
||||
raise ValueError("Nano Banana 2 的 512px 分辨率仅支持特价计费")
|
||||
模型 = "nano-banana-2-0.5k"
|
||||
else:
|
||||
模型 = f"{base_model_id}-{res_key}"
|
||||
if is_official:
|
||||
模型 += "-official"
|
||||
|
||||
|
||||
try:
|
||||
@@ -767,7 +912,7 @@ class BatchNanoBananaPro:
|
||||
# 验证:至少需要填写一个文件夹路径
|
||||
has_any_folder = any(
|
||||
f and f.strip()
|
||||
for f in [文件夹1, 文件夹2, 文件夹3, 文件夹4, 文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9]
|
||||
for f in [文件夹1, 文件夹2, 文件夹3, 文件夹4, 文件夹5]
|
||||
)
|
||||
if not has_any_folder:
|
||||
raise ValueError("请至少填写一个文件夹路径,该节点专为批量文件夹处理设计")
|
||||
@@ -782,26 +927,16 @@ class BatchNanoBananaPro:
|
||||
|
||||
# 校验宽高比与模型的兼容性
|
||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
||||
if supported_ratios and 宽高比 not in supported_ratios:
|
||||
if 宽高比 != "智能" and supported_ratios and 宽高比 not in supported_ratios:
|
||||
raise ValueError(
|
||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容!\n"
|
||||
f"该模型支持的宽高比:{', '.join(supported_ratios)}"
|
||||
)
|
||||
|
||||
# 校验图片搜索(联网)与模型的兼容性
|
||||
# 仅 nano-banana-2-限时特价 和 gemini-3.1-flash-image-preview 支持图片搜索
|
||||
IMAGE_SEARCH_UNSUPPORTED_MODELS = ["nano-banana-pro-次卡", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview"]
|
||||
if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS:
|
||||
raise ValueError(
|
||||
f"模型 \"{模型}\" 不支持【图片搜索(联网)】功能!"
|
||||
f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用"
|
||||
)
|
||||
|
||||
# 加载文件夹图片
|
||||
print("BatchNanoBananaPro: 开始加载图片...")
|
||||
image_lists = self._load_folders(
|
||||
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9
|
||||
文件夹1, 文件夹2, 文件夹3, 文件夹4, 文件夹5
|
||||
)
|
||||
|
||||
# 验证文件夹是否有可用图片
|
||||
@@ -811,7 +946,7 @@ class BatchNanoBananaPro:
|
||||
|
||||
# 处理独立的参考图输入
|
||||
manual_images = []
|
||||
for i in range(1, 10): # 1-9
|
||||
for i in range(1, 6): # 1-5
|
||||
key = f"参考图{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
pil_images = tensor_to_pil(kwargs[key])
|
||||
@@ -848,17 +983,15 @@ class BatchNanoBananaPro:
|
||||
total_tasks = len(pairs)
|
||||
|
||||
# 打印首行概览
|
||||
# 图片搜索(联网)开启时隐含谷歌搜索接地,与客户端请求逻辑保持一致
|
||||
grounding_str = ""
|
||||
if enable_image_search:
|
||||
grounding_str = " | 谷歌图片搜索接地"
|
||||
elif enable_grounding:
|
||||
if enable_grounding:
|
||||
grounding_str = " | 谷歌搜索接地"
|
||||
thinking_str = f" | 思考:{thinking_level}" if thinking_level else ""
|
||||
|
||||
if batch_prompts:
|
||||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}")
|
||||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}{thinking_str}")
|
||||
else:
|
||||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}")
|
||||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}{thinking_str}")
|
||||
|
||||
# 创建 ComfyUI 原生进度条
|
||||
pbar = None
|
||||
@@ -890,17 +1023,9 @@ class BatchNanoBananaPro:
|
||||
except Exception as e:
|
||||
raise ValueError(f"保存路径无效或无写入权限: {保存路径} - {str(e)}")
|
||||
|
||||
# 初始化 API 客户端
|
||||
if self.client is None:
|
||||
try:
|
||||
self.client = GeminiAPIClient()
|
||||
except ValueError as e:
|
||||
raise ValueError(f"初始化 API 客户端失败: {str(e)}")
|
||||
|
||||
# 注入代理设置(每次执行都刷新,支持用户中途修改端口)
|
||||
self.client.proxy_url = GeminiAPIClient.build_proxy_url(proxy_port)
|
||||
if self.client.proxy_url:
|
||||
print(f"BatchNanoBananaPro: 已启用代理加速 → {self.client.proxy_url}")
|
||||
# 获取 API 密钥和基础 URL
|
||||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
base_url = get_base_url_by_route(网络)
|
||||
|
||||
# 判断是否使用默认 output 目录
|
||||
original_save_path = kwargs.get('保存路径', '')
|
||||
@@ -920,11 +1045,13 @@ class BatchNanoBananaPro:
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
output_folder=保存路径,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
pbar=pbar,
|
||||
prompts_per_task=prompts_per_task,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
image_format=image_format,
|
||||
image_format=图片格式,
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -1046,15 +1173,16 @@ class BatchNanoBananaPro:
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
except Exception as e:
|
||||
raise type(e)(str(e)) from None
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
finally:
|
||||
# 查询余额
|
||||
if self.client is not None:
|
||||
try:
|
||||
balance_data = self.client.query_balance_sync()
|
||||
balance_info = self.client.format_balance_info(balance_data)
|
||||
print(f"BatchNanaBananaPro: {balance_info}")
|
||||
client = GeminiAPIClient()
|
||||
client.base_url = get_base_url_by_route(网络)
|
||||
balance_data = client.query_balance_sync()
|
||||
balance_info = client.format_balance_info(balance_data)
|
||||
print(f"BatchNanoBananaPro: {balance_info}")
|
||||
print("=" * 60)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -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,)
|
||||
+582
-20
@@ -3,9 +3,23 @@ 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
|
||||
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
|
||||
@@ -15,6 +29,68 @@ except ImportError:
|
||||
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:
|
||||
"""
|
||||
@@ -52,30 +128,33 @@ class O1keyGPTImage:
|
||||
], {
|
||||
"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)",
|
||||
"1365x1024(1K 横版 4:3)",
|
||||
"1024x1365(1K 竖版 3:4)",
|
||||
"1820x1024(1K 横版 16:9)",
|
||||
"1024x1820(1K 竖版 9:16)",
|
||||
"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)",
|
||||
"2732x2048(2K 横版 4:3)",
|
||||
"2048x2732(2K 竖版 3:4)",
|
||||
"3640x2048(2K 横版 16:9)",
|
||||
"2048x3640(2K 竖版 9:16)",
|
||||
"2736x2048(2K 横版 4:3)",
|
||||
"2048x2736(2K 竖版 3:4)",
|
||||
"3648x2048(2K 横版 16:9)",
|
||||
"2048x3648(2K 竖版 9:16)",
|
||||
# ── 4K ──
|
||||
"3840x3840(4K 正方形 1:1)",
|
||||
"3840x2560(4K 横版 3:2)",
|
||||
"2560x3840(4K 竖版 2:3)",
|
||||
"3840x2880(4K 横版 4:3)",
|
||||
"2880x3840(4K 竖版 3:4)",
|
||||
"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)",
|
||||
], {
|
||||
@@ -94,6 +173,10 @@ class O1keyGPTImage:
|
||||
"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,
|
||||
@@ -128,8 +211,10 @@ class O1keyGPTImage:
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str = "gpt-image-2-次卡",
|
||||
网络: str = "全球加速",
|
||||
分辨率: str = "auto",
|
||||
质量: str = "自动",
|
||||
输出格式: str = "jpeg",
|
||||
生图数量: int = 1,
|
||||
seed: int = 0,
|
||||
遮罩=None,
|
||||
@@ -160,15 +245,20 @@ class O1keyGPTImage:
|
||||
raise ValueError("提供了遮罩但未提供图片,请同时提供图片和遮罩")
|
||||
|
||||
# ── 2. 解析分辨率显示值 → API 参数值 ──────────────────────────────────
|
||||
size = "auto" if 分辨率 == "智能" else 分辨率.split("(")[0].strip()
|
||||
size = _resolve_async_size(分辨率)
|
||||
|
||||
# ── 2b. 解析质量显示值 → API 参数值 ───────────────────────────────────
|
||||
# ── 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] 请联系作者授权后方可使用!")
|
||||
@@ -181,6 +271,8 @@ class O1keyGPTImage:
|
||||
|
||||
# ── 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:
|
||||
# 批量模式:逐条提示词调用
|
||||
@@ -191,15 +283,17 @@ class O1keyGPTImage:
|
||||
print("[o1key GPT Image] 用户取消,已中断批量生成")
|
||||
raise InterruptProcessingException()
|
||||
try:
|
||||
pil_images = client.run_sync(
|
||||
pil_images = client.generate_image_async_sync(
|
||||
prompt=p,
|
||||
model=模型,
|
||||
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 "")
|
||||
@@ -210,20 +304,24 @@ class O1keyGPTImage:
|
||||
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.run_sync(
|
||||
pil_images = client.generate_image_async_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
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:
|
||||
@@ -260,3 +358,467 @@ class O1keyGPTImage:
|
||||
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",
|
||||
}
|
||||
@@ -8,6 +8,7 @@ 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
|
||||
|
||||
@@ -121,6 +122,7 @@ class KlingVideo:
|
||||
"required": {
|
||||
"提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"反向提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"模型版本": (["v3", "v2-6"], {"default": "v3"}),
|
||||
"时长": ([5, 10, 15],),
|
||||
"分辨率": (["1080p", "720p"],),
|
||||
@@ -244,6 +246,7 @@ class KlingVideo:
|
||||
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"
|
||||
|
||||
@@ -251,6 +254,7 @@ class KlingVideo:
|
||||
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="kling_")
|
||||
|
||||
client = KlingClient()
|
||||
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
|
||||
|
||||
# ── 进度条 ────────────────────────────────────────────────────
|
||||
try:
|
||||
@@ -307,6 +311,7 @@ class KlingFirstLastFrame:
|
||||
"首帧": ("IMAGE",),
|
||||
"尾帧": ("IMAGE",),
|
||||
"提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"模型": (["v3", "v2-6"], {"default": "v3"}),
|
||||
"分辨率": (["1080p", "720p"],),
|
||||
"时长": ([5, 10, 15],),
|
||||
@@ -391,6 +396,7 @@ class KlingFirstLastFrame:
|
||||
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="kling_")
|
||||
|
||||
client = KlingClient()
|
||||
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
|
||||
|
||||
# 进度条:0~100 步
|
||||
try:
|
||||
@@ -454,6 +460,7 @@ class KlingMotionControlTest:
|
||||
"提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"参考图片": ("IMAGE",),
|
||||
"参考视频": ("VIDEO",),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
},
|
||||
"optional": {
|
||||
"模型": (["v3", "v2-6"], {"default": "v3"}),
|
||||
@@ -560,6 +567,7 @@ class KlingMotionControlTest:
|
||||
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="kling_motion_")
|
||||
|
||||
client = KlingClient()
|
||||
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
|
||||
|
||||
# ── 进度条 ────────────────────────────────────────────────────
|
||||
try:
|
||||
|
||||
@@ -1,165 +0,0 @@
|
||||
"""
|
||||
多分辨率图像预览节点
|
||||
ComfyUI 自定义节点,支持同时预览多张不同分辨率的图像
|
||||
|
||||
背景:
|
||||
ComfyUI 原生「预览图像」节点要求 batch 内所有图片分辨率相同(因为它们被
|
||||
stack 成一个 [B, H, W, C] tensor)。当 API 返回多张不同尺寸的图片时
|
||||
(例如 nano-banana-2 同时返回 1K + 2K),原生节点会报错。
|
||||
|
||||
解决方案:
|
||||
声明 INPUT_IS_LIST = True,ComfyUI 会将连入的所有图像作为
|
||||
Python list[Tensor] 传入,而不是强行 stack 成单个 tensor。
|
||||
节点逐张单独保存为临时 PNG,再通过 ui.images 列表返回给前端并列展示,
|
||||
完全不受分辨率一致性的限制。
|
||||
"""
|
||||
|
||||
import os
|
||||
import uuid
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
|
||||
def _get_temp_dir() -> str:
|
||||
"""获取 ComfyUI temp 目录,不可用时回退到系统临时目录"""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
return folder_paths.get_temp_directory()
|
||||
import tempfile
|
||||
return tempfile.gettempdir()
|
||||
|
||||
|
||||
def _tensor_to_pil(tensor) -> list:
|
||||
"""
|
||||
将单个 IMAGE tensor 转换为 PIL Image 列表。
|
||||
|
||||
ComfyUI IMAGE tensor 格式:[B, H, W, C],float32,值域 [0, 1]
|
||||
支持:
|
||||
- 单张图 tensor: shape [H, W, C] 或 [1, H, W, C]
|
||||
- batch tensor: shape [B, H, W, C](B 张相同尺寸图)
|
||||
"""
|
||||
import torch
|
||||
if not isinstance(tensor, torch.Tensor):
|
||||
return []
|
||||
|
||||
if tensor.ndim == 3:
|
||||
tensor = tensor.unsqueeze(0)
|
||||
|
||||
results = []
|
||||
for i in range(tensor.shape[0]):
|
||||
img_np = tensor[i].cpu().numpy()
|
||||
img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)
|
||||
results.append(Image.fromarray(img_np))
|
||||
return results
|
||||
|
||||
|
||||
class MultiResPreview:
|
||||
"""
|
||||
多分辨率图像预览节点
|
||||
|
||||
功能:
|
||||
- 单个「图像」输入端口,支持接入批次图像
|
||||
- INPUT_IS_LIST = True:ComfyUI 将每张图作为独立 tensor 传入,
|
||||
不强制要求尺寸相同,彻底解决不同分辨率无法共存的问题
|
||||
- 每张图像独立保存为临时 PNG,在节点上并列展示所有图像
|
||||
|
||||
用法:
|
||||
将 Nano Banana 节点的输出直接连入「图像」端口即可,
|
||||
无论返回几张、分辨率是否相同,都能正确展示。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"图像": ("IMAGE",),
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
},
|
||||
}
|
||||
|
||||
# 关键:告知 ComfyUI 以 list[Tensor] 而非 stacked Tensor 传入图像
|
||||
# 这样不同分辨率的图片可以共存于同一个输入中
|
||||
INPUT_IS_LIST = True
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "preview"
|
||||
CATEGORY = "image"
|
||||
|
||||
DESCRIPTION = (
|
||||
"多分辨率图像预览节点。\n"
|
||||
"单个图像输入端口,支持任意数量、任意分辨率的批次图像。\n"
|
||||
"解决了原生「预览图像」节点要求 batch 内图片尺寸相同的限制。\n"
|
||||
"常用场景:nano-banana-2 同时返回 1K + 2K 图时,直接连入本节点即可。"
|
||||
)
|
||||
|
||||
def preview(self, 图像, prompt=None, extra_pnginfo=None) -> dict:
|
||||
"""
|
||||
逐张将图像保存到 temp 目录,返回 ui.images 供前端展示。
|
||||
|
||||
Args:
|
||||
图像: list[Tensor],每个元素是一张或一批图(INPUT_IS_LIST)
|
||||
prompt: ComfyUI 注入的 prompt 元数据(可选)
|
||||
extra_pnginfo: ComfyUI 注入的额外 PNG 信息(可选)
|
||||
|
||||
Returns:
|
||||
{"ui": {"images": [...]}} 格式,每项对应一张图
|
||||
"""
|
||||
temp_dir = _get_temp_dir()
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
# 构建 PNG 元数据(与原生预览节点行为一致)
|
||||
metadata = PngInfo()
|
||||
# INPUT_IS_LIST 时 hidden 值也会被包装成 list,取第一个元素
|
||||
_prompt = prompt[0] if isinstance(prompt, list) else prompt
|
||||
_extra = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo
|
||||
if _prompt is not None:
|
||||
try:
|
||||
metadata.add_text("prompt", json.dumps(_prompt))
|
||||
except Exception:
|
||||
pass
|
||||
if _extra is not None:
|
||||
try:
|
||||
for k, v in _extra.items():
|
||||
metadata.add_text(k, json.dumps(v))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
saved = []
|
||||
total_input = 0
|
||||
total_saved = 0
|
||||
|
||||
# 图像 是 list[Tensor],逐个处理(每个 Tensor 可能自身是个 batch)
|
||||
for tensor in 图像:
|
||||
pil_images = _tensor_to_pil(tensor)
|
||||
total_input += len(pil_images)
|
||||
|
||||
for pil_img in pil_images:
|
||||
try:
|
||||
filename = f"multi_res_preview_{uuid.uuid4().hex[:12]}.png"
|
||||
filepath = os.path.join(temp_dir, filename)
|
||||
pil_img.save(filepath, pnginfo=metadata, compress_level=1)
|
||||
|
||||
saved.append({
|
||||
"filename": filename,
|
||||
"subfolder": "",
|
||||
"type": "temp",
|
||||
})
|
||||
total_saved += 1
|
||||
except Exception as e:
|
||||
print(f"多分辨率预览: ⚠️ 保存图像失败 - {e}")
|
||||
|
||||
if total_input == 0:
|
||||
print("多分辨率预览: ⚠️ 没有接收到任何图像")
|
||||
|
||||
return {"ui": {"images": saved}}
|
||||
@@ -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,831 +0,0 @@
|
||||
"""
|
||||
Nano Banana Pro 节点
|
||||
ComfyUI 自定义节点,用于调用 Gemini 模型生成图像
|
||||
"""
|
||||
|
||||
import time
|
||||
import math
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||||
from ..utils.file_utils import ImageInfo, generate_timestamp_filename, save_image
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
from ..models_config import (
|
||||
get_enabled_models, get_model_description,
|
||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
||||
get_model_supported_resolutions, get_all_supported_resolutions
|
||||
)
|
||||
|
||||
# 检查 folder_paths 是否可用
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
# 导入 ComfyUI 原生进度条
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ NanoBananaPro: comfy.utils.ProgressBar 不可用,将只使用终端进度显示")
|
||||
|
||||
# 内存监控(可选)
|
||||
try:
|
||||
import psutil
|
||||
MEMORY_MONITOR_AVAILABLE = True
|
||||
except ImportError:
|
||||
MEMORY_MONITOR_AVAILABLE = False
|
||||
print("⚠️ NanoBananaPro: psutil 不可用,内存监控功能禁用")
|
||||
|
||||
# ============================================================================
|
||||
# 调试日志配置
|
||||
# ============================================================================
|
||||
# 是否启用调试日志(打印完整的 API 响应内容)
|
||||
# 设置为 True 以启用调试日志,False 以禁用
|
||||
DEBUG_LOG_ENABLED = True
|
||||
# 是否启用请求体日志(打印发送给 API 的请求体,base64 图片数据将自动截断)
|
||||
# 设置为 True 以启用请求体日志,False 以禁用
|
||||
REQUEST_LOG_ENABLED = True
|
||||
# ============================================================================
|
||||
|
||||
_NODE = "Nano Banana Pro"
|
||||
|
||||
|
||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
||||
"""
|
||||
将 PIL Image 列表转换为 ComfyUI tensor,安全处理多张不同尺寸的情况。
|
||||
|
||||
策略:
|
||||
- 以像素数最大的图尺寸为基准
|
||||
- 只输出与最大尺寸相同的图,其余较小的图丢弃
|
||||
"""
|
||||
if not images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
return pil_to_tensor([placeholder])
|
||||
|
||||
base_size = max(images, key=lambda img: img.size[0] * img.size[1]).size
|
||||
matched = [img for img in images if img.size == base_size]
|
||||
skipped = [img for img in images if img.size != base_size]
|
||||
|
||||
if skipped:
|
||||
sizes_str = ", ".join(f"{img.size[0]}×{img.size[1]}" for img in skipped)
|
||||
print(
|
||||
f"{node_label}: 丢弃 {len(skipped)} 张较小尺寸的图 ({sizes_str}),"
|
||||
f"仅输出最大尺寸 {base_size[0]}×{base_size[1]} 的 {len(matched)} 张"
|
||||
)
|
||||
|
||||
return pil_to_tensor(matched)
|
||||
|
||||
|
||||
class NanoBananaPro:
|
||||
"""
|
||||
Nano Banana Pro 节点
|
||||
|
||||
功能:
|
||||
- 文生图:基于提示词生成图像
|
||||
- 图生图:基于输入图像和提示词生成新图像
|
||||
- 批量生成:支持并发生成多张图像
|
||||
|
||||
注意:
|
||||
- 支持的模型列表从 models_config.py 动态加载
|
||||
- 要添加/禁用模型,请编辑 models_config.py 文件
|
||||
"""
|
||||
|
||||
# 支持的模型列表(从配置文件动态加载)
|
||||
MODELS = None # 将在 INPUT_TYPES 中动态获取
|
||||
|
||||
# 支持的宽高比列表(全量:所有启用模型的并集,动态加载)
|
||||
# 实际渲染时通过 get_all_supported_aspect_ratios() 获取
|
||||
ASPECT_RATIOS = [
|
||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
||||
"1:4", "4:1", "1:8", "8:1"
|
||||
]
|
||||
|
||||
# 支持的分辨率列表(全量兜底,实际由 get_all_supported_resolutions() 动态生成)
|
||||
RESOLUTIONS = ["512px", "1K", "2K", "4K"]
|
||||
|
||||
def __init__(self):
|
||||
"""初始化节点"""
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
|
||||
ComfyUI 节点规范:
|
||||
- required: 必选参数
|
||||
- optional: 可选参数
|
||||
"""
|
||||
# 从配置文件动态获取启用的模型列表
|
||||
enabled_models = get_enabled_models()
|
||||
|
||||
# 如果没有启用的模型,使用空列表(会导致节点不可用,提示用户配置)
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
||||
|
||||
# 动态获取所有启用模型支持的宽高比(去重合并)
|
||||
all_aspect_ratios = get_all_supported_aspect_ratios()
|
||||
if not all_aspect_ratios:
|
||||
all_aspect_ratios = cls.ASPECT_RATIOS
|
||||
|
||||
# 动态获取所有启用模型支持的分辨率(去重合并)
|
||||
all_resolutions = get_all_supported_resolutions()
|
||||
if not all_resolutions:
|
||||
all_resolutions = cls.RESOLUTIONS
|
||||
|
||||
# 创建9个独立的图像输入
|
||||
optional_inputs = {}
|
||||
for i in range(1, 10): # 1-9
|
||||
optional_inputs[f"参考图{i}"] = ("IMAGE",)
|
||||
|
||||
optional_inputs["代理端口(如7897)"] = ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"placeholder": "本地代理端口,如 7897(Clash Verge)或 10808(v2rayN),留空不使用"
|
||||
})
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "一个中国女子的OOTD",
|
||||
"multiline": True
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0]
|
||||
}),
|
||||
"宽高比": (all_aspect_ratios, {
|
||||
"default": "1:1"
|
||||
}),
|
||||
"分辨率": (all_resolutions, {
|
||||
"default": "2K"
|
||||
}),
|
||||
"生图数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 1000,
|
||||
"step": 1
|
||||
}),
|
||||
"谷歌搜索(联网)": (["关闭", "打开"], {
|
||||
"default": "关闭"
|
||||
}),
|
||||
"图片搜索(联网)": (["关闭", "打开"], {
|
||||
"default": "关闭"
|
||||
}),
|
||||
"返回格式": (["url", "base64"], {
|
||||
"default": "url"
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff
|
||||
})
|
||||
},
|
||||
"optional": optional_inputs
|
||||
}
|
||||
|
||||
# 返回值类型
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("输出图像",)
|
||||
|
||||
# 导入 ComfyUI 的文件夹路径管理
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
# 执行函数名
|
||||
FUNCTION = "generate"
|
||||
|
||||
# 节点分类
|
||||
CATEGORY = "image/generation"
|
||||
|
||||
def resize_to_megapixels(
|
||||
self,
|
||||
image: Image.Image,
|
||||
target_megapixels: float
|
||||
) -> Image.Image:
|
||||
"""
|
||||
将图像缩放到指定的总像素数,保持纵横比
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
target_megapixels: 目标像素数(百万像素)
|
||||
|
||||
Returns:
|
||||
缩放后的 PIL Image
|
||||
|
||||
Example:
|
||||
>>> resized = self.resize_to_megapixels(img, 2.0) # 缩放到2百万像素
|
||||
"""
|
||||
# 计算当前像素数
|
||||
current_pixels = image.width * image.height
|
||||
target_pixels = int(target_megapixels * 1_000_000)
|
||||
|
||||
# 如果当前像素数已经接近目标,则不缩放
|
||||
if abs(current_pixels - target_pixels) / target_pixels < 0.05:
|
||||
return image
|
||||
|
||||
# 计算缩放比例
|
||||
scale = (target_pixels / current_pixels) ** 0.5
|
||||
|
||||
# 计算新尺寸
|
||||
new_width = int(image.width * scale)
|
||||
new_height = int(image.height * scale)
|
||||
|
||||
# 确保至少为1像素
|
||||
new_width = max(1, new_width)
|
||||
new_height = max(1, new_height)
|
||||
|
||||
# 使用 Lanczos 重采样
|
||||
resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
return resized_image
|
||||
|
||||
def validate_inputs(
|
||||
self,
|
||||
images: Optional[torch.Tensor],
|
||||
batch_size: int
|
||||
) -> None:
|
||||
"""
|
||||
验证输入参数
|
||||
|
||||
Args:
|
||||
images: 输入图像张量(可选)
|
||||
batch_size: 批次大小
|
||||
|
||||
Raises:
|
||||
ValueError: 如果输入参数不合法
|
||||
"""
|
||||
# 检查图像数量
|
||||
if images is not None:
|
||||
num_images = images.shape[0]
|
||||
if num_images > 14:
|
||||
raise ValueError(
|
||||
f"输入图像数量 {num_images} 超过限制 14 张,请减少输入图像数量"
|
||||
)
|
||||
|
||||
# 检查批次大小
|
||||
if batch_size < 1 or batch_size > 1000:
|
||||
raise ValueError(
|
||||
f"批次大小 {batch_size} 超出范围 [1, 1000]"
|
||||
)
|
||||
|
||||
async def _generate_single_task(
|
||||
self,
|
||||
session: aiohttp.ClientSession,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: List[Image.Image],
|
||||
output_folder: str,
|
||||
global_task_index: int,
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
save_to_disk: bool = True,
|
||||
image_format: str = "url",
|
||||
) -> dict:
|
||||
"""执行单个生成任务"""
|
||||
result = {
|
||||
"global_task_index": global_task_index,
|
||||
"prompt": prompt,
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"saved_files": [],
|
||||
"output_images": [],
|
||||
"error": None
|
||||
}
|
||||
|
||||
try:
|
||||
gen_result = await self.client.generate_single_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images if images else None,
|
||||
session=session,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
image_format=image_format,
|
||||
)
|
||||
if gen_result:
|
||||
images_list, _ = gen_result
|
||||
if save_to_disk:
|
||||
for gen_img in images_list:
|
||||
output_path = generate_timestamp_filename(
|
||||
output_folder=output_folder,
|
||||
extension=".png"
|
||||
)
|
||||
save_image(gen_img, output_path)
|
||||
result["saved_files"].append(output_path)
|
||||
gen_img = None
|
||||
else:
|
||||
result["output_images"] = images_list
|
||||
|
||||
result["success"] = True
|
||||
result["generated_count"] = len(images_list)
|
||||
except Exception as e:
|
||||
result["error"] = str(e)
|
||||
|
||||
return result
|
||||
|
||||
async def _process_batch_async(
|
||||
self,
|
||||
prompts: List[str],
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images_per_prompt: int,
|
||||
input_images: List[Image.Image],
|
||||
output_folder: str,
|
||||
pbar=None,
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
save_to_disk: bool = True,
|
||||
image_format: str = "url",
|
||||
) -> List[dict]:
|
||||
"""异步批量处理:每个提示词独立调用 API"""
|
||||
# 构建任务列表:(prompt, sub_index) 用于 images_per_prompt > 1 的情况
|
||||
tasks_def = []
|
||||
for p_idx, prompt in enumerate(prompts):
|
||||
for sub_idx in range(images_per_prompt):
|
||||
tasks_def.append((p_idx, sub_idx, prompt))
|
||||
|
||||
total_tasks = len(tasks_def)
|
||||
num_prompts = len(prompts)
|
||||
|
||||
max_concurrent = 50
|
||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
||||
|
||||
all_results = []
|
||||
completed = 0
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
for batch_idx in range(num_batches):
|
||||
start_idx = batch_idx * max_concurrent
|
||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||||
|
||||
tasks = []
|
||||
for i in range(start_idx, end_idx):
|
||||
_, _, prompt = tasks_def[i]
|
||||
task = asyncio.create_task(
|
||||
self._generate_single_task(
|
||||
session=session,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_images,
|
||||
output_folder=output_folder,
|
||||
global_task_index=i,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
save_to_disk=save_to_disk,
|
||||
image_format=image_format,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
batch_results = []
|
||||
for coro in asyncio.as_completed(tasks):
|
||||
result_data = None
|
||||
try:
|
||||
result = await coro
|
||||
if isinstance(result, Exception):
|
||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": [], "prompt": ""}
|
||||
else:
|
||||
result_data = result
|
||||
batch_results.append(result_data)
|
||||
except Exception as e:
|
||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": [], "prompt": ""}
|
||||
batch_results.append(result_data)
|
||||
|
||||
completed += 1
|
||||
prompt_snippet = (result_data.get("prompt", "") or "")[:30]
|
||||
|
||||
if result_data and result_data.get("success", False):
|
||||
success_count += 1
|
||||
count = result_data.get("generated_count", 1)
|
||||
print(f"Nano Banana Pro: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✓成功({count}张)")
|
||||
else:
|
||||
fail_count += 1
|
||||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
||||
print(f"Nano Banana Pro: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
all_results.extend(batch_results)
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
return all_results
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str,
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
生图数量: int,
|
||||
seed: int,
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""
|
||||
生成图像
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
模型: 模型名称
|
||||
宽高比: 宽高比
|
||||
分辨率: 分辨率
|
||||
生图数量: 批次大小
|
||||
seed: 随机种子
|
||||
**kwargs: 搜索开关(谷歌搜索(联网)/ 图片搜索(联网))及动态参考图输入 (参考图1-9)
|
||||
注:两个搜索参数名含全角括号,不能作为 Python 形参,从 kwargs 中提取
|
||||
|
||||
注意:
|
||||
调试日志功能已移至文件顶部配置,通过修改 DEBUG_LOG_ENABLED 常量控制
|
||||
|
||||
Returns:
|
||||
生成的图像张量 (IMAGE,)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# 从 kwargs 提取搜索参数(界面显示为「关闭/打开」,转为 bool 供调用)
|
||||
enable_grounding: bool = (kwargs.pop("谷歌搜索(联网)", "关闭") == "打开")
|
||||
enable_image_search: bool = (kwargs.pop("图片搜索(联网)", "关闭") == "打开")
|
||||
proxy_port: str = kwargs.pop("代理端口(如7897)", "")
|
||||
image_format: str = kwargs.pop("返回格式", "url")
|
||||
|
||||
# 创建 ComfyUI 原生进度条
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(生图数量)
|
||||
|
||||
try:
|
||||
# 设置随机种子(用于本地随机操作)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed % (2**32))
|
||||
|
||||
# 内存监控初始化
|
||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
||||
import psutil
|
||||
process = psutil.Process()
|
||||
initial_memory = process.memory_info().rss / 1024 / 1024
|
||||
print(f"Nano Banana Pro: 初始内存使用: {initial_memory:.1f} MB")
|
||||
|
||||
# 初始化 API 客户端
|
||||
if self.client is None:
|
||||
try:
|
||||
self.client = GeminiAPIClient()
|
||||
except ValueError as e:
|
||||
raise ValueError(f"初始化失败: {str(e)}")
|
||||
|
||||
# 注入代理设置(每次执行都刷新,支持用户中途修改端口)
|
||||
self.client.proxy_url = GeminiAPIClient.build_proxy_url(proxy_port)
|
||||
if self.client.proxy_url:
|
||||
print(f"Nano Banana Pro: 已启用代理加速 → {self.client.proxy_url}")
|
||||
|
||||
# 校验分辨率与模型的兼容性
|
||||
supported_resolutions = get_model_supported_resolutions(模型)
|
||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
||||
raise ValueError(
|
||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容!\n"
|
||||
f"该模型支持的分辨率:{', '.join(supported_resolutions)}"
|
||||
)
|
||||
|
||||
# 校验宽高比与模型的兼容性
|
||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
||||
if supported_ratios and 宽高比 not in supported_ratios:
|
||||
raise ValueError(
|
||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容!\n"
|
||||
f"该模型支持的宽高比:{', '.join(supported_ratios)}"
|
||||
)
|
||||
|
||||
# 校验图片搜索(联网)与模型的兼容性
|
||||
# 仅 nano-banana-2-限时特价 和 gemini-3.1-flash-image-preview 支持图片搜索
|
||||
IMAGE_SEARCH_UNSUPPORTED_MODELS = ["nano-banana-pro-次卡", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview"]
|
||||
if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS:
|
||||
raise ValueError(
|
||||
f"模型 \"{模型}\" 不支持【图片搜索(联网)】功能!"
|
||||
f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用"
|
||||
)
|
||||
|
||||
# 收集独立输入的参考图
|
||||
input_images = []
|
||||
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 张,请减少输入图像数量"
|
||||
)
|
||||
|
||||
# 解析批量提示词
|
||||
batch_prompts = parse_batch_prompts(prompt)
|
||||
|
||||
# 打印首行概览
|
||||
# 图片搜索(联网)开启时隐含谷歌搜索接地,与客户端请求逻辑保持一致
|
||||
grounding_str = ""
|
||||
if enable_image_search:
|
||||
grounding_str = " | 谷歌图片搜索接地"
|
||||
elif enable_grounding:
|
||||
grounding_str = " | 谷歌搜索接地"
|
||||
|
||||
if batch_prompts:
|
||||
# 批量提示词模式
|
||||
num_prompts = len(batch_prompts)
|
||||
total_images = num_prompts * 生图数量
|
||||
mode_str = f"批量提示词模式 ({num_prompts}个提示词)"
|
||||
if input_images:
|
||||
mode_str += f" (输入{len(input_images)}张)"
|
||||
print(f"Nano Banana Pro: {mode_str} | {分辨率} {宽高比} | 共{total_images}张{grounding_str}")
|
||||
|
||||
# 大批量警告
|
||||
if total_images > 100:
|
||||
print(f"⚠️ Nano Banana Pro: 警告!批量生成 {total_images} 张图片,内存占用可能较高")
|
||||
print(f"⚠️ 建议:分批执行或减少生图数量")
|
||||
else:
|
||||
# 单提示词模式
|
||||
mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式"
|
||||
print(f"Nano Banana Pro: {mode_str} | {分辨率} {宽高比} | {生图数量}张{grounding_str}")
|
||||
|
||||
# 大批量警告
|
||||
if 生图数量 > 100:
|
||||
print(f"⚠️ Nano Banana Pro: 警告!批量生成 {生图数量} 张图片,内存占用可能较高")
|
||||
print(f"⚠️ 建议:分批执行或减少生图数量")
|
||||
|
||||
# 统计变量
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
# 进度回调 - 打印错误信息并更新进度条,添加内存监控
|
||||
def progress_callback(current, total, success, error_msg=None):
|
||||
nonlocal success_count, fail_count
|
||||
if success:
|
||||
success_count += 1
|
||||
else:
|
||||
fail_count += 1
|
||||
|
||||
# 更新 ComfyUI 原生进度条
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
# 内存监控(每完成10个任务检查一次)
|
||||
if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0:
|
||||
import gc
|
||||
gc.collect() # 强制垃圾回收
|
||||
current_memory = process.memory_info().rss / 1024 / 1024
|
||||
memory_increase = current_memory - initial_memory
|
||||
print(f"Nano Banana Pro: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)")
|
||||
|
||||
# 内存警告阈值(2GB)
|
||||
if current_memory > 2000:
|
||||
print(f"⚠️ Nano Banana Pro: 内存使用过高!建议减少生图数量或分批执行")
|
||||
|
||||
# 根据是否有批量提示词选择生成模式
|
||||
if batch_prompts:
|
||||
num_prompts = len(batch_prompts)
|
||||
total_images = num_prompts * 生图数量
|
||||
|
||||
# ===== 批量提示词模式:异步并发,内存输出 =====
|
||||
if pbar is not None:
|
||||
pbar = ProgressBar(total_images)
|
||||
|
||||
def run_async_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
self._process_batch_async(
|
||||
prompts=batch_prompts,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images_per_prompt=生图数量,
|
||||
input_images=input_images,
|
||||
output_folder="",
|
||||
pbar=pbar,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
save_to_disk=False,
|
||||
image_format=image_format,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async_in_thread)
|
||||
try:
|
||||
results = future.result(timeout=900)
|
||||
except TimeoutError:
|
||||
raise RuntimeError("任务执行超时(900秒),请减少提示词数量或检查网络连接")
|
||||
|
||||
# 统计结果
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||||
|
||||
print(f"完成!总耗时 {time_str} | 成功: {success_count}/{total_images} | 失败: {fail_count}")
|
||||
|
||||
# 失败详情
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
if failed_results:
|
||||
for fr in failed_results:
|
||||
idx = fr.get("global_task_index", -1) + 1
|
||||
prompt_snippet = (fr.get("prompt", "") or "")[:30]
|
||||
error_msg = fr.get("error", "未知错误")
|
||||
print(f" 失败 #{idx}: {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → {error_msg}")
|
||||
|
||||
# 收集内存中的图像
|
||||
output_images = []
|
||||
for r in results:
|
||||
output_images.extend(r.get("output_images", []))
|
||||
|
||||
if not output_images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
output_images = [placeholder]
|
||||
|
||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
else:
|
||||
# 单提示词模式
|
||||
if 生图数量 == 1:
|
||||
# 单张:同步生成
|
||||
generated_images = self.client.generate_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
batch_size=1,
|
||||
images=input_images,
|
||||
progress_callback=progress_callback,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
image_format=image_format,
|
||||
)
|
||||
else:
|
||||
# 多张:异步并发,内存输出
|
||||
|
||||
if pbar is not None:
|
||||
pbar = ProgressBar(生图数量)
|
||||
|
||||
def run_async_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
self._process_batch_async(
|
||||
prompts=[prompt],
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images_per_prompt=生图数量,
|
||||
input_images=input_images,
|
||||
output_folder="",
|
||||
pbar=pbar,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
save_to_disk=False,
|
||||
image_format=image_format,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async_in_thread)
|
||||
try:
|
||||
results = future.result(timeout=900)
|
||||
except TimeoutError:
|
||||
raise RuntimeError("任务执行超时(900秒),请减少生图数量或检查网络连接")
|
||||
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||||
print(f"完成!总耗时 {time_str} | 成功: {success_count}/{生图数量} | 失败: {fail_count}")
|
||||
|
||||
# 失败详情
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
if failed_results:
|
||||
for fr in failed_results:
|
||||
idx = fr.get("global_task_index", -1) + 1
|
||||
error_msg = fr.get("error", "未知错误")
|
||||
print(f" 失败 #{idx}: {prompt[:30]}{'...' if len(prompt) >= 30 else ''} → {error_msg}")
|
||||
|
||||
# 收集内存中的图像
|
||||
output_images = []
|
||||
for r in results:
|
||||
output_images.extend(r.get("output_images", []))
|
||||
|
||||
if not output_images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
output_images = [placeholder]
|
||||
|
||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
|
||||
|
||||
# 优化:限制输出图片数量,避免内存爆炸
|
||||
max_output_images = 20 # 最多输出20张图片到ComfyUI
|
||||
|
||||
if len(generated_images) > max_output_images:
|
||||
print(f"Nano Banana Pro: 生成 {len(generated_images)} 张图片,限制输出前 {max_output_images} 张到ComfyUI")
|
||||
output_images = generated_images[:max_output_images]
|
||||
else:
|
||||
output_images = generated_images
|
||||
|
||||
# 转换输出图像
|
||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||
|
||||
# 计算耗时并打印最终统计
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed < 1:
|
||||
time_str = f"{elapsed:.3f}s"
|
||||
else:
|
||||
time_str = f"{elapsed:.2f}s"
|
||||
|
||||
# 打印最终汇总
|
||||
if fail_count > 0:
|
||||
print(f"完成!总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张")
|
||||
else:
|
||||
print(f"完成!总耗时 {time_str} | 成功 {len(generated_images)}张")
|
||||
|
||||
# 最终内存清理
|
||||
import gc
|
||||
gc.collect()
|
||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
||||
final_memory = process.memory_info().rss / 1024 / 1024
|
||||
print(f"Nano Banana Pro: 最终内存使用: {final_memory:.1f} MB")
|
||||
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as e:
|
||||
# 检测是否为授权错误
|
||||
if str(e) == "未授权!":
|
||||
print("请联系作者授权后方可使用!")
|
||||
raise ValueError("未授权!") from None
|
||||
raise ValueError(str(e)) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
except Exception as e:
|
||||
raise type(e)(str(e)) from None
|
||||
|
||||
finally:
|
||||
# 查询余额
|
||||
if self.client is not None:
|
||||
try:
|
||||
balance_data = self.client.query_balance_sync()
|
||||
balance_info = self.client.format_balance_info(balance_data)
|
||||
print(f"Nano Banana Pro: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 最终内存清理
|
||||
import gc
|
||||
gc.collect()
|
||||
+26
-27
@@ -30,7 +30,8 @@ from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||||
from ..utils.file_utils import load_images_from_folder, pair_images_by_name, pair_images_cartesian
|
||||
from ..utils.config import get_api_key_or_raise
|
||||
from ..utils.config import get_api_key_or_raise, NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.http_error import async_request_with_retry, extract_structured_error_message, get_friendly_message
|
||||
from ..models_config import (
|
||||
get_enabled_async_models,
|
||||
get_model_provider,
|
||||
@@ -169,15 +170,16 @@ class NanoBananaV2:
|
||||
"default": "一个中国女子的OOTD",
|
||||
"multiline": True
|
||||
}),
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
|
||||
"模型": (models, {"default": models[0]}),
|
||||
"宽高比": (all_aspect_ratios, {"default": "1:1"}),
|
||||
"宽高比": (["智能"] + all_aspect_ratios, {"default": "智能"}),
|
||||
"分辨率": (all_resolutions, {"default": "2K"}),
|
||||
"生图数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 9,
|
||||
"step": 1
|
||||
})
|
||||
}),
|
||||
},
|
||||
"optional": optional
|
||||
}
|
||||
@@ -239,6 +241,10 @@ class NanoBananaV2:
|
||||
@staticmethod
|
||||
def _friendly_error(error_msg: str) -> str:
|
||||
"""将上游错误转化为用户友好的提示"""
|
||||
structured_message = extract_structured_error_message(error_msg)
|
||||
if structured_message and structured_message != error_msg:
|
||||
error_msg = structured_message
|
||||
|
||||
if "No available channel for model" in error_msg:
|
||||
return (
|
||||
"当前分组下模型不可用,请检查分组是否正确。"
|
||||
@@ -299,13 +305,11 @@ class NanoBananaV2:
|
||||
print(f"[异步提交] URL: {url}")
|
||||
print(f"[异步提交] 请求体: {json.dumps(_log_body, ensure_ascii=False)[:500]}")
|
||||
|
||||
async with session.post(url, json=request_body, headers=headers, proxy=provider.proxy_url) as response:
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
if not error_text.strip():
|
||||
error_text = "(服务器未返回错误详情)"
|
||||
raise RuntimeError(f"提交任务失败 ({response.status}): {error_text}")
|
||||
data = await response.json()
|
||||
resp = await async_request_with_retry(
|
||||
session, "POST", url, json=request_body, headers=headers,
|
||||
proxy=provider.proxy_url, prefix="异步提交: "
|
||||
)
|
||||
data = await resp.json()
|
||||
|
||||
if DEBUG_LOG_ENABLED:
|
||||
import json
|
||||
@@ -343,16 +347,16 @@ class NanoBananaV2:
|
||||
error_text = await response.text()
|
||||
if not error_text.strip():
|
||||
error_text = "(服务器未返回错误详情)"
|
||||
raise RuntimeError(f"查询任务失败 ({response.status}): {error_text}")
|
||||
raise RuntimeError(f"查询任务失败: {get_friendly_message(response.status, error_text)}")
|
||||
|
||||
result = await response.json()
|
||||
status = provider.extract_status(result)
|
||||
|
||||
# 提取进度并回调(封顶 1.0 防止异常值导致进度条溢出)
|
||||
if on_progress and status in ("SUBMITTED", "IN_PROGRESS"):
|
||||
# 提取进度并回调;运行中状态不显示 100%,只有 SUCCESS 才补满。
|
||||
if on_progress and status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"):
|
||||
p = provider.extract_progress(result)
|
||||
if p is not None:
|
||||
p = min(p, 1.0)
|
||||
p = min(p, 0.99)
|
||||
if p > last_progress:
|
||||
on_progress(p - last_progress)
|
||||
last_progress = p
|
||||
@@ -374,7 +378,7 @@ class NanoBananaV2:
|
||||
print(f"{self.NODE_LABEL}: [轮询] FAILURE 但无错误信息,原始响应: {json.dumps(result, ensure_ascii=False)[:500]}")
|
||||
friendly_msg = self._friendly_error(error_msg)
|
||||
raise RuntimeError(f"任务失败: {friendly_msg}")
|
||||
elif status in ("SUBMITTED", "IN_PROGRESS"):
|
||||
elif status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"):
|
||||
# 分段 sleep,每 0.1 秒检查一次取消信号
|
||||
sleep_iterations = int(_POLL_INTERVAL / _INTERRUPT_CHECK_INTERVAL)
|
||||
for _ in range(sleep_iterations):
|
||||
@@ -406,10 +410,7 @@ class NanoBananaV2:
|
||||
"error": None,
|
||||
}
|
||||
|
||||
contributed = [0.0] # mutable container,追踪本任务已贡献的 pbar 进度
|
||||
|
||||
def _track_progress(delta):
|
||||
contributed[0] += delta
|
||||
if on_progress:
|
||||
on_progress(delta)
|
||||
|
||||
@@ -435,15 +436,9 @@ class NanoBananaV2:
|
||||
result["request_time"] = request_time
|
||||
result["download_time"] = download_time
|
||||
except InterruptProcessingException:
|
||||
# 用户取消:补齐进度后向上传播,不吞掉
|
||||
if contributed[0] < 1.0 and on_progress:
|
||||
on_progress(1.0 - contributed[0])
|
||||
raise
|
||||
except Exception as e:
|
||||
result["error"] = str(e) or f"{type(e).__name__}(无错误详情)"
|
||||
# 失败也补齐 1.0 进度,保证进度条总数正确
|
||||
if contributed[0] < 1.0 and on_progress:
|
||||
on_progress(1.0 - contributed[0])
|
||||
|
||||
return result
|
||||
|
||||
@@ -541,6 +536,7 @@ class NanoBananaV2:
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
生图数量: int,
|
||||
网络线路: str = "全球加速",
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""生成图像(异步模式)"""
|
||||
@@ -557,6 +553,7 @@ class NanoBananaV2:
|
||||
proxy_url = BaseAsyncImageProvider.build_proxy_url(proxy_port)
|
||||
provider = self._get_provider(模型, proxy_url=proxy_url, api_key_override=api_key_override)
|
||||
provider.image_compression = "webp" if 图片质量 == "日常" else None
|
||||
provider._route_base_url = get_base_url_by_route(网络线路)
|
||||
|
||||
if proxy_url:
|
||||
print(f"{self.NODE_LABEL}: 已启用代理加速 -> {proxy_url}")
|
||||
@@ -590,7 +587,7 @@ class NanoBananaV2:
|
||||
|
||||
# 运行时验证宽高比
|
||||
supported_ratios = provider.get_model_aspect_ratios(模型)
|
||||
if supported_ratios and 宽高比 not in supported_ratios:
|
||||
if 宽高比 != "智能" and supported_ratios and 宽高比 not in supported_ratios:
|
||||
raise ValueError(
|
||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容!\n"
|
||||
f"该模型支持的宽高比:{', '.join(supported_ratios)}"
|
||||
@@ -702,7 +699,7 @@ class NanoBananaV2:
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
except Exception as e:
|
||||
raise type(e)(str(e)) from None
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
finally:
|
||||
# 查询并打印余额
|
||||
@@ -977,6 +974,7 @@ class NanoBananaV2Batch(NanoBananaV2):
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
生图数量: int = 1,
|
||||
网络线路: str = "全球加速",
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""生成图像(异步模式 - 批量版:全并发 + 即时落盘)"""
|
||||
@@ -997,6 +995,7 @@ class NanoBananaV2Batch(NanoBananaV2):
|
||||
proxy_url = BaseAsyncImageProvider.build_proxy_url(proxy_port)
|
||||
provider = self._get_provider(模型, proxy_url=proxy_url, api_key_override=api_key_override)
|
||||
provider.image_compression = "webp" if 图片质量 == "日常" else None
|
||||
provider._route_base_url = get_base_url_by_route(网络线路)
|
||||
|
||||
if proxy_url:
|
||||
print(f"{self.NODE_LABEL}: 已启用代理加速 -> {proxy_url}")
|
||||
@@ -1219,7 +1218,7 @@ class NanoBananaV2Batch(NanoBananaV2):
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
except Exception as e:
|
||||
raise type(e)(str(e)) from None
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
finally:
|
||||
try:
|
||||
|
||||
@@ -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,)
|
||||
+2
-185
@@ -1,79 +1,18 @@
|
||||
"""
|
||||
图像元数据去除节点
|
||||
替代 ComfyUI 原生"保存图像"节点,保存时不写入提示词、工作流等 AI 元数据
|
||||
|
||||
提供两种节点:
|
||||
1. SaveCleanImage - 接收 IMAGE 张量,去除元数据后直接保存到 output 目录
|
||||
2. BatchCleanMetadata - 指定文件夹路径,批量去除已有图片中的元数据
|
||||
提供批量去除已有图片中元数据的功能
|
||||
"""
|
||||
|
||||
import os
|
||||
from datetime import datetime
|
||||
import random
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..utils.file_utils import _get_port_suffix
|
||||
|
||||
# 尝试导入 ComfyUI 的 folder_paths
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
# 支持的图片格式
|
||||
SUPPORTED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.webp', '.bmp', '.tiff', '.tif'}
|
||||
|
||||
|
||||
def _get_output_dir() -> str:
|
||||
"""
|
||||
获取 ComfyUI output 目录
|
||||
|
||||
Returns:
|
||||
output 目录的绝对路径
|
||||
"""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
return folder_paths.get_output_directory()
|
||||
# fallback: 相对于插件目录推断
|
||||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
return os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output")
|
||||
|
||||
|
||||
def _get_next_counter(directory: str, prefix: str) -> int:
|
||||
"""
|
||||
扫描目录,获取下一个可用的文件计数器
|
||||
|
||||
Args:
|
||||
directory: 目标目录
|
||||
prefix: 文件名前缀
|
||||
|
||||
Returns:
|
||||
下一个计数器值
|
||||
"""
|
||||
if not os.path.exists(directory):
|
||||
return 1
|
||||
|
||||
if prefix:
|
||||
pattern = re.compile(rf'^{re.escape(prefix)}_(\d+)')
|
||||
else:
|
||||
pattern = re.compile(rf'^(\d+)\.')
|
||||
max_counter = 0
|
||||
|
||||
for f in os.listdir(directory):
|
||||
m = pattern.match(f)
|
||||
if m:
|
||||
counter = int(m.group(1))
|
||||
max_counter = max(max_counter, counter)
|
||||
|
||||
return max_counter + 1
|
||||
|
||||
|
||||
def _save_image_clean(image: Image.Image, path: str, fmt: str = None, quality: int = 95) -> None:
|
||||
"""
|
||||
保存图像,不包含任何元数据
|
||||
@@ -127,129 +66,7 @@ def _save_image_clean(image: Image.Image, path: str, fmt: str = None, quality: i
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 节点 1:保存干净图像
|
||||
# ============================================================================
|
||||
|
||||
class SaveCleanImage:
|
||||
"""
|
||||
保存干净图像节点(不含元数据)
|
||||
|
||||
功能:
|
||||
- 接收 IMAGE 张量(支持单图和批次)
|
||||
- 去除所有元数据后保存到 ComfyUI/output 目录
|
||||
- 文件名自动添加 nometa 标识,方便辨认
|
||||
- 支持 PNG/JPEG/WEBP 格式
|
||||
- 作为终端节点,替代 ComfyUI 原生"保存图像"节点
|
||||
|
||||
使用场景:
|
||||
- 生图完成后,直接保存不含 AI 元数据的干净图像
|
||||
- 分享图像时不暴露提示词和工作流
|
||||
"""
|
||||
|
||||
SAVE_FORMATS = ["PNG", "JPEG", "WEBP"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
|
||||
Returns:
|
||||
输入参数配置字典
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"图像": ("IMAGE",),
|
||||
"文件名前缀": ("STRING", {"default": "ComfyUI_nometa"}),
|
||||
"保存格式": (cls.SAVE_FORMATS, {"default": "PNG"}),
|
||||
},
|
||||
"optional": {
|
||||
"JPEG/WEBP质量": ("INT", {
|
||||
"default": 95,
|
||||
"min": 1,
|
||||
"max": 100,
|
||||
"step": 1
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "save_clean"
|
||||
CATEGORY = "image"
|
||||
|
||||
DESCRIPTION = (
|
||||
"保存干净图像(不含元数据)。\n"
|
||||
"替代 ComfyUI 原生'保存图像'节点,保存时不写入提示词、工作流等 AI 元数据。\n"
|
||||
"文件保存到 ComfyUI/output 目录。"
|
||||
)
|
||||
|
||||
def save_clean(
|
||||
self,
|
||||
图像: torch.Tensor,
|
||||
文件名前缀: str = "ComfyUI_nometa",
|
||||
保存格式: str = "PNG",
|
||||
**kwargs
|
||||
) -> dict:
|
||||
"""
|
||||
去除元数据并保存图像
|
||||
|
||||
Args:
|
||||
图像: ComfyUI 图像张量 [B, H, W, C]
|
||||
文件名前缀: 保存文件名前缀
|
||||
保存格式: 图像格式(PNG/JPEG/WEBP)
|
||||
**kwargs: 可选参数(JPEG/WEBP质量)
|
||||
|
||||
Returns:
|
||||
UI 结果字典,包含保存的图像信息用于前端预览
|
||||
"""
|
||||
quality = kwargs.get("JPEG/WEBP质量", 95)
|
||||
|
||||
output_dir = _get_output_dir()
|
||||
port_suffix = _get_port_suffix()
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# 格式与扩展名映射
|
||||
ext_map = {"PNG": ".png", "JPEG": ".jpg", "WEBP": ".webp"}
|
||||
ext = ext_map.get(保存格式, ".png")
|
||||
|
||||
# 转换为 PIL 图像
|
||||
pil_images = tensor_to_pil(图像)
|
||||
|
||||
results = []
|
||||
saved_paths = []
|
||||
for img in pil_images:
|
||||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
ms = random.randint(0, 999)
|
||||
|
||||
while True:
|
||||
if 文件名前缀:
|
||||
filename = f"{文件名前缀}_{ts}_{ms:03d}{port_suffix}{ext}"
|
||||
else:
|
||||
filename = f"{ts}_{ms:03d}{port_suffix}{ext}"
|
||||
filepath = os.path.join(output_dir, filename)
|
||||
if not os.path.exists(filepath):
|
||||
break
|
||||
ms = (ms + 1) % 1000
|
||||
|
||||
_save_image_clean(img, filepath, fmt=保存格式, quality=quality)
|
||||
|
||||
results.append({
|
||||
"filename": filename,
|
||||
"subfolder": "",
|
||||
"type": "output"
|
||||
})
|
||||
saved_paths.append(filepath)
|
||||
|
||||
# 打印详细日志,方便用户定位保存的文件
|
||||
print(f"保存干净图像: 已保存 {len(pil_images)} 张无元数据图像 (格式: {保存格式})")
|
||||
for p in saved_paths:
|
||||
print(f" → {p}")
|
||||
|
||||
return {"ui": {"images": results}}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 节点 2:批量去除元数据
|
||||
# 批量去除元数据
|
||||
# ============================================================================
|
||||
|
||||
class BatchCleanMetadata:
|
||||
|
||||
@@ -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,)
|
||||
+55
-8
@@ -4,7 +4,9 @@ Seedance 视频生成节点
|
||||
- Seedance: 文生视频 / 图生视频 / 首尾帧生视频(根据图片输入自动切换模式)
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
@@ -13,8 +15,9 @@ import torch
|
||||
|
||||
from ..clients.seedance_client import SeedanceClient
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64, pil_to_tensor
|
||||
from ..utils.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
|
||||
|
||||
@@ -27,6 +30,9 @@ _MODELS = [
|
||||
|
||||
_RESOLUTIONS = ["720p", "1080p", "480p"]
|
||||
|
||||
_MAX_IMAGE_BYTES = 30 * 1024 * 1024
|
||||
_MAX_REQUEST_BODY_BYTES = 64 * 1024 * 1024
|
||||
|
||||
|
||||
# ── 模型能力判断 ──────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -37,13 +43,45 @@ def _supports_camera_fixed(model: str) -> bool:
|
||||
|
||||
# ── 工具函数 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _tensor_to_base64_url(tensor) -> str:
|
||||
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)
|
||||
b64 = encode_image_to_base64(pil_images[0], format="PNG")
|
||||
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"""
|
||||
@@ -116,6 +154,7 @@ class Seedance:
|
||||
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"],
|
||||
@@ -199,7 +238,7 @@ class Seedance:
|
||||
}
|
||||
|
||||
elif mode == "i2v":
|
||||
first_url = _tensor_to_base64_url(first_image)
|
||||
first_url = _tensor_to_base64_url(first_image, "首帧图片")
|
||||
metadata["content"] = [
|
||||
{
|
||||
"type": "image_url",
|
||||
@@ -216,8 +255,8 @@ class Seedance:
|
||||
}
|
||||
|
||||
else: # flipflop
|
||||
first_url = _tensor_to_base64_url(first_image)
|
||||
last_url = _tensor_to_base64_url(last_image)
|
||||
first_url = _tensor_to_base64_url(first_image, "首帧图片")
|
||||
last_url = _tensor_to_base64_url(last_image, "尾帧图片")
|
||||
metadata["content"] = [
|
||||
{
|
||||
"type": "image_url",
|
||||
@@ -238,10 +277,13 @@ class Seedance:
|
||||
"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)
|
||||
|
||||
@@ -268,6 +310,7 @@ class SeedanceMultiModal:
|
||||
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"],
|
||||
@@ -311,6 +354,7 @@ class SeedanceMultiModal:
|
||||
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)
|
||||
@@ -340,11 +384,11 @@ class SeedanceMultiModal:
|
||||
imgs = ref_images[:9]
|
||||
if len(ref_images) > 9:
|
||||
print(f"[SeedanceMultiModal] 参考图片超过9张,仅取前9张(共{len(ref_images)}张)")
|
||||
for img_tensor in imgs:
|
||||
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)
|
||||
url = _tensor_to_base64_url(img_tensor, f"参考图片{idx}")
|
||||
content.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": url},
|
||||
@@ -410,10 +454,13 @@ class SeedanceMultiModal:
|
||||
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)
|
||||
|
||||
|
||||
+10
-3
@@ -17,7 +17,7 @@ import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
from ..utils.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
|
||||
|
||||
# ============================================================================
|
||||
@@ -26,10 +26,12 @@ from ..utils.file_types import FileList
|
||||
|
||||
SUPPORTED_MODELS = [
|
||||
"gpt-5.5",
|
||||
"gemini-3.1-flash-lite-preview",
|
||||
"gemini-3.1-pro-preview",
|
||||
"deepseek-v4-pro",
|
||||
"doubao-seed-2-0-pro-260215",
|
||||
"claude-opus-4-7",
|
||||
"claude-opus-4-6",
|
||||
"gemini-3.5-flash",
|
||||
"doubao-seed-2.0-pro",
|
||||
]
|
||||
|
||||
# 图片缩放最大尺寸
|
||||
@@ -64,6 +66,9 @@ class UniversalLLMChat:
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"网络线路": (NETWORK_ROUTE_OPTIONS, {
|
||||
"default": "全球加速"
|
||||
}),
|
||||
"模型": (SUPPORTED_MODELS, {
|
||||
"default": SUPPORTED_MODELS[0]
|
||||
}),
|
||||
@@ -368,6 +373,7 @@ class UniversalLLMChat:
|
||||
self,
|
||||
模型: str,
|
||||
提示词: str,
|
||||
网络线路: str = "全球加速",
|
||||
图片: Optional[torch.Tensor] = None,
|
||||
视频=None,
|
||||
文件: Optional[FileList] = None,
|
||||
@@ -378,6 +384,7 @@ class UniversalLLMChat:
|
||||
|
||||
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
|
||||
|
||||
@@ -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()
|
||||
@@ -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")
|
||||
@@ -21,6 +21,14 @@ DEFAULT_API_BASE_URL = "https://api.o1key.com"
|
||||
# 可通过环境变量 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]:
|
||||
"""
|
||||
@@ -136,3 +144,10 @@ def get_async_api_base_url() -> str:
|
||||
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,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}")
|
||||
+176
-1
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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,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)
|
||||
@@ -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))
|
||||
@@ -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,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);
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user