""" 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()