Initial commit: Comfyui_o1key v1.10.0
This commit is contained in:
o1key
@@ -0,0 +1,381 @@
|
||||
"""
|
||||
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)}")
|
||||
Reference in New Issue
Block a user