""" Nano Banana 节点 (V3) ComfyUI 自定义节点，用于调用异步生图模型 使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动 """ import time import math import random import asyncio import aiohttp from concurrent.futures import ThreadPoolExecutor from typing import 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 ( NETWORK_ROUTE_OPTIONS, get_base_url_by_route, get_api_key_or_raise, ) from ..utils.nano_banana_async import generate_nano_banana_async from ..clients.gemini_client import GeminiAPIClient 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 = False _NODE = "Nano Banana" _REQUEST_TIMEOUT = 900 _INTERRUPT_CHECK_INTERVAL = 0.2 _client_instance = None def _get_client(): global _client_instance if _client_instance is None: _client_instance = GeminiAPIClient() 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) 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 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, progress_callback: Optional[Callable[[float], None]] = None, ) -> List[Image.Image]: 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, enable_grounding=enable_grounding, thinking_level=thinking_level, node_label="Nano Banana", request_log_enabled=REQUEST_LOG_ENABLED, check_interrupt=_check_interrupt, progress_callback=progress_callback, ) return result_images, timing["task_ms"], timing["parse_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, progress_callback: Optional[Callable[[float], None]] = None, ) -> dict: result = { "global_task_index": global_task_index, "prompt": prompt, "success": False, "generated_count": 0, "output_images": [], "error": None, } try: gen_images, task_ms, parse_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, progress_callback=progress_callback, ) del task_ms, parse_ms result["output_images"] = gen_images result["success"] = True result["generated_count"] = len(gen_images) 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, 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): _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, global_task_index=i, enable_grounding=enable_grounding, thinking_level=thinking_level, progress_callback=_make_progress_callback(pbar), ) ) 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 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}") 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() was_interrupted = False 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 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( _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, 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=_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(): 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, progress_callback=_make_progress_callback(pbar), ) return loop.run_until_complete(_run_with_interrupt(_do())) finally: loop.close() with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(run_single) generated_images, task_ms, parse_ms = 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"{task_ms/1000:.2f}s" parse_str = f"{parse_ms/1000:.2f}s" print(f"完成！总耗时 {time_str} | 异步任务 {task_str} | 解析 {parse_str} | 成功 {len(generated_images)}张") 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()