751 lines
27 KiB
Python
751 lines
27 KiB
Python
"""
|
|
批量 Nano Banana Pro 节点
|
|
ComfyUI 自定义节点，用于批量处理图像生成任务
|
|
支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存
|
|
"""
|
|
|
|
import time
|
|
import math
|
|
import random
|
|
import asyncio
|
|
import aiohttp
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Optional, Tuple, List
|
|
from PIL import Image
|
|
|
|
import torch
|
|
import numpy as np
|
|
|
|
from ..utils.image_utils import tensor_to_pil, pil_to_tensor
|
|
from ..utils.file_utils import (
|
|
ImageInfo,
|
|
load_images_from_folder,
|
|
pair_images_indexed,
|
|
pair_images_cartesian,
|
|
generate_output_filename,
|
|
save_image
|
|
)
|
|
from ..clients.gemini_client import GeminiAPIClient
|
|
from ..models_config import get_enabled_models
|
|
|
|
# 导入 ComfyUI 原生进度条
|
|
try:
|
|
from comfy.utils import ProgressBar
|
|
PROGRESS_BAR_AVAILABLE = True
|
|
except ImportError:
|
|
PROGRESS_BAR_AVAILABLE = False
|
|
print("⚠️ BatchNanoBananaPro: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")
|
|
|
|
|
|
class BatchNanoBananaPro:
|
|
"""
|
|
批量 Nano Banana Pro 节点
|
|
|
|
功能：
|
|
- 从多个文件夹加载图片
|
|
- 支持三种配对模式：
|
|
* 1:1 - 索引配对（文件夹之间按位置配对）
|
|
* 1*N - 笛卡尔积配对（所有可能组合）
|
|
* 不配对 - 固定参考图模式（文件夹图片依次与所有参考图组合）
|
|
- 批量调用 API 生成图像
|
|
- 智能命名保存（保留原始文件名）
|
|
- 并发控制（默认最大 100）
|
|
|
|
注意：
|
|
- 「不配对」模式只支持单个文件夹
|
|
- 支持的模型列表从 models_config.py 动态加载
|
|
- 要添加/禁用模型，请编辑 models_config.py 文件
|
|
"""
|
|
|
|
# 支持的模型列表（从配置文件动态加载）
|
|
MODELS = None # 将在 INPUT_TYPES 中动态获取
|
|
|
|
# 支持的宽高比列表
|
|
ASPECT_RATIOS = [
|
|
"1:1", "4:3", "3:4", "16:9", "9:16",
|
|
"2:3", "3:2", "4:5", "5:4", "21:9"
|
|
]
|
|
|
|
# 支持的分辨率列表
|
|
RESOLUTIONS = ["1K", "2K", "4K"]
|
|
|
|
# 配对模式
|
|
PAIRING_MODES = ["1:1", "1*N", "不配对"]
|
|
|
|
def __init__(self):
|
|
"""初始化节点"""
|
|
self.client = None
|
|
|
|
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
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
定义输入参数
|
|
|
|
ComfyUI 节点规范：
|
|
- required: 必选参数
|
|
- optional: 可选参数
|
|
"""
|
|
# 从配置文件动态获取启用的模型列表
|
|
enabled_models = get_enabled_models()
|
|
|
|
# 如果没有启用的模型，使用空列表（会导致节点不可用，提示用户配置）
|
|
if not enabled_models:
|
|
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
|
|
|
# 创建9个独立的图像输入
|
|
optional_inputs = {}
|
|
for i in range(1, 10): # 1-9
|
|
optional_inputs[f"参考图{i}"] = ("IMAGE",)
|
|
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {
|
|
"default": "一个中国女子的OOTD",
|
|
"multiline": True
|
|
}),
|
|
"模型": (enabled_models, {
|
|
"default": enabled_models[0]
|
|
}),
|
|
"宽高比": (cls.ASPECT_RATIOS, {
|
|
"default": "1:1"
|
|
}),
|
|
"分辨率": (cls.RESOLUTIONS, {
|
|
"default": "2K"
|
|
}),
|
|
"像素缩放": ("BOOLEAN", {
|
|
"default": False
|
|
}),
|
|
"分辨率像素": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.1,
|
|
"max": 100.0,
|
|
"step": 0.1,
|
|
"display": "number"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff
|
|
}),
|
|
"文件夹1": ("STRING", {
|
|
"default": "",
|
|
"multiline": False
|
|
}),
|
|
"文件夹2": ("STRING", {
|
|
"default": "",
|
|
"multiline": False
|
|
}),
|
|
"文件夹3": ("STRING", {
|
|
"default": "",
|
|
"multiline": False
|
|
}),
|
|
"文件夹4": ("STRING", {
|
|
"default": "",
|
|
"multiline": False
|
|
}),
|
|
"保存路径": ("STRING", {
|
|
"default": "",
|
|
"multiline": False
|
|
}),
|
|
"图片配对模式": (cls.PAIRING_MODES, {
|
|
"default": "不配对"
|
|
})
|
|
},
|
|
"optional": optional_inputs
|
|
}
|
|
|
|
# 返回值类型
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("输出图像",)
|
|
|
|
# 执行函数名
|
|
FUNCTION = "process_batch"
|
|
|
|
# 节点分类
|
|
CATEGORY = "image/batch"
|
|
|
|
def _load_folders(
|
|
self,
|
|
folder1: str,
|
|
folder2: Optional[str],
|
|
folder3: Optional[str],
|
|
folder4: Optional[str],
|
|
enable_scaling: bool,
|
|
target_megapixels: float
|
|
) -> List[List[ImageInfo]]:
|
|
"""
|
|
加载所有文件夹中的图片
|
|
|
|
Args:
|
|
folder1-4: 文件夹路径
|
|
enable_scaling: 是否启用像素缩放
|
|
target_megapixels: 目标像素数（百万像素）
|
|
|
|
Returns:
|
|
图片列表的列表
|
|
"""
|
|
folders = [folder1, folder2, folder3, folder4]
|
|
all_images = []
|
|
|
|
for i, folder in enumerate(folders, 1):
|
|
if folder and folder.strip():
|
|
try:
|
|
images = load_images_from_folder(folder)
|
|
if images:
|
|
# 应用像素缩放
|
|
if enable_scaling:
|
|
scaled_images = []
|
|
for img_info in images:
|
|
scaled_img = self.resize_to_megapixels(
|
|
img_info.image,
|
|
target_megapixels
|
|
)
|
|
# 创建新的 ImageInfo，保留其他元数据
|
|
scaled_info = ImageInfo(
|
|
image=scaled_img,
|
|
filename=img_info.filename,
|
|
extension=img_info.extension,
|
|
source_path=img_info.source_path
|
|
)
|
|
scaled_images.append(scaled_info)
|
|
images = scaled_images
|
|
|
|
all_images.append(images)
|
|
print(f"BatchNanoBananaPro: 文件夹{i} 加载了 {len(images)} 张图片")
|
|
else:
|
|
print(f"BatchNanoBananaPro: 文件夹{i} 为空或没有有效图片")
|
|
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: 配对模式 (1:1, 1*N, 不配对)
|
|
manual_images: 手动输入的参考图
|
|
|
|
Returns:
|
|
配对后的元组列表
|
|
|
|
Raises:
|
|
ValueError: 不配对模式下填入多个文件夹时
|
|
"""
|
|
# === 新模式：不配对 ===
|
|
if pairing_mode == "不配对":
|
|
# 验证：只支持单个文件夹
|
|
if len(image_lists) > 1:
|
|
raise ValueError("「不配对」模式只支持单个文件夹，请清空其他文件夹路径")
|
|
|
|
# 场景1：有文件夹 + 有参考图
|
|
if image_lists and manual_images:
|
|
folder_images = image_lists[0]
|
|
# 每张文件夹图片 + 所有参考图
|
|
pairs = []
|
|
for img in folder_images:
|
|
pair = (img,) + tuple(manual_images)
|
|
pairs.append(pair)
|
|
return pairs
|
|
|
|
# 场景2：有文件夹 + 无参考图
|
|
elif image_lists:
|
|
# 每张图片单独成组
|
|
return [(img,) for img in image_lists[0]]
|
|
|
|
# 场景3：无文件夹 + 有参考图
|
|
elif manual_images:
|
|
# 每张参考图单独成组
|
|
return [(img,) for img in manual_images]
|
|
|
|
else:
|
|
return []
|
|
|
|
# === 原有逻辑：1:1 和 1*N ===
|
|
# 如果有手动参考图，添加到列表中（所有参考图作为一个列表）
|
|
if manual_images:
|
|
image_lists.append(manual_images)
|
|
|
|
if not image_lists:
|
|
return []
|
|
|
|
# 如果只有一个列表，直接返回每个图片作为单元素元组
|
|
if len(image_lists) == 1:
|
|
return [(img,) for img in image_lists[0]]
|
|
|
|
# 根据配对模式选择配对函数
|
|
if pairing_mode == "1:1":
|
|
pairs = pair_images_indexed(*image_lists)
|
|
else: # 1*N
|
|
pairs = pair_images_cartesian(*image_lists)
|
|
|
|
return pairs
|
|
|
|
async def _generate_single_task(
|
|
self,
|
|
client: GeminiAPIClient,
|
|
session: aiohttp.ClientSession,
|
|
prompt: str,
|
|
model: str,
|
|
resolution: str,
|
|
aspect_ratio: str,
|
|
images: List[ImageInfo],
|
|
output_folder: str,
|
|
task_index: int
|
|
) -> dict:
|
|
"""
|
|
执行单个生成任务
|
|
|
|
Args:
|
|
client: API 客户端
|
|
session: aiohttp 会话
|
|
prompt: 提示词
|
|
model: 模型名称
|
|
resolution: 分辨率
|
|
aspect_ratio: 宽高比
|
|
images: 输入图片列表
|
|
output_folder: 输出文件夹
|
|
task_index: 任务索引
|
|
|
|
Returns:
|
|
包含结果信息的字典
|
|
"""
|
|
result = {
|
|
"task_index": task_index,
|
|
"success": False,
|
|
"generated_count": 0,
|
|
"saved_files": [],
|
|
"error": None
|
|
}
|
|
|
|
try:
|
|
# 准备输入图片
|
|
input_pil_images = [info.image for info in images]
|
|
|
|
# 调用 API 生成图片（固定生成1次）
|
|
generated_images = []
|
|
try:
|
|
gen_result = await client.generate_single_async(
|
|
prompt=prompt,
|
|
model=model,
|
|
resolution=resolution,
|
|
aspect_ratio=aspect_ratio,
|
|
images=input_pil_images,
|
|
session=session
|
|
)
|
|
if gen_result:
|
|
generated_images.extend(gen_result)
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
print(f"BatchNanoBananaPro: 任务 {task_index + 1} 生成失败 - {error_msg}")
|
|
result["error"] = error_msg
|
|
|
|
# 保存生成的图片
|
|
for i, gen_img in enumerate(generated_images):
|
|
# 使用任务索引作为唯一标识，确保并发安全
|
|
output_path = generate_output_filename(
|
|
source_images=list(images),
|
|
batch_index=i,
|
|
output_folder=output_folder,
|
|
extension=".png",
|
|
task_id=f"task{task_index}"
|
|
)
|
|
save_image(gen_img, output_path)
|
|
result["saved_files"].append(output_path)
|
|
|
|
# 只有生成了图片才标记为成功
|
|
if len(generated_images) > 0:
|
|
result["success"] = True
|
|
result["generated_count"] = len(generated_images)
|
|
|
|
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,
|
|
output_folder: str,
|
|
pbar=None
|
|
) -> List[dict]:
|
|
"""
|
|
异步批量处理所有任务
|
|
|
|
Args:
|
|
pairs: 配对后的图片组合
|
|
prompt: 提示词
|
|
model: 模型名称
|
|
resolution: 分辨率
|
|
aspect_ratio: 宽高比
|
|
output_folder: 输出文件夹
|
|
|
|
Returns:
|
|
所有任务的结果列表
|
|
"""
|
|
if self.client is None:
|
|
self.client = GeminiAPIClient()
|
|
|
|
# 固定最大并发数为 100
|
|
max_concurrent = 100
|
|
|
|
total_tasks = len(pairs)
|
|
all_results = []
|
|
completed = 0
|
|
success_count = 0
|
|
fail_count = 0
|
|
|
|
# 计算分批数量
|
|
num_batches = math.ceil(total_tasks / max_concurrent)
|
|
|
|
# 进度打印配置：任务数 >= 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} 批执行")
|
|
|
|
connector = aiohttp.TCPConnector(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)
|
|
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):
|
|
task = asyncio.create_task(
|
|
self._generate_single_task(
|
|
client=self.client,
|
|
session=session,
|
|
prompt=prompt,
|
|
model=model,
|
|
resolution=resolution,
|
|
aspect_ratio=aspect_ratio,
|
|
images=list(pair),
|
|
output_folder=output_folder,
|
|
task_index=start_idx + i
|
|
)
|
|
)
|
|
tasks.append(task)
|
|
|
|
# 使用 as_completed 实时获取完成的任务
|
|
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": []
|
|
}
|
|
all_results.append(result_data)
|
|
else:
|
|
result_data = result
|
|
all_results.append(result)
|
|
except Exception as e:
|
|
result_data = {
|
|
"success": False,
|
|
"error": str(e),
|
|
"generated_count": 0,
|
|
"saved_files": []
|
|
}
|
|
all_results.append(result_data)
|
|
|
|
completed += 1
|
|
|
|
# 根据成功/失败状态打印不同信息
|
|
if result_data and result_data.get("success", False):
|
|
success_count += 1
|
|
print(f"BatchNanoBananaPro: 任务 {completed}/{total_tasks} 成功 ✓")
|
|
else:
|
|
fail_count += 1
|
|
# 提取错误信息的第一行
|
|
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
|
# 截取第一行或前50个字符
|
|
if '\n' in error_msg:
|
|
error_msg = error_msg.split('\n')[0]
|
|
if len(error_msg) > 50:
|
|
error_msg = error_msg[:50] + "..."
|
|
print(f"BatchNanoBananaPro: 任务 {completed}/{total_tasks} 失败 ✗ - {error_msg}")
|
|
|
|
# 更新 ComfyUI 原生进度条
|
|
if pbar is not None:
|
|
pbar.update(1)
|
|
|
|
# 大任务额外显示百分比里程碑
|
|
if show_milestone and milestone_index < len(milestones):
|
|
progress = completed / total_tasks
|
|
if progress >= milestones[milestone_index]:
|
|
percentage = int(milestones[milestone_index] * 100)
|
|
print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<")
|
|
milestone_index += 1
|
|
|
|
return all_results
|
|
|
|
def process_batch(
|
|
self,
|
|
prompt: str,
|
|
文件夹1: str,
|
|
文件夹2: str,
|
|
文件夹3: str,
|
|
文件夹4: str,
|
|
像素缩放: bool,
|
|
分辨率像素: float,
|
|
seed: int,
|
|
保存路径: str,
|
|
图片配对模式: str,
|
|
模型: str,
|
|
宽高比: str,
|
|
分辨率: str,
|
|
**kwargs
|
|
) -> Tuple[torch.Tensor]:
|
|
"""
|
|
批量处理图像生成任务
|
|
|
|
Args:
|
|
prompt: 提示词
|
|
文件夹1-4: 图片文件夹路径
|
|
像素缩放: 是否启用像素缩放
|
|
分辨率像素: 目标像素数（百万像素）
|
|
seed: 随机种子
|
|
保存路径: 输出保存路径
|
|
图片配对模式: 1:1 或 1*N
|
|
模型: 模型名称
|
|
宽高比: 输出宽高比
|
|
分辨率: 输出分辨率
|
|
**kwargs: 动态参考图输入 (参考图1-9)
|
|
|
|
Returns:
|
|
输出图像张量
|
|
"""
|
|
start_time = time.time()
|
|
|
|
try:
|
|
# 设置随机种子（用于本地随机操作）
|
|
random.seed(seed)
|
|
np.random.seed(seed % (2**32))
|
|
# 验证保存路径
|
|
if not 保存路径 or not 保存路径.strip():
|
|
raise ValueError("请提供保存路径")
|
|
|
|
# 加载文件夹图片
|
|
print("BatchNanoBananaPro: 开始加载图片...")
|
|
image_lists = self._load_folders(
|
|
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
|
像素缩放, 分辨率像素
|
|
)
|
|
|
|
# 处理独立的参考图输入
|
|
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):
|
|
# 如果启用像素缩放，也对参考图进行缩放
|
|
if 像素缩放:
|
|
img = self.resize_to_megapixels(img, 分辨率像素)
|
|
|
|
manual_images.append(
|
|
ImageInfo(
|
|
image=img,
|
|
filename=f"manual_{i}_{j}",
|
|
extension=".png",
|
|
source_path=""
|
|
)
|
|
)
|
|
|
|
if manual_images:
|
|
print(f"BatchNanoBananaPro: 加载了 {len(manual_images)} 张参考图")
|
|
|
|
# 验证是否有图片
|
|
total_folder_images = sum(len(lst) for lst in image_lists)
|
|
total_manual_images = len(manual_images)
|
|
|
|
if total_folder_images == 0 and total_manual_images == 0:
|
|
raise ValueError("未找到任何图片，请检查文件夹路径或提供参考图")
|
|
|
|
# 创建配对
|
|
print(f"BatchNanoBananaPro: 使用 {图片配对模式} 模式创建配对...")
|
|
pairs = self._create_pairs(image_lists, 图片配对模式, manual_images if manual_images else None)
|
|
|
|
if not pairs:
|
|
raise ValueError("配对结果为空，请检查输入")
|
|
|
|
total_tasks = len(pairs)
|
|
print(f"BatchNanoBananaPro: 共 {total_tasks} 组配对")
|
|
|
|
# 创建 ComfyUI 原生进度条
|
|
pbar = None
|
|
if PROGRESS_BAR_AVAILABLE:
|
|
pbar = ProgressBar(total_tasks)
|
|
|
|
# 初始化 API 客户端
|
|
if self.client is None:
|
|
try:
|
|
self.client = GeminiAPIClient()
|
|
except ValueError as e:
|
|
raise ValueError(f"初始化 API 客户端失败: {str(e)}")
|
|
|
|
# 执行批量生成
|
|
print("BatchNanoBananaPro: 开始批量生成...")
|
|
|
|
# 在新线程中运行异步代码，避免事件循环冲突
|
|
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(
|
|
pairs=pairs,
|
|
prompt=prompt,
|
|
model=模型,
|
|
resolution=分辨率,
|
|
aspect_ratio=宽高比,
|
|
output_folder=保存路径,
|
|
pbar=pbar
|
|
)
|
|
)
|
|
finally:
|
|
loop.close()
|
|
|
|
# 使用线程池在新线程中运行事件循环
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
future = executor.submit(run_async_in_thread)
|
|
results = future.result()
|
|
|
|
# 统计结果
|
|
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
|
|
|
|
# 精简统计信息
|
|
print("=" * 50)
|
|
print(f"BatchNanoBananaPro 处理完成 | 总耗时: {elapsed:.2f}s | 成功: {success_count}/{total_tasks} | 生成: {total_generated}张")
|
|
print(f"保存路径: {保存路径}")
|
|
|
|
# 失败详情（如果有）
|
|
failed_results = [r for r in results if not r.get("success", False)]
|
|
if failed_results:
|
|
# 收集失败任务的索引
|
|
failed_indices = [str(r.get('task_index', '?') + 1) for r in failed_results[:5]]
|
|
failed_str = ",".join(failed_indices)
|
|
if len(failed_results) > 5:
|
|
failed_str += f"... (共{len(failed_results)}个)"
|
|
# 显示第一个失败原因作为示例
|
|
first_error = failed_results[0].get('error', '未知错误')
|
|
print(f"失败 {len(failed_results)}个: 任务{failed_str} - {first_error}")
|
|
|
|
# 收集所有生成的图片
|
|
output_images = []
|
|
for file_path in all_saved_files:
|
|
try:
|
|
img = Image.open(file_path)
|
|
output_images.append(img)
|
|
except Exception as e:
|
|
print(f"BatchNanoBananaPro: 无法加载图片 {file_path} - {e}")
|
|
|
|
# 如果没有生成成功的图片，创建一个占位图
|
|
if not output_images:
|
|
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
|
output_images = [placeholder]
|
|
|
|
# 转换为张量
|
|
output_tensor = pil_to_tensor(output_images)
|
|
|
|
return (output_tensor,)
|
|
|
|
except ValueError as e:
|
|
# 检测是否为授权错误
|
|
if str(e) == "未授权！":
|
|
print("请联系作者授权后方可使用！")
|
|
else:
|
|
print(f"BatchNanoBananaPro: 输入错误 - {str(e)}")
|
|
raise
|
|
|
|
except RuntimeError as e:
|
|
print(f"BatchNanoBananaPro: 运行时错误 - {str(e)}")
|
|
raise
|
|
|
|
except Exception as e:
|
|
print(f"BatchNanoBananaPro: 未知错误 - {str(e)}")
|
|
raise
|
|
|
|
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"{balance_info}")
|
|
print("=" * 50)
|
|
except Exception as e:
|
|
print(f"⚠️ 余额查询失败 - {str(e)}")
|
|
print("=" * 50) |