Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
723 lines
27 KiB
Python
723 lines
27 KiB
Python
"""
|
||
Nano Banana 节点 (V3)
|
||
ComfyUI 自定义节点，用于调用异步生图模型
|
||
使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动
|
||
"""
|
||
|
||
import time
|
||
import math
|
||
import random
|
||
import asyncio
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from typing import Any, Callable, List, Optional
|
||
|
||
import torch
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
from comfy_api.latest import io
|
||
|
||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||
from ..utils.config import (
|
||
get_base_url_by_route,
|
||
get_api_key_or_raise,
|
||
get_runtime_config_signature,
|
||
)
|
||
from ..utils.nano_banana_async import (
|
||
generate_nano_banana_async,
|
||
VERBOSE_LOG_ENABLED,
|
||
)
|
||
from ..utils.http2_client import create_http_client
|
||
from ..clients.gemini_client import GeminiAPIClient
|
||
from ..utils.nano_banana_models import (
|
||
NANO_BANANA_MODEL_OPTIONS,
|
||
NANO_BANANA_ROUTE_OPTIONS,
|
||
resolve_nano_banana_model,
|
||
)
|
||
|
||
try:
|
||
from comfy.utils import ProgressBar
|
||
PROGRESS_BAR_AVAILABLE = True
|
||
except ImportError:
|
||
PROGRESS_BAR_AVAILABLE = False
|
||
|
||
try:
|
||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||
INTERRUPT_AVAILABLE = True
|
||
except ImportError:
|
||
INTERRUPT_AVAILABLE = False
|
||
InterruptProcessingException = RuntimeError
|
||
processing_interrupted = lambda: False
|
||
|
||
# 完整原始报文仅在显式开启详细日志时打印。
|
||
REQUEST_LOG_ENABLED = VERBOSE_LOG_ENABLED
|
||
|
||
_NODE = "Nano Banana"
|
||
_REQUEST_TIMEOUT = 900
|
||
_INTERRUPT_CHECK_INTERVAL = 0.2
|
||
MAX_REFERENCE_IMAGES = 14
|
||
_MAX_GENERATION_CONCURRENCY = 12
|
||
_MAX_DOWNLOAD_CONCURRENCY = 6
|
||
_HTTP_MAX_CONNECTIONS = 32
|
||
_HTTP_MAX_KEEPALIVE_CONNECTIONS = 16
|
||
|
||
_client_instance = None
|
||
_client_config_signature = None
|
||
|
||
|
||
def _get_client():
|
||
global _client_instance, _client_config_signature
|
||
config_signature = get_runtime_config_signature()
|
||
if _client_instance is None or config_signature != _client_config_signature:
|
||
_client_instance = GeminiAPIClient()
|
||
_client_config_signature = config_signature
|
||
return _client_instance
|
||
|
||
|
||
async def _poll_interrupt():
|
||
while True:
|
||
await asyncio.sleep(_INTERRUPT_CHECK_INTERVAL)
|
||
if INTERRUPT_AVAILABLE and processing_interrupted():
|
||
return
|
||
|
||
|
||
async def _run_with_interrupt(coro):
|
||
if not INTERRUPT_AVAILABLE:
|
||
return await coro
|
||
|
||
request_task = asyncio.ensure_future(coro)
|
||
interrupt_task = asyncio.ensure_future(_poll_interrupt())
|
||
|
||
done, pending = await asyncio.wait(
|
||
[request_task, interrupt_task],
|
||
return_when=asyncio.FIRST_COMPLETED,
|
||
)
|
||
|
||
for task in pending:
|
||
task.cancel()
|
||
try:
|
||
await task
|
||
except (asyncio.CancelledError, Exception):
|
||
pass
|
||
|
||
if interrupt_task in done and request_task not in done:
|
||
raise InterruptProcessingException()
|
||
|
||
return request_task.result()
|
||
|
||
|
||
def _check_interrupt():
|
||
if INTERRUPT_AVAILABLE and processing_interrupted():
|
||
raise InterruptProcessingException()
|
||
|
||
|
||
def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]:
|
||
if pbar is None:
|
||
return None
|
||
|
||
last_progress = [0.0]
|
||
|
||
def _on_progress(progress: float) -> None:
|
||
try:
|
||
progress = max(0.0, min(float(progress), 1.0))
|
||
except (TypeError, ValueError):
|
||
return
|
||
if progress <= last_progress[0]:
|
||
return
|
||
pbar.update(progress - last_progress[0])
|
||
last_progress[0] = progress
|
||
|
||
return _on_progress
|
||
|
||
|
||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
||
if not images:
|
||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||
return pil_to_tensor([placeholder])
|
||
|
||
base_size = max(images, key=lambda img: img.size[0] * img.size[1]).size
|
||
matched = [img for img in images if img.size == base_size]
|
||
skipped = [img for img in images if img.size != base_size]
|
||
|
||
if skipped:
|
||
sizes_str = ", ".join(f"{img.size[0]}x{img.size[1]}" for img in skipped)
|
||
print(
|
||
f"{node_label}: 丢弃 {len(skipped)} 张较小尺寸的图 ({sizes_str})，"
|
||
f"仅输出最大尺寸 {base_size[0]}x{base_size[1]} 的 {len(matched)} 张"
|
||
)
|
||
|
||
return pil_to_tensor(matched)
|
||
|
||
|
||
def _collect_autogrow_inputs(value) -> list:
|
||
"""Collect connected Autogrow slots while tolerating a single legacy value."""
|
||
if value is None:
|
||
return []
|
||
if isinstance(value, dict):
|
||
return [item for item in value.values() if item is not None]
|
||
return [value]
|
||
|
||
|
||
THINKING_LEVEL_MAP = {
|
||
"低": "minimal",
|
||
"高": "high",
|
||
}
|
||
|
||
|
||
def _build_model_id(model_name: str, resolution: str, route: str) -> str:
|
||
"""兼容旧调用签名；新模型名只由主模型和线路决定。"""
|
||
del resolution
|
||
return resolve_nano_banana_model(model_name, route)
|
||
|
||
|
||
async def _generate_single(
|
||
session: Any,
|
||
base_url: str,
|
||
api_key: str,
|
||
prompt: str,
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
images: Optional[List[Image.Image]] = None,
|
||
image_urls: Optional[List[str]] = None,
|
||
thinking_level: Optional[str] = None,
|
||
progress_callback: Optional[Callable[[float], None]] = None,
|
||
node_label: str = "Nano Banana",
|
||
result_url_callback: Optional[Callable[[str], None]] = None,
|
||
log_task_success: bool = True,
|
||
upload_cache: Optional[dict] = None,
|
||
download_semaphore: Optional[asyncio.Semaphore] = None,
|
||
resize_mode: str = "不缩放",
|
||
google_search: bool = False,
|
||
) -> tuple[List[Image.Image], dict]:
|
||
result_images, timing = await generate_nano_banana_async(
|
||
session=session,
|
||
base_url=base_url,
|
||
api_key=api_key,
|
||
prompt=prompt,
|
||
model=model,
|
||
resolution=resolution,
|
||
aspect_ratio=aspect_ratio,
|
||
images=images,
|
||
image_urls=image_urls,
|
||
upload_cache=upload_cache,
|
||
download_semaphore=download_semaphore,
|
||
thinking_level=thinking_level,
|
||
google_search=google_search,
|
||
node_label=node_label,
|
||
request_log_enabled=REQUEST_LOG_ENABLED,
|
||
check_interrupt=_check_interrupt,
|
||
progress_callback=progress_callback,
|
||
result_url_callback=result_url_callback,
|
||
log_task_success=log_task_success,
|
||
resize_mode=resize_mode,
|
||
)
|
||
return result_images, timing
|
||
|
||
|
||
async def _generate_single_task(
|
||
session: Any,
|
||
base_url: str,
|
||
api_key: str,
|
||
prompt: str,
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
images: Optional[List[Image.Image]],
|
||
image_urls: Optional[List[str]],
|
||
global_task_index: int,
|
||
thinking_level: Optional[str] = None,
|
||
progress_callback: Optional[Callable[[float], None]] = None,
|
||
upload_cache: Optional[dict] = None,
|
||
download_semaphore: Optional[asyncio.Semaphore] = None,
|
||
resize_mode: str = "不缩放",
|
||
google_search: bool = False,
|
||
) -> dict:
|
||
result = {
|
||
"global_task_index": global_task_index,
|
||
"prompt": prompt,
|
||
"success": False,
|
||
"generated_count": 0,
|
||
"output_images": [],
|
||
"error": None,
|
||
}
|
||
task_started = time.time()
|
||
try:
|
||
gen_images, timing = 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,
|
||
image_urls=image_urls if image_urls else None,
|
||
thinking_level=thinking_level,
|
||
google_search=google_search,
|
||
progress_callback=progress_callback,
|
||
node_label=f"Nano Banana#{global_task_index + 1}",
|
||
upload_cache=upload_cache,
|
||
download_semaphore=download_semaphore,
|
||
resize_mode=resize_mode,
|
||
)
|
||
result["output_images"] = gen_images
|
||
result["success"] = True
|
||
result["generated_count"] = len(gen_images)
|
||
print(
|
||
f"Nano Banana#{global_task_index + 1}: 完成 ✓ | "
|
||
f"生成={len(gen_images)} 张 | 耗时={time.time() - task_started:.1f}s"
|
||
)
|
||
except InterruptProcessingException:
|
||
raise
|
||
except Exception as e:
|
||
result["error"] = str(e)
|
||
return result
|
||
|
||
|
||
async def _process_batch_async(
|
||
base_url: str,
|
||
api_key: str,
|
||
prompts: List[str],
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
images_per_prompt: int,
|
||
input_images: Optional[List[Image.Image]],
|
||
pbar=None,
|
||
thinking_level: Optional[str] = None,
|
||
resize_mode: str = "不缩放",
|
||
unlimited_downloads: bool = False,
|
||
google_search: bool = False,
|
||
) -> 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 = _MAX_GENERATION_CONCURRENCY
|
||
num_batches = math.ceil(total_tasks / max_concurrent)
|
||
|
||
all_results = []
|
||
completed = 0
|
||
success_count = 0
|
||
fail_count = 0
|
||
|
||
upload_cache = {}
|
||
download_semaphore = (
|
||
None
|
||
if unlimited_downloads
|
||
else asyncio.Semaphore(_MAX_DOWNLOAD_CONCURRENCY)
|
||
)
|
||
|
||
async with create_http_client(
|
||
http2=True,
|
||
max_connections=_HTTP_MAX_CONNECTIONS,
|
||
max_keepalive_connections=_HTTP_MAX_KEEPALIVE_CONNECTIONS,
|
||
) as session:
|
||
for batch_idx in range(num_batches):
|
||
_check_interrupt()
|
||
start_idx = batch_idx * max_concurrent
|
||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||
|
||
tasks = []
|
||
for i in range(start_idx, end_idx):
|
||
_check_interrupt()
|
||
_, _, prompt = tasks_def[i]
|
||
task = asyncio.create_task(
|
||
_generate_single_task(
|
||
session=session,
|
||
base_url=base_url,
|
||
api_key=api_key,
|
||
prompt=prompt,
|
||
model=model,
|
||
resolution=resolution,
|
||
aspect_ratio=aspect_ratio,
|
||
images=input_images,
|
||
image_urls=None,
|
||
global_task_index=i,
|
||
thinking_level=thinking_level,
|
||
google_search=google_search,
|
||
progress_callback=_make_progress_callback(pbar),
|
||
upload_cache=upload_cache,
|
||
download_semaphore=download_semaphore,
|
||
resize_mode=resize_mode,
|
||
)
|
||
)
|
||
tasks.append(task)
|
||
|
||
batch_results = []
|
||
for coro in asyncio.as_completed(tasks):
|
||
_check_interrupt()
|
||
result_data = None
|
||
try:
|
||
result = await coro
|
||
if isinstance(result, Exception):
|
||
result_data = {"success": False, "error": str(result), "generated_count": 0, "output_images": [], "prompt": ""}
|
||
else:
|
||
result_data = result
|
||
except InterruptProcessingException:
|
||
for task in tasks:
|
||
task.cancel()
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
raise
|
||
except Exception as e:
|
||
result_data = {"success": False, "error": str(e), "generated_count": 0, "output_images": [], "prompt": ""}
|
||
|
||
batch_results.append(result_data)
|
||
completed += 1
|
||
task_num = result_data.get("global_task_index", "?")
|
||
if isinstance(task_num, int):
|
||
task_num += 1
|
||
|
||
if result_data and result_data.get("success", False):
|
||
success_count += 1
|
||
else:
|
||
fail_count += 1
|
||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
||
print(f"Nano Banana#{task_num}: 失败 | error={error_msg}")
|
||
|
||
all_results.extend(batch_results)
|
||
print(
|
||
f"Nano Banana: 批次 {batch_idx + 1}/{num_batches} 完成 "
|
||
f"| 成功={success_count} | 失败={fail_count} "
|
||
f"| 总进度={completed}/{total_tasks}"
|
||
)
|
||
import gc; gc.collect()
|
||
await asyncio.sleep(0.1)
|
||
|
||
return all_results
|
||
|
||
|
||
class NanoBanana(io.ComfyNode):
|
||
|
||
@staticmethod
|
||
def _validate_model_config(model_name: str, aspect_ratio: str, resolution: str):
|
||
"""验证模型配置是否合法"""
|
||
# 验证模型名称
|
||
valid_models = set(NANO_BANANA_MODEL_OPTIONS)
|
||
if model_name not in valid_models:
|
||
raise ValueError(
|
||
f"模型 '{model_name}' 无效，支持的模型：{', '.join(sorted(valid_models))}"
|
||
)
|
||
|
||
# 验证宽高比
|
||
valid_aspect_ratios = {
|
||
"智能", "1:1", "1:4", "1:8", "2:3", "3:2", "3:4",
|
||
"4:1", "4:3", "4:5", "5:4", "8:1",
|
||
"9:16", "16:9", "21:9",
|
||
}
|
||
if aspect_ratio not in valid_aspect_ratios:
|
||
raise ValueError(
|
||
f"宽高比 '{aspect_ratio}' 无效，支持的宽高比：{', '.join(sorted(valid_aspect_ratios))}"
|
||
)
|
||
|
||
# 验证分辨率
|
||
# “智能”仅由统一生图节点传入；独立节点仍保持原有下拉选项。
|
||
valid_resolutions = {"智能", "1K", "2K", "4K"}
|
||
if resolution not in valid_resolutions:
|
||
raise ValueError(
|
||
f"分辨率 '{resolution}' 无效，支持的分辨率：{', '.join(sorted(valid_resolutions))}"
|
||
)
|
||
|
||
# Nano Banana 2 系列特有的宽高比
|
||
nano_2_exclusive_ratios = {"1:4", "1:8", "4:1", "8:1"}
|
||
# 其他模型使用了 Nano Banana 2 系列专属宽高比
|
||
if model_name not in {"Nano Banana 2", "Nano Banana 2 Lite"} and aspect_ratio in nano_2_exclusive_ratios:
|
||
raise ValueError(
|
||
f"宽高比 {aspect_ratio} 仅支持 Nano Banana 2 系列模型，"
|
||
f"当前模型 {model_name} 不支持此宽高比"
|
||
)
|
||
|
||
# Nano Banana 只支持 1K
|
||
if model_name == "Nano Banana" and resolution not in {"智能", "1K"}:
|
||
raise ValueError(
|
||
f"Nano Banana 模型仅支持 1K 分辨率，"
|
||
f"当前选择的 {resolution} 不支持"
|
||
)
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
reference_images = io.Autogrow.Input(
|
||
"参考图组",
|
||
template=io.Autogrow.TemplateNames(
|
||
input=io.Image.Input("参考图"),
|
||
names=[f"参考图{i}" for i in range(1, MAX_REFERENCE_IMAGES + 1)],
|
||
min=0,
|
||
),
|
||
tooltip=f"连接后自动增加输入端口，合计最多 {MAX_REFERENCE_IMAGES} 张参考图。",
|
||
)
|
||
return io.Schema(
|
||
node_id="NanoBanana",
|
||
display_name="Nano Banana",
|
||
category="image/generation",
|
||
inputs=[
|
||
io.String.Input(
|
||
"prompt",
|
||
default="",
|
||
multiline=True,
|
||
),
|
||
io.Combo.Input("模型", options=NANO_BANANA_MODEL_OPTIONS, default="Nano Banana 2"),
|
||
io.Combo.Input(
|
||
"模型线路",
|
||
options=NANO_BANANA_ROUTE_OPTIONS,
|
||
default="畅速",
|
||
),
|
||
io.Combo.Input(
|
||
"思考等级",
|
||
options=["低", "高"],
|
||
default="高",
|
||
tooltip="仅 Nano Banana 2 生效：低=minimal，高=high。",
|
||
),
|
||
io.Combo.Input("分辨率", options=["1K", "2K", "4K"], default="2K"),
|
||
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=["1", "2", "4", "9"],
|
||
default="1",
|
||
tooltip="选择本次生成的图像数量。",
|
||
),
|
||
io.Int.Input(
|
||
"seed",
|
||
default=0,
|
||
min=0,
|
||
max=0xFFFFFFFFFFFFFFFF,
|
||
),
|
||
reference_images,
|
||
io.Combo.Input(
|
||
"缩放图片",
|
||
options=["不缩放", "智能缩放"],
|
||
default="不缩放",
|
||
tooltip="请求体超过 18 MiB 时，智能缩放会等比缩小占用最大的参考图。",
|
||
),
|
||
],
|
||
outputs=[
|
||
io.Image.Output(display_name="输出图像"),
|
||
],
|
||
# Accept the former 参考图1～参考图9 keys when executing workflows
|
||
# saved before the Autogrow migration.
|
||
accept_all_inputs=True,
|
||
)
|
||
|
||
@classmethod
|
||
def execute(
|
||
cls,
|
||
prompt,
|
||
模型,
|
||
分辨率,
|
||
宽高比,
|
||
生图数量,
|
||
模型线路="畅速",
|
||
seed=0,
|
||
思考等级="高",
|
||
缩放图片="不缩放",
|
||
**kwargs,
|
||
) -> io.NodeOutput:
|
||
start_time = time.time()
|
||
was_interrupted = False
|
||
生图数量 = int(生图数量)
|
||
|
||
model_name = 模型
|
||
# 兼容旧工作流/外部调用传入的“计费”字段。
|
||
模型线路 = kwargs.pop("计费", 模型线路)
|
||
unlimited_downloads = kwargs.pop("_o1key_unlimited_downloads", False) is True
|
||
requested_google_search = kwargs.pop("_o1key_google_search", False) is True
|
||
google_search = model_name == "Nano Banana 2" and requested_google_search
|
||
resize_mode = str(缩放图片)
|
||
if resize_mode not in {"不缩放", "智能缩放"}:
|
||
raise ValueError("缩放图片参数无效")
|
||
# 验证模型和宽高比、分辨率的组合是否合法
|
||
cls._validate_model_config(model_name, 宽高比, 分辨率)
|
||
|
||
if 思考等级 not in THINKING_LEVEL_MAP:
|
||
raise ValueError("思考等级无效，仅支持：低、高")
|
||
thinking_level = (
|
||
THINKING_LEVEL_MAP[思考等级]
|
||
if model_name == "Nano Banana 2"
|
||
else None
|
||
)
|
||
|
||
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))
|
||
|
||
reference_inputs = _collect_autogrow_inputs(kwargs.get("参考图组"))
|
||
if not reference_inputs:
|
||
# Keep execution compatibility with workflows created before
|
||
# the Autogrow input replaced the nine fixed image sockets.
|
||
reference_inputs = [
|
||
kwargs[f"参考图{i}"]
|
||
for i in range(1, 10)
|
||
if kwargs.get(f"参考图{i}") is not None
|
||
]
|
||
|
||
input_images = []
|
||
for image_input in reference_inputs:
|
||
input_images.extend(tensor_to_pil(image_input))
|
||
|
||
if len(input_images) > MAX_REFERENCE_IMAGES:
|
||
raise ValueError(
|
||
f"输入图像数量 {len(input_images)} 超过限制 {MAX_REFERENCE_IMAGES} 张"
|
||
)
|
||
|
||
batch_prompts = parse_batch_prompts(prompt)
|
||
|
||
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}张")
|
||
else:
|
||
mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式"
|
||
print(f"Nano Banana: {mode_str} | {分辨率} {宽高比} | {生图数量}张")
|
||
|
||
if batch_prompts or 生图数量 > 1:
|
||
prompts = batch_prompts if batch_prompts else [prompt]
|
||
images_per_prompt = 生图数量
|
||
total_tasks = len(prompts) * images_per_prompt
|
||
|
||
if pbar is not None:
|
||
pbar = ProgressBar(total_tasks)
|
||
|
||
def run_async_in_thread():
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
return loop.run_until_complete(
|
||
_run_with_interrupt(_process_batch_async(
|
||
base_url=base_url,
|
||
api_key=api_key,
|
||
prompts=prompts,
|
||
model=actual_model,
|
||
resolution=分辨率,
|
||
aspect_ratio=宽高比,
|
||
images_per_prompt=images_per_prompt,
|
||
input_images=input_images,
|
||
pbar=pbar,
|
||
thinking_level=thinking_level,
|
||
google_search=google_search,
|
||
resize_mode=resize_mode,
|
||
unlimited_downloads=unlimited_downloads,
|
||
))
|
||
)
|
||
finally:
|
||
asyncio.set_event_loop(None)
|
||
loop.close()
|
||
|
||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||
future = executor.submit(run_async_in_thread)
|
||
try:
|
||
results = future.result(timeout=_REQUEST_TIMEOUT)
|
||
except TimeoutError:
|
||
raise RuntimeError(f"任务执行超时（{_REQUEST_TIMEOUT}秒）")
|
||
|
||
success_count = sum(1 for r in results if r.get("success", False))
|
||
fail_count = len(results) - success_count
|
||
elapsed = time.time() - start_time
|
||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||
print(f"完成！总耗时 {time_str} | 成功: {success_count}/{total_tasks} | 失败: {fail_count}")
|
||
|
||
output_images = []
|
||
for r in results:
|
||
output_images.extend(r.get("output_images", []))
|
||
if not output_images:
|
||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||
output_images = [placeholder]
|
||
|
||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||
import gc; gc.collect()
|
||
return io.NodeOutput(output_tensor)
|
||
|
||
else:
|
||
def run_single():
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
async def _do():
|
||
async with create_http_client(
|
||
http2=True,
|
||
max_connections=_HTTP_MAX_CONNECTIONS,
|
||
max_keepalive_connections=_HTTP_MAX_KEEPALIVE_CONNECTIONS,
|
||
) 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,
|
||
thinking_level=thinking_level,
|
||
google_search=google_search,
|
||
download_semaphore=(
|
||
None
|
||
if unlimited_downloads
|
||
else asyncio.Semaphore(_MAX_DOWNLOAD_CONCURRENCY)
|
||
),
|
||
progress_callback=_make_progress_callback(pbar),
|
||
resize_mode=resize_mode,
|
||
)
|
||
return loop.run_until_complete(_run_with_interrupt(_do()))
|
||
finally:
|
||
asyncio.set_event_loop(None)
|
||
loop.close()
|
||
|
||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||
future = executor.submit(run_single)
|
||
generated_images, timing = future.result(timeout=_REQUEST_TIMEOUT)
|
||
|
||
output_tensor = _images_to_tensor_safe(generated_images, _NODE)
|
||
elapsed = time.time() - start_time
|
||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||
task_str = f"{timing['task_ms']/1000:.2f}s"
|
||
download_str = f"{timing['download_ms']/1000:.2f}s"
|
||
parse_str = f"{max(0, timing['parse_ms'] - timing['download_ms'])/1000:.2f}s"
|
||
inline_suffix = f" | 内联={timing['inline_images']}张" if timing['inline_images'] else ""
|
||
print(
|
||
f"Nano Banana: 完成 ✓ | task_id={timing['task_id']} | "
|
||
f"生成={len(generated_images)} 张 | 耗时={time_str} "
|
||
f"(生成 {task_str} | 下载 {download_str} | 解析 {parse_str}{inline_suffix})"
|
||
)
|
||
import gc; gc.collect()
|
||
return io.NodeOutput(output_tensor)
|
||
|
||
except InterruptProcessingException:
|
||
was_interrupted = True
|
||
print("Nano Banana: 用户取消")
|
||
raise
|
||
except ValueError as e:
|
||
if str(e) == "未授权！":
|
||
print("请联系作者授权后方可使用！")
|
||
raise ValueError("未授权！") from None
|
||
raise ValueError(str(e)) from None
|
||
except RuntimeError as e:
|
||
raise RuntimeError(str(e)) from None
|
||
except Exception as e:
|
||
raise RuntimeError(str(e)) from None
|
||
finally:
|
||
if not was_interrupted:
|
||
try:
|
||
client = _get_client()
|
||
client.base_url = base_url
|
||
balance_data = client.query_balance_sync()
|
||
balance_info = client.format_balance_info(balance_data)
|
||
print(f"Nano Banana: {balance_info}")
|
||
except Exception:
|
||
pass
|
||
import gc; gc.collect()
|