feat: 新增 Grok 图像节点、前端 UI 增强、重构 nano-banana 系列

- 新增 Grok Image 节点及客户端
- 新增 save_image_format 节点
- 新增前端 JS 扩展:画笔工具、点阵网格、侧边栏隐藏、资源切换、重命名等
- 重构 nano-banana 节点,移除 pro 版本
- 移除 multi_res_preview 节点
- 新增 http_error 工具模块
- 各客户端和节点优化改进

Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Jony
2026-05-24 22:46:55 +08:00
co-authored by Claude Opus 4.6
parent c491731c99
commit 69279c654d
40 changed files with 3406 additions and 1599 deletions
+148
View File
@@ -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()
"""
+360 -8
View File
@@ -11,9 +11,17 @@ Comfyui_o1key - ComfyUI 自定义节点集合
import ssl
import logging
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)
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, O1keyGrokImage, KVideoFirstLast, KVideoImage2Video
from .nodes import K3Video, K3VideoFirstLast, K3MotionControl, K3MotionVideoCheck, NanoBananaV2, NanoBananaV2Batch, SaveImageFormat
# 报错弹框友好文案(不修改原节点代码,仅在外层统一处理)
_MSG_TIMEOUT = "API 请求超时,请稍后重试或检查网络。"
@@ -52,12 +60,12 @@ _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,
@@ -67,13 +75,14 @@ NODE_CLASS_MAPPINGS = {
"KlingFirstLastFrame": KlingFirstLastFrame,
"KlingMotionControlTest": KlingMotionControlTest,
"AspectRatioPreset": AspectRatioPreset,
"MultiResPreview": MultiResPreview,
"BatchImagesO1key": BatchImagesO1key,
"Seedance": Seedance,
"SeedanceMultiModal": SeedanceMultiModal,
"StreamPreview": StreamPreview,
"DoubaoImage": DoubaoImage,
"O1keyGPTImage": O1keyGPTImage,
"O1keyGrokImage": O1keyGrokImage,
"KVideoFirstLast": KVideoFirstLast,
"KVideoImage2Video": KVideoImage2Video,
"K3Video": K3Video,
@@ -82,15 +91,16 @@ NODE_CLASS_MAPPINGS = {
"K3MotionVideoCheck": K3MotionVideoCheck,
"NanoBananaV2": NanoBananaV2,
"NanoBananaV2Batch": NanoBananaV2Batch,
"SaveImageFormat": SaveImageFormat,
}
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",
@@ -100,13 +110,14 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"KlingFirstLastFrame": "首尾帧生视频 自研模型",
"KlingMotionControlTest": "动作控制 自研模型",
"AspectRatioPreset": "图片宽高比预设",
"MultiResPreview": "预览图像(v2",
"BatchImagesO1key": "加载图像(批量)",
"Seedance": "Seedance 视频生成",
"SeedanceMultiModal": "Seedance 多模态参考生视频",
"StreamPreview": "流式文本预览",
"DoubaoImage": "豆包生图",
"O1keyGPTImage": "o1key GPT Image",
"O1keyGrokImage": "Grok Image",
"KVideoFirstLast": "K26 图生视频(首尾帧)",
"KVideoImage2Video": "K26 图生视频",
"K3Video": "K3 图生视频 自研",
@@ -115,6 +126,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"K3MotionVideoCheck": "视频时长检测 K3",
"NanoBananaV2": "Nano Banana V2",
"NanoBananaV2Batch": "Nano Banana V2(批量)",
"SaveImageFormat": "保存图像(格式转换)",
}
WEB_DIRECTORY = "./web"
@@ -126,11 +138,351 @@ try:
from aiohttp import web
from server import PromptServer
import folder_paths
from .utils.config import CONFIG_FILE, load_config, NETWORK_ROUTES
@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("全球加速", "https://api.o1key.cn")
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, uuid, 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 = os.path.join(output_dir, ".o1key_history.json")
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 os.listdir(output_dir):
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 = {}
ungrouped = []
for f in all_files:
m = meta.get(f["name"], {})
wid = m.get("workflow_id")
if wid:
groups.setdefault(wid, []).append((f, m))
else:
ungrouped.append((f, {}))
# 构建 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,
})
# 无元数据的文件各自作为独立 job
for f, m in ungrouped:
mtime_ms = int(f["mtime"] * 1000)
job_id = str(uuid.uuid5(uuid.NAMESPACE_URL, f["name"]))
jobs.append({
"id": job_id,
"status": "completed",
"create_time": mtime_ms,
"execution_start_time": mtime_ms,
"execution_end_time": mtime_ms,
"preview_output": {
"filename": f["name"],
"subfolder": "",
"type": "output",
"nodeId": "0",
"mediaType": f["media"],
},
"outputs_count": 1,
"execution_error": None,
"workflow_id": None,
})
# 按时间倒序排列,分页
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 (workflow_id 或 uuid5) 返回含工作流的 job 详情"""
import os, uuid, struct, json as _json
job_id = request.match_info["job_id"]
output_dir = os.path.abspath(folder_paths.get_output_directory())
meta_file = os.path.join(output_dir, ".o1key_history.json")
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 的所有文件(按 workflow_id 或 uuid5 匹配)
matched_files = []
for fname in os.listdir(output_dir):
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)
elif str(uuid.uuid5(uuid.NAMESPACE_URL, fname)) == 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)
# === 执行事件 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 = _os.path.join(output_dir, ".o1key_history.json")
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
+53 -40
View File
@@ -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
class BaseAPIClient(ABC):
@@ -206,7 +208,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,31 +234,55 @@ class BaseAPIClient(ABC):
return
try:
if _interrupt_available:
request_task = asyncio.ensure_future(_do_request())
interrupt_task = asyncio.ensure_future(_poll_interrupt())
last_error_status = None
last_error_text = ""
done, pending = await asyncio.wait(
[request_task, interrupt_task],
return_when=asyncio.FIRST_COMPLETED
)
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())
# 取消未完成的任务
for t in pending:
t.cancel()
try:
await t
except (asyncio.CancelledError, Exception):
pass
done, pending = await asyncio.wait(
[request_task, interrupt_task],
return_when=asyncio.FIRST_COMPLETED
)
# 判断是哪个先完成
if interrupt_task in done and request_task not in done:
raise InterruptProcessingException()
for t in pending:
t.cancel()
try:
await t
except (asyncio.CancelledError, Exception):
pass
# 请求完成,取出结果(可能含异常)
return request_task.result()
else:
return await _do_request()
if interrupt_task in done and request_task not in done:
raise InterruptProcessingException()
result = request_task.result()
else:
result = await _do_request()
if isinstance(result, dict) and result.get("_error"):
status = result["_status"]
error_text = result["_text"]
last_error_status = status
last_error_text = error_text
if status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
friendly = HTTP_ERROR_MESSAGES.get(status, f"请求失败 ({status})")
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
print(f"{friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
await asyncio.sleep(delay)
continue
if status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[status])
raise RuntimeError(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(last_error_text)
except InterruptProcessingException:
raise
@@ -352,30 +379,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"
+46 -29
View File
@@ -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,39 +246,55 @@ class DoubaoImageClient:
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
# 3. 发送 POST 请求
t0 = time.time()
async with session.post(
url,
json=body,
headers=self._headers(),
) as resp:
elapsed_req = time.time() - t0
text = await resp.text()
# 3. 发送 POST 请求(带退避重试)
last_status = None
for attempt in range(DEFAULT_MAX_RETRIES + 1):
t0 = time.time()
async with session.post(
url,
json=body,
headers=self._headers(),
) as resp:
elapsed_req = time.time() - t0
text = await resp.text()
if resp.status != 200:
last_status = resp.status
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
friendly = HTTP_ERROR_MESSAGES.get(resp.status)
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
print(f"[豆包生图] {friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
await asyncio.sleep(delay)
continue
if resp.status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
try:
err_json = json.loads(text)
err_obj = err_json.get("error", {})
if isinstance(err_obj, dict):
msg = (
err_obj.get("message")
or err_obj.get("msg")
or text
)
else:
msg = str(err_obj) or text
except Exception:
msg = text
raise RuntimeError(
f"请求失败 HTTP {resp.status}: {msg}"
)
if resp.status != 200:
# 尝试解析错误信息
try:
err_json = json.loads(text)
err_obj = err_json.get("error", {})
if isinstance(err_obj, dict):
msg = (
err_obj.get("message")
or err_obj.get("msg")
or text
)
else:
msg = str(err_obj) or text
resp_json = json.loads(text)
except Exception:
msg = text
raise RuntimeError(
f"请求失败 HTTP {resp.status}: {msg}"
)
raise RuntimeError(f"响应 JSON 解析失败,原始内容:{text[:500]}")
try:
resp_json = json.loads(text)
except Exception:
raise RuntimeError(f"响应 JSON 解析失败,原始内容:{text[:500]}")
break
else:
if last_status and last_status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
raise RuntimeError(f"请求失败: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
print(f"[豆包生图] API 响应耗时 {elapsed_req:.1f}s,开始下载图像...")
+3
View File
@@ -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]}"
+1 -1
View File
@@ -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(
+23 -8
View File
@@ -187,6 +187,7 @@ class GeminiAPIClient(BaseAPIClient):
enable_grounding: bool = False,
enable_image_search: bool = False,
image_compression: str = None,
thinking_level: str = None,
**kwargs
) -> Dict[str, Any]:
"""
@@ -227,15 +228,15 @@ class GeminiAPIClient(BaseAPIClient):
})
# 估算完整 body 大小(不含工具字段,工具字段很小可忽略)
est_image_config = {"imageSize": resolution}
if aspect_ratio and aspect_ratio != "智能":
est_image_config["aspectRatio"] = aspect_ratio
estimated = self._estimate_body_size(
parts + img_parts,
{
"generationConfig": {
"responseModalities": ["IMAGE"],
"imageConfig": {
"aspectRatio": aspect_ratio,
"imageSize": resolution
}
"imageConfig": est_image_config
}
}
)
@@ -267,6 +268,10 @@ class GeminiAPIClient(BaseAPIClient):
parts.extend(img_parts)
# 构建请求体
image_config = {"imageSize": resolution}
if aspect_ratio and aspect_ratio != "智能":
image_config["aspectRatio"] = aspect_ratio
request_body = {
"contents": [
{
@@ -276,13 +281,17 @@ class GeminiAPIClient(BaseAPIClient):
],
"generationConfig": {
"responseModalities": ["IMAGE"],
"imageConfig": {
"aspectRatio": aspect_ratio,
"imageSize": resolution
}
"imageConfig": image_config
}
}
# 添加思考深度配置
if thinking_level:
request_body["generationConfig"]["thinkingConfig"] = {
"thinkingLevel": thinking_level,
"includeThoughts": True
}
# 添加图片压缩参数
if image_compression:
request_body["image_compression"] = image_compression
@@ -564,6 +573,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 +612,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 +748,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 +813,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 +886,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 +920,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)
+38 -20
View File
@@ -26,6 +26,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
try:
from comfy.model_management import processing_interrupted, InterruptProcessingException
@@ -330,31 +331,46 @@ class GptImageClient:
async def _do_request():
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
t0 = time.time()
async with session.post(url, json=body, headers=self._json_headers()) as resp:
elapsed = time.time() - t0
text = await resp.text()
last_status = None
for attempt in range(DEFAULT_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 != 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"[o1key GPT Image] {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", {})
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}")
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
)
resp_json = json.loads(text)
except Exception:
msg = text
raise RuntimeError(f"请求失败 HTTP {resp.status}: {msg}")
raise RuntimeError(f"响应 JSON 解析失败,原始内容:{text[:500]}")
try:
resp_json = json.loads(text)
except Exception:
raise RuntimeError(f"响应 JSON 解析失败,原始内容:{text[:500]}")
print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
return await self._parse_response(resp_json, session)
print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
return await self._parse_response(resp_json, session)
if last_status and last_status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
raise RuntimeError(f"请求失败: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
return await self._run_with_interrupt(_do_request())
@@ -446,6 +462,8 @@ class GptImageClient:
text = await resp.text()
if resp.status != 200:
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", {})
+355
View File
@@ -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}"
+11 -16
View File
@@ -10,6 +10,7 @@ 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
class KlingClient:
@@ -49,11 +50,11 @@ 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:
text = await resp.text()
if resp.status != 200:
raise RuntimeError(f"提交失败 ({resp.status}): {text}")
return json.loads(text)
resp = await async_request_with_retry(
session, "POST", url, json=body, headers=self._headers(), prefix="Kling 提交: "
)
text = await resp.text()
return json.loads(text)
# ── 轮询状态 ──────────────────────────────────────────────────────
@@ -203,17 +204,11 @@ class KlingClient:
if on_stage:
on_stage("submitting")
create_url = f"{self.base_url}{self.NEW_API_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
# 尝试提取友好错误信息
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"动作控制提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
resp = await async_request_with_retry(
session, "POST", create_url, json=body, headers=headers, prefix="Kling 动作控制提交: "
)
text = await resp.text()
create_resp = json.loads(text)
video_id = create_resp.get("id")
if not video_id:
+3 -1
View File
@@ -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]:
+6 -12
View File
@@ -11,6 +11,7 @@ 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
class SeedanceClient:
@@ -47,18 +48,11 @@ 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:
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)
resp = await async_request_with_retry(
session, "POST", url, json=body, headers=self._headers(), prefix="Seedance 提交: "
)
text = await resp.text()
data = json.loads(text)
# new-api 返回字段:id / task_id
task_id = data.get("id") or data.get("task_id")
+11 -16
View File
@@ -13,9 +13,10 @@ 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
try:
from comfy_api.latest import InputImpl
@@ -106,6 +107,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 +125,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",
@@ -207,20 +209,13 @@ class K3MotionControl:
# 1. 提交任务
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(
create_url,
resp = await async_request_with_retry(
session, "POST", create_url,
data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
headers=headers,
) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"K3 动作控制提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
headers=headers, prefix="K3 动作控制提交: "
)
text = await resp.text()
create_resp = json.loads(text)
# task_id 兼容扁平结构和 data 嵌套结构
task_id = (
+10 -13
View File
@@ -11,8 +11,9 @@ 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
try:
from comfy_api.latest import InputImpl
@@ -113,6 +114,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 +139,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",
@@ -246,16 +248,11 @@ class K3Video:
# 1. 提交
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"K3 提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
resp = await async_request_with_retry(
session, "POST", create_url, json=body, headers=headers, prefix="K3 提交: "
)
text = await resp.text()
create_resp = json.loads(text)
task_id = (
create_resp.get("task_id")
+10 -13
View File
@@ -10,8 +10,9 @@ 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
try:
from comfy_api.latest import InputImpl
@@ -94,6 +95,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 +111,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",
@@ -185,16 +187,11 @@ class K3VideoFirstLast:
# 1. 提交
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"K3 首尾帧提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
resp = await async_request_with_retry(
session, "POST", create_url, json=body, headers=headers, prefix="K3 首尾帧提交: "
)
text = await resp.text()
create_resp = json.loads(text)
task_id = (
create_resp.get("task_id")
+10 -13
View File
@@ -10,8 +10,9 @@ 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
try:
from comfy_api.latest import InputImpl
@@ -55,6 +56,7 @@ class KVideoFirstLast:
"模式": (["1080p"],),
"时长": ([5, 10],),
"生成音频": (["关闭", "打开"], {"default": "关闭"}),
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
"seed": ("INT", {
"default": 0, "min": 0, "max": 2147483647,
"tooltip": "seed 仅控制节点是否重新运行,结果本身不可复现。",
@@ -70,9 +72,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",
@@ -154,16 +156,11 @@ class KVideoFirstLast:
# 1. 提交
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"K26 提交失败 ({resp.status}): {msg}")
create_resp = json.loads(text)
resp = await async_request_with_retry(
session, "POST", create_url, json=body, headers=headers, prefix="K26 提交: "
)
text = await resp.text()
create_resp = json.loads(text)
task_id = (
create_resp.get("task_id")
+9 -8
View File
@@ -11,8 +11,9 @@ 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
try:
from comfy_api.latest import InputImpl
@@ -56,6 +57,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 +70,9 @@ class KVideoImage2Video:
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/KVideo"
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", seed=0):
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 网络线路="全球加速", 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",
@@ -150,11 +152,10 @@ class KVideoImage2Video:
# 1. 提交
_stage("submitting")
create_url = f"{base_url}{_ENDPOINT_CREATE}"
async with session.post(create_url, json=body, headers=headers) as resp:
if resp.status != 200:
err_text = await resp.text()
raise RuntimeError(f"提交失败 ({resp.status}): {err_text}")
sr = await resp.json()
resp = await async_request_with_retry(
session, "POST", create_url, json=body, headers=headers, prefix="K26 图生视频提交: "
)
sr = await resp.json()
task_id = sr.get("task_id") or sr.get("id")
if not task_id:
+7 -5
View File
@@ -4,27 +4,29 @@
"""
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 .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 .grok_image import O1keyGrokImage
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
__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', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGrokImage', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat']
@@ -1,11 +1,15 @@
"""
批量 Nano Banana Pro 节点
批量 Nano Banana 节点
ComfyUI 自定义节点用于批量处理图像生成任务
支持多文件夹加载1:1/笛卡尔积配对智能命名保存
"""
import io as _io
import re
import json
import time
import math
import base64
import random
import asyncio
import aiohttp
@@ -16,7 +20,7 @@ from PIL import Image
import torch
import numpy as np
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts, encode_image_to_base64, encode_image_to_base64_limited
from ..utils.file_utils import (
ImageInfo,
load_images_from_folder,
@@ -25,9 +29,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.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
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 +72,146 @@ DEBUG_LOG_ENABLED = False
REQUEST_LOG_ENABLED = False
# ============================================================================
_NODE = "Nano Banana Pro"
_NODE = "Nano Banana"
_ENDPOINT = "/v1/chat/completions"
_IMAGE_RE = re.compile(r"!\[.*?\]\(data:image/(\w+);base64,([A-Za-z0-9+/=]+)\)")
def _get_headers(api_key: str) -> dict:
return {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"X-Accel-Buffering": "no",
"Cache-Control": "no-cache, no-transform",
}
def _build_request_body(
prompt: str,
model: str,
aspect_ratio: str,
resolution: str,
images: Optional[List[Image.Image]] = None,
enable_grounding: bool = False,
) -> dict:
content_parts = [{"type": "text", "text": prompt}]
if images:
for img in images:
b64 = encode_image_to_base64_limited(img, format="PNG")
content_parts.append({
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{b64}"}
})
body = {
"model": model,
"stream": True,
"messages": [{"role": "user", "content": content_parts}],
}
google_config = {
"image_config": {
"image_size": resolution,
}
}
if aspect_ratio and aspect_ratio != "智能":
google_config["image_config"]["aspect_ratio"] = aspect_ratio
body["extra_body"] = {"google": google_config}
if enable_grounding:
body["extra_body"]["google_search"] = True
return body
async def _generate_single_openai(
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,
) -> List[Image.Image]:
url = f"{base_url}{_ENDPOINT}"
headers = _get_headers(api_key)
body = _build_request_body(
prompt=prompt,
model=model,
aspect_ratio=aspect_ratio,
resolution=resolution,
images=images,
enable_grounding=enable_grounding,
)
if REQUEST_LOG_ENABLED:
extra = json.dumps(body.get("extra_body", {}), ensure_ascii=False)
print(f"[请求] POST {url} | model={model} | extra_body={extra}")
last_status = None
for attempt in range(DEFAULT_MAX_RETRIES + 1):
resp = await session.post(url, headers=headers, json=body)
if resp.status == 200:
break
last_status = resp.status
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
friendly = HTTP_ERROR_MESSAGES.get(resp.status, f"请求失败 ({resp.status})")
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
resp.close()
await asyncio.sleep(delay)
continue
error_text = await resp.text()
resp.close()
if resp.status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
try:
err_json = json.loads(error_text)
msg = err_json.get("error", {}).get("message", error_text[:200])
except Exception:
msg = error_text[:200]
raise RuntimeError(f"API 错误 ({resp.status}): {msg}")
else:
if last_status and last_status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
raise RuntimeError(f"API 错误: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
full_content = ""
buffer = ""
async for raw_chunk in resp.content.iter_any():
buffer += raw_chunk.decode("utf-8")
while "\n" in buffer:
line_str, buffer = buffer.split("\n", 1)
line_str = line_str.strip()
if not line_str or not line_str.startswith("data:"):
continue
data_str = line_str[5:].strip()
if data_str == "[DONE]":
break
try:
chunk = json.loads(data_str)
delta = chunk.get("choices", [{}])[0].get("delta", {})
if "content" in delta:
full_content += delta["content"]
except (json.JSONDecodeError, IndexError):
continue
resp.close()
if not full_content:
raise RuntimeError("API 未返回有效内容")
matches = list(_IMAGE_RE.finditer(full_content))
if not matches:
raise RuntimeError(f"响应中未找到图片: {full_content[:100]}")
last_match = matches[-1]
img_data = base64.b64decode(last_match.group(2))
final_image = Image.open(_io.BytesIO(img_data)).convert("RGB")
return [final_image]
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
@@ -98,7 +242,7 @@ def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.
class BatchNanoBananaPro:
"""
批量 Nano Banana Pro 节点
批量 Nano Banana 节点
功能
- 从多个文件夹加载图片
@@ -116,26 +260,32 @@ 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
def resize_to_megapixels(
self,
@@ -181,33 +331,17 @@ 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",)
# 图片配对模式移到可选参数
@@ -215,35 +349,29 @@ class BatchNanoBananaPro:
"default": "不配对"
})
optional_inputs["代理端口(如7897"] = ("STRING", {
"default": "",
"multiline": False,
"placeholder": "本地代理端口,如 7897Clash Verge)或 10808v2rayN),留空不使用"
})
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 +398,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 +515,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 +525,12 @@ 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 = "原始",
) -> dict:
"""
执行单个生成任务
Args:
client: API 客户端
session: aiohttp 会话
prompt: 提示词
model: 模型名称
resolution: 分辨率
aspect_ratio: 宽高比
images: 输入图片列表
output_folder: 输出文件夹
task_index: 任务索引
Returns:
包含结果信息的字典
执行单个生成任务OpenAI 兼容接口
"""
result = {
"task_index": task_index,
@@ -440,7 +538,7 @@ class BatchNanoBananaPro:
"success": False,
"generated_count": 0,
"saved_files": [],
"output_images": [], # 无保存路径时存储内存图片
"output_images": [],
"error": None
}
@@ -448,26 +546,21 @@ class BatchNanoBananaPro:
# 准备输入图片
input_pil_images = [info.image for info in images]
# 调用 API 生成图片(固定生成1次)
# 调用 OpenAI 兼容接口生成图片
generated_images = []
try:
gen_result = await client.generate_single_async(
gen_images = await _generate_single_openai(
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,
)
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 +583,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
)
save_image(gen_img, output_path)
# 保存时不压缩
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 +651,20 @@ 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 = "原始",
) -> List[dict]:
"""
异步批量处理所有任务 - 改进版支持分批保存
Args:
pairs: 配对后的图片组合
prompt: 提示词单提示词模式时使用
model: 模型名称
resolution: 分辨率
aspect_ratio: 宽高比
output_folder: 输出文件夹
pbar: ComfyUI 进度条
prompts_per_task: 每个任务对应的提示词列表批量提示词模式时使用
Returns:
所有任务的结果列表
异步批量处理所有任务OpenAI 兼容接口
"""
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 +717,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,7 +728,6 @@ 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,
)
@@ -721,15 +821,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 +837,14 @@ class BatchNanoBananaPro:
Args:
prompt: 提示词
文件夹1-9: 图片文件夹路径
文件夹1-5: 图片文件夹路径
seed: 随机种子
保存路径: 输出保存路径
图片配对模式: 1:1 1*N
模型: 模型名称
宽高比: 输出宽高比
分辨率: 输出分辨率
**kwargs: 动态参考图输入 (参考图1-9)
**kwargs: 动态参考图输入 (参考图1-5)
Returns:
输出图像张量
@@ -753,10 +852,27 @@ 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 = False
# 拼接实际模型 ID
base_model_id = self.MODEL_ID_MAP.get(模型, "nano-banana-pro")
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 +883,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 +898,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 +917,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,11 +954,8 @@ class BatchNanoBananaPro:
total_tasks = len(pairs)
# 打印首行概览
# 图片搜索(联网)开启时隐含谷歌搜索接地,与客户端请求逻辑保持一致
grounding_str = ""
if enable_image_search:
grounding_str = " | 谷歌图片搜索接地"
elif enable_grounding:
if enable_grounding:
grounding_str = " | 谷歌搜索接地"
if batch_prompts:
@@ -890,17 +993,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 +1015,12 @@ 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=图片格式,
)
)
except Exception as e:
@@ -1050,14 +1146,15 @@ class BatchNanoBananaPro:
finally:
# 查询余额
if self.client is not None:
try:
balance_data = self.client.query_balance_sync()
balance_info = self.client.format_balance_info(balance_data)
print(f"BatchNanaBananaPro: {balance_info}")
print("=" * 60)
except Exception:
pass
try:
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
# 最终内存清理
import gc
+19 -13
View File
@@ -6,6 +6,7 @@ o1key GPT Image 节点
import time
from ..clients.gpt_image_client import GptImageClient
from ..utils.image_utils import parse_batch_prompts
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
try:
from comfy.model_management import processing_interrupted, InterruptProcessingException
@@ -52,30 +53,33 @@ class O1keyGPTImage:
], {
"default": "gpt-image-2-次卡",
})
optional_inputs["网络"] = (NETWORK_ROUTE_OPTIONS, {
"default": "全球加速",
})
optional_inputs["分辨率"] = ([
"智能",
# ── 1K ──
"1024x10241K 正方形 1:1",
"1536x10241K 横版 3:2",
"1024x15361K 竖版 2:3",
"1365x10241K 横版 4:3",
"1024x13651K 竖版 3:4",
"1820x10241K 横版 16:9",
"1024x18201K 竖版 9:16",
"1360x10241K 横版 4:3",
"1024x13601K 竖版 3:4",
"1824x10241K 横版 16:9",
"1024x18241K 竖版 9:16",
# ── 2K ──
"2048x20482K 正方形 1:1",
"3072x20482K 横版 3:2",
"2048x30722K 竖版 2:3",
"2732x20482K 横版 4:3",
"2048x27322K 竖版 3:4",
"3640x20482K 横版 16:9",
"2048x36402K 竖版 9:16",
"2736x20482K 横版 4:3",
"2048x27362K 竖版 3:4",
"3648x20482K 横版 16:9",
"2048x36482K 竖版 9:16",
# ── 4K ──
"3840x38404K 正方形 1:1",
"3840x25604K 横版 3:2",
"2560x38404K 竖版 2:3",
"3840x28804K 横版 4:3",
"2880x38404K 竖版 3:4",
"2880x28804K 正方形 1:1",
"3504x23364K 横版 3:2",
"2336x35044K 竖版 2:3",
"3264x24484K 横版 4:3",
"2448x32644K 竖版 3:4",
"3840x21604K 横版 16:9",
"2160x38404K 竖版 9:16",
], {
@@ -128,6 +132,7 @@ class O1keyGPTImage:
self,
prompt: str,
模型: str = "gpt-image-2-次卡",
网络: str = "全球加速",
分辨率: str = "auto",
质量: str = "自动",
生图数量: int = 1,
@@ -169,6 +174,7 @@ class O1keyGPTImage:
# ── 3. 创建客户端 ─────────────────────────────────────────────────────
try:
client = GptImageClient()
client.base_url = get_base_url_by_route(网络)
except ValueError as e:
if str(e) == "未授权!":
print("[o1key GPT Image] 请联系作者授权后方可使用!")
+165
View File
@@ -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
+9 -1
View File
@@ -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,13 +246,15 @@ class KlingVideo:
body["image"] = _tensor_to_base64(start_frame)
endpoint_type = "image2video"
else:
body["metadata"] = {"aspect_ratio": aspect_ratio}
if aspect_ratio and aspect_ratio != "智能":
body["metadata"] = {"aspect_ratio": aspect_ratio}
endpoint_type = "text2video"
# ── 保存路径(临时文件,避免与下游保存节点重复落盘)──────────────────
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="kling_")
client = KlingClient()
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
# ── 进度条 ────────────────────────────────────────────────────
try:
@@ -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:
-165
View File
@@ -1,165 +0,0 @@
"""
多分辨率图像预览节点
ComfyUI 自定义节点,支持同时预览多张不同分辨率的图像
背景:
ComfyUI 原生「预览图像」节点要求 batch 内所有图片分辨率相同(因为它们被
stack 成一个 [B, H, W, C] tensor)。当 API 返回多张不同尺寸的图片时
(例如 nano-banana-2 同时返回 1K + 2K),原生节点会报错。
解决方案:
声明 INPUT_IS_LIST = TrueComfyUI 会将连入的所有图像作为
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 = TrueComfyUI 将每张图作为独立 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}}
+636
View File
@@ -0,0 +1,636 @@
"""
Nano Banana 节点 (V3)
ComfyUI 自定义节点,用于调用生图模型(OpenAI 兼容接口)
使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动
"""
import io as _io
import re
import json
import time
import math
import base64
import random
import asyncio
import aiohttp
from concurrent.futures import ThreadPoolExecutor
from typing import 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, encode_image_to_base64, encode_image_to_base64_limited
from ..utils.config import (
NETWORK_ROUTE_OPTIONS,
get_base_url_by_route,
get_api_key_or_raise,
)
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
from ..clients.gemini_client import GeminiAPIClient
try:
import folder_paths
FOLDER_PATHS_AVAILABLE = True
except ImportError:
FOLDER_PATHS_AVAILABLE = False
try:
from comfy.utils import ProgressBar
PROGRESS_BAR_AVAILABLE = True
except ImportError:
PROGRESS_BAR_AVAILABLE = False
try:
import psutil
MEMORY_MONITOR_AVAILABLE = True
except ImportError:
MEMORY_MONITOR_AVAILABLE = False
DEBUG_LOG_ENABLED = True
REQUEST_LOG_ENABLED = False
_NODE = "Nano Banana"
_ENDPOINT = "/v1/chat/completions"
_client_instance = None
def _get_client():
global _client_instance
if _client_instance is None:
_client_instance = GeminiAPIClient()
return _client_instance
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
_IMAGE_RE = re.compile(r"!\[.*?\]\(data:image/(\w+);base64,([A-Za-z0-9+/=]+)\)")
def _get_headers(api_key: str) -> dict:
return {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"X-Accel-Buffering": "no",
"Cache-Control": "no-cache, no-transform",
}
def _build_request_body(
prompt: str,
model: str,
aspect_ratio: str,
resolution: str,
images: Optional[List[Image.Image]] = None,
enable_grounding: bool = False,
thinking_level: Optional[str] = None,
) -> dict:
content_parts = [{"type": "text", "text": prompt}]
if images:
for img in images:
b64 = encode_image_to_base64_limited(img, format="PNG")
content_parts.append({
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{b64}"}
})
body = {
"model": model,
"stream": True,
"messages": [{"role": "user", "content": content_parts}],
}
google_config = {
"image_config": {
"image_size": resolution,
}
}
if aspect_ratio and aspect_ratio != "智能":
google_config["image_config"]["aspect_ratio"] = aspect_ratio
if thinking_level:
google_config["thinking_config"] = {
"thinking_level": thinking_level.lower(),
"include_thoughts": True,
}
body["extra_body"] = {"google": google_config}
if enable_grounding:
body["extra_body"]["google_search"] = True
return body
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,
) -> List[Image.Image]:
url = f"{base_url}{_ENDPOINT}"
headers = _get_headers(api_key)
body = _build_request_body(
prompt=prompt,
model=model,
aspect_ratio=aspect_ratio,
resolution=resolution,
images=images,
enable_grounding=enable_grounding,
thinking_level=thinking_level,
)
if REQUEST_LOG_ENABLED:
extra = json.dumps(body.get("extra_body", {}), ensure_ascii=False)
print(f"[请求] POST {url} | model={model} | extra_body={extra}")
last_status = None
for attempt in range(DEFAULT_MAX_RETRIES + 1):
resp = await session.post(url, headers=headers, json=body)
if resp.status == 200:
break
last_status = resp.status
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
friendly = HTTP_ERROR_MESSAGES.get(resp.status, f"请求失败 ({resp.status})")
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
print(f"Nano Banana: {friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
resp.close()
await asyncio.sleep(delay)
continue
error_text = await resp.text()
resp.close()
if resp.status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
try:
err_json = json.loads(error_text)
msg = err_json.get("error", {}).get("message", error_text[:200])
except Exception:
msg = error_text[:200]
raise RuntimeError(f"API 错误 ({resp.status}): {msg}")
else:
if last_status and last_status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
raise RuntimeError(f"API 错误: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
full_content = ""
buffer = ""
t_request = time.time()
t_first_token = None
async for raw_chunk in resp.content.iter_any():
if t_first_token is None:
t_first_token = time.time()
buffer += raw_chunk.decode("utf-8")
while "\n" in buffer:
line_str, buffer = buffer.split("\n", 1)
line_str = line_str.strip()
if not line_str or not line_str.startswith("data:"):
continue
data_str = line_str[5:].strip()
if data_str == "[DONE]":
break
try:
chunk = json.loads(data_str)
delta = chunk.get("choices", [{}])[0].get("delta", {})
if "content" in delta:
full_content += delta["content"]
except (json.JSONDecodeError, IndexError):
continue
t_done = time.time()
resp.close()
if not full_content:
raise RuntimeError("API 未返回有效内容")
# 思考模型可能输出多张临时图片,最终图片始终是最后一张
matches = list(_IMAGE_RE.finditer(full_content))
if not matches:
raise RuntimeError(f"响应中未找到图片: {full_content[:100]}")
last_match = matches[-1]
img_data = base64.b64decode(last_match.group(2))
final_image = Image.open(_io.BytesIO(img_data)).convert("RGB")
first_token_ms = (t_first_token - t_request) * 1000 if t_first_token else 0
download_ms = (t_done - t_first_token) * 1000 if t_first_token else 0
return [final_image], first_token_ms, download_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,
) -> dict:
result = {
"global_task_index": global_task_index,
"prompt": prompt,
"success": False,
"generated_count": 0,
"output_images": [],
"error": None,
}
try:
gen_images, first_token_ms, download_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,
)
result["output_images"] = gen_images
result["success"] = True
result["generated_count"] = len(gen_images)
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):
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(
_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,
)
)
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, "output_images": [], "prompt": ""}
else:
result_data = result
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}")
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
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()
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 input_images and 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(
_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=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
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,
)
return loop.run_until_complete(_do())
finally:
loop.close()
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(run_single)
generated_images, first_token_ms, download_ms = future.result(timeout=900)
if pbar is not None:
pbar.update(1)
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"
ft_str = f"{first_token_ms/1000:.2f}s"
dl_str = f"{download_ms/1000:.2f}s"
print(f"完成!总耗时 {time_str} | 首字 {ft_str} | 下载 {dl_str} | 成功 {len(generated_images)}")
import gc; gc.collect()
return io.NodeOutput(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:
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()
-831
View File
@@ -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": "本地代理端口,如 7897Clash Verge)或 10808v2rayN),留空不使用"
})
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()
+15 -11
View File
@@ -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
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
}
@@ -299,13 +301,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
@@ -541,6 +541,7 @@ class NanoBananaV2:
宽高比: str,
分辨率: str,
生图数量: int,
网络线路: str = "全球加速",
**kwargs
) -> Tuple[torch.Tensor]:
"""生成图像(异步模式)"""
@@ -557,6 +558,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 +592,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)}"
@@ -977,6 +979,7 @@ class NanoBananaV2Batch(NanoBananaV2):
宽高比: str,
分辨率: str,
生图数量: int = 1,
网络线路: str = "全球加速",
**kwargs
) -> Tuple[torch.Tensor]:
"""生成图像(异步模式 - 批量版:全并发 + 即时落盘)"""
@@ -997,6 +1000,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}")
+2 -185
View File
@@ -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:
+91
View File
@@ -0,0 +1,91 @@
"""保存图像节点 - 支持 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": {
"images": ("IMAGE",),
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
"format": (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, images, filename_prefix="ComfyUI", format="PNG",
prompt=None, extra_pnginfo=None):
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}}
+5
View File
@@ -15,6 +15,7 @@ from ..clients.seedance_client import SeedanceClient
from ..clients.gemini_client import GeminiAPIClient
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64, pil_to_tensor
from ..utils.r2_uploader import upload_video, upload_audio
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
from comfy_api.latest import InputImpl
@@ -116,6 +117,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"],
@@ -242,6 +244,7 @@ class Seedance:
_, 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 +271,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"],
@@ -414,6 +418,7 @@ class SeedanceMultiModal:
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="seedance_mm_")
client = SeedanceClient()
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
pbar = _make_pbar()
on_stage, on_prog = _make_callbacks("Seedance多模态", pbar)
+6 -1
View File
@@ -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
# ============================================================================
@@ -66,6 +66,9 @@ class UniversalLLMChat:
def INPUT_TYPES(cls):
return {
"required": {
"网络线路": (NETWORK_ROUTE_OPTIONS, {
"default": "全球加速"
}),
"模型": (SUPPORTED_MODELS, {
"default": SUPPORTED_MODELS[0]
}),
@@ -370,6 +373,7 @@ class UniversalLLMChat:
self,
模型: str,
提示词: str,
网络线路: str = "全球加速",
图片: Optional[torch.Tensor] = None,
视频=None,
文件: Optional[FileList] = None,
@@ -380,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
+13
View File
@@ -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,8 @@ 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 垫底"""
return NETWORK_ROUTES.get(route, get_api_base_url())
+125
View File
@@ -0,0 +1,125 @@
"""
统一 HTTP 错误处理 & 退避重试模块
使用方式:
1. 对于 aiohttp 请求,用 async_request_with_retry() 包裹 POST/GET 调用
2. 对于已拿到 status code 的场景,调用 raise_for_status() 抛出友好错误
新增生图/视频节点时,请统一使用本模块处理 HTTP 错误。
"""
import asyncio
import random
from typing import Optional
import aiohttp
# ═══════════════════════════════════════════════════════════════════════════════
# 状态码 → 用户友好文案
# ═══════════════════════════════════════════════════════════════════════════════
HTTP_ERROR_MESSAGES = {
429: "模型速率超限或额度不足!",
502: "网关超时。请重试或将网络切换为美国直连",
503: "模型超载。请稍后重试!",
504: "网关超时。请稍后重试。",
}
# 可退避重试的状态码
RETRYABLE_STATUS_CODES = {429, 502, 503, 504}
# 退避重试默认参数
DEFAULT_MAX_RETRIES = 3
DEFAULT_BASE_DELAY = 2.0 # 首次重试等待秒数
DEFAULT_MAX_DELAY = 30.0 # 最大等待秒数
DEFAULT_BACKOFF_FACTOR = 2.0 # 指数退避因子
def get_friendly_message(status_code: int, raw_message: str = "") -> str:
"""根据状态码返回友好文案,未匹配则返回原始信息"""
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/503/504) 进行重试。
超过最大重试次数后抛出友好 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, prefix=prefix)
raw_msg = last_message[:200] if last_message else ""
raise RuntimeError(f"{prefix}请求失败 ({last_status}): {raw_msg}")
+51
View File
@@ -110,6 +110,57 @@ def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
return base64.b64encode(img_bytes).decode('utf-8')
_MAX_IMAGE_BYTES = 10 * 1024 * 1024 # 10MB base64 上限
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
+92
View File
@@ -0,0 +1,92 @@
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 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;
}
};
},
});
+41
View File
@@ -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);
}
},
});
+215
View File
@@ -0,0 +1,215 @@
import { app } from "../../../scripts/app.js";
app.registerExtension({
name: "o1key.hideSidebarItems",
async setup() {
const hide = () => {
// 隐藏侧边栏的"说明"、"应用"、"模型"、"节点"按钮
const hiddenLabels = ["说明", "帮助", "Help", "应用", "Apps", "模型", "Models", "节点", "Nodes"];
document.querySelectorAll(".side-bar-button, [class*='sidebar'] button, .p-togglebutton").forEach(btn => {
const label = btn.getAttribute("aria-label") || btn.textContent || "";
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);
},
});
+1
View File
File diff suppressed because one or more lines are too long
+483
View File
@@ -0,0 +1,483 @@
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 .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);
}
// --- 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());
}
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], () => {
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") 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) return;
state.drawing = false;
shape = null;
});
}
// --- 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, 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, () => {
// 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);
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();
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="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-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.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 {
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;
}
};
// 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)
});
};
}
});
+29
View File
@@ -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);
},
});
+18
View File
@@ -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 }
);
},
});