""" Nano Banana Pro 节点 ComfyUI 自定义节点，用于调用 Gemini 3 Pro 模型生成图像 """ import time import random from typing import Optional, Tuple import torch import numpy as np from PIL import Image from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts from ..clients.gemini_client import GeminiAPIClient from ..models_config import get_enabled_models, get_model_description # 导入 ComfyUI 原生进度条 try: from comfy.utils import ProgressBar PROGRESS_BAR_AVAILABLE = True except ImportError: PROGRESS_BAR_AVAILABLE = False print("⚠️ NanoBananaPro: comfy.utils.ProgressBar 不可用，将只使用终端进度显示") class NanoBananaPro: """ Nano Banana Pro 节点 功能： - 文生图：基于提示词生成图像 - 图生图：基于输入图像和提示词生成新图像 - 批量生成：支持并发生成多张图像 注意： - 支持的模型列表从 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"] def __init__(self): """初始化节点""" self.client = None @classmethod def INPUT_TYPES(cls): """ 定义输入参数 ComfyUI 节点规范： - required: 必选参数 - optional: 可选参数 """ # 从配置文件动态获取启用的模型列表 enabled_models = get_enabled_models() # 如果没有启用的模型，使用空列表（会导致节点不可用，提示用户配置） if not enabled_models: enabled_models = ["请在 models_config.py 中启用至少一个模型"] # 创建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" }), "生图数量": ("INT", { "default": 1, "min": 1, "max": 1000, "step": 1 }), "像素缩放": ("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 }) }, "optional": optional_inputs } # 返回值类型 RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("输出图像",) # 执行函数名 FUNCTION = "generate" # 节点分类 CATEGORY = "image/generation" def resize_to_megapixels( self, image: Image.Image, target_megapixels: float ) -> Image.Image: """ 将图像缩放到指定的总像素数，保持纵横比 Args: image: PIL Image 对象 target_megapixels: 目标像素数（百万像素） Returns: 缩放后的 PIL Image Example: >>> resized = self.resize_to_megapixels(img, 2.0) # 缩放到2百万像素 """ # 计算当前像素数 current_pixels = image.width * image.height target_pixels = int(target_megapixels * 1_000_000) # 如果当前像素数已经接近目标，则不缩放 if abs(current_pixels - target_pixels) / target_pixels < 0.05: return image # 计算缩放比例 scale = (target_pixels / current_pixels) ** 0.5 # 计算新尺寸 new_width = int(image.width * scale) new_height = int(image.height * scale) # 确保至少为1像素 new_width = max(1, new_width) new_height = max(1, new_height) # 使用 Lanczos 重采样 resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) return resized_image def validate_inputs( self, images: Optional[torch.Tensor], batch_size: int ) -> None: """ 验证输入参数 Args: images: 输入图像张量（可选） batch_size: 批次大小 Raises: ValueError: 如果输入参数不合法 """ # 检查图像数量 if images is not None: num_images = images.shape[0] if num_images > 14: raise ValueError( f"输入图像数量 {num_images} 超过限制 14 张，请减少输入图像数量" ) # 检查批次大小 if batch_size < 1 or batch_size > 1000: raise ValueError( f"批次大小 {batch_size} 超出范围 [1, 1000]" ) def generate( self, prompt: str, 模型: str, 宽高比: str, 分辨率: str, 生图数量: int, 像素缩放: bool, 分辨率像素: float, seed: int, **kwargs ) -> Tuple[torch.Tensor]: """ 生成图像 Args: prompt: 提示词 模型: 模型名称 宽高比: 宽高比 分辨率: 分辨率 生图数量: 批次大小 像素缩放: 是否启用像素缩放 分辨率像素: 目标像素数（百万像素） seed: 随机种子 **kwargs: 动态参考图输入 (参考图1-9) Returns: 生成的图像张量 (IMAGE,) """ start_time = time.time() # 创建 ComfyUI 原生进度条 pbar = None if PROGRESS_BAR_AVAILABLE: pbar = ProgressBar(生图数量) try: # 设置随机种子（用于本地随机操作） random.seed(seed) np.random.seed(seed % (2**32)) # 初始化 API 客户端 if self.client is None: try: self.client = GeminiAPIClient() except ValueError as e: raise ValueError(f"初始化失败: {str(e)}") # 收集独立输入的参考图 input_images = [] for i in range(1, 10): # 1-9 key = f"参考图{i}" if key in kwargs and kwargs[key] is not None: pil_imgs = tensor_to_pil(kwargs[key]) input_images.extend(pil_imgs) # 验证输入图像数量 if input_images: if len(input_images) > 14: raise ValueError( f"输入图像数量 {len(input_images)} 超过限制 14 张，请减少输入图像数量" ) # 应用像素缩放（如果启用） if input_images and 像素缩放: scaled_images = [] for img in input_images: scaled = self.resize_to_megapixels(img, 分辨率像素) scaled_images.append(scaled) input_images = scaled_images print(f"Nano Banana Pro: 已缩放 {len(scaled_images)} 张图像到 {分辨率像素}M 像素") # 转换为 API 所需的格式 if input_images: print(f"Nano Banana Pro: 图生图模式 (输入 {len(input_images)} 张图像)") # 解析批量提示词 batch_prompts = parse_batch_prompts(prompt) # 统计变量 success_count = 0 fail_count = 0 # 进度回调 - 实时显示每个任务的完成状态，并更新 ComfyUI 进度条 def progress_callback(current, total, success, error_msg=None): nonlocal success_count, fail_count if success: success_count += 1 print(f"Nano Banana Pro: ✓ [{current}/{total}] 第 {success_count} 张生成成功") else: fail_count += 1 error_brief = error_msg[:50] + "..." if error_msg and len(error_msg) > 50 else error_msg print(f"Nano Banana Pro: ✗ [{current}/{total}] 生成失败 - {error_brief}") # 更新 ComfyUI 原生进度条 if pbar is not None: pbar.update(1) # 根据是否有批量提示词选择生成模式 if batch_prompts: # 批量提示词模式 num_prompts = len(batch_prompts) total_images = num_prompts * 生图数量 print(f"Nano Banana Pro: 批量提示词模式 ({num_prompts} 个提示词 × {生图数量} 张/提示词 = {total_images} 张图)") print(f"Nano Banana Pro: 发送请求") print(f"Nano Banana Pro: 生图中...") # 重新创建进度条以匹配实际总数 if pbar is not None: pbar = ProgressBar(total_images) generated_images = self.client.generate_multi_prompts_sync( prompts=batch_prompts, model=模型, resolution=分辨率, aspect_ratio=宽高比, images_per_prompt=生图数量, images=input_images, progress_callback=progress_callback ) if fail_count > 0: print(f"Nano Banana Pro: 生图完成 (成功: {success_count}, 失败: {fail_count})") else: print(f"Nano Banana Pro: 全部生图成功！") else: # 单提示词模式 print(f"Nano Banana Pro: {'图生图' if input_images else '文生图'}模式") print(f"Nano Banana Pro: 发送请求") print(f"Nano Banana Pro: 生图中...") generated_images = self.client.generate_sync( prompt=prompt, model=模型, resolution=分辨率, aspect_ratio=宽高比, batch_size=生图数量, images=input_images, progress_callback=progress_callback ) if fail_count > 0: print(f"Nano Banana Pro: 生图完成 (成功: {success_count}, 失败: {fail_count})") else: print(f"Nano Banana Pro: 全部生图成功！") # 转换输出图像 output_tensor = pil_to_tensor(generated_images) # 计算耗时 elapsed = time.time() - start_time print(f"Nano Banana Pro: 完成生图 (耗时: {elapsed:.2f}s, 成功生成 {len(generated_images)} 张图像)") return (output_tensor,) except ValueError as e: # 检测是否为授权错误 if str(e) == "未授权！": print("请联系作者授权后方可使用！") else: # 用户输入错误 print(f"Nano Banana Pro: 输入错误 - {str(e)}") raise except RuntimeError as e: # API 或网络错误 print(f"Nano Banana Pro: API 错误 - {str(e)}") raise except Exception as e: # 其他未知错误 print(f"Nano Banana Pro: 未知错误 - {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"Nano Banana Pro: {balance_info}") except Exception as e: print(f"Nano Banana Pro: ⚠️ 余额查询失败 - {str(e)}")