576 lines
21 KiB
Python
576 lines
21 KiB
Python
"""
|
||
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()
|