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
@@ -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()
|
||||
Reference in New Issue
Block a user