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.
1481 lines
58 KiB
Python
1481 lines
58 KiB
Python
"""
|
||
批量 Nano Banana 节点
|
||
ComfyUI 自定义节点，用于批量处理图像生成任务
|
||
支持多文件夹加载、同序号/同名/全匹配配对、智能命名保存
|
||
"""
|
||
|
||
import time
|
||
import math
|
||
import asyncio
|
||
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
|
||
from threading import Event
|
||
from typing import Any, Callable, Optional, Tuple, List
|
||
from PIL import Image
|
||
|
||
import torch
|
||
|
||
from comfy_api.latest import io
|
||
|
||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||
from ..utils.file_utils import (
|
||
ImageInfo,
|
||
load_images_from_folder,
|
||
pair_images_indexed,
|
||
pair_images_by_name,
|
||
pair_images_cartesian,
|
||
generate_timestamp_filename,
|
||
save_image,
|
||
)
|
||
from ..utils.config import get_base_url_by_route, get_api_key_or_raise
|
||
from ..utils.nano_banana_async import (
|
||
generate_nano_banana_async,
|
||
DEBUG_LOG_ENABLED as _ASYNC_DEBUG_LOG,
|
||
VERBOSE_LOG_ENABLED as _ASYNC_VERBOSE_LOG,
|
||
)
|
||
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,
|
||
)
|
||
|
||
# 导入 ComfyUI 原生进度条
|
||
try:
|
||
from comfy.utils import ProgressBar
|
||
PROGRESS_BAR_AVAILABLE = True
|
||
except ImportError:
|
||
PROGRESS_BAR_AVAILABLE = False
|
||
print("⚠️ BatchNanoBananaPro: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")
|
||
|
||
try:
|
||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||
INTERRUPT_AVAILABLE = True
|
||
except ImportError:
|
||
INTERRUPT_AVAILABLE = False
|
||
InterruptProcessingException = RuntimeError
|
||
processing_interrupted = lambda: False
|
||
|
||
# 导入 ComfyUI 的文件夹路径管理
|
||
try:
|
||
import folder_paths
|
||
FOLDER_PATHS_AVAILABLE = True
|
||
except ImportError:
|
||
FOLDER_PATHS_AVAILABLE = False
|
||
print("⚠️ BatchNanoBananaPro: folder_paths 不可用，将无法使用默认保存路径")
|
||
|
||
# 内存监控（可选）
|
||
try:
|
||
import psutil
|
||
MEMORY_MONITOR_AVAILABLE = True
|
||
except ImportError:
|
||
MEMORY_MONITOR_AVAILABLE = False
|
||
print("⚠️ BatchNanoBananaPro: psutil 不可用，内存监控功能禁用")
|
||
|
||
# ============================================================================
|
||
# 调试日志配置
|
||
# ============================================================================
|
||
# O1KEY_DEBUG_LOG 控制精简诊断；O1KEY_VERBOSE_LOG=1 才打印完整原始报文。
|
||
DEBUG_LOG_ENABLED = _ASYNC_DEBUG_LOG
|
||
REQUEST_LOG_ENABLED = _ASYNC_VERBOSE_LOG
|
||
|
||
THINKING_LEVEL_MAP = {
|
||
"低": "minimal",
|
||
"高": "high",
|
||
}
|
||
# ============================================================================
|
||
|
||
_NODE = "Nano Banana"
|
||
_BATCH_SIZE = 50
|
||
_DOWNLOAD_CONCURRENCY = 6
|
||
_HTTP_MAX_CONNECTIONS = 32
|
||
_HTTP_MAX_KEEPALIVE_CONNECTIONS = 16
|
||
_PER_BATCH_TIMEOUT_SECONDS = 1320
|
||
_BATCH_TIMEOUT_GRACE_SECONDS = 30
|
||
_THREAD_STOP_GRACE_SECONDS = 10
|
||
_INTERRUPT_CHECK_INTERVAL = 0.2
|
||
_MAX_FOLDER_INPUTS = 5
|
||
_MAX_FIXED_REFERENCE_IMAGES = 9
|
||
class _BatchStopRequested(RuntimeError):
|
||
pass
|
||
|
||
|
||
def _check_comfy_interrupt() -> None:
|
||
if INTERRUPT_AVAILABLE and processing_interrupted():
|
||
raise InterruptProcessingException()
|
||
|
||
|
||
def _make_interrupt_checker(stop_event: Event) -> Callable[[], None]:
|
||
def _check() -> None:
|
||
if stop_event.is_set():
|
||
raise _BatchStopRequested("Batch stop requested")
|
||
_check_comfy_interrupt()
|
||
|
||
return _check
|
||
|
||
|
||
async def _wait_for_interrupt(stop_event: Event) -> str:
|
||
while True:
|
||
if stop_event.is_set():
|
||
return "stop"
|
||
if INTERRUPT_AVAILABLE and processing_interrupted():
|
||
return "comfy"
|
||
await asyncio.sleep(_INTERRUPT_CHECK_INTERVAL)
|
||
|
||
|
||
async def _run_with_interrupt(coro, stop_event: Event):
|
||
request_task = asyncio.create_task(coro)
|
||
interrupt_task = asyncio.create_task(_wait_for_interrupt(stop_event))
|
||
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:
|
||
reason = interrupt_task.result()
|
||
request_task.cancel()
|
||
await asyncio.gather(request_task, return_exceptions=True)
|
||
if reason == "comfy":
|
||
raise InterruptProcessingException()
|
||
raise _BatchStopRequested("Batch stop requested")
|
||
|
||
for task in pending:
|
||
task.cancel()
|
||
await asyncio.gather(*pending, return_exceptions=True)
|
||
return request_task.result()
|
||
|
||
|
||
def _batch_timeout_seconds(total_tasks: int) -> int:
|
||
batch_count = max(1, math.ceil(total_tasks / _BATCH_SIZE))
|
||
return batch_count * _PER_BATCH_TIMEOUT_SECONDS + _BATCH_TIMEOUT_GRACE_SECONDS
|
||
|
||
|
||
def _normalize_output_format(fmt: Optional[str]) -> Optional[str]:
|
||
if not fmt:
|
||
return None
|
||
normalized = str(fmt).upper()
|
||
if normalized == "JPG":
|
||
return "JPEG"
|
||
if normalized == "WEBP":
|
||
return "WebP"
|
||
return normalized
|
||
|
||
|
||
def _format_extension(fmt: Optional[str]) -> str:
|
||
normalized = _normalize_output_format(fmt)
|
||
if normalized == "JPEG":
|
||
return ".jpg"
|
||
if normalized == "WebP":
|
||
return ".webp"
|
||
return ".png"
|
||
|
||
|
||
def _image_original_format(image: Image.Image) -> Optional[str]:
|
||
return _normalize_output_format(
|
||
getattr(image, "_o1key_original_format", None) or getattr(image, "format", None)
|
||
)
|
||
|
||
|
||
def _normalize_image_quality(value) -> int:
|
||
"""Return a 1-100 lossy image quality, defaulting to 95 for invalid input."""
|
||
try:
|
||
quality = int(value)
|
||
except (TypeError, ValueError):
|
||
quality = 95
|
||
return quality if 1 <= quality <= 100 else 95
|
||
|
||
|
||
def _save_generated_image(
|
||
image: Image.Image,
|
||
output_path: str,
|
||
save_format: str,
|
||
keep_original: bool,
|
||
image_quality: int = 95,
|
||
) -> None:
|
||
normalized = _normalize_output_format(save_format) or "PNG"
|
||
image_quality = _normalize_image_quality(image_quality)
|
||
|
||
if keep_original:
|
||
original_bytes = getattr(image, "_o1key_original_bytes", None)
|
||
original_format = _image_original_format(image)
|
||
if original_bytes is not None and original_format == normalized:
|
||
with open(output_path, "wb") as f:
|
||
f.write(original_bytes)
|
||
return
|
||
|
||
if normalized == "JPEG":
|
||
working = image if image.mode == "RGB" else image.convert("RGB")
|
||
working.save(output_path, format="JPEG", quality=image_quality, subsampling=0)
|
||
elif normalized == "WebP":
|
||
image.save(output_path, format="WEBP", lossless=False, quality=image_quality, method=6)
|
||
elif normalized == "PNG":
|
||
image.save(output_path, format="PNG", compress_level=0)
|
||
else:
|
||
image.save(output_path, format=normalized)
|
||
|
||
|
||
def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]:
|
||
if pbar is None:
|
||
return None
|
||
|
||
last_progress = [0.0]
|
||
|
||
def _on_progress(progress: float) -> None:
|
||
try:
|
||
progress = max(0.0, min(float(progress), 1.0))
|
||
except (TypeError, ValueError):
|
||
return
|
||
if progress <= last_progress[0]:
|
||
return
|
||
pbar.update(progress - last_progress[0])
|
||
last_progress[0] = progress
|
||
|
||
return _on_progress
|
||
|
||
|
||
async def _generate_single_async(
|
||
session: Any,
|
||
base_url: str,
|
||
api_key: str,
|
||
prompt: str,
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
thinking_level: Optional[str] = None,
|
||
images: Optional[List[Image.Image]] = None,
|
||
image_urls: Optional[List[str]] = None,
|
||
progress_callback: Optional[Callable[[float], None]] = None,
|
||
node_label: str = "BatchNanoBananaPro",
|
||
check_interrupt: Optional[Callable[[], None]] = None,
|
||
download_semaphore: Optional[asyncio.Semaphore] = None,
|
||
result_url_callback: Optional[Callable[[str], None]] = None,
|
||
log_task_success: bool = True,
|
||
log_downloads: bool = True,
|
||
upload_cache: Optional[dict] = None,
|
||
resize_mode: str = "不缩放",
|
||
) -> List[Image.Image]:
|
||
result_images, _ = await generate_nano_banana_async(
|
||
session=session,
|
||
base_url=base_url,
|
||
api_key=api_key,
|
||
prompt=prompt,
|
||
model=model,
|
||
resolution=resolution,
|
||
aspect_ratio=aspect_ratio,
|
||
thinking_level=thinking_level,
|
||
images=images,
|
||
image_urls=image_urls,
|
||
upload_cache=upload_cache,
|
||
node_label=node_label,
|
||
request_log_enabled=REQUEST_LOG_ENABLED,
|
||
check_interrupt=check_interrupt,
|
||
progress_callback=progress_callback,
|
||
download_semaphore=download_semaphore,
|
||
result_url_callback=result_url_callback,
|
||
log_task_success=log_task_success,
|
||
log_downloads=log_downloads,
|
||
resize_mode=resize_mode,
|
||
)
|
||
return result_images
|
||
|
||
|
||
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)
|
||
|
||
|
||
def _collect_autogrow_inputs(group) -> list:
|
||
"""按端口顺序提取 Autogrow 值，并兼容单值输入。"""
|
||
if isinstance(group, dict):
|
||
return [value for value in group.values() if value is not None]
|
||
if group is None:
|
||
return []
|
||
return [group]
|
||
|
||
|
||
def _path_count_from_label(value) -> int:
|
||
"""解析“1个路径”或旧版“1个文件夹”，异常值回退为 1。"""
|
||
try:
|
||
return max(1, min(int(str(value).split("个", 1)[0]), _MAX_FOLDER_INPUTS))
|
||
except (TypeError, ValueError):
|
||
return 1
|
||
|
||
|
||
def _path_input_name(index: int) -> str:
|
||
return "参考图1（主图）" if index == 1 else f"参考图{index}"
|
||
|
||
|
||
def _path_option(count: int):
|
||
"""构建指定数量图片路径的 DynamicCombo 分支。"""
|
||
inputs = [
|
||
io.String.Input(
|
||
_path_input_name(index),
|
||
default="",
|
||
placeholder="填写图片文件夹路径",
|
||
tooltip=(
|
||
"主图文件夹路径；请求时作为第 1 张参考图优先提交。"
|
||
if index == 1
|
||
else f"第 {index} 个参考图文件夹路径。"
|
||
),
|
||
)
|
||
for index in range(1, count + 1)
|
||
]
|
||
if count >= 2:
|
||
inputs.append(io.Combo.Input(
|
||
"图片配对模式",
|
||
options=BatchNanoBananaPro.PAIRING_MODES,
|
||
default="不配对",
|
||
tooltip="可按相同文件名、同序号 1:1、全部组合进行配对，或选择不配对。",
|
||
))
|
||
return io.DynamicCombo.Option(f"{count}个路径", inputs)
|
||
|
||
|
||
class BatchNanoBananaPro(io.ComfyNode):
|
||
"""
|
||
批量 Nano Banana 节点
|
||
|
||
功能：
|
||
- 从多个文件夹加载图片
|
||
- 支持四种配对模式：
|
||
* 相同文件名 - 仅匹配文件名相同的图片
|
||
* 同序号 - 按文件夹内顺序进行 1:1 配对
|
||
* 全匹配 - 笛卡尔积配对（所有可能组合）
|
||
* 不配对 - 固定参考图模式（文件夹图片依次与所有参考图组合）
|
||
- 批量调用 API 生成图像
|
||
- 智能命名保存（保留原始文件名）
|
||
- 并发控制（默认最大 100）
|
||
|
||
注意：
|
||
- 「不配对」模式只支持单个文件夹
|
||
- 支持的模型列表从 models_config.py 动态加载
|
||
- 要添加/禁用模型，请编辑 models_config.py 文件
|
||
"""
|
||
|
||
# 模型展示名到基础 ID 的映射
|
||
MODEL_DISPLAY_NAMES = NANO_BANANA_MODEL_OPTIONS
|
||
|
||
# 配对模式
|
||
PAIRING_MODES = ["相同文件名", "同序号", "全匹配", "不配对"]
|
||
PAIRING_MODE_ALIASES = {
|
||
"按相同图片命名": "相同文件名",
|
||
"1*N": "全匹配",
|
||
}
|
||
|
||
@classmethod
|
||
def _normalize_pairing_mode(cls, value: str) -> str:
|
||
normalized = cls.PAIRING_MODE_ALIASES.get(value, value)
|
||
if normalized not in cls.PAIRING_MODES:
|
||
raise ValueError(
|
||
f"图片配对模式“{value}”无效，可选：{', '.join(cls.PAIRING_MODES)}"
|
||
)
|
||
return normalized
|
||
|
||
@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 != "1K":
|
||
raise ValueError(
|
||
f"Nano Banana 模型仅支持 1K 分辨率，"
|
||
f"当前选择的 {resolution} 不支持"
|
||
)
|
||
|
||
def __init__(self):
|
||
pass
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
reference_images = io.Autogrow.Input(
|
||
"参考图组",
|
||
template=io.Autogrow.TemplateNames(
|
||
input=io.Image.Input("参考图"),
|
||
names=[
|
||
f"参考图{index}"
|
||
for index in range(1, _MAX_FIXED_REFERENCE_IMAGES + 1)
|
||
],
|
||
min=0,
|
||
),
|
||
tooltip="连接后自动增加输入端参考图，编号会接在图片路径之后；最多 9 张，并在请求末尾提交。",
|
||
)
|
||
|
||
return io.Schema(
|
||
node_id="BatchNanoBananaPro",
|
||
display_name="Nano Banana 批量跑图",
|
||
category="image/batch",
|
||
inputs=[
|
||
io.String.Input(
|
||
"prompt",
|
||
default="一个中国女子的OOTD",
|
||
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="高"),
|
||
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.DynamicCombo.Input(
|
||
"图片路径数量",
|
||
options=[_path_option(count) for count in range(1, _MAX_FOLDER_INPUTS + 1)],
|
||
tooltip="选择后仅显示对应数量的参考图文件夹路径；最多 5 个。",
|
||
),
|
||
reference_images,
|
||
io.Combo.Input(
|
||
"缩放图片",
|
||
options=["不缩放", "智能缩放"],
|
||
default="不缩放",
|
||
tooltip="请求体超过 18 MiB 时，智能缩放会等比缩小占用最大的参考图。",
|
||
),
|
||
io.Combo.Input(
|
||
"图片输出格式",
|
||
options=["原始", "JPEG", "PNG", "WebP"],
|
||
default="原始",
|
||
),
|
||
io.Int.Input(
|
||
"图片质量",
|
||
default=95,
|
||
min=1,
|
||
max=100,
|
||
step=1,
|
||
display_mode=io.NumberDisplay.number,
|
||
tooltip="仅 JPEG 和 WebP 输出生效。",
|
||
),
|
||
io.Combo.Input(
|
||
"图片保存命名规则",
|
||
options=["和原始图片名保持一致", "自然数字"],
|
||
default="和原始图片名保持一致",
|
||
),
|
||
io.String.Input(
|
||
"图片保存路径",
|
||
default="",
|
||
placeholder="留空时保存到 ComfyUI output 目录",
|
||
),
|
||
io.Int.Input(
|
||
"seed",
|
||
default=0,
|
||
min=0,
|
||
max=0xFFFFFFFFFFFFFFFF,
|
||
),
|
||
],
|
||
outputs=[
|
||
io.Image.Output(display_name="输出图像"),
|
||
],
|
||
accept_all_inputs=True,
|
||
)
|
||
|
||
@classmethod
|
||
def execute(
|
||
cls,
|
||
prompt,
|
||
模型,
|
||
思考等级="高",
|
||
分辨率="2K",
|
||
宽高比="智能",
|
||
图片路径数量=None,
|
||
模型线路="畅速",
|
||
seed=0,
|
||
图片输出格式="原始",
|
||
图片质量=95,
|
||
图片保存命名规则="和原始图片名保持一致",
|
||
图片保存路径="",
|
||
缩放图片="不缩放",
|
||
**kwargs,
|
||
) -> io.NodeOutput:
|
||
model_name = 模型
|
||
# 兼容旧工作流/外部调用传入的“计费”字段。
|
||
模型线路 = kwargs.pop("计费", 模型线路)
|
||
|
||
legacy_group = kwargs.get("图片文件夹数量")
|
||
nested_path_group = (
|
||
图片路径数量
|
||
if isinstance(图片路径数量, dict)
|
||
else legacy_group if isinstance(legacy_group, dict) else None
|
||
)
|
||
path_inputs = nested_path_group or kwargs
|
||
selected_count = _path_count_from_label(
|
||
path_inputs.get(
|
||
"图片路径数量",
|
||
path_inputs.get("图片文件夹数量", 图片路径数量 or legacy_group),
|
||
)
|
||
)
|
||
image_paths = [
|
||
path_inputs.get(
|
||
_path_input_name(index),
|
||
path_inputs.get(f"图片路径{index}", kwargs.get(f"图片路径{index}", "")),
|
||
)
|
||
for index in range(1, _MAX_FOLDER_INPUTS + 1)
|
||
]
|
||
# DynamicCombo 只定义当前数量的路径；旧工作流直接调用仍保留全部编号路径。
|
||
if nested_path_group is not None:
|
||
image_paths = image_paths[:selected_count] + [""] * (_MAX_FOLDER_INPUTS - selected_count)
|
||
pairing_mode = path_inputs.get("图片配对模式", kwargs.get("图片配对模式", "不配对"))
|
||
pairing_mode = cls._normalize_pairing_mode(pairing_mode)
|
||
|
||
reference_inputs = _collect_autogrow_inputs(kwargs.get("参考图组"))
|
||
if not reference_inputs:
|
||
reference_inputs = [
|
||
kwargs[f"参考图{index}"]
|
||
for index in range(1, _MAX_FIXED_REFERENCE_IMAGES + 1)
|
||
if kwargs.get(f"参考图{index}") is not None
|
||
]
|
||
reference_kwargs = {
|
||
f"参考图{index}": value
|
||
for index, value in enumerate(reference_inputs, start=1)
|
||
}
|
||
|
||
node = cls()
|
||
output_tensor, = node.process_batch(
|
||
prompt=prompt,
|
||
文件夹1=image_paths[0],
|
||
文件夹2=image_paths[1],
|
||
文件夹3=image_paths[2],
|
||
文件夹4=image_paths[3],
|
||
文件夹5=image_paths[4],
|
||
seed=seed,
|
||
图片配对模式=pairing_mode,
|
||
模型=model_name,
|
||
思考等级=思考等级,
|
||
模型线路=模型线路,
|
||
宽高比=宽高比,
|
||
分辨率=分辨率,
|
||
图片格式=图片输出格式,
|
||
图片质量=图片质量,
|
||
保存路径=图片保存路径,
|
||
命名规则=图片保存命名规则,
|
||
缩放图片=缩放图片,
|
||
**reference_kwargs,
|
||
)
|
||
return io.NodeOutput(output_tensor)
|
||
|
||
def _load_folders(
|
||
self,
|
||
folder1: str,
|
||
folder2: Optional[str],
|
||
folder3: Optional[str],
|
||
folder4: Optional[str],
|
||
folder5: Optional[str] = None,
|
||
folder6: Optional[str] = None,
|
||
folder7: Optional[str] = None,
|
||
folder8: Optional[str] = None,
|
||
folder9: Optional[str] = None,
|
||
) -> List[List[ImageInfo]]:
|
||
"""
|
||
加载所有文件夹中的图片
|
||
|
||
Args:
|
||
folder1-9: 文件夹路径
|
||
|
||
Returns:
|
||
图片列表的列表
|
||
"""
|
||
folders = [folder1, folder2, folder3, folder4, folder5, folder6, folder7, folder8, folder9]
|
||
all_images = []
|
||
|
||
for i, folder in enumerate(folders, 1):
|
||
if folder and folder.strip():
|
||
try:
|
||
images = load_images_from_folder(folder)
|
||
if images:
|
||
all_images.append(images)
|
||
except ValueError as e:
|
||
print(f"BatchNanoBananaPro: 文件夹{i} 加载失败 - {e}")
|
||
|
||
return all_images
|
||
|
||
def _create_pairs(
|
||
self,
|
||
image_lists: List[List[ImageInfo]],
|
||
pairing_mode: str,
|
||
manual_images: Optional[List[ImageInfo]] = None,
|
||
) -> List[Tuple[ImageInfo, ...]]:
|
||
"""
|
||
根据配对模式创建图片组合
|
||
|
||
Args:
|
||
image_lists: 参与配对的文件夹图片列表
|
||
pairing_mode: 配对模式 (相同文件名, 同序号, 全匹配, 不配对)
|
||
manual_images: 手动输入的参考图(固定追加到每组末尾)
|
||
|
||
Returns:
|
||
配对后的元组列表
|
||
|
||
Raises:
|
||
ValueError: 不配对模式下填入多个配对文件夹时
|
||
"""
|
||
pairing_mode = self._normalize_pairing_mode(pairing_mode)
|
||
def _apply_extras(pairs: List[Tuple[ImageInfo, ...]]) -> List[Tuple[ImageInfo, ...]]:
|
||
"""为每组追加固定参考图。"""
|
||
manual_tuple = tuple(manual_images) if manual_images else ()
|
||
return [pair + manual_tuple for pair in pairs]
|
||
|
||
# === 不配对模式 ===
|
||
if pairing_mode == "不配对":
|
||
# 验证：不配对模式只支持单个文件夹
|
||
if len(image_lists) > 1:
|
||
raise ValueError("「不配对」模式只支持单个配对图片路径，请清空多余的图片路径或改用其他配对模式")
|
||
|
||
if image_lists:
|
||
base_pairs = [(img,) for img in image_lists[0]]
|
||
else:
|
||
return []
|
||
|
||
return _apply_extras(base_pairs)
|
||
|
||
# === 相同文件名、同序号和全匹配模式 ===
|
||
if not image_lists:
|
||
return []
|
||
|
||
# 文件夹图片配对
|
||
if len(image_lists) == 1:
|
||
base_pairs = [(img,) for img in image_lists[0]]
|
||
elif pairing_mode == "相同文件名":
|
||
base_pairs = list(pair_images_by_name(*image_lists))
|
||
elif pairing_mode == "同序号":
|
||
base_pairs = list(pair_images_indexed(*image_lists))
|
||
else: # 全匹配
|
||
base_pairs = list(pair_images_cartesian(*image_lists))
|
||
|
||
return _apply_extras(base_pairs)
|
||
|
||
async def _generate_single_task(
|
||
self,
|
||
session: Any,
|
||
base_url: str,
|
||
api_key: str,
|
||
prompt: str,
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
thinking_level: Optional[str],
|
||
images: List[ImageInfo],
|
||
output_folder: str,
|
||
task_index: int,
|
||
base_filename: str = None,
|
||
base_extension: str = None,
|
||
image_format: str = "原始",
|
||
image_quality: int = 95,
|
||
progress_callback: Optional[Callable[[float], None]] = None,
|
||
naming_rule: str = "和原始图片名保持一致",
|
||
check_interrupt: Optional[Callable[[], None]] = None,
|
||
download_semaphore: Optional[asyncio.Semaphore] = None,
|
||
upload_cache: Optional[dict] = None,
|
||
resize_mode: str = "不缩放",
|
||
) -> dict:
|
||
"""
|
||
执行单个生成任务（异步生图接口）
|
||
"""
|
||
result = {
|
||
"task_index": task_index,
|
||
"prompt": prompt,
|
||
"success": False,
|
||
"generated_count": 0,
|
||
"saved_files": [],
|
||
"output_images": [],
|
||
"error": None
|
||
}
|
||
task_started = time.time()
|
||
|
||
try:
|
||
if check_interrupt:
|
||
check_interrupt()
|
||
|
||
# 准备输入图片
|
||
input_pil_images = [info.image for info in images]
|
||
|
||
# 调用异步生图接口生成图片
|
||
generated_images = []
|
||
try:
|
||
gen_images = await _generate_single_async(
|
||
session=session,
|
||
base_url=base_url,
|
||
api_key=api_key,
|
||
prompt=prompt,
|
||
model=model,
|
||
resolution=resolution,
|
||
aspect_ratio=aspect_ratio,
|
||
thinking_level=thinking_level,
|
||
images=input_pil_images if input_pil_images else None,
|
||
progress_callback=progress_callback,
|
||
node_label=f"BatchNanoBananaPro#{task_index + 1}",
|
||
check_interrupt=check_interrupt,
|
||
download_semaphore=download_semaphore,
|
||
log_task_success=True,
|
||
upload_cache=upload_cache,
|
||
resize_mode=resize_mode,
|
||
)
|
||
generated_images.extend(gen_images)
|
||
except (InterruptProcessingException, _BatchStopRequested, asyncio.CancelledError):
|
||
raise
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
print(
|
||
f"BatchNanoBananaPro#{task_index + 1}: 失败 "
|
||
f"| model={model} | resolution={resolution} | error={error_msg}"
|
||
)
|
||
if REQUEST_LOG_ENABLED:
|
||
import traceback
|
||
print(traceback.format_exc())
|
||
result["error"] = error_msg
|
||
|
||
# 保存生成的图片到磁盘（始终保存）
|
||
import os
|
||
|
||
for i, gen_img in enumerate(generated_images):
|
||
if check_interrupt:
|
||
check_interrupt()
|
||
if image_format == "原始":
|
||
save_format = _image_original_format(gen_img) or "PNG"
|
||
else:
|
||
save_format = image_format
|
||
|
||
save_ext = _format_extension(save_format)
|
||
if save_format in ("JPEG", "WebP") and gen_img.mode in ("RGBA", "LA", "P"):
|
||
gen_img = gen_img.convert("RGB")
|
||
|
||
# 根据命名规则确定文件名
|
||
if naming_rule == "自然数字":
|
||
# 以自然数字命名 (1, 2, 3, ...)，按任务提交顺序（task_index）命名
|
||
# 如果一个任务生成多张图，则为 task_index_0, task_index_1...
|
||
if i == 0:
|
||
# 第一张图直接使用任务索引
|
||
num = task_index + 1
|
||
filename = f"{num}{save_ext}"
|
||
else:
|
||
# 同一任务的后续图片加后缀
|
||
num = task_index + 1
|
||
filename = f"{num}_{i}{save_ext}"
|
||
output_path = os.path.join(output_folder, filename)
|
||
|
||
# 如果文件已存在，添加额外后缀避免覆盖
|
||
if os.path.exists(output_path):
|
||
counter = 1
|
||
base_num = num
|
||
while True:
|
||
if i == 0:
|
||
filename = f"{base_num}_dup{counter}{save_ext}"
|
||
else:
|
||
filename = f"{base_num}_{i}_dup{counter}{save_ext}"
|
||
output_path = os.path.join(output_folder, filename)
|
||
if not os.path.exists(output_path):
|
||
break
|
||
counter += 1
|
||
elif base_filename:
|
||
# 和原始图片名保持一致，重名则在后面加 1,2,3,4...
|
||
base_name = base_filename
|
||
counter = 0
|
||
while True:
|
||
if counter == 0:
|
||
filename = f"{base_name}{save_ext}"
|
||
else:
|
||
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:
|
||
output_path = generate_timestamp_filename(
|
||
output_folder=output_folder,
|
||
extension=save_ext
|
||
)
|
||
|
||
# 保存时不做额外压缩
|
||
_save_generated_image(
|
||
gen_img,
|
||
output_path,
|
||
save_format=save_format,
|
||
keep_original=(image_format == "原始"),
|
||
image_quality=image_quality,
|
||
)
|
||
|
||
result["saved_files"].append(output_path)
|
||
gen_img = None
|
||
|
||
# 只有生成了图片才标记为成功
|
||
if len(generated_images) > 0:
|
||
result["success"] = True
|
||
result["generated_count"] = len(generated_images)
|
||
if result["success"]:
|
||
elapsed = time.time() - task_started
|
||
print(
|
||
f"BatchNanoBananaPro#{task_index + 1}: 完成 ✓ | "
|
||
f"生成={len(generated_images)} 张 | 耗时={elapsed:.1f}s"
|
||
)
|
||
|
||
except (InterruptProcessingException, _BatchStopRequested, asyncio.CancelledError):
|
||
raise
|
||
except Exception as e:
|
||
result["error"] = str(e)
|
||
|
||
return result
|
||
|
||
async def _process_batch_async(
|
||
self,
|
||
pairs: List[Tuple[ImageInfo, ...]],
|
||
prompt: str,
|
||
model: str,
|
||
resolution: str,
|
||
aspect_ratio: str,
|
||
thinking_level: Optional[str],
|
||
output_folder: str,
|
||
base_url: str,
|
||
api_key: str,
|
||
pbar=None,
|
||
prompts_per_task: Optional[List[str]] = None,
|
||
image_format: str = "原始",
|
||
image_quality: int = 95,
|
||
naming_rule: str = "和原始图片名保持一致",
|
||
check_interrupt: Optional[Callable[[], None]] = None,
|
||
resize_mode: str = "不缩放",
|
||
) -> List[dict]:
|
||
"""
|
||
异步批量处理所有任务（异步生图接口）
|
||
"""
|
||
total_tasks = len(pairs)
|
||
|
||
max_concurrent = _BATCH_SIZE
|
||
download_semaphore = asyncio.Semaphore(_DOWNLOAD_CONCURRENCY)
|
||
|
||
print(
|
||
f"BatchNanoBananaPro: 开始 | 任务={total_tasks} | "
|
||
f"本批并发={min(total_tasks, max_concurrent)}/{max_concurrent}"
|
||
)
|
||
|
||
all_results = []
|
||
completed = 0
|
||
success_count = 0
|
||
fail_count = 0
|
||
|
||
# 计算生成批次数量
|
||
num_batches = math.ceil(total_tasks / max_concurrent)
|
||
|
||
# 内存监控初始化
|
||
if MEMORY_MONITOR_AVAILABLE and total_tasks > 50:
|
||
import psutil
|
||
process = psutil.Process()
|
||
initial_memory = process.memory_info().rss / 1024 / 1024
|
||
print(f"BatchNanoBananaPro: 初始内存使用: {initial_memory:.1f} MB")
|
||
|
||
# 进度打印配置：任务数 >= 50 时，额外显示百分比里程碑
|
||
show_milestone = total_tasks >= 50
|
||
milestones = [0.2, 0.4, 0.6, 0.8, 1.0] # 20%, 40%, 60%, 80%, 100%
|
||
milestone_index = 0
|
||
|
||
if num_batches > 1:
|
||
print(f"BatchNanoBananaPro: 任务数 {total_tasks} 超过并发上限 {max_concurrent}，将分 {num_batches} 批执行")
|
||
|
||
upload_cache = {}
|
||
|
||
async with create_http_client(
|
||
http2=True,
|
||
max_connections=_HTTP_MAX_CONNECTIONS,
|
||
max_keepalive_connections=_HTTP_MAX_KEEPALIVE_CONNECTIONS,
|
||
) as session:
|
||
# 分批处理：每批最多 50 个任务
|
||
for batch_idx in range(num_batches):
|
||
if check_interrupt:
|
||
check_interrupt()
|
||
start_idx = batch_idx * max_concurrent
|
||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||
batch_pairs = pairs[start_idx:end_idx]
|
||
|
||
if num_batches > 1:
|
||
print(f"BatchNanoBananaPro: 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...")
|
||
|
||
# 创建当前批次的任务
|
||
tasks = []
|
||
for i, pair in enumerate(batch_pairs):
|
||
if check_interrupt:
|
||
check_interrupt()
|
||
# 批量提示词模式时，每个任务使用对应的提示词；否则使用统一提示词
|
||
task_prompt = prompts_per_task[start_idx + i] if prompts_per_task else prompt
|
||
|
||
# 提取文件夹1图片的名称作为保存文件名
|
||
base_filename = None
|
||
base_extension = None
|
||
if pair and len(pair) > 0:
|
||
first_image = pair[0]
|
||
if hasattr(first_image, 'filename'):
|
||
base_filename = first_image.filename
|
||
if hasattr(first_image, 'extension'):
|
||
base_extension = first_image.extension
|
||
|
||
task = asyncio.create_task(
|
||
self._generate_single_task(
|
||
session=session,
|
||
base_url=base_url,
|
||
api_key=api_key,
|
||
prompt=task_prompt,
|
||
model=model,
|
||
resolution=resolution,
|
||
aspect_ratio=aspect_ratio,
|
||
thinking_level=thinking_level,
|
||
images=list(pair),
|
||
output_folder=output_folder,
|
||
task_index=start_idx + i,
|
||
base_filename=base_filename,
|
||
base_extension=base_extension,
|
||
image_format=image_format,
|
||
image_quality=image_quality,
|
||
progress_callback=_make_progress_callback(pbar),
|
||
naming_rule=naming_rule,
|
||
check_interrupt=check_interrupt,
|
||
download_semaphore=download_semaphore,
|
||
upload_cache=upload_cache,
|
||
resize_mode=resize_mode,
|
||
)
|
||
)
|
||
tasks.append(task)
|
||
|
||
# 收集当前批次的结果（使用 gather 保持任务提交顺序）
|
||
try:
|
||
batch_results = await asyncio.gather(*tasks, return_exceptions=True)
|
||
except (InterruptProcessingException, _BatchStopRequested, asyncio.CancelledError):
|
||
for task in tasks:
|
||
task.cancel()
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
raise
|
||
|
||
stop_error = next(
|
||
(
|
||
result for result in batch_results
|
||
if isinstance(result, (InterruptProcessingException, _BatchStopRequested))
|
||
),
|
||
None,
|
||
)
|
||
if stop_error is not None:
|
||
raise stop_error
|
||
|
||
# 规范化结果，并使用真实任务编号而不是完成顺序。
|
||
normalized_batch_results = []
|
||
for idx, result in enumerate(batch_results):
|
||
if isinstance(result, Exception):
|
||
result_data = {
|
||
"task_index": start_idx + idx,
|
||
"success": False,
|
||
"error": str(result),
|
||
"generated_count": 0,
|
||
"saved_files": []
|
||
}
|
||
elif isinstance(result, dict):
|
||
result_data = result
|
||
else:
|
||
result_data = {
|
||
"task_index": start_idx + idx,
|
||
"success": False,
|
||
"error": "Unknown result type",
|
||
"generated_count": 0,
|
||
"saved_files": []
|
||
}
|
||
|
||
result_data.setdefault("task_index", start_idx + idx)
|
||
normalized_batch_results.append(result_data)
|
||
completed += 1
|
||
|
||
if result_data and result_data.get("success", False):
|
||
success_count += 1
|
||
else:
|
||
fail_count += 1
|
||
|
||
# 大任务额外显示百分比里程碑
|
||
if show_milestone and milestone_index < len(milestones):
|
||
if pbar is not None and getattr(pbar, "total", 0):
|
||
progress = pbar.current / pbar.total
|
||
else:
|
||
progress = success_count / total_tasks
|
||
if progress >= milestones[milestone_index]:
|
||
percentage = int(milestones[milestone_index] * 100)
|
||
print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<")
|
||
milestone_index += 1
|
||
|
||
# 当前批次完成后汇总结果并清理内存。
|
||
all_results.extend(normalized_batch_results)
|
||
|
||
# 统计当前批次的结果
|
||
batch_success = sum(1 for r in normalized_batch_results if r.get("success", False))
|
||
batch_fail = len(normalized_batch_results) - batch_success
|
||
batch_generated = sum(r.get("generated_count", 0) for r in normalized_batch_results)
|
||
|
||
print(
|
||
f"BatchNanoBananaPro: 批次 {batch_idx + 1}/{num_batches} 完成 "
|
||
f"| 成功={batch_success} | 失败={batch_fail} "
|
||
f"| 生成={batch_generated} | 总进度={completed}/{total_tasks}"
|
||
)
|
||
|
||
# 强制垃圾回收，释放内存
|
||
import gc
|
||
gc.collect()
|
||
|
||
# 内存监控
|
||
if MEMORY_MONITOR_AVAILABLE and total_tasks > 50:
|
||
current_memory = process.memory_info().rss / 1024 / 1024
|
||
memory_increase = current_memory - initial_memory
|
||
print(f"BatchNanoBananaPro: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)")
|
||
|
||
# 内存警告阈值（2GB）
|
||
if current_memory > 2000:
|
||
print(f"⚠️ BatchNanoBananaPro: 内存使用过高！但图片已分批保存，即使崩溃也不会丢失已完成的任务")
|
||
|
||
# 短暂暂停，让系统有时间处理文件I/O
|
||
await asyncio.sleep(0.5)
|
||
if check_interrupt:
|
||
check_interrupt()
|
||
|
||
return all_results
|
||
|
||
def process_batch(
|
||
self,
|
||
prompt: str,
|
||
文件夹1: str,
|
||
文件夹2: str,
|
||
文件夹3: str,
|
||
文件夹4: str,
|
||
文件夹5: str,
|
||
seed: int,
|
||
图片配对模式: str,
|
||
模型: str,
|
||
模型线路: str = "畅速",
|
||
宽高比: str = "智能",
|
||
分辨率: str = "2K",
|
||
图片格式: str = "原始",
|
||
图片质量: int = 95,
|
||
思考等级: str = "高",
|
||
保存路径: str = "",
|
||
命名规则: str = "和原始图片名保持一致",
|
||
缩放图片: str = "不缩放",
|
||
**kwargs
|
||
) -> Tuple[torch.Tensor]:
|
||
"""
|
||
批量处理图像生成任务
|
||
|
||
Args:
|
||
prompt: 提示词
|
||
文件夹1-5: 图片文件夹路径
|
||
seed: 随机种子
|
||
保存路径: 输出保存路径
|
||
图片配对模式: 相同文件名、同序号、全匹配或不配对
|
||
模型: 模型名称
|
||
宽高比: 输出宽高比
|
||
分辨率: 输出分辨率
|
||
**kwargs: 动态参考图输入 (参考图1-5)
|
||
|
||
Returns:
|
||
输出图像张量
|
||
"""
|
||
图片配对模式 = self._normalize_pairing_mode(图片配对模式)
|
||
模型线路 = kwargs.pop("计费", 模型线路)
|
||
图片质量 = _normalize_image_quality(图片质量)
|
||
if 缩放图片 not in {"不缩放", "智能缩放"}:
|
||
raise ValueError("缩放图片参数无效")
|
||
start_time = time.time()
|
||
was_interrupted = False
|
||
skip_balance_query = False
|
||
|
||
# 只有 Nano Banana 2 支持思考等级；其他模型不向接口传该字段。
|
||
if 思考等级 not in THINKING_LEVEL_MAP:
|
||
raise ValueError(f"不支持的思考等级: {思考等级}")
|
||
thinking_level = THINKING_LEVEL_MAP[思考等级] if 模型 == "Nano Banana 2" else None
|
||
|
||
# 验证模型配置是否合法（使用与 NanoBanana 相同的验证逻辑）
|
||
self._validate_model_config(模型, 宽高比, 分辨率)
|
||
|
||
# 模型线路与主模型共同决定实际 API 模型名。
|
||
display_model = 模型
|
||
模型 = resolve_nano_banana_model(display_model, 模型线路)
|
||
|
||
|
||
try:
|
||
# 收集 5 个图片路径(按编号 1-5)
|
||
path_list = [文件夹1, 文件夹2, 文件夹3, 文件夹4, 文件夹5]
|
||
filled_paths = [f for f in path_list if f and f.strip()]
|
||
if not filled_paths:
|
||
raise ValueError("请至少填写一个图片路径，该节点专为批量文件夹处理设计")
|
||
|
||
pairing_path_entries = [
|
||
(idx, path) for idx, path in enumerate(path_list, 1)
|
||
if path and path.strip()
|
||
]
|
||
|
||
# 校验：图片配对模式仅在至少 2 个路径下生效
|
||
if len(pairing_path_entries) < 2 and 图片配对模式 != "不配对":
|
||
raise ValueError(
|
||
f"当前参与配对的图片路径仅 {len(pairing_path_entries)} 个，无法使用「{图片配对模式}」配对模式。\n"
|
||
f"图片配对模式仅在至少 2 个图片路径时起作用，"
|
||
f"请将「图片配对模式」改为「不配对」，或增加配对图片路径。"
|
||
)
|
||
|
||
# 加载配对路径图片
|
||
print("BatchNanoBananaPro: 开始加载图片...")
|
||
image_lists = []
|
||
for idx, p in pairing_path_entries:
|
||
try:
|
||
imgs = load_images_from_folder(p)
|
||
if imgs:
|
||
image_lists.append(imgs)
|
||
else:
|
||
print(f"BatchNanoBananaPro: 图片路径{idx} 未找到图片，已跳过")
|
||
except ValueError as e:
|
||
print(f"BatchNanoBananaPro: 图片路径{idx} 加载失败 - {e}")
|
||
|
||
# 验证是否有可用图片
|
||
total_folder_images = sum(len(lst) for lst in image_lists)
|
||
if total_folder_images == 0:
|
||
raise ValueError("图片路径中未找到任何图片，请检查路径是否正确")
|
||
|
||
# 处理独立的参考图输入
|
||
manual_images = []
|
||
for i in range(1, 10): # 1-9
|
||
key = f"参考图{i}"
|
||
if key in kwargs and kwargs[key] is not None:
|
||
pil_images = tensor_to_pil(kwargs[key])
|
||
for j, img in enumerate(pil_images):
|
||
manual_images.append(
|
||
ImageInfo(
|
||
image=img,
|
||
filename=f"manual_{i}_{j}",
|
||
extension=".png",
|
||
source_path=""
|
||
)
|
||
)
|
||
|
||
# 创建配对并追加固定参考图
|
||
pairs = self._create_pairs(
|
||
image_lists,
|
||
图片配对模式,
|
||
manual_images if manual_images else None,
|
||
)
|
||
|
||
if not pairs:
|
||
raise ValueError("配对结果为空，请检查输入")
|
||
|
||
# 校验单任务参考图数量上限（与单节点一致，模型最多接受 14 张）
|
||
max_images_per_task = max(len(pair) for pair in pairs)
|
||
if max_images_per_task > 14:
|
||
raise ValueError(
|
||
f"单个任务的参考图数量 {max_images_per_task} 超过限制 14 张"
|
||
f"（配对图片 + 固定参考图）。请减少配对路径或固定参考图的数量。"
|
||
)
|
||
|
||
# 解析批量提示词（使用 --- 分隔多个提示词）
|
||
batch_prompts = parse_batch_prompts(prompt)
|
||
prompts_per_task = None
|
||
if batch_prompts:
|
||
# 展开 pairs × prompts：每个图片组合 × 每个提示词 = 一个任务
|
||
expanded_pairs = []
|
||
expanded_prompts = []
|
||
for pair in pairs:
|
||
for bp in batch_prompts:
|
||
expanded_pairs.append(pair)
|
||
expanded_prompts.append(bp)
|
||
pairs = expanded_pairs
|
||
prompts_per_task = expanded_prompts
|
||
|
||
total_tasks = len(pairs)
|
||
|
||
if batch_prompts:
|
||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务")
|
||
else:
|
||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务")
|
||
|
||
# 创建 ComfyUI 原生进度条
|
||
pbar = None
|
||
if PROGRESS_BAR_AVAILABLE:
|
||
pbar = ProgressBar(total_tasks)
|
||
|
||
# 检查保存路径（重要！）
|
||
has_save_path = bool(保存路径 and 保存路径.strip())
|
||
if not has_save_path:
|
||
# 使用 ComfyUI 默认 output 目录作为保存路径
|
||
if FOLDER_PATHS_AVAILABLE:
|
||
保存路径 = folder_paths.get_output_directory()
|
||
has_save_path = True
|
||
print(f"BatchNanoBananaPro: 未设置保存路径，将使用 ComfyUI 默认 output 目录: {保存路径}")
|
||
else:
|
||
print("BatchNanoBananaPro: 未设置保存路径，图片将输出到节点")
|
||
|
||
if has_save_path:
|
||
# 验证保存路径
|
||
import os
|
||
try:
|
||
os.makedirs(保存路径, exist_ok=True)
|
||
# 测试写入权限
|
||
test_file = os.path.join(保存路径, ".write_test")
|
||
with open(test_file, 'w') as f:
|
||
f.write("test")
|
||
os.remove(test_file)
|
||
print(f"BatchNanoBananaPro: 保存路径验证通过: {保存路径}")
|
||
except Exception as e:
|
||
raise ValueError(f"保存路径无效或无写入权限: {保存路径} - {str(e)}")
|
||
|
||
# 获取 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('保存路径', '')
|
||
user_set_save_path = bool(original_save_path and original_save_path.strip())
|
||
|
||
# 执行批量生成
|
||
# 在新线程中运行异步代码，避免事件循环冲突
|
||
stop_event = Event()
|
||
check_interrupt = _make_interrupt_checker(stop_event)
|
||
batch_timeout = _batch_timeout_seconds(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(
|
||
self._process_batch_async(
|
||
pairs=pairs,
|
||
prompt=prompt,
|
||
model=模型,
|
||
resolution=分辨率,
|
||
aspect_ratio=宽高比,
|
||
thinking_level=thinking_level,
|
||
output_folder=保存路径,
|
||
base_url=base_url,
|
||
api_key=api_key,
|
||
pbar=pbar,
|
||
prompts_per_task=prompts_per_task,
|
||
image_format=图片格式,
|
||
image_quality=图片质量,
|
||
naming_rule=命名规则,
|
||
check_interrupt=check_interrupt,
|
||
resize_mode=缩放图片,
|
||
),
|
||
stop_event,
|
||
)
|
||
)
|
||
except (InterruptProcessingException, _BatchStopRequested):
|
||
raise
|
||
except Exception as e:
|
||
# 即使崩溃，也记录错误
|
||
print(f"BatchNanoBananaPro: 任务执行异常: {str(e)}")
|
||
raise
|
||
finally:
|
||
pending = asyncio.all_tasks(loop)
|
||
for task in pending:
|
||
task.cancel()
|
||
if pending:
|
||
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
|
||
loop.run_until_complete(loop.shutdown_asyncgens())
|
||
asyncio.set_event_loop(None)
|
||
loop.close()
|
||
|
||
# Keep the worker cancellable; a ThreadPoolExecutor context manager would
|
||
# wait indefinitely for a stuck worker when leaving the timeout handler.
|
||
executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="nano-banana-batch")
|
||
future = executor.submit(run_async_in_thread)
|
||
try:
|
||
if DEBUG_LOG_ENABLED:
|
||
print(
|
||
f"BatchNanoBananaPro: 总超时={batch_timeout}s "
|
||
f"| 批次数={math.ceil(total_tasks / _BATCH_SIZE)} | 每批={_BATCH_SIZE}"
|
||
)
|
||
results = future.result(timeout=batch_timeout)
|
||
except FuturesTimeoutError:
|
||
skip_balance_query = True
|
||
stop_event.set()
|
||
try:
|
||
future.result(timeout=_THREAD_STOP_GRACE_SECONDS)
|
||
except BaseException:
|
||
pass
|
||
print(f"BatchNanoBananaPro: 任务执行超时（{batch_timeout}秒），已请求取消后台任务")
|
||
raise RuntimeError(
|
||
f"批量任务执行超时（{batch_timeout}秒），后台网络任务已取消"
|
||
) from None
|
||
except InterruptProcessingException:
|
||
was_interrupted = True
|
||
skip_balance_query = True
|
||
stop_event.set()
|
||
raise
|
||
except Exception:
|
||
# 即使失败，也尝试返回部分结果
|
||
if 'all_saved_files' in locals():
|
||
print(f"BatchNanoBananaPro: 部分保存的图片: {len(all_saved_files)} 张")
|
||
raise
|
||
finally:
|
||
stop_event.set()
|
||
executor.shutdown(wait=False, cancel_futures=True)
|
||
|
||
# 统计结果
|
||
success_count = sum(1 for r in results if r.get("success", False))
|
||
fail_count = len(results) - success_count
|
||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||
all_saved_files = []
|
||
for r in results:
|
||
all_saved_files.extend(r.get("saved_files", []))
|
||
|
||
elapsed = time.time() - start_time
|
||
|
||
# 格式化时间
|
||
if elapsed < 1:
|
||
time_str = f"{elapsed:.3f}s"
|
||
else:
|
||
time_str = f"{elapsed:.2f}s"
|
||
|
||
# 计算平均耗时
|
||
avg_time = elapsed / success_count if success_count > 0 else 0
|
||
avg_time_str = f"{avg_time:.1f}s/张" if success_count > 0 else "N/A"
|
||
|
||
# 精简统计信息
|
||
has_save_path = bool(保存路径 and 保存路径.strip())
|
||
is_default_path = not bool(kwargs.get('保存路径', '').strip() if '保存路径' in locals() else False)
|
||
print("=" * 60)
|
||
print(f"完成！总耗时 {time_str} | 成功: {success_count}/{total_tasks} | 生成 {total_generated} 张 | 平均 {avg_time_str}")
|
||
if has_save_path:
|
||
if is_default_path:
|
||
print(f"保存路径: {保存路径} (ComfyUI 默认 output 目录)")
|
||
else:
|
||
print(f"保存路径: {保存路径}")
|
||
else:
|
||
print("保存路径: 未设置（仅输出到节点）")
|
||
|
||
# 失败详情（如果有）
|
||
failed_results = [r for r in results if not r.get("success", False)]
|
||
if failed_results:
|
||
print(f"失败任务汇总 ({len(failed_results)} 个):")
|
||
for failed in failed_results[:3]:
|
||
task_num = failed.get('task_index', '?') + 1
|
||
error_msg = failed.get('error', '未知错误')
|
||
print(f" #{task_num}: {error_msg}")
|
||
|
||
if len(failed_results) > 3:
|
||
remaining = [str(r.get('task_index', '?') + 1) for r in failed_results[3:]]
|
||
print(f" 其他失败任务编号: {', '.join(remaining)}")
|
||
|
||
# 收集最后几张图片用于 ComfyUI 节点输出
|
||
output_images = []
|
||
max_output_images = 10
|
||
|
||
if all_saved_files:
|
||
# 从磁盘加载最近的图片
|
||
recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):]
|
||
for file_path in recent_files:
|
||
try:
|
||
img = Image.open(file_path)
|
||
output_images.append(img)
|
||
except Exception as e:
|
||
print(f"BatchNanoBananaPro: 无法加载图片 {file_path} - {e}")
|
||
|
||
# 策略3：如果还是没有图片，创建一个占位图
|
||
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()
|
||
|
||
# 打印最终统计信息
|
||
total_saved = len(all_saved_files)
|
||
print(f"BatchNanoBananaPro: 任务完成！共保存 {total_saved} 张图片到磁盘")
|
||
if total_saved > 0:
|
||
print(f"BatchNanoBananaPro: 最新保存的文件: {all_saved_files[-1]}")
|
||
|
||
return (output_tensor,)
|
||
|
||
except InterruptProcessingException:
|
||
was_interrupted = True
|
||
skip_balance_query = True
|
||
print("BatchNanoBananaPro: 用户取消")
|
||
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 skip_balance_query and not was_interrupted:
|
||
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
|
||
gc.collect()
|
||
if DEBUG_LOG_ENABLED:
|
||
print("BatchNanoBananaPro: 最终内存清理完成")
|