feat: sync latest local version as authoritative codebase
Complete rewrite/sync of comfyui_o1key custom nodes. Treat this commit as the current canonical version. Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
This commit is contained in:
o1key
@@ -0,0 +1,24 @@
|
||||
"""
|
||||
节点模块
|
||||
包含所有 ComfyUI 自定义节点的实现
|
||||
"""
|
||||
|
||||
from .nano_banana_pro import NanoBananaPro
|
||||
from .batch_nano_banana_pro import BatchNanoBananaPro
|
||||
from .google_gemini import GoogleGemini
|
||||
from .load_file import LoadFile
|
||||
from .image_stitch_pro import ImageStitchPro
|
||||
from .remove_metadata import SaveCleanImage, BatchCleanMetadata
|
||||
from .video_preview import VideoPreview
|
||||
from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset
|
||||
from .veo_video import GoogleVeo
|
||||
from .flux_edit import FluxImageEdit
|
||||
from .universal_llm import UniversalLLMChat
|
||||
from .quan_neng_sheng_tu import QuanNengShengTu
|
||||
from .batch_quan_neng_sheng_tu import BatchQuanNengShengTu
|
||||
from .multi_res_preview import MultiResPreview
|
||||
from .batch_images_o1key import BatchImagesO1key
|
||||
from .nano_banana_v2 import NanaBananaV2
|
||||
from .batch_nano_banana_v2 import BatchNanaBananaV2
|
||||
|
||||
__all__ = ['NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'QuanNengShengTu', 'BatchQuanNengShengTu', 'MultiResPreview', 'BatchImagesO1key', 'NanaBananaV2', 'BatchNanaBananaV2']
|
||||
@@ -0,0 +1,80 @@
|
||||
"""
|
||||
批量图像（o1key）节点
|
||||
复刻 ComfyUI 原生「批量图像」节点的动态输入行为：
|
||||
|
||||
- 默认显示 2 个图像输入端口（图1, 图2）
|
||||
- 当最后一个端口连上图像后，自动追加新端口
|
||||
- 断开连线后，多余的端口自动消失，最少保留 2 个
|
||||
|
||||
与原生节点的区别：
|
||||
原生节点会把所有图像强制 resize 到第一张的分辨率再合并为单一 tensor。
|
||||
本节点保留每张图的原始分辨率，以 list[Tensor] 形式输出（is_output_list）。
|
||||
下游节点（如「多分辨率图像预览」）需开启 INPUT_IS_LIST 才能正确接收。
|
||||
|
||||
实现方式：使用 V3 API 的 io.Autogrow.TemplateNames，
|
||||
框架原生支持动态 slot 增减，无需编写任何 JS 扩展。
|
||||
"""
|
||||
|
||||
import torch
|
||||
from comfy_api.latest import io
|
||||
|
||||
# 预生成 50 个端口名：图1, 图2, ..., 图50
|
||||
_SLOT_NAMES = [f"图{i}" for i in range(1, 51)]
|
||||
|
||||
|
||||
class BatchImagesO1key(io.ComfyNode):
|
||||
"""
|
||||
批量图像（o1key）
|
||||
|
||||
- 动态输入端口（默认 2 个，最多 50 个），端口名为 图1、图2、图3...
|
||||
- 连接最后一个端口时自动增加新端口
|
||||
- 断开后自动减少，保持界面整洁
|
||||
- 保留每张图的原始分辨率，不做任何 resize / 裁剪
|
||||
- 输出为图像列表，可直接接入「多分辨率图像预览」节点
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
autogrow_template = io.Autogrow.TemplateNames(
|
||||
input=io.Image.Input("image"),
|
||||
names=_SLOT_NAMES,
|
||||
min=2,
|
||||
)
|
||||
return io.Schema(
|
||||
node_id="BatchImagesO1key",
|
||||
display_name="加载图像（批量）",
|
||||
category="image",
|
||||
description=(
|
||||
"将多个独立图像收集为图像列表输出，保留每张图的原始分辨率。\n"
|
||||
"• 默认显示 2 个输入端口（图1、图2），连接最后一个后自动追加新端口\n"
|
||||
"• 断开连线后端口自动减少，最少保留 2 个\n"
|
||||
"• 不做任何 resize / 裁剪，原图尺寸原样输出\n"
|
||||
"• 输出为图像列表，可直接接入「多分辨率图像预览」节点"
|
||||
),
|
||||
search_aliases=["批量图像", "batch images", "合并图像", "图像合并", "stack images"],
|
||||
inputs=[
|
||||
io.Autogrow.Input("images", template=autogrow_template)
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="图像", is_output_list=True),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images: io.Autogrow.Type) -> io.NodeOutput:
|
||||
# images 是 dict，key 为 "图1", "图2", ... ；未连接的 slot 值为 None
|
||||
tensors = [v for v in images.values() if v is not None]
|
||||
|
||||
if not tensors:
|
||||
raise ValueError("批量图像（o1key）：请至少连接一张图像")
|
||||
|
||||
for i, t in enumerate(tensors):
|
||||
h, w = t.shape[1], t.shape[2]
|
||||
print(f"批量图像（o1key）：图{i + 1} → {w}×{h}，shape={list(t.shape)}")
|
||||
|
||||
print(f"批量图像（o1key）：共收集 {len(tensors)} 张，原始分辨率原样输出")
|
||||
|
||||
# 以 list[Tensor] 形式返回，每张图保持自身分辨率
|
||||
return io.NodeOutput(tensors)
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,781 @@
|
||||
"""
|
||||
批量 Nano Banana v2 节点
|
||||
BatchNanoBananaPro 的完全复刻，唯一改动：
|
||||
|
||||
将原来 9 个独立「参考图1~9」输入端
|
||||
改为 1 个「参考图」输入端（可选），配合「加载图像（批量）」节点使用。
|
||||
|
||||
「加载图像（批量）」输出 is_output_list=True（list[Tensor]），
|
||||
本节点声明 INPUT_IS_LIST = True 来整体接收该列表，
|
||||
然后在 process_batch() 开头对所有参数统一解包，其余业务逻辑与原节点完全一致。
|
||||
"""
|
||||
|
||||
import os
|
||||
import gc
|
||||
import time
|
||||
import math
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Optional, Tuple, List
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||||
from ..utils.file_utils import (
|
||||
ImageInfo,
|
||||
load_images_from_folder,
|
||||
pair_images_by_name,
|
||||
pair_images_cartesian,
|
||||
generate_timestamp_filename,
|
||||
save_image,
|
||||
)
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
from ..models_config import (
|
||||
get_enabled_models,
|
||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
||||
get_model_supported_resolutions, get_all_supported_resolutions
|
||||
)
|
||||
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import psutil
|
||||
MEMORY_MONITOR_AVAILABLE = True
|
||||
except ImportError:
|
||||
MEMORY_MONITOR_AVAILABLE = False
|
||||
|
||||
DEBUG_LOG_ENABLED = False
|
||||
REQUEST_LOG_ENABLED = False
|
||||
|
||||
_NODE = "BatchNanoBananaV2"
|
||||
|
||||
|
||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
||||
"""
|
||||
将 PIL Image 列表转换为 ComfyUI tensor，安全处理多张不同尺寸的情况。
|
||||
|
||||
ComfyUI 的 IMAGE tensor 格式为 [B, H, W, C]，要求 batch 内所有图尺寸相同。
|
||||
当 API 返回多张不同分辨率的图时（主图 + 附图），直接 stack 会崩溃。
|
||||
|
||||
策略：
|
||||
- 所有图均已按原始分辨率保存到磁盘（调用此函数前已完成）
|
||||
- 以第一张图的尺寸为基准，只将尺寸相同的图纳入 tensor 输出
|
||||
- 尺寸不同的图跳过（不 resize、不丢弃磁盘文件），并打印日志提示
|
||||
- 若没有任何图与第一张尺寸相同（极罕见），则只输出第一张
|
||||
"""
|
||||
if not images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
return pil_to_tensor([placeholder])
|
||||
|
||||
base_size = images[0].size # PIL size = (W, H)
|
||||
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]}×{img.size[1]}" for img in skipped)
|
||||
print(
|
||||
f"{node_label}: API 额外返回了 {len(skipped)} 张不同尺寸的图 ({sizes_str})，"
|
||||
f"已按原始分辨率保存到磁盘，tensor 输出仅包含与主图尺寸相同的 {len(matched)} 张 "
|
||||
f"({base_size[0]}×{base_size[1]})"
|
||||
)
|
||||
|
||||
return pil_to_tensor(matched if matched else [images[0]])
|
||||
|
||||
|
||||
class BatchNanaBananaV2:
|
||||
"""
|
||||
批量 Nano Banana v2
|
||||
|
||||
与 BatchNanoBananaPro 完全一致，参考图输入方式不同：
|
||||
- 原版：9 个独立可选端口（参考图1~9）
|
||||
- v2：1 个可选端口「参考图」，配合「加载图像（批量）」可传入任意数量图片
|
||||
"""
|
||||
|
||||
ASPECT_RATIOS = [
|
||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
||||
"1:4", "4:1", "1:8", "8:1"
|
||||
]
|
||||
RESOLUTIONS = ["512", "1K", "2K", "4K"]
|
||||
PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"]
|
||||
|
||||
def __init__(self):
|
||||
self.client = None
|
||||
|
||||
def resize_to_megapixels(self, image: Image.Image, target_megapixels: float) -> Image.Image:
|
||||
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 = max(1, int(image.width * scale))
|
||||
new_height = max(1, int(image.height * scale))
|
||||
return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
enabled_models = get_enabled_models()
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
||||
|
||||
all_aspect_ratios = get_all_supported_aspect_ratios() or cls.ASPECT_RATIOS
|
||||
all_resolutions = get_all_supported_resolutions() or cls.RESOLUTIONS
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {"default": "一个中国女子的OOTD", "multiline": True}),
|
||||
"模型": (enabled_models, {"default": enabled_models[0]}),
|
||||
"宽高比": (all_aspect_ratios, {"default": "1:1"}),
|
||||
"分辨率": (all_resolutions, {"default": "2K"}),
|
||||
"像素缩放": ("BOOLEAN", {"default": False, "label_on": "打开", "label_off": "关闭"}),
|
||||
"分辨率像素": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.1, "display": "number"}),
|
||||
"谷歌搜索（联网）": (["关闭", "打开"], {"default": "关闭"}),
|
||||
"图片搜索（联网）": (["关闭", "打开"], {"default": "关闭"}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"文件夹1": ("STRING", {"default": "", "multiline": False}),
|
||||
"文件夹2": ("STRING", {"default": "", "multiline": False}),
|
||||
"文件夹3": ("STRING", {"default": "", "multiline": False}),
|
||||
"文件夹4": ("STRING", {"default": "", "multiline": False}),
|
||||
"文件夹5": ("STRING", {"default": "", "multiline": False}),
|
||||
"文件夹6": ("STRING", {"default": "", "multiline": False}),
|
||||
"文件夹7": ("STRING", {"default": "", "multiline": False}),
|
||||
"文件夹8": ("STRING", {"default": "", "multiline": False}),
|
||||
"文件夹9": ("STRING", {"default": "", "multiline": False}),
|
||||
"保存路径": ("STRING", {"default": "", "multiline": False}),
|
||||
},
|
||||
"optional": {
|
||||
# 单个参考图端口，接受普通 IMAGE 或「加载图像（批量）」输出的列表
|
||||
"参考图": ("IMAGE",),
|
||||
"图片配对模式": (cls.PAIRING_MODES, {"default": "不配对"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("输出图像",)
|
||||
FUNCTION = "process_batch"
|
||||
CATEGORY = "image/batch"
|
||||
|
||||
# 声明 INPUT_IS_LIST，使 ComfyUI 将「加载图像（批量）」的 list[Tensor]
|
||||
# 整体传入而非逐张迭代执行，同时其余所有参数也会被包进 list，需解包。
|
||||
INPUT_IS_LIST = True
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 以下方法与 BatchNanoBananaPro 完全相同
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _load_folders(
|
||||
self,
|
||||
folder1, folder2, folder3, folder4,
|
||||
enable_scaling, target_megapixels,
|
||||
folder5=None, folder6=None, folder7=None, folder8=None, folder9=None,
|
||||
) -> List[List[ImageInfo]]:
|
||||
folders = [folder1, folder2, folder3, folder4,
|
||||
folder5, folder6, folder7, folder8, folder9]
|
||||
all_images = []
|
||||
for i, folder in enumerate(folders, 1):
|
||||
if folder and folder.strip():
|
||||
try:
|
||||
images = load_images_from_folder(folder)
|
||||
if images:
|
||||
if enable_scaling:
|
||||
scaled = []
|
||||
for info in images:
|
||||
scaled_img = self.resize_to_megapixels(info.image, target_megapixels)
|
||||
scaled.append(ImageInfo(
|
||||
image=scaled_img,
|
||||
filename=info.filename,
|
||||
extension=info.extension,
|
||||
source_path=info.source_path
|
||||
))
|
||||
images = scaled
|
||||
all_images.append(images)
|
||||
except ValueError as e:
|
||||
print(f"{_NODE}: 文件夹{i} 加载失败 - {e}")
|
||||
return all_images
|
||||
|
||||
def _create_pairs(
|
||||
self,
|
||||
image_lists: List[List[ImageInfo]],
|
||||
pairing_mode: str,
|
||||
manual_images: Optional[List[ImageInfo]] = None
|
||||
) -> List[Tuple[ImageInfo, ...]]:
|
||||
if pairing_mode == "不配对":
|
||||
if len(image_lists) > 1:
|
||||
raise ValueError("「不配对」模式只支持单个文件夹，请清空其他文件夹路径")
|
||||
if image_lists and manual_images:
|
||||
return [(img,) + tuple(manual_images) for img in image_lists[0]]
|
||||
elif image_lists:
|
||||
return [(img,) for img in image_lists[0]]
|
||||
else:
|
||||
return []
|
||||
|
||||
if not image_lists:
|
||||
return []
|
||||
|
||||
if len(image_lists) == 1:
|
||||
base_pairs = [(img,) for img in image_lists[0]]
|
||||
elif pairing_mode == "按相同图片命名":
|
||||
base_pairs = list(pair_images_by_name(*image_lists))
|
||||
else:
|
||||
base_pairs = list(pair_images_cartesian(*image_lists))
|
||||
|
||||
if manual_images:
|
||||
manual_tuple = tuple(manual_images)
|
||||
base_pairs = [pair + manual_tuple for pair in base_pairs]
|
||||
|
||||
return base_pairs
|
||||
|
||||
async def _generate_single_task(
|
||||
self,
|
||||
client: GeminiAPIClient,
|
||||
session: aiohttp.ClientSession,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: List[ImageInfo],
|
||||
output_folder: str,
|
||||
task_index: int,
|
||||
enable_grounding: bool = True,
|
||||
enable_image_search: bool = False,
|
||||
base_filename: str = None,
|
||||
) -> dict:
|
||||
result = {
|
||||
"task_index": task_index,
|
||||
"prompt": prompt,
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"saved_files": [],
|
||||
"output_images": [],
|
||||
"error": None
|
||||
}
|
||||
try:
|
||||
input_pil_images = [info.image for info in images]
|
||||
generated_images = []
|
||||
try:
|
||||
gen_result = await client.generate_single_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_pil_images,
|
||||
session=session,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
if gen_result:
|
||||
images_list, timing_info = gen_result
|
||||
generated_images.extend(images_list)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_msg = str(e)
|
||||
error_traceback = traceback.format_exc()
|
||||
print(f"=" * 80)
|
||||
print(f"🔍 【原始报错信息展示】")
|
||||
print(f"=" * 80)
|
||||
print(f"任务编号: {task_index + 1}")
|
||||
print(f"失败时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
print(f"模型: {model}")
|
||||
print(f"分辨率: {resolution}")
|
||||
print(f"宽高比: {aspect_ratio}")
|
||||
print(f"-" * 80)
|
||||
print(f"错误信息: {error_msg}")
|
||||
print(f"-" * 80)
|
||||
print(f"完整堆栈追踪:")
|
||||
print(error_traceback)
|
||||
print(f"=" * 80)
|
||||
result["error"] = error_msg
|
||||
|
||||
for i, gen_img in enumerate(generated_images):
|
||||
if base_filename:
|
||||
base_name = base_filename
|
||||
counter = 0
|
||||
while True:
|
||||
filename = f"{base_name}.png" if counter == 0 else f"{base_name}+{counter}.png"
|
||||
output_path = os.path.join(output_folder, filename)
|
||||
if not os.path.exists(output_path):
|
||||
break
|
||||
counter += 1
|
||||
else:
|
||||
output_path = generate_timestamp_filename(
|
||||
output_folder=output_folder, extension=".png"
|
||||
)
|
||||
save_image(gen_img, output_path)
|
||||
result["saved_files"].append(output_path)
|
||||
gen_img = None
|
||||
|
||||
if len(generated_images) > 0:
|
||||
result["success"] = True
|
||||
result["generated_count"] = len(generated_images)
|
||||
|
||||
except Exception as e:
|
||||
result["error"] = str(e)
|
||||
|
||||
return result
|
||||
|
||||
async def _process_batch_async(
|
||||
self,
|
||||
pairs: List[Tuple[ImageInfo, ...]],
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
output_folder: str,
|
||||
pbar=None,
|
||||
prompts_per_task: Optional[List[str]] = None,
|
||||
enable_grounding: bool = True,
|
||||
enable_image_search: bool = False,
|
||||
) -> List[dict]:
|
||||
if self.client is None:
|
||||
self.client = GeminiAPIClient()
|
||||
|
||||
total_tasks = len(pairs)
|
||||
max_concurrent = 10
|
||||
|
||||
print(f"{_NODE}: 检测到 {total_tasks} 个任务")
|
||||
|
||||
all_results = []
|
||||
completed = 0
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
||||
|
||||
if MEMORY_MONITOR_AVAILABLE and total_tasks > 50:
|
||||
process = psutil.Process()
|
||||
initial_memory = process.memory_info().rss / 1024 / 1024
|
||||
print(f"{_NODE}: 初始内存使用: {initial_memory:.1f} MB")
|
||||
|
||||
show_milestone = total_tasks >= 50
|
||||
milestones = [0.2, 0.4, 0.6, 0.8, 1.0]
|
||||
milestone_index = 0
|
||||
|
||||
if num_batches > 1:
|
||||
print(f"{_NODE}: 任务数 {total_tasks} 超过并发上限 {max_concurrent}，将分 {num_batches} 批执行")
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
for batch_idx in range(num_batches):
|
||||
start_idx = batch_idx * max_concurrent
|
||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||||
batch_pairs = pairs[start_idx:end_idx]
|
||||
|
||||
if num_batches > 1:
|
||||
print(f"{_NODE}: 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...")
|
||||
|
||||
tasks = []
|
||||
for i, pair in enumerate(batch_pairs):
|
||||
task_prompt = prompts_per_task[start_idx + i] if prompts_per_task else prompt
|
||||
base_filename = None
|
||||
if pair and len(pair) > 0:
|
||||
first_image = pair[0]
|
||||
if hasattr(first_image, 'filename'):
|
||||
base_filename = first_image.filename
|
||||
|
||||
task = asyncio.create_task(
|
||||
self._generate_single_task(
|
||||
client=self.client,
|
||||
session=session,
|
||||
prompt=task_prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=list(pair),
|
||||
output_folder=output_folder,
|
||||
task_index=start_idx + i,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
base_filename=base_filename,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
batch_results = []
|
||||
for coro in asyncio.as_completed(tasks):
|
||||
result_data = None
|
||||
try:
|
||||
result = await coro
|
||||
if isinstance(result, Exception):
|
||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": []}
|
||||
batch_results.append(result_data)
|
||||
else:
|
||||
result_data = result
|
||||
batch_results.append(result)
|
||||
except Exception as e:
|
||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": []}
|
||||
batch_results.append(result_data)
|
||||
|
||||
completed += 1
|
||||
if result_data and result_data.get("success", False):
|
||||
success_count += 1
|
||||
print(f"{_NODE}: 任务 {completed}/{total_tasks} 成功 ✓")
|
||||
else:
|
||||
fail_count += 1
|
||||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
||||
print(f"{_NODE}: 任务 {completed}/{total_tasks} 失败 ✗")
|
||||
print(f"=" * 80)
|
||||
print(f"🔍 【原始报错信息展示】")
|
||||
print(f"=" * 80)
|
||||
print(f"任务编号: {completed}/{total_tasks}")
|
||||
print(f"失败时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
|
||||
print(f"-" * 80)
|
||||
print(f"错误详情:")
|
||||
print(error_msg)
|
||||
print(f"=" * 80)
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
if show_milestone and milestone_index < len(milestones):
|
||||
progress = completed / total_tasks
|
||||
if progress >= milestones[milestone_index]:
|
||||
percentage = int(milestones[milestone_index] * 100)
|
||||
print(f"{_NODE}: >>> 进度 {percentage}% <<<")
|
||||
milestone_index += 1
|
||||
|
||||
all_results.extend(batch_results)
|
||||
print(f"{_NODE}: 第 {batch_idx + 1} 批完成，开始分批保存...")
|
||||
|
||||
batch_success = sum(1 for r in batch_results if r.get("success", False))
|
||||
batch_fail = len(batch_results) - batch_success
|
||||
batch_generated = sum(r.get("generated_count", 0) for r in batch_results)
|
||||
print(f"{_NODE}: 本批结果 - 成功: {batch_success}/{len(batch_results)}，生成: {batch_generated} 张")
|
||||
|
||||
gc.collect()
|
||||
|
||||
if MEMORY_MONITOR_AVAILABLE and total_tasks > 50:
|
||||
current_memory = process.memory_info().rss / 1024 / 1024
|
||||
memory_increase = current_memory - initial_memory
|
||||
print(f"{_NODE}: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)")
|
||||
if current_memory > 2000:
|
||||
print(f"⚠️ {_NODE}: 内存使用过高！但图片已分批保存，即使崩溃也不会丢失已完成的任务")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
return all_results
|
||||
|
||||
def process_batch(
|
||||
self,
|
||||
prompt,
|
||||
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9,
|
||||
像素缩放,
|
||||
分辨率像素,
|
||||
seed,
|
||||
模型,
|
||||
宽高比,
|
||||
分辨率,
|
||||
保存路径,
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
|
||||
# ----------------------------------------------------------------
|
||||
# INPUT_IS_LIST=True 时，所有参数均为 list，先统一解包为标量
|
||||
# ----------------------------------------------------------------
|
||||
def _unpack(v):
|
||||
return v[0] if isinstance(v, list) else v
|
||||
|
||||
prompt = _unpack(prompt)
|
||||
文件夹1 = _unpack(文件夹1)
|
||||
文件夹2 = _unpack(文件夹2)
|
||||
文件夹3 = _unpack(文件夹3)
|
||||
文件夹4 = _unpack(文件夹4)
|
||||
文件夹5 = _unpack(文件夹5)
|
||||
文件夹6 = _unpack(文件夹6)
|
||||
文件夹7 = _unpack(文件夹7)
|
||||
文件夹8 = _unpack(文件夹8)
|
||||
文件夹9 = _unpack(文件夹9)
|
||||
像素缩放 = _unpack(像素缩放)
|
||||
分辨率像素 = _unpack(分辨率像素)
|
||||
seed = _unpack(seed)
|
||||
模型 = _unpack(模型)
|
||||
宽高比 = _unpack(宽高比)
|
||||
分辨率 = _unpack(分辨率)
|
||||
保存路径 = _unpack(保存路径)
|
||||
|
||||
# 含全角括号的参数名无法作为形参，从 kwargs 中提取
|
||||
enable_grounding: bool = (_unpack(kwargs.pop("谷歌搜索（联网）", "关闭"))) == "打开"
|
||||
enable_image_search: bool = (_unpack(kwargs.pop("图片搜索（联网）", "关闭"))) == "打开"
|
||||
|
||||
# 图片配对模式（可选参数）
|
||||
图片配对模式 = _unpack(kwargs.pop("图片配对模式", "不配对"))
|
||||
|
||||
# ----------------------------------------------------------------
|
||||
# 收集参考图：兼容两种来源
|
||||
# 1. 「加载图像（批量）」→ is_output_list=True → list[Tensor]
|
||||
# INPUT_IS_LIST 下传入的是 list[list[Tensor]] 或 list[Tensor]，需展平
|
||||
# 2. 普通 IMAGE 端口（单 tensor 或 batch tensor）→ list 中只有 1 个元素
|
||||
# ----------------------------------------------------------------
|
||||
ref_raw = kwargs.pop("参考图", None)
|
||||
manual_images: List[ImageInfo] = []
|
||||
|
||||
if ref_raw is not None:
|
||||
items = ref_raw if isinstance(ref_raw, list) else [ref_raw]
|
||||
idx = 0
|
||||
for item in items:
|
||||
if item is None:
|
||||
continue
|
||||
if isinstance(item, list):
|
||||
sub_tensors = item
|
||||
elif isinstance(item, torch.Tensor):
|
||||
sub_tensors = [item]
|
||||
else:
|
||||
continue
|
||||
for tensor in sub_tensors:
|
||||
if tensor is None or not isinstance(tensor, torch.Tensor):
|
||||
continue
|
||||
pil_images = tensor_to_pil(tensor)
|
||||
for j, img in enumerate(pil_images):
|
||||
if 像素缩放:
|
||||
img = self.resize_to_megapixels(img, 分辨率像素)
|
||||
manual_images.append(ImageInfo(
|
||||
image=img,
|
||||
filename=f"manual_{idx}_{j}",
|
||||
extension=".png",
|
||||
source_path=""
|
||||
))
|
||||
idx += 1
|
||||
|
||||
# ----------------------------------------------------------------
|
||||
# 以下逻辑与 BatchNanoBananaPro.process_batch() 完全一致
|
||||
# ----------------------------------------------------------------
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed % (2 ** 32))
|
||||
|
||||
has_any_folder = any(
|
||||
f and f.strip()
|
||||
for f in [文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9]
|
||||
)
|
||||
if not has_any_folder:
|
||||
raise ValueError("请至少填写一个文件夹路径，该节点专为批量文件夹处理设计")
|
||||
|
||||
supported_resolutions = get_model_supported_resolutions(模型)
|
||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
||||
raise ValueError(
|
||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的分辨率：{', '.join(supported_resolutions)}"
|
||||
)
|
||||
|
||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
||||
if supported_ratios and 宽高比 not in supported_ratios:
|
||||
raise ValueError(
|
||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的宽高比：{', '.join(supported_ratios)}"
|
||||
)
|
||||
|
||||
IMAGE_SEARCH_UNSUPPORTED_MODELS = [
|
||||
"nano-banana-pro-限时特价", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview"
|
||||
]
|
||||
if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS:
|
||||
raise ValueError(
|
||||
f"模型 \"{模型}\" 不支持【图片搜索（联网）】功能！"
|
||||
f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用"
|
||||
)
|
||||
|
||||
print(f"{_NODE}: 开始加载图片...")
|
||||
image_lists = self._load_folders(
|
||||
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
||||
像素缩放, 分辨率像素,
|
||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9
|
||||
)
|
||||
|
||||
total_folder_images = sum(len(lst) for lst in image_lists)
|
||||
if total_folder_images == 0:
|
||||
raise ValueError("文件夹中未找到任何图片，请检查文件夹路径是否正确")
|
||||
|
||||
pairs = self._create_pairs(image_lists, 图片配对模式, manual_images if manual_images else None)
|
||||
|
||||
if not pairs:
|
||||
raise ValueError("配对结果为空，请检查输入")
|
||||
|
||||
batch_prompts = parse_batch_prompts(prompt)
|
||||
prompts_per_task = None
|
||||
if batch_prompts:
|
||||
expanded_pairs = []
|
||||
expanded_prompts = []
|
||||
for pair in pairs:
|
||||
for bp in batch_prompts:
|
||||
expanded_pairs.append(pair)
|
||||
expanded_prompts.append(bp)
|
||||
pairs = expanded_pairs
|
||||
prompts_per_task = expanded_prompts
|
||||
|
||||
total_tasks = len(pairs)
|
||||
|
||||
grounding_str = ""
|
||||
if enable_image_search:
|
||||
grounding_str = " | 谷歌图片搜索接地"
|
||||
elif enable_grounding:
|
||||
grounding_str = " | 谷歌搜索接地"
|
||||
|
||||
if batch_prompts:
|
||||
print(f"{_NODE}: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}")
|
||||
else:
|
||||
print(f"{_NODE}: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}")
|
||||
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(total_tasks)
|
||||
|
||||
has_save_path = bool(保存路径 and 保存路径.strip())
|
||||
if not has_save_path:
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
保存路径 = folder_paths.get_output_directory()
|
||||
has_save_path = True
|
||||
print(f"{_NODE}: 未设置保存路径，将使用 ComfyUI 默认 output 目录: {保存路径}")
|
||||
else:
|
||||
print(f"{_NODE}: 未设置保存路径，图片将输出到节点")
|
||||
|
||||
if has_save_path:
|
||||
try:
|
||||
os.makedirs(保存路径, exist_ok=True)
|
||||
test_file = os.path.join(保存路径, ".write_test")
|
||||
with open(test_file, 'w') as f:
|
||||
f.write("test")
|
||||
os.remove(test_file)
|
||||
print(f"{_NODE}: 保存路径验证通过: {保存路径}")
|
||||
except Exception as e:
|
||||
raise ValueError(f"保存路径无效或无写入权限: {保存路径} - {str(e)}")
|
||||
|
||||
if self.client is None:
|
||||
try:
|
||||
self.client = GeminiAPIClient()
|
||||
except ValueError as e:
|
||||
raise ValueError(f"初始化 API 客户端失败: {str(e)}")
|
||||
|
||||
def run_async_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
self._process_batch_async(
|
||||
pairs=pairs,
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
output_folder=保存路径,
|
||||
pbar=pbar,
|
||||
prompts_per_task=prompts_per_task,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"{_NODE}: 异步任务执行异常: {str(e)}")
|
||||
raise
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async_in_thread)
|
||||
try:
|
||||
results = future.result(timeout=3600)
|
||||
except TimeoutError:
|
||||
print(f"{_NODE}: 任务执行超时（1小时）")
|
||||
raise RuntimeError("任务执行超时，请减少任务数量或检查网络连接")
|
||||
except Exception as e:
|
||||
print(f"{_NODE}: 任务执行失败: {str(e)}")
|
||||
raise
|
||||
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||
all_saved_files = [f for r in results for f in r.get("saved_files", [])]
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||||
avg_time = elapsed / success_count if success_count > 0 else 0
|
||||
avg_time_str = f"{avg_time:.1f}s/张" if success_count > 0 else "N/A"
|
||||
|
||||
print("=" * 60)
|
||||
print(f"完成！总耗时 {time_str} | 成功: {success_count}/{total_tasks} | 生成 {total_generated} 张 | 平均 {avg_time_str}")
|
||||
if has_save_path:
|
||||
print(f"保存路径: {保存路径}")
|
||||
else:
|
||||
print("保存路径: 未设置（仅输出到节点）")
|
||||
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
if failed_results:
|
||||
print(f"-" * 60)
|
||||
print(f"❌ 失败任务汇总: {len(failed_results)} 个")
|
||||
print(f"-" * 60)
|
||||
for idx, failed in enumerate(failed_results[:3], 1):
|
||||
task_num = failed.get('task_index', '?') + 1
|
||||
error_msg = failed.get('error', '未知错误')
|
||||
print(f"\n【失败任务 #{task_num}】")
|
||||
print(f"错误信息: {error_msg}")
|
||||
if len(failed_results) > 3:
|
||||
remaining = [str(r.get('task_index', '?') + 1) for r in failed_results[3:]]
|
||||
print(f"\n其他失败任务编号: {', '.join(remaining)}")
|
||||
print(f"-" * 60)
|
||||
|
||||
output_images = []
|
||||
if all_saved_files:
|
||||
for fp in all_saved_files[-min(10, len(all_saved_files)):]:
|
||||
try:
|
||||
output_images.append(Image.open(fp))
|
||||
except Exception as e:
|
||||
print(f"{_NODE}: 无法加载图片 {fp} - {e}")
|
||||
|
||||
if not output_images:
|
||||
output_images = [Image.new('RGB', (512, 512), color=(128, 128, 128))]
|
||||
|
||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||
gc.collect()
|
||||
|
||||
total_saved = len(all_saved_files)
|
||||
print(f"{_NODE}: 任务完成！共保存 {total_saved} 张图片到磁盘")
|
||||
if total_saved > 0:
|
||||
print(f"{_NODE}: 最新保存的文件: {all_saved_files[-1]}")
|
||||
|
||||
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as e:
|
||||
if str(e) == "未授权！":
|
||||
print("请联系作者授权后方可使用！")
|
||||
raise ValueError("未授权！") from None
|
||||
error_msg = str(e)
|
||||
print(f"{_NODE}: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
error_full = str(e)
|
||||
print(f"{_NODE}: ❌ {error_full}")
|
||||
raise RuntimeError(error_full) from None
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"{_NODE}: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
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"{_NODE}: {balance_info}")
|
||||
print("=" * 60)
|
||||
except Exception:
|
||||
pass
|
||||
gc.collect()
|
||||
@@ -0,0 +1,642 @@
|
||||
"""
|
||||
全能生图（批量）节点
|
||||
ComfyUI 自定义节点，用于批量处理图像生成任务
|
||||
支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存
|
||||
"""
|
||||
|
||||
import time
|
||||
import math
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Optional, Tuple, List
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||||
from ..utils.file_utils import (
|
||||
ImageInfo,
|
||||
load_images_from_folder,
|
||||
pair_images_by_name,
|
||||
pair_images_cartesian,
|
||||
generate_timestamp_filename,
|
||||
save_image,
|
||||
)
|
||||
from ..clients.openai_client import OpenAIAPIClient
|
||||
from ..models_config import (
|
||||
get_enabled_models,
|
||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
||||
get_model_supported_resolutions, get_all_supported_resolutions
|
||||
)
|
||||
|
||||
# 导入 ComfyUI 原生进度条
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ 全能生图（批量）: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")
|
||||
|
||||
# 导入 ComfyUI 的文件夹路径管理
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
print("⚠️ 全能生图（批量）: folder_paths 不可用，将无法使用默认保存路径")
|
||||
|
||||
# 内存监控（可选）
|
||||
try:
|
||||
import psutil
|
||||
MEMORY_MONITOR_AVAILABLE = True
|
||||
except ImportError:
|
||||
MEMORY_MONITOR_AVAILABLE = False
|
||||
print("⚠️ 全能生图（批量）: psutil 不可用，内存监控功能禁用")
|
||||
|
||||
# ============================================================================
|
||||
# 调试日志配置
|
||||
# ============================================================================
|
||||
DEBUG_LOG_ENABLED = False
|
||||
REQUEST_LOG_ENABLED = False
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class BatchQuanNengShengTu:
|
||||
"""
|
||||
全能生图（批量）节点
|
||||
|
||||
功能：
|
||||
- 从多个文件夹加载图片
|
||||
- 支持三种配对模式：
|
||||
* 按相同图片命名 - 索引配对（文件夹之间按位置配对）
|
||||
* 1*N - 笛卡尔积配对（所有可能组合）
|
||||
* 不配对 - 固定参考图模式（文件夹图片依次与所有参考图组合）
|
||||
- 批量调用 API 生成图像
|
||||
- 智能命名保存（保留原始文件名）
|
||||
- 并发控制（默认最大 10）
|
||||
|
||||
注意：
|
||||
- 「不配对」模式只支持单个文件夹
|
||||
- 支持的模型列表从 models_config.py 动态加载
|
||||
"""
|
||||
|
||||
MODELS = None
|
||||
ASPECT_RATIOS = [
|
||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
||||
"1:4", "4:1", "1:8", "8:1"
|
||||
]
|
||||
RESOLUTIONS = ["512", "1K", "2K", "4K"]
|
||||
PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"]
|
||||
|
||||
def __init__(self):
|
||||
"""初始化节点"""
|
||||
self.client = None
|
||||
|
||||
def resize_to_megapixels(
|
||||
self,
|
||||
image: Image.Image,
|
||||
target_megapixels: float
|
||||
) -> Image.Image:
|
||||
"""将图像缩放到指定的总像素数，保持纵横比"""
|
||||
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 = max(1, int(image.width * scale))
|
||||
new_height = max(1, int(image.height * scale))
|
||||
|
||||
return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""定义输入参数"""
|
||||
enabled_models = get_enabled_models()
|
||||
enabled_models = [m for m in enabled_models if "限时特价" not in m]
|
||||
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
||||
|
||||
all_aspect_ratios = get_all_supported_aspect_ratios()
|
||||
if not all_aspect_ratios:
|
||||
all_aspect_ratios = cls.ASPECT_RATIOS
|
||||
|
||||
all_resolutions = get_all_supported_resolutions()
|
||||
if not all_resolutions:
|
||||
all_resolutions = cls.RESOLUTIONS
|
||||
|
||||
optional_inputs = {}
|
||||
for i in range(1, 10):
|
||||
optional_inputs[f"参考图{i}"] = ("IMAGE",)
|
||||
|
||||
optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, {
|
||||
"default": "不配对"
|
||||
})
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"提示词": ("STRING", {
|
||||
"default": "一个中国女子的OOTD",
|
||||
"multiline": True
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0]
|
||||
}),
|
||||
"宽高比": (all_aspect_ratios, {
|
||||
"default": "1:1"
|
||||
}),
|
||||
"分辨率": (all_resolutions, {
|
||||
"default": "2K"
|
||||
}),
|
||||
"像素缩放": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"label_on": "打开",
|
||||
"label_off": "关闭"
|
||||
}),
|
||||
"分辨率像素": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.1,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"display": "number"
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff
|
||||
}),
|
||||
"文件夹1": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹2": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹3": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹4": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹5": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹6": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹7": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹8": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"文件夹9": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
}),
|
||||
"保存路径": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
})
|
||||
},
|
||||
"optional": optional_inputs
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("输出图像",)
|
||||
FUNCTION = "process_batch"
|
||||
CATEGORY = "image/batch"
|
||||
|
||||
def _load_folders(
|
||||
self,
|
||||
folder1: str,
|
||||
folder2: Optional[str],
|
||||
folder3: Optional[str],
|
||||
folder4: Optional[str],
|
||||
enable_scaling: bool,
|
||||
target_megapixels: float,
|
||||
folder5: Optional[str] = None,
|
||||
folder6: Optional[str] = None,
|
||||
folder7: Optional[str] = None,
|
||||
folder8: Optional[str] = None,
|
||||
folder9: Optional[str] = None,
|
||||
) -> List[List[ImageInfo]]:
|
||||
"""加载所有文件夹中的图片"""
|
||||
folders = [folder1, folder2, folder3, folder4, folder5, folder6, folder7, folder8, folder9]
|
||||
all_images = []
|
||||
|
||||
for i, folder in enumerate(folders, 1):
|
||||
if folder and folder.strip():
|
||||
try:
|
||||
images = load_images_from_folder(folder)
|
||||
if images:
|
||||
if enable_scaling:
|
||||
scaled_images = []
|
||||
for img_info in images:
|
||||
scaled_img = self.resize_to_megapixels(
|
||||
img_info.image,
|
||||
target_megapixels
|
||||
)
|
||||
scaled_info = ImageInfo(
|
||||
image=scaled_img,
|
||||
filename=img_info.filename,
|
||||
extension=img_info.extension,
|
||||
source_path=img_info.source_path
|
||||
)
|
||||
scaled_images.append(scaled_info)
|
||||
images = scaled_images
|
||||
all_images.append(images)
|
||||
except ValueError as e:
|
||||
print(f"全能生图（批量）: 文件夹{i} 加载失败 - {e}")
|
||||
|
||||
return all_images
|
||||
|
||||
def _create_pairs(
|
||||
self,
|
||||
image_lists: List[List[ImageInfo]],
|
||||
pairing_mode: str,
|
||||
manual_images: Optional[List[ImageInfo]] = None
|
||||
) -> List[Tuple[ImageInfo, ...]]:
|
||||
"""根据配对模式创建图片组合"""
|
||||
if pairing_mode == "不配对":
|
||||
if len(image_lists) > 1:
|
||||
raise ValueError("「不配对」模式只支持单个文件夹，请清空其他文件夹路径")
|
||||
|
||||
if image_lists and manual_images:
|
||||
folder_images = image_lists[0]
|
||||
pairs = []
|
||||
for img in folder_images:
|
||||
pair = (img,) + tuple(manual_images)
|
||||
pairs.append(pair)
|
||||
return pairs
|
||||
elif image_lists:
|
||||
return [(img,) for img in image_lists[0]]
|
||||
else:
|
||||
return []
|
||||
|
||||
if not image_lists:
|
||||
return []
|
||||
|
||||
if len(image_lists) == 1:
|
||||
base_pairs = [(img,) for img in image_lists[0]]
|
||||
elif pairing_mode == "按相同图片命名":
|
||||
base_pairs = list(pair_images_by_name(*image_lists))
|
||||
else:
|
||||
base_pairs = list(pair_images_cartesian(*image_lists))
|
||||
|
||||
if manual_images:
|
||||
manual_tuple = tuple(manual_images)
|
||||
base_pairs = [pair + manual_tuple for pair in base_pairs]
|
||||
|
||||
return base_pairs
|
||||
|
||||
async def _generate_single_task(
|
||||
self,
|
||||
session: aiohttp.ClientSession,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: List[ImageInfo],
|
||||
output_folder: str,
|
||||
task_index: int,
|
||||
base_filename: str = None,
|
||||
) -> dict:
|
||||
"""执行单个生成任务"""
|
||||
result = {
|
||||
"task_index": task_index,
|
||||
"prompt": prompt,
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"saved_files": [],
|
||||
"error": None
|
||||
}
|
||||
|
||||
try:
|
||||
input_pil_images = [info.image for info in images]
|
||||
|
||||
gen_result = await self.client.generate_single_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_pil_images,
|
||||
session=session,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=False,
|
||||
enable_image_search=False
|
||||
)
|
||||
|
||||
if gen_result:
|
||||
images_list, _ = gen_result
|
||||
|
||||
import os
|
||||
for gen_img in images_list:
|
||||
if base_filename:
|
||||
base_name = base_filename
|
||||
counter = 0
|
||||
while True:
|
||||
filename = f"{base_name}.png" if counter == 0 else f"{base_name}+{counter}.png"
|
||||
output_path = os.path.join(output_folder, filename)
|
||||
if not os.path.exists(output_path):
|
||||
break
|
||||
counter += 1
|
||||
else:
|
||||
output_path = generate_timestamp_filename(
|
||||
output_folder=output_folder,
|
||||
extension=".png"
|
||||
)
|
||||
save_image(gen_img, output_path)
|
||||
result["saved_files"].append(output_path)
|
||||
gen_img = None
|
||||
|
||||
result["success"] = True
|
||||
result["generated_count"] = len(images_list)
|
||||
|
||||
except Exception as e:
|
||||
result["error"] = str(e)
|
||||
|
||||
return result
|
||||
|
||||
async def _process_batch_async(
|
||||
self,
|
||||
pairs: List[Tuple[ImageInfo, ...]],
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
output_folder: str,
|
||||
pbar=None,
|
||||
prompts_per_task: Optional[List[str]] = None,
|
||||
) -> List[dict]:
|
||||
"""异步批量处理所有任务"""
|
||||
if self.client is None:
|
||||
self.client = OpenAIAPIClient()
|
||||
|
||||
total_tasks = len(pairs)
|
||||
max_concurrent = 10
|
||||
|
||||
print(f"全能生图（批量）: 检测到 {total_tasks} 个任务")
|
||||
|
||||
all_results = []
|
||||
completed = 0
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
||||
|
||||
if num_batches > 1:
|
||||
print(f"全能生图（批量）: 任务数 {total_tasks} 超过并发上限 {max_concurrent}，将分 {num_batches} 批执行")
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
for batch_idx in range(num_batches):
|
||||
start_idx = batch_idx * max_concurrent
|
||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||||
batch_pairs = pairs[start_idx:end_idx]
|
||||
|
||||
if num_batches > 1:
|
||||
print(f"全能生图（批量）: 执行第 {batch_idx + 1}/{num_batches} 批 ({start_idx + 1}-{end_idx})...")
|
||||
|
||||
tasks = []
|
||||
for i, pair in enumerate(batch_pairs):
|
||||
task_prompt = prompts_per_task[start_idx + i] if prompts_per_task else prompt
|
||||
|
||||
base_filename = None
|
||||
if pair and len(pair) > 0:
|
||||
first_image = pair[0]
|
||||
if hasattr(first_image, 'filename'):
|
||||
base_filename = first_image.filename
|
||||
|
||||
task = asyncio.create_task(
|
||||
self._generate_single_task(
|
||||
session=session,
|
||||
prompt=task_prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=list(pair),
|
||||
output_folder=output_folder,
|
||||
task_index=start_idx + i,
|
||||
base_filename=base_filename,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
batch_results = []
|
||||
for coro in asyncio.as_completed(tasks):
|
||||
result_data = None
|
||||
try:
|
||||
result = await coro
|
||||
if isinstance(result, Exception):
|
||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": []}
|
||||
else:
|
||||
result_data = result
|
||||
batch_results.append(result_data)
|
||||
except Exception as e:
|
||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": []}
|
||||
batch_results.append(result_data)
|
||||
|
||||
completed += 1
|
||||
|
||||
if result_data and result_data.get("success", False):
|
||||
success_count += 1
|
||||
print(f"全能生图（批量）: 任务 {completed}/{total_tasks} 成功 ✓")
|
||||
else:
|
||||
fail_count += 1
|
||||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
||||
print(f"全能生图（批量）: 任务 {completed}/{total_tasks} 失败 ✗ - {error_msg}")
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
all_results.extend(batch_results)
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
return all_results
|
||||
|
||||
def process_batch(
|
||||
self,
|
||||
提示词: str,
|
||||
模型: str,
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
像素缩放: bool,
|
||||
分辨率像素: float,
|
||||
seed: int,
|
||||
文件夹1: str,
|
||||
文件夹2: str,
|
||||
文件夹3: str,
|
||||
文件夹4: str,
|
||||
文件夹5: str,
|
||||
文件夹6: str,
|
||||
文件夹7: str,
|
||||
文件夹8: str,
|
||||
文件夹9: str,
|
||||
保存路径: str,
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""批量处理图像生成"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed % (2**32))
|
||||
|
||||
if self.client is None:
|
||||
self.client = OpenAIAPIClient()
|
||||
|
||||
supported_resolutions = get_model_supported_resolutions(模型)
|
||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
||||
raise ValueError(
|
||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的分辨率：{', '.join(supported_resolutions)}"
|
||||
)
|
||||
|
||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
||||
if supported_ratios and 宽高比 not in supported_ratios:
|
||||
raise ValueError(
|
||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的宽高比：{', '.join(supported_ratios)}"
|
||||
)
|
||||
|
||||
manual_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])
|
||||
for pil_img in pil_imgs:
|
||||
manual_images.append(ImageInfo(
|
||||
image=pil_img,
|
||||
filename=f"manual_{i}",
|
||||
extension=".png",
|
||||
source_path=""
|
||||
))
|
||||
|
||||
if 像素缩放 and manual_images:
|
||||
scaled_manual = []
|
||||
for img_info in manual_images:
|
||||
scaled_img = self.resize_to_megapixels(img_info.image, 分辨率像素)
|
||||
scaled_manual.append(ImageInfo(
|
||||
image=scaled_img,
|
||||
filename=img_info.filename,
|
||||
extension=img_info.extension,
|
||||
source_path=img_info.source_path
|
||||
))
|
||||
manual_images = scaled_manual
|
||||
|
||||
folder_images = self._load_folders(
|
||||
文件夹1, 文件夹2, 文件夹3, 文件夹4,
|
||||
像素缩放, 分辨率像素,
|
||||
文件夹5, 文件夹6, 文件夹7, 文件夹8, 文件夹9
|
||||
)
|
||||
|
||||
pairing_mode = kwargs.get("图片配对模式", "不配对")
|
||||
pairs = self._create_pairs(folder_images, pairing_mode, manual_images if manual_images else None)
|
||||
|
||||
if not pairs:
|
||||
raise ValueError("没有可处理的图片组合，请检查文件夹路径和参考图输入")
|
||||
|
||||
batch_prompts = parse_batch_prompts(提示词)
|
||||
prompts_per_task = None
|
||||
|
||||
if batch_prompts:
|
||||
if len(batch_prompts) != len(pairs):
|
||||
raise ValueError(
|
||||
f"批量提示词数量 ({len(batch_prompts)}) 与任务数量 ({len(pairs)}) 不匹配！\n"
|
||||
f"请确保提示词数量与图片组合数量一致"
|
||||
)
|
||||
prompts_per_task = batch_prompts
|
||||
print(f"全能生图（批量）: 批量提示词模式 - {len(batch_prompts)} 个提示词")
|
||||
|
||||
output_folder = 保存路径.strip() if 保存路径 else ""
|
||||
if not output_folder and FOLDER_PATHS_AVAILABLE:
|
||||
output_folder = folder_paths.get_output_directory()
|
||||
|
||||
if not output_folder:
|
||||
raise ValueError("无法确定保存路径，请指定保存路径或确保 folder_paths 可用")
|
||||
|
||||
import os
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
print(f"全能生图（批量）: 保存路径 → {output_folder}")
|
||||
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(len(pairs))
|
||||
|
||||
def run_async():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
self._process_batch_async(
|
||||
pairs=pairs,
|
||||
prompt=提示词,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
output_folder=output_folder,
|
||||
pbar=pbar,
|
||||
prompts_per_task=prompts_per_task,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async)
|
||||
results = future.result(timeout=3600)
|
||||
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
all_saved_files = []
|
||||
for r in results:
|
||||
all_saved_files.extend(r.get("saved_files", []))
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"全能生图（批量）: 完成！总耗时 {elapsed:.2f}s | 成功: {success_count}/{len(pairs)} | 失败: {fail_count}")
|
||||
|
||||
output_images = []
|
||||
max_output = 10
|
||||
recent_files = all_saved_files[-min(max_output, len(all_saved_files)):]
|
||||
for file_path in recent_files:
|
||||
try:
|
||||
img = Image.open(file_path)
|
||||
output_images.append(img)
|
||||
except Exception as e:
|
||||
print(f"全能生图（批量）: 无法加载 {file_path} - {e}")
|
||||
|
||||
if not output_images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
output_images = [placeholder]
|
||||
|
||||
output_tensor = pil_to_tensor(output_images)
|
||||
print(f"全能生图（批量）: 共保存 {len(all_saved_files)} 张图片，节点输出最后 {len(output_images)} 张")
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
|
||||
except Exception as e:
|
||||
print(f"全能生图（批量）: ❌ {str(e)}")
|
||||
raise
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
"""
|
||||
Flux2 图像编辑节点
|
||||
通过 vip.o1key.com 调用 Flux2 + SeedVR2 远程服务进行图像编辑和超分辨率
|
||||
|
||||
功能：
|
||||
- 接收主图和参考图
|
||||
- 上传到远程服务器执行图像编辑
|
||||
- 轮询等待 SeedVR2 超分辨率结果
|
||||
- 返回最终放大后的图像
|
||||
"""
|
||||
|
||||
import time
|
||||
from io import BytesIO
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor
|
||||
from ..clients.flux_edit_client import FluxEditClient
|
||||
|
||||
|
||||
class FluxImageEdit:
|
||||
"""
|
||||
Flux2 图像编辑节点
|
||||
|
||||
通过远程 API 将主图与参考图结合，按照提示词进行图像编辑，
|
||||
并经 SeedVR2 超分辨率放大后返回最终结果。
|
||||
"""
|
||||
|
||||
SIZES = ["2K", "4K"]
|
||||
|
||||
def __init__(self):
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"主图": ("IMAGE",),
|
||||
"参考图": ("IMAGE",),
|
||||
"提示词": ("STRING", {
|
||||
"default": "Replace the woman's underwear in Figure 1 with the strapless bra in Figure 2",
|
||||
"multiline": True,
|
||||
}),
|
||||
"分辨率": (cls.SIZES, {
|
||||
"default": "4K",
|
||||
}),
|
||||
"轮询间隔": ("INT", {
|
||||
"default": 15,
|
||||
"min": 5,
|
||||
"max": 60,
|
||||
"step": 5,
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff,
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("输出图像",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "image/edit"
|
||||
|
||||
def _image_to_jpeg_bytes(self, image: Image.Image, quality: int = 92) -> bytes:
|
||||
"""将 PIL Image 转为 JPEG 二进制"""
|
||||
if image.mode in ("RGBA", "P", "LA"):
|
||||
image = image.convert("RGB")
|
||||
buf = BytesIO()
|
||||
image.save(buf, format="JPEG", quality=quality)
|
||||
return buf.getvalue()
|
||||
|
||||
def generate(
|
||||
self,
|
||||
主图: torch.Tensor,
|
||||
参考图: torch.Tensor,
|
||||
提示词: str,
|
||||
分辨率: str,
|
||||
轮询间隔: int,
|
||||
seed: int,
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""
|
||||
执行图像编辑
|
||||
|
||||
Args:
|
||||
主图: 要编辑的原始图像 (ComfyUI tensor, [B, H, W, C])
|
||||
参考图: 参考/风格图像 (ComfyUI tensor, [B, H, W, C])
|
||||
提示词: 编辑指令
|
||||
分辨率: 超分辨率目标 ("2K" 或 "4K"，会自动映射为 2048/4096)
|
||||
轮询间隔: 轮询秒数
|
||||
seed: 随机种子
|
||||
|
||||
Returns:
|
||||
输出图像 tensor (IMAGE,)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 初始化客户端
|
||||
if self.client is None:
|
||||
self.client = FluxEditClient()
|
||||
|
||||
# Tensor → PIL（取第一张）
|
||||
main_pils = tensor_to_pil(主图)
|
||||
ref_pils = tensor_to_pil(参考图)
|
||||
|
||||
if not main_pils:
|
||||
raise ValueError("主图不能为空")
|
||||
if not ref_pils:
|
||||
raise ValueError("参考图不能为空")
|
||||
|
||||
main_img = main_pils[0]
|
||||
ref_img = ref_pils[0]
|
||||
|
||||
# PIL → JPEG bytes
|
||||
main_bytes = self._image_to_jpeg_bytes(main_img)
|
||||
ref_bytes = self._image_to_jpeg_bytes(ref_img)
|
||||
|
||||
print(f"Flux Edit: 开始处理 | 主图 {main_img.size} | 参考图 {ref_img.size} | 分辨率 {分辨率} | seed {seed}")
|
||||
|
||||
# 进度回调
|
||||
def progress_callback(status_str: str):
|
||||
print(f"Flux Edit: {status_str}")
|
||||
|
||||
# 提交任务并等待结果
|
||||
result_bytes = self.client.submit_and_wait(
|
||||
image_bytes=main_bytes,
|
||||
mask_bytes=ref_bytes,
|
||||
prompt=提示词,
|
||||
size=分辨率,
|
||||
poll_interval=轮询间隔,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
# 解码结果
|
||||
result_img = Image.open(BytesIO(result_bytes))
|
||||
if result_img.mode != "RGB":
|
||||
result_img = result_img.convert("RGB")
|
||||
|
||||
print(f"Flux Edit: 结果图像尺寸 {result_img.size}")
|
||||
|
||||
# 转为 tensor
|
||||
output_tensor = pil_to_tensor([result_img])
|
||||
|
||||
# 打印耗时
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed < 60:
|
||||
time_str = f"{elapsed:.1f}s"
|
||||
else:
|
||||
minutes = int(elapsed // 60)
|
||||
seconds = elapsed % 60
|
||||
time_str = f"{minutes}m {seconds:.0f}s"
|
||||
print(f"Flux Edit: 完成！总耗时 {time_str}")
|
||||
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as e:
|
||||
if str(e) == "未授权！":
|
||||
print("请联系作者授权后方可使用！")
|
||||
raise ValueError("未授权！") from None
|
||||
print(f"Flux Edit: ❌ {e}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"Flux Edit: ❌ {error_msg}")
|
||||
raise RuntimeError(error_msg) from None
|
||||
@@ -0,0 +1,755 @@
|
||||
"""
|
||||
Google Gemini 节点
|
||||
ComfyUI 自定义节点,用于调用 Gemini Flash 模型进行多模态文本生成
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
import tempfile
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from io import BytesIO
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
|
||||
from ..utils.file_types import FileData
|
||||
from ..clients.gemini_flash_client import GeminiFlashClient
|
||||
from ..models_config import get_enabled_flash_models
|
||||
|
||||
# 文件大小限制（20MB）
|
||||
MAX_FILE_SIZE = 20 * 1024 * 1024
|
||||
|
||||
# 图片缩放后最大尺寸（1K分辨率 = 1024像素）
|
||||
MAX_IMAGE_DIMENSION = 1024
|
||||
|
||||
# 视频压缩目标大小（1-10MB）
|
||||
TARGET_VIDEO_SIZE_MIN = 1 * 1024 * 1024
|
||||
TARGET_VIDEO_SIZE_MAX = 10 * 1024 * 1024
|
||||
|
||||
|
||||
# 支持的视频 MIME 类型映射
|
||||
VIDEO_MIME_TYPES = {
|
||||
".mp4": "video/mp4",
|
||||
".mpeg": "video/mpeg",
|
||||
".mpg": "video/mpg",
|
||||
".mov": "video/quicktime",
|
||||
".avi": "video/x-msvideo",
|
||||
".flv": "video/x-flv",
|
||||
".webm": "video/webm",
|
||||
".wmv": "video/x-ms-wmv",
|
||||
".3gp": "video/3gpp",
|
||||
".3gpp": "video/3gpp"
|
||||
}
|
||||
|
||||
# 尝试导入视频处理库
|
||||
try:
|
||||
import cv2
|
||||
CV2_AVAILABLE = True
|
||||
except ImportError:
|
||||
CV2_AVAILABLE = False
|
||||
print("⚠️ Google Gemini: OpenCV (cv2) 不可用，视频压缩功能将受限")
|
||||
|
||||
try:
|
||||
import subprocess
|
||||
FFMPEG_AVAILABLE = True
|
||||
except ImportError:
|
||||
FFMPEG_AVAILABLE = False
|
||||
|
||||
|
||||
class GoogleGemini:
|
||||
"""
|
||||
Google Gemini 节点
|
||||
|
||||
功能：
|
||||
- 支持多个 Gemini Flash 模型
|
||||
- 支持图片、视频和文件输入
|
||||
- 支持不同思考等级（不思考/低/中/高）- 通过 thinkingConfig.thinkingLevel 控制
|
||||
- 输出生成的文本内容（主要内容 + 思考内容）
|
||||
"""
|
||||
|
||||
# 支持的思考等级选项
|
||||
THINKING_LEVELS = ["不思考", "低", "中", "高"]
|
||||
|
||||
def __init__(self):
|
||||
"""初始化节点"""
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
"""
|
||||
# 从配置获取启用的模型列表
|
||||
enabled_models = get_enabled_flash_models()
|
||||
default_model = enabled_models[0] if enabled_models else "gemini-3-flash-preview"
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"模型": (enabled_models, {
|
||||
"default": default_model
|
||||
}),
|
||||
"提示词": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True
|
||||
}),
|
||||
"思考等级": (cls.THINKING_LEVELS, {
|
||||
"default": "不思考"
|
||||
})
|
||||
},
|
||||
"optional": {
|
||||
"图片": ("IMAGE",),
|
||||
"视频": ("VIDEO",),
|
||||
"文件": ("FILE",)
|
||||
}
|
||||
}
|
||||
|
||||
# 返回值类型
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("主要内容",)
|
||||
|
||||
# 执行函数名
|
||||
FUNCTION = "generate"
|
||||
|
||||
# 节点分类
|
||||
CATEGORY = "text/generation"
|
||||
|
||||
# 允许输出到 UI
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def _resize_image_if_needed(self, img: Image.Image) -> Image.Image:
|
||||
"""
|
||||
如果图片过大，缩放到1K分辨率
|
||||
|
||||
Args:
|
||||
img: PIL Image 对象
|
||||
|
||||
Returns:
|
||||
缩放后的 PIL Image
|
||||
"""
|
||||
width, height = img.size
|
||||
max_dim = max(width, height)
|
||||
|
||||
if max_dim > MAX_IMAGE_DIMENSION:
|
||||
# 计算缩放比例
|
||||
scale = MAX_IMAGE_DIMENSION / max_dim
|
||||
new_width = int(width * scale)
|
||||
new_height = int(height * scale)
|
||||
|
||||
print(f"Google Gemini: 图片尺寸 {width}x{height} 超过限制，缩放至 {new_width}x{new_height}")
|
||||
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
return img
|
||||
|
||||
def _check_and_compress_image(self, img: Image.Image) -> str:
|
||||
"""
|
||||
检查图片大小，如果超过20MB则进行压缩
|
||||
|
||||
Args:
|
||||
img: PIL Image 对象
|
||||
|
||||
Returns:
|
||||
base64 编码的字符串
|
||||
"""
|
||||
# 先进行尺寸缩放（如果需要）
|
||||
img = self._resize_image_if_needed(img)
|
||||
|
||||
# 尝试不同的压缩质量
|
||||
qualities = [95, 85, 75, 65, 55, 45]
|
||||
|
||||
for quality in qualities:
|
||||
buffer = BytesIO()
|
||||
# 转换为RGB模式（去除alpha通道）以减小体积
|
||||
if img.mode in ('RGBA', 'P'):
|
||||
img_rgb = img.convert('RGB')
|
||||
else:
|
||||
img_rgb = img
|
||||
|
||||
img_rgb.save(buffer, format='JPEG', quality=quality, optimize=True)
|
||||
buffer.seek(0)
|
||||
data = buffer.getvalue()
|
||||
|
||||
if len(data) <= MAX_FILE_SIZE:
|
||||
print(f"Google Gemini: 图片压缩后大小 {len(data) / 1024 / 1024:.2f}MB (质量{quality})")
|
||||
return base64.b64encode(data).decode('utf-8')
|
||||
|
||||
# 如果所有质量都无法满足，使用最低质量
|
||||
print(f"Google Gemini: 警告 - 即使最低质量仍超过20MB，将使用最低质量发送")
|
||||
return base64.b64encode(data).decode('utf-8')
|
||||
|
||||
def _prepare_image_data(
|
||||
self,
|
||||
images: Optional[torch.Tensor]
|
||||
) -> Optional[List[Dict[str, str]]]:
|
||||
"""
|
||||
准备图片数据
|
||||
|
||||
如果图片超过20MB，会自动进行缩放和压缩
|
||||
|
||||
Args:
|
||||
images: ComfyUI 图片张量 [B, H, W, C]
|
||||
|
||||
Returns:
|
||||
图片数据列表，每个元素包含 mime_type 和 data
|
||||
"""
|
||||
if images is None:
|
||||
return None
|
||||
|
||||
pil_images = tensor_to_pil(images)
|
||||
if not pil_images:
|
||||
return None
|
||||
|
||||
# 将所有图片转为 RGB PIL Image 并首次编码
|
||||
processed = [] # [(pil_img_rgb, b64_data, mime_type)]
|
||||
for img in pil_images:
|
||||
buffer = BytesIO()
|
||||
img.save(buffer, format='PNG')
|
||||
original_size = buffer.tell()
|
||||
buffer.close()
|
||||
|
||||
if original_size > MAX_FILE_SIZE:
|
||||
print(f"Google Gemini: 检测到图片过大 ({original_size / 1024 / 1024:.2f}MB)，正在进行压缩...")
|
||||
img_rgb = img.convert('RGB') if img.mode != 'RGB' else img.copy()
|
||||
b64_str = self._check_and_compress_image(img_rgb)
|
||||
processed.append((img_rgb, b64_str, "image/jpeg"))
|
||||
else:
|
||||
b64_str = encode_image_to_base64(img)
|
||||
processed.append((None, b64_str, "image/png"))
|
||||
|
||||
# 多图总体积控制
|
||||
def calc_total_bytes():
|
||||
return sum(len(base64.b64decode(item[1])) for item in processed)
|
||||
|
||||
total = calc_total_bytes()
|
||||
if total > MAX_FILE_SIZE and len(processed) > 1:
|
||||
print(f"Google Gemini: 图片总体积 {total / 1024 / 1024:.2f}MB 超过 {MAX_FILE_SIZE // 1024 // 1024}MB 限制，正在压缩...")
|
||||
|
||||
# 降质量
|
||||
for quality in range(70, 19, -10):
|
||||
new_processed = []
|
||||
for pil_img, _, _ in processed:
|
||||
if pil_img is None:
|
||||
# PNG 原图需要转 RGB
|
||||
continue
|
||||
buf = BytesIO()
|
||||
pil_img.save(buf, format='JPEG', quality=quality, optimize=True)
|
||||
data = buf.getvalue()
|
||||
new_processed.append((pil_img, base64.b64encode(data).decode('utf-8'), "image/jpeg"))
|
||||
if not new_processed:
|
||||
break
|
||||
processed = new_processed
|
||||
total = calc_total_bytes()
|
||||
if total <= MAX_FILE_SIZE:
|
||||
print(f"Google Gemini: 图片压缩完成，总体积 {total / 1024 / 1024:.2f}MB ({len(processed)}张图片，质量{quality})")
|
||||
break
|
||||
|
||||
# 降分辨率
|
||||
if total > MAX_FILE_SIZE:
|
||||
for scale in [0.75, 0.5, 0.35]:
|
||||
new_processed = []
|
||||
for pil_img, _, _ in processed:
|
||||
if pil_img is None:
|
||||
continue
|
||||
w, h = pil_img.size
|
||||
resized = pil_img.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
|
||||
buf = BytesIO()
|
||||
resized.save(buf, format='JPEG', quality=20, optimize=True)
|
||||
data = buf.getvalue()
|
||||
new_processed.append((resized, base64.b64encode(data).decode('utf-8'), "image/jpeg"))
|
||||
if not new_processed:
|
||||
break
|
||||
processed = new_processed
|
||||
total = calc_total_bytes()
|
||||
if total <= MAX_FILE_SIZE:
|
||||
print(f"Google Gemini: 图片压缩完成，总体积 {total / 1024 / 1024:.2f}MB ({len(processed)}张图片，缩放{int(scale*100)}%)")
|
||||
break
|
||||
|
||||
if total > MAX_FILE_SIZE:
|
||||
print(f"Google Gemini: 无法将 {len(processed)} 张图片压缩到 {MAX_FILE_SIZE // 1024 // 1024}MB 以内，请减少图片数量或降低分辨率")
|
||||
raise ValueError(f"图片总体积 {total / 1024 / 1024:.2f}MB 超过限制，无法压缩到 {MAX_FILE_SIZE // 1024 // 1024}MB 以内")
|
||||
|
||||
image_data = [{"mime_type": mt, "data": b64} for _, b64, mt in processed]
|
||||
return image_data
|
||||
|
||||
def _compress_video_with_ffmpeg(self, input_path: str, output_path: str, target_size: int) -> bool:
|
||||
"""
|
||||
使用 FFmpeg 压缩视频到目标大小
|
||||
|
||||
Args:
|
||||
input_path: 输入视频路径
|
||||
output_path: 输出视频路径
|
||||
target_size: 目标文件大小（字节）
|
||||
|
||||
Returns:
|
||||
是否压缩成功
|
||||
"""
|
||||
try:
|
||||
# 获取视频时长（秒）
|
||||
probe_cmd = ['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
|
||||
'-of', 'default=noprint_wrappers=1:nokey=1', input_path]
|
||||
duration = float(subprocess.check_output(probe_cmd).decode().strip())
|
||||
|
||||
# 计算目标比特率（bit/s），预留一些余量
|
||||
target_bitrate = int((target_size * 8) / duration * 0.9)
|
||||
|
||||
# 使用 FFmpeg 压缩视频
|
||||
# -c:v libx264: 使用 H.264 编码器
|
||||
# -b:v: 视频比特率
|
||||
# -maxrate 和 -bufsize: 控制码率波动
|
||||
# -c:a aac: 音频使用 AAC 编码
|
||||
# -b:a 128k: 音频比特率 128k
|
||||
# -movflags +faststart: 优化网络播放
|
||||
cmd = [
|
||||
'ffmpeg', '-y', '-i', input_path,
|
||||
'-c:v', 'libx264',
|
||||
'-b:v', f'{target_bitrate}',
|
||||
'-maxrate', f'{int(target_bitrate * 1.5)}',
|
||||
'-bufsize', f'{target_bitrate * 2}',
|
||||
'-c:a', 'aac',
|
||||
'-b:a', '128k',
|
||||
'-movflags', '+faststart',
|
||||
'-preset', 'fast',
|
||||
output_path
|
||||
]
|
||||
|
||||
print(f"Google Gemini: 正在压缩视频到 {target_size / 1024 / 1024:.1f}MB...")
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
|
||||
if result.returncode == 0 and os.path.exists(output_path):
|
||||
final_size = os.path.getsize(output_path)
|
||||
print(f"Google Gemini: 视频压缩完成，最终大小 {final_size / 1024 / 1024:.2f}MB")
|
||||
return True
|
||||
else:
|
||||
print(f"Google Gemini: FFmpeg 压缩失败: {result.stderr}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Google Gemini: 视频压缩异常: {str(e)}")
|
||||
return False
|
||||
|
||||
def _compress_video_with_opencv(self, input_path: str, output_path: str, scale: float = 0.5) -> bool:
|
||||
"""
|
||||
使用 OpenCV 压缩视频（备用方案）
|
||||
|
||||
Args:
|
||||
input_path: 输入视频路径
|
||||
output_path: 输出视频路径
|
||||
scale: 尺寸缩放比例
|
||||
|
||||
Returns:
|
||||
是否压缩成功
|
||||
"""
|
||||
if not CV2_AVAILABLE:
|
||||
return False
|
||||
|
||||
try:
|
||||
cap = cv2.VideoCapture(input_path)
|
||||
if not cap.isOpened():
|
||||
return False
|
||||
|
||||
# 获取原视频参数
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
|
||||
# 计算新尺寸
|
||||
new_width = int(width * scale)
|
||||
new_height = int(height * scale)
|
||||
|
||||
# 创建视频写入器
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
out = cv2.VideoWriter(output_path, fourcc, fps, (new_width, new_height))
|
||||
|
||||
print(f"Google Gemini: 使用 OpenCV 压缩视频，分辨率 {width}x{height} -> {new_width}x{new_height}")
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
# 缩放帧
|
||||
resized = cv2.resize(frame, (new_width, new_height))
|
||||
out.write(resized)
|
||||
|
||||
cap.release()
|
||||
out.release()
|
||||
|
||||
if os.path.exists(output_path):
|
||||
final_size = os.path.getsize(output_path)
|
||||
print(f"Google Gemini: 视频压缩完成，最终大小 {final_size / 1024 / 1024:.2f}MB")
|
||||
return True
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Google Gemini: OpenCV 压缩失败: {str(e)}")
|
||||
return False
|
||||
|
||||
def _compress_video(self, video_path: str) -> str:
|
||||
"""
|
||||
压缩视频到 1-10MB 之间
|
||||
|
||||
Args:
|
||||
video_path: 原视频路径
|
||||
|
||||
Returns:
|
||||
压缩后的视频路径（临时文件）
|
||||
"""
|
||||
original_size = os.path.getsize(video_path)
|
||||
print(f"Google Gemini: 视频文件过大 ({original_size / 1024 / 1024:.2f}MB)，正在压缩...")
|
||||
|
||||
# 创建临时文件
|
||||
temp_dir = tempfile.gettempdir()
|
||||
_, ext = os.path.splitext(video_path)
|
||||
output_path = os.path.join(temp_dir, f"compressed_{int(time.time())}{ext}")
|
||||
|
||||
# 确定目标大小（优先尝试 10MB，如果不行再降低）
|
||||
target_sizes = [
|
||||
TARGET_VIDEO_SIZE_MAX, # 10MB
|
||||
int(TARGET_VIDEO_SIZE_MAX * 0.8), # 8MB
|
||||
int(TARGET_VIDEO_SIZE_MAX * 0.6), # 6MB
|
||||
int(TARGET_VIDEO_SIZE_MAX * 0.5), # 5MB
|
||||
TARGET_VIDEO_SIZE_MIN * 5, # 5MB
|
||||
TARGET_VIDEO_SIZE_MIN * 3, # 3MB
|
||||
TARGET_VIDEO_SIZE_MIN * 2, # 2MB
|
||||
]
|
||||
|
||||
# 优先尝试 FFmpeg
|
||||
if FFMPEG_AVAILABLE:
|
||||
for target_size in target_sizes:
|
||||
if self._compress_video_with_ffmpeg(video_path, output_path, target_size):
|
||||
# 检查最终大小
|
||||
final_size = os.path.getsize(output_path)
|
||||
if TARGET_VIDEO_SIZE_MIN <= final_size <= MAX_FILE_SIZE:
|
||||
return output_path
|
||||
# 如果仍然太大，继续降低目标
|
||||
os.remove(output_path)
|
||||
|
||||
# FFmpeg 失败或不可用，尝试 OpenCV
|
||||
if CV2_AVAILABLE:
|
||||
scales = [0.7, 0.5, 0.4, 0.3, 0.25]
|
||||
for scale in scales:
|
||||
if self._compress_video_with_opencv(video_path, output_path, scale):
|
||||
final_size = os.path.getsize(output_path)
|
||||
if final_size <= MAX_FILE_SIZE:
|
||||
return output_path
|
||||
# 如果仍然太大，继续降低分辨率
|
||||
os.remove(output_path)
|
||||
|
||||
# 所有压缩方法都失败
|
||||
raise ValueError(
|
||||
f"视频文件过大 ({original_size / 1024 / 1024:.2f}MB) 且无法压缩到 20MB 以下。"
|
||||
f"请安装 FFmpeg 以获得更好的压缩效果，或手动压缩视频。"
|
||||
)
|
||||
|
||||
def _prepare_video_data(
|
||||
self,
|
||||
video
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
准备视频数据
|
||||
|
||||
ComfyUI VIDEO 类型包含视频文件路径信息。
|
||||
读取视频文件并转换为 base64。
|
||||
如果视频超过 20MB，会自动进行压缩。
|
||||
|
||||
Args:
|
||||
video: ComfyUI VIDEO 类型数据
|
||||
|
||||
Returns:
|
||||
视频数据字典，包含 mime_type 和 data
|
||||
"""
|
||||
if video is None:
|
||||
return None
|
||||
|
||||
# VIDEO 类型处理：支持多种格式
|
||||
video_path = None
|
||||
temp_compressed_path = None
|
||||
|
||||
if isinstance(video, dict):
|
||||
# 字典格式：尝试常见的键名
|
||||
video_path = video.get("video") or video.get("path") or video.get("file") or video.get("filename")
|
||||
# 如果还是找不到，遍历所有键找到有效路径
|
||||
if not video_path:
|
||||
for key, val in video.items():
|
||||
if isinstance(val, str) and os.path.exists(val):
|
||||
video_path = val
|
||||
break
|
||||
elif isinstance(video, str):
|
||||
# 字符串格式：直接作为路径
|
||||
video_path = video
|
||||
else:
|
||||
# 对象格式：尝试常见属性
|
||||
# 1. 尝试 __file 属性（VideoFromFile 对象）
|
||||
if hasattr(video, "__file"):
|
||||
video_path = video.__file
|
||||
# 2. 尝试其他常见属性
|
||||
elif hasattr(video, "video"):
|
||||
video_path = video.video
|
||||
elif hasattr(video, "path"):
|
||||
video_path = video.path
|
||||
elif hasattr(video, "filename"):
|
||||
video_path = video.filename
|
||||
# 3. 尝试从 __dict__ 中查找路径（支持私有属性如 _VideoFromFile__file）
|
||||
elif hasattr(video, "__dict__"):
|
||||
for attr_name, attr_value in video.__dict__.items():
|
||||
# 查找字符串类型的属性，且包含 file 或 path 关键字
|
||||
if isinstance(attr_value, str):
|
||||
if "file" in attr_name.lower() or "path" in attr_name.lower():
|
||||
# 验证路径是否有效
|
||||
if os.path.exists(attr_value):
|
||||
video_path = attr_value
|
||||
break
|
||||
# 如果属性值本身看起来像文件路径，也尝试使用
|
||||
elif os.path.exists(attr_value) and os.path.isfile(attr_value):
|
||||
video_path = attr_value
|
||||
break
|
||||
|
||||
if not video_path or not os.path.exists(video_path):
|
||||
print(f"Google Gemini: 视频文件不存在或路径无效: {video_path}")
|
||||
return None
|
||||
|
||||
# 获取文件扩展名和 MIME 类型
|
||||
_, ext = os.path.splitext(video_path)
|
||||
ext = ext.lower()
|
||||
|
||||
mime_type = VIDEO_MIME_TYPES.get(ext, "video/mp4")
|
||||
|
||||
try:
|
||||
# 检查文件大小
|
||||
file_size = os.path.getsize(video_path)
|
||||
|
||||
# 如果超过 20MB，进行压缩
|
||||
if file_size > MAX_FILE_SIZE:
|
||||
video_path = self._compress_video(video_path)
|
||||
temp_compressed_path = video_path
|
||||
# 压缩后统一使用 mp4 格式
|
||||
mime_type = "video/mp4"
|
||||
|
||||
# 读取并编码视频
|
||||
with open(video_path, "rb") as f:
|
||||
video_bytes = f.read()
|
||||
|
||||
b64_str = base64.b64encode(video_bytes).decode("utf-8")
|
||||
|
||||
# 清理临时文件
|
||||
if temp_compressed_path and os.path.exists(temp_compressed_path):
|
||||
try:
|
||||
os.remove(temp_compressed_path)
|
||||
print(f"Google Gemini: 临时压缩文件已清理")
|
||||
except:
|
||||
pass
|
||||
|
||||
return {
|
||||
"mime_type": mime_type,
|
||||
"data": b64_str
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
# 清理临时文件
|
||||
if temp_compressed_path and os.path.exists(temp_compressed_path):
|
||||
try:
|
||||
os.remove(temp_compressed_path)
|
||||
except:
|
||||
pass
|
||||
|
||||
print(f"Google Gemini: 处理视频文件失败 - {str(e)}")
|
||||
return None
|
||||
|
||||
def _prepare_file_data(
|
||||
self,
|
||||
file: Optional[FileData]
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
准备文件数据
|
||||
|
||||
从 FILE 类型提取文件数据
|
||||
|
||||
Args:
|
||||
file: FileData 对象（来自 LoadFile 节点）
|
||||
|
||||
Returns:
|
||||
文件数据字典，包含 mime_type 和 data
|
||||
"""
|
||||
if file is None:
|
||||
return None
|
||||
|
||||
return {
|
||||
"mime_type": file.mime_type,
|
||||
"data": file.data
|
||||
}
|
||||
|
||||
def _parse_dual_output(self, raw_response: Dict) -> Tuple[str, str]:
|
||||
"""
|
||||
解析包含思考内容和主要内容的响应
|
||||
|
||||
Args:
|
||||
raw_response: API 原始响应字典
|
||||
|
||||
Returns:
|
||||
(主要内容, 思考内容)
|
||||
"""
|
||||
candidates = raw_response.get("candidates", [])
|
||||
if not candidates:
|
||||
return ("", "")
|
||||
|
||||
parts = candidates[0].get("content", {}).get("parts", [])
|
||||
|
||||
thought_text = ""
|
||||
main_text = ""
|
||||
|
||||
for part in parts:
|
||||
if part.get("thought") is True:
|
||||
# 思考部分
|
||||
thought_text = part.get("text", "")
|
||||
elif "thoughtSignature" in part or "text" in part:
|
||||
# 主要内容
|
||||
main_text = part.get("text", "")
|
||||
|
||||
return main_text
|
||||
|
||||
def generate(
|
||||
self,
|
||||
模型: str,
|
||||
提示词: str,
|
||||
思考等级: str,
|
||||
图片: Optional[torch.Tensor] = None,
|
||||
视频=None,
|
||||
文件: Optional[FileData] = None
|
||||
) -> Tuple[str]:
|
||||
"""
|
||||
生成文本
|
||||
|
||||
Args:
|
||||
模型: 使用的模型名称
|
||||
提示词: 用户提示词
|
||||
思考等级: 思考等级选项
|
||||
图片: 输入图片
|
||||
视频: 输入视频
|
||||
文件: 输入文件（PDF/TXT）
|
||||
|
||||
Returns:
|
||||
(主要内容, 思考内容)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 初始化 API 客户端
|
||||
if self.client is None:
|
||||
try:
|
||||
self.client = GeminiFlashClient()
|
||||
except ValueError as e:
|
||||
raise ValueError(f"初始化失败: {str(e)}")
|
||||
|
||||
# 准备图片数据
|
||||
image_data = self._prepare_image_data(图片)
|
||||
if image_data:
|
||||
print(f"Google Gemini: 输入 {len(image_data)} 张图片")
|
||||
|
||||
# 准备视频数据
|
||||
video_data = self._prepare_video_data(视频)
|
||||
if video_data:
|
||||
print(f"Google Gemini: 输入视频 ({video_data['mime_type']})")
|
||||
|
||||
# 准备文件数据
|
||||
document_data = self._prepare_file_data(文件)
|
||||
if document_data:
|
||||
file_type = "PDF" if document_data['mime_type'] == "application/pdf" else "TXT"
|
||||
print(f"Google Gemini: 输入文件 ({file_type})")
|
||||
|
||||
# 构建输入描述
|
||||
input_desc = []
|
||||
if 提示词:
|
||||
input_desc.append("文本")
|
||||
if image_data:
|
||||
input_desc.append(f"{len(image_data)}张图片")
|
||||
if video_data:
|
||||
input_desc.append("视频")
|
||||
if document_data:
|
||||
input_desc.append("文件")
|
||||
|
||||
print(f"Google Gemini: 模型 = {模型}")
|
||||
print(f"Google Gemini: 多模态输入 ({', '.join(input_desc)})")
|
||||
print(f"Google Gemini: 思考等级 = {思考等级}")
|
||||
|
||||
# 获取端点和构建请求体
|
||||
endpoint = self.client.get_endpoint(model=模型)
|
||||
request_body = self.client.build_request_body(
|
||||
prompt=提示词,
|
||||
model=模型,
|
||||
thinking_level=思考等级,
|
||||
image_data=image_data,
|
||||
video_data=video_data,
|
||||
document_data=document_data
|
||||
)
|
||||
|
||||
print(f"Google Gemini: 发送请求...")
|
||||
|
||||
# 调用底层 API 获取原始响应
|
||||
async def get_raw_response():
|
||||
return await self.client.request_async(
|
||||
endpoint,
|
||||
request_body,
|
||||
session=None
|
||||
)
|
||||
|
||||
# 在独立线程中执行异步请求
|
||||
raw_response = self.client.run_async_in_thread(get_raw_response())
|
||||
|
||||
# 计算耗时
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# 解析响应，分离主要内容和思考内容
|
||||
main_text = self._parse_dual_output(raw_response)
|
||||
|
||||
# 打印响应 token 用量
|
||||
usage = raw_response.get("usageMetadata", {})
|
||||
prompt_tokens = usage.get("promptTokenCount", 0)
|
||||
candidates_tokens = usage.get("candidatesTokenCount", 0)
|
||||
thoughts_tokens = usage.get("thoughtsTokenCount", 0)
|
||||
total_tokens = usage.get("totalTokenCount", 0)
|
||||
finish_reason = ""
|
||||
candidates = raw_response.get("candidates", [])
|
||||
if candidates:
|
||||
finish_reason = candidates[0].get("finishReason", "")
|
||||
|
||||
print(f"Google Gemini: 生成完成 (耗时: {elapsed:.2f}s)")
|
||||
print(f"Google Gemini: finishReason = {finish_reason}")
|
||||
print(f"Google Gemini: Token 用量 — 输入: {prompt_tokens}, 输出: {candidates_tokens}, 思考: {thoughts_tokens}, 合计: {total_tokens}")
|
||||
print(f"Google Gemini: 主要内容长度: {len(main_text)} 字符")
|
||||
|
||||
# 输出预览
|
||||
if main_text:
|
||||
preview = main_text[:100] + "..." if len(main_text) > 100 else main_text
|
||||
print(f"Google Gemini: 主要内容预览: {preview}")
|
||||
|
||||
return (main_text,)
|
||||
if str(e) == "未授权！":
|
||||
print("请联系作者授权后方可使用！")
|
||||
raise ValueError("未授权！") from None
|
||||
else:
|
||||
# 用户输入错误 - 只显示简洁信息
|
||||
error_msg = str(e).split('\n')[0] # 只取第一行
|
||||
print(f"Google Gemini: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
# 日志只打第一行；报错框展示完整多行
|
||||
error_full = str(e)
|
||||
print(f"Google Gemini: ❌ {error_full.split('\n')[0]}")
|
||||
raise RuntimeError(error_full) from None
|
||||
|
||||
except Exception as e:
|
||||
# 其他未知错误 - 只显示简洁信息
|
||||
error_msg = str(e).split('\n')[0]
|
||||
print(f"Google Gemini: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
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"Google Gemini: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,241 @@
|
||||
"""
|
||||
高级图像拼接节点
|
||||
支持最多 10 张图像按指定方向（上、下、左、右）依次拼接，
|
||||
支持调整图像大小匹配和添加间隔。
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor
|
||||
from ..utils.file_utils import load_images_from_folder
|
||||
|
||||
|
||||
# 间隔颜色映射
|
||||
SPACING_COLOR_MAP = {
|
||||
"white": (255, 255, 255),
|
||||
"black": (0, 0, 0),
|
||||
"red": (255, 0, 0),
|
||||
"green": (0, 255, 0),
|
||||
"blue": (0, 0, 255),
|
||||
}
|
||||
|
||||
|
||||
def _resize_to_match(img: Image.Image, ref: Image.Image, direction: str) -> Image.Image:
|
||||
"""
|
||||
按拼接方向将 img 缩放，使其与 ref 在垂直于拼接轴的尺寸上一致。
|
||||
|
||||
- 水平拼接 (right/left)：统一高度
|
||||
- 垂直拼接 (down/up)：统一宽度
|
||||
"""
|
||||
ref_w, ref_h = ref.size
|
||||
img_w, img_h = img.size
|
||||
|
||||
if direction in ("right", "left"):
|
||||
if img_h != ref_h:
|
||||
scale = ref_h / img_h
|
||||
new_w = max(1, int(img_w * scale))
|
||||
img = img.resize((new_w, ref_h), Image.LANCZOS)
|
||||
else:
|
||||
if img_w != ref_w:
|
||||
scale = ref_w / img_w
|
||||
new_h = max(1, int(img_h * scale))
|
||||
img = img.resize((ref_w, new_h), Image.LANCZOS)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def _make_spacer(ref: Image.Image, spacing_width: int,
|
||||
direction: str, color: Tuple[int, int, int]) -> Image.Image:
|
||||
"""创建间隔色块"""
|
||||
if direction in ("right", "left"):
|
||||
return Image.new("RGB", (spacing_width, ref.size[1]), color)
|
||||
else:
|
||||
return Image.new("RGB", (ref.size[0], spacing_width), color)
|
||||
|
||||
|
||||
def _stitch_two(img_a: Image.Image, img_b: Image.Image,
|
||||
direction: str, match_size: bool,
|
||||
spacing_width: int, spacing_color: Tuple[int, int, int]) -> Image.Image:
|
||||
"""
|
||||
将两张 PIL 图像按指定方向拼接。
|
||||
img_a 为基准图像，img_b 拼接在 img_a 的指定方向侧。
|
||||
direction="right" → img_b 在 img_a 右侧
|
||||
direction="left" → img_b 在 img_a 左侧
|
||||
direction="down" → img_b 在 img_a 下方
|
||||
direction="up" → img_b 在 img_a 上方
|
||||
"""
|
||||
if img_a.mode != "RGB":
|
||||
img_a = img_a.convert("RGB")
|
||||
if img_b.mode != "RGB":
|
||||
img_b = img_b.convert("RGB")
|
||||
|
||||
if match_size:
|
||||
img_b = _resize_to_match(img_b, img_a, direction)
|
||||
|
||||
if direction == "right":
|
||||
pieces = [img_a, img_b]
|
||||
elif direction == "left":
|
||||
pieces = [img_b, img_a]
|
||||
elif direction == "down":
|
||||
pieces = [img_a, img_b]
|
||||
else: # up
|
||||
pieces = [img_b, img_a]
|
||||
|
||||
if spacing_width > 0:
|
||||
interleaved: List[Image.Image] = []
|
||||
for idx, piece in enumerate(pieces):
|
||||
interleaved.append(piece)
|
||||
if idx < len(pieces) - 1:
|
||||
interleaved.append(_make_spacer(piece, spacing_width, direction, spacing_color))
|
||||
pieces = interleaved
|
||||
|
||||
if direction in ("right", "left"):
|
||||
total_w = sum(p.size[0] for p in pieces)
|
||||
max_h = max(p.size[1] for p in pieces)
|
||||
canvas = Image.new("RGB", (total_w, max_h), spacing_color)
|
||||
x = 0
|
||||
for piece in pieces:
|
||||
canvas.paste(piece, (x, 0))
|
||||
x += piece.size[0]
|
||||
else:
|
||||
max_w = max(p.size[0] for p in pieces)
|
||||
total_h = sum(p.size[1] for p in pieces)
|
||||
canvas = Image.new("RGB", (max_w, total_h), spacing_color)
|
||||
y = 0
|
||||
for piece in pieces:
|
||||
canvas.paste(piece, (0, y))
|
||||
y += piece.size[1]
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def _natural_sort_key(filename: str):
|
||||
"""按数字优先的文件名排序，使 1, 2, 3, 10 而非 1, 10, 2, 3"""
|
||||
try:
|
||||
return (0, int(filename))
|
||||
except ValueError:
|
||||
return (1, filename.lower())
|
||||
|
||||
|
||||
class ImageStitchPro:
|
||||
"""
|
||||
高级图像拼接节点
|
||||
|
||||
在 ComfyUI 原生拼接节点基础上扩展，支持同时输入最多 10 张图像，
|
||||
按指定方向依次拼接，并可在图像间添加任意颜色的间隔。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"方向": (["right", "down", "left", "up"], {"default": "down"}),
|
||||
"匹配图像尺寸": ("BOOLEAN", {"default": True}),
|
||||
"间距宽度": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 2}),
|
||||
"间距颜色": (["white", "black", "red", "green", "blue"], {"default": "white"}),
|
||||
},
|
||||
"optional": {
|
||||
"图1": ("IMAGE",),
|
||||
"图2": ("IMAGE",),
|
||||
"图3": ("IMAGE",),
|
||||
"图4": ("IMAGE",),
|
||||
"图5": ("IMAGE",),
|
||||
"图6": ("IMAGE",),
|
||||
"图7": ("IMAGE",),
|
||||
"图8": ("IMAGE",),
|
||||
"图9": ("IMAGE",),
|
||||
"图10": ("IMAGE",),
|
||||
"图11": ("IMAGE",),
|
||||
"图12": ("IMAGE",),
|
||||
"图片路径（可选）": ("STRING", {"default": "", "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("拼接图像",)
|
||||
FUNCTION = "stitch"
|
||||
CATEGORY = "image"
|
||||
|
||||
DESCRIPTION = (
|
||||
"高级图像拼接节点，支持最多 12 张图像按指定方向（右/下/左/上）依次拼接。\n"
|
||||
"可选择是否将后续图像缩放以匹配第一张图像的尺寸，并可在图像间添加彩色间隔。\n"
|
||||
"可选填「图片路径」：仅处理该文件夹内图片，按文件名顺序依次拼接；与输入端图片不可同时使用。"
|
||||
)
|
||||
|
||||
def stitch(
|
||||
self,
|
||||
方向: str = "down",
|
||||
匹配图像尺寸: bool = True,
|
||||
间距宽度: int = 0,
|
||||
间距颜色: str = "white",
|
||||
图1: Optional[torch.Tensor] = None,
|
||||
图2: Optional[torch.Tensor] = None,
|
||||
图3: Optional[torch.Tensor] = None,
|
||||
图4: Optional[torch.Tensor] = None,
|
||||
图5: Optional[torch.Tensor] = None,
|
||||
图6: Optional[torch.Tensor] = None,
|
||||
图7: Optional[torch.Tensor] = None,
|
||||
图8: Optional[torch.Tensor] = None,
|
||||
图9: Optional[torch.Tensor] = None,
|
||||
图10: Optional[torch.Tensor] = None,
|
||||
图11: Optional[torch.Tensor] = None,
|
||||
图12: Optional[torch.Tensor] = None,
|
||||
**kwargs: object,
|
||||
) -> Tuple[torch.Tensor]:
|
||||
|
||||
color = SPACING_COLOR_MAP.get(间距颜色, (255, 255, 255))
|
||||
raw_tensors = [图1, 图2, 图3, 图4, 图5, 图6, 图7, 图8, 图9, 图10, 图11, 图12]
|
||||
tensors = [t for t in raw_tensors if t is not None]
|
||||
has_input_images = len(tensors) > 0
|
||||
image_folder = (kwargs.get("图片路径（可选）") or "").strip()
|
||||
|
||||
if image_folder and has_input_images:
|
||||
raise ValueError("不可同时使用「图片路径（可选）」与输入端图片，请二选一。")
|
||||
|
||||
if image_folder:
|
||||
infos = load_images_from_folder(image_folder)
|
||||
if not infos:
|
||||
raise ValueError(f"文件夹中未找到可用的图片，或路径无效: {image_folder}")
|
||||
infos.sort(key=lambda x: _natural_sort_key(x.filename))
|
||||
pil_list = [info.image for info in infos]
|
||||
if len(pil_list) == 1:
|
||||
return (pil_to_tensor(pil_list),)
|
||||
base = pil_list[0]
|
||||
for next_img in pil_list[1:]:
|
||||
base = _stitch_two(
|
||||
base, next_img,
|
||||
direction=方向,
|
||||
match_size=匹配图像尺寸,
|
||||
spacing_width=间距宽度,
|
||||
spacing_color=color,
|
||||
)
|
||||
return (pil_to_tensor([base]),)
|
||||
else:
|
||||
if not has_input_images:
|
||||
raise ValueError("请至少接入一张图片，或填写「图片路径（可选）」中的文件夹路径。")
|
||||
|
||||
if len(tensors) == 1:
|
||||
return (tensors[0],)
|
||||
|
||||
pil_batches: List[List[Image.Image]] = [tensor_to_pil(t) for t in tensors]
|
||||
|
||||
batch_size = min(len(b) for b in pil_batches)
|
||||
result_images: List[Image.Image] = []
|
||||
|
||||
for i in range(batch_size):
|
||||
frames = [batch[i] for batch in pil_batches]
|
||||
base = frames[0]
|
||||
for next_img in frames[1:]:
|
||||
base = _stitch_two(
|
||||
base, next_img,
|
||||
direction=方向,
|
||||
match_size=匹配图像尺寸,
|
||||
spacing_width=间距宽度,
|
||||
spacing_color=color,
|
||||
)
|
||||
result_images.append(base)
|
||||
|
||||
return (pil_to_tensor(result_images),)
|
||||
@@ -0,0 +1,736 @@
|
||||
"""
|
||||
Kling 3.0 Video Nodes
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
|
||||
from ..clients.kling_client import KlingClient
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
|
||||
|
||||
from comfy_api.latest import InputImpl
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
|
||||
def _get_video_output_dir() -> str:
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
base = folder_paths.get_output_directory()
|
||||
else:
|
||||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
base = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output")
|
||||
video_dir = os.path.join(base, "video")
|
||||
os.makedirs(video_dir, exist_ok=True)
|
||||
return video_dir
|
||||
|
||||
|
||||
def _get_next_counter(directory: str, prefix: str) -> int:
|
||||
if not os.path.exists(directory):
|
||||
return 1
|
||||
pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)")
|
||||
max_counter = 0
|
||||
for f in os.listdir(directory):
|
||||
m = pattern.match(f)
|
||||
if m:
|
||||
max_counter = max(max_counter, int(m.group(1)))
|
||||
return max_counter + 1
|
||||
|
||||
|
||||
def _tensor_to_base64(tensor) -> str:
|
||||
"""ComfyUI IMAGE tensor → base64 PNG 字符串"""
|
||||
pil_images = tensor_to_pil(tensor)
|
||||
return encode_image_to_base64(pil_images[0], format="PNG")
|
||||
|
||||
|
||||
def _validate_prompt(prompt: str, *, required: bool = True) -> None:
|
||||
"""校验单条提示词。
|
||||
|
||||
Args:
|
||||
prompt: 提示词字符串。
|
||||
required: 为 True 时不允许为空（多镜头关闭或 shot_type 为 intelligence 时适用）。
|
||||
"""
|
||||
if required and not prompt.strip():
|
||||
raise ValueError("提示词不能为空（非多镜头模式下必填）。")
|
||||
if len(prompt) > 2500:
|
||||
raise ValueError(
|
||||
f"提示词长度 ({len(prompt)}) 超过上限 2500 个字符，请缩短后重试。"
|
||||
)
|
||||
|
||||
|
||||
def _validate_multi_prompt(multi_prompt_list: list, total_duration: int) -> None:
|
||||
"""校验多镜头分镜列表。
|
||||
|
||||
规则：
|
||||
- 分镜数量：1 ~ 6；
|
||||
- 每个分镜提示词不超过 512 个字符；
|
||||
- 每个分镜时长 ≥ 1 且 ≤ total_duration；
|
||||
- 所有分镜时长之和必须等于 total_duration。
|
||||
"""
|
||||
count = len(multi_prompt_list)
|
||||
if count < 1 or count > 6:
|
||||
raise ValueError(
|
||||
f"多镜头分镜数量须在 1~6 之间，当前为 {count}。"
|
||||
)
|
||||
|
||||
duration_sum = 0
|
||||
for entry in multi_prompt_list:
|
||||
idx = entry["index"]
|
||||
p = entry.get("prompt", "")
|
||||
dur = entry.get("duration", 0)
|
||||
|
||||
if len(p) > 512:
|
||||
raise ValueError(
|
||||
f"镜头 {idx} 提示词长度 ({len(p)}) 超过上限 512 个字符。"
|
||||
)
|
||||
if dur < 1:
|
||||
raise ValueError(
|
||||
f"镜头 {idx} 时长 ({dur}s) 不能小于 1 秒。"
|
||||
)
|
||||
if dur > total_duration:
|
||||
raise ValueError(
|
||||
f"镜头 {idx} 时长 ({dur}s) 超过任务总时长 ({total_duration}s)。"
|
||||
)
|
||||
duration_sum += dur
|
||||
|
||||
if duration_sum != total_duration:
|
||||
raise ValueError(
|
||||
f"所有分镜时长之和 ({duration_sum}s) 必须等于任务总时长 ({total_duration}s)。"
|
||||
)
|
||||
|
||||
|
||||
def _validate_image(tensor, label: str = "图片") -> None:
|
||||
"""校验图片张量。
|
||||
|
||||
规则：
|
||||
- 文件大小（PNG）不超过 10MB；
|
||||
- 宽、高均不小于 300px；
|
||||
- 宽高比介于 1:2.5 ~ 2.5:1 之间（即 ratio ∈ [0.4, 2.5]）。
|
||||
"""
|
||||
import io
|
||||
|
||||
pil_images = tensor_to_pil(tensor)
|
||||
img = pil_images[0]
|
||||
w, h = img.size
|
||||
|
||||
# ── 最小尺寸 ──────────────────────────────────────────────────────
|
||||
if w < 300 or h < 300:
|
||||
raise ValueError(
|
||||
f"{label} 宽高不得小于 300px，当前为 {w}×{h}px。"
|
||||
)
|
||||
|
||||
# ── 宽高比 ────────────────────────────────────────────────────────
|
||||
ratio = w / h
|
||||
if ratio < 1 / 2.5 or ratio > 2.5:
|
||||
raise ValueError(
|
||||
f"{label} 宽高比须在 1:2.5 ~ 2.5:1 之间，"
|
||||
f"当前为 {w}:{h}（比值 {ratio:.2f}）。"
|
||||
)
|
||||
|
||||
# ── 文件大小 ──────────────────────────────────────────────────────
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
size_mb = buf.tell() / (1024 * 1024)
|
||||
if size_mb > 10:
|
||||
raise ValueError(
|
||||
f"{label} PNG 大小 ({size_mb:.1f}MB) 超过上限 10MB。"
|
||||
)
|
||||
|
||||
|
||||
class KlingVideo:
|
||||
"""Kling 3.0 视频生成节点（支持多镜头）"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"反向提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"时长": ([5, 10, 15],),
|
||||
"分辨率": (["1080p", "720p"],),
|
||||
"宽高比": (["智能", "16:9", "9:16", "1:1"], {"default": "智能"}),
|
||||
"生成音频": (["打开", "关闭"], {"default": "打开"}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}),
|
||||
},
|
||||
"optional": {
|
||||
"起始帧": ("IMAGE",),
|
||||
"镜头1_提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"镜头1_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}),
|
||||
"镜头2_提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"镜头2_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}),
|
||||
"镜头3_提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"镜头3_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}),
|
||||
"镜头4_提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"镜头4_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}),
|
||||
"镜头5_提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"镜头5_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}),
|
||||
"镜头6_提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"镜头6_时长": ("INT", {"default": 5, "min": 1, "max": 15, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("视频",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/Kling"
|
||||
|
||||
async def generate(self, **kwargs):
|
||||
"""生成视频（支持多镜头）"""
|
||||
prompt = kwargs["提示词"]
|
||||
negative_prompt = kwargs["反向提示词"]
|
||||
duration = kwargs["时长"]
|
||||
resolution = kwargs["分辨率"]
|
||||
aspect_ratio = kwargs["宽高比"]
|
||||
generate_audio = kwargs["生成音频"]
|
||||
start_frame = kwargs.get("起始帧", None)
|
||||
seed = kwargs.get("seed", 0) # noqa: F841 — 触发 ComfyUI 缓存刷新
|
||||
|
||||
mode = "pro" if resolution == "1080p" else "std"
|
||||
voice = "voice" if generate_audio == "打开" else "novoice"
|
||||
|
||||
# ── 多镜头检测 ────────────────────────────────────────────────
|
||||
multi_prompt_list = []
|
||||
for i in range(1, 7):
|
||||
sb_prompt = kwargs.get(f"镜头{i}_提示词", "").strip()
|
||||
if sb_prompt:
|
||||
sb_duration = kwargs.get(f"镜头{i}_时长", 5)
|
||||
multi_prompt_list.append({
|
||||
"index": i,
|
||||
"prompt": sb_prompt,
|
||||
"duration": sb_duration,
|
||||
})
|
||||
|
||||
multi_shot_enabled = len(multi_prompt_list) > 0
|
||||
|
||||
if multi_shot_enabled:
|
||||
total_duration = sum(e["duration"] for e in multi_prompt_list)
|
||||
if total_duration < 3 or total_duration > 15:
|
||||
raise ValueError(
|
||||
f"多镜头总时长 ({total_duration}s) 必须在 3~15 秒之间。"
|
||||
)
|
||||
_validate_multi_prompt(multi_prompt_list, total_duration)
|
||||
duration = total_duration
|
||||
else:
|
||||
_validate_prompt(prompt, required=True)
|
||||
|
||||
# ── 构建模型名 & 请求体 ───────────────────────────────────────
|
||||
import json, base64, copy
|
||||
model_name = f"kling-v3-{mode}-{duration}s-{voice}"
|
||||
|
||||
body = {
|
||||
"model": model_name,
|
||||
"mode": mode,
|
||||
"duration": duration,
|
||||
}
|
||||
|
||||
sound = "on" if generate_audio == "打开" else "off"
|
||||
|
||||
if multi_shot_enabled or sound == "on":
|
||||
ms_payload = {}
|
||||
ms_payload["prompt"] = prompt
|
||||
|
||||
if sound == "on":
|
||||
ms_payload["sound"] = "on"
|
||||
|
||||
if multi_shot_enabled:
|
||||
ms_payload["multi_shot"] = True
|
||||
ms_payload["shot_type"] = "customize"
|
||||
ms_payload["multi_prompt"] = multi_prompt_list
|
||||
|
||||
encoded = base64.b64encode(
|
||||
json.dumps(ms_payload, ensure_ascii=False).encode("utf-8")
|
||||
).decode("utf-8")
|
||||
body["prompt"] = f"__MS__:{encoded}"
|
||||
else:
|
||||
body["prompt"] = prompt
|
||||
|
||||
if negative_prompt.strip():
|
||||
body["negative_prompt"] = negative_prompt
|
||||
|
||||
if start_frame is not None:
|
||||
_validate_image(start_frame, "起始帧")
|
||||
body["image"] = _tensor_to_base64(start_frame)
|
||||
endpoint_type = "image2video"
|
||||
else:
|
||||
body["metadata"] = {"aspect_ratio": aspect_ratio}
|
||||
endpoint_type = "text2video"
|
||||
|
||||
# ── 保存路径 ──────────────────────────────────────────────────
|
||||
video_dir = _get_video_output_dir()
|
||||
counter = _get_next_counter(video_dir, "kling")
|
||||
save_path = os.path.join(video_dir, f"kling_{counter:05d}.mp4")
|
||||
|
||||
client = KlingClient()
|
||||
|
||||
# ── 进度条 ────────────────────────────────────────────────────
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
pbar = ProgressBar(100)
|
||||
except Exception:
|
||||
pbar = None
|
||||
|
||||
def on_stage(stage: str):
|
||||
if stage == "submitting":
|
||||
print("[视频生成] 提交中...")
|
||||
if pbar: pbar.update_absolute(0, 100)
|
||||
elif stage.startswith("submitted:"):
|
||||
print(f"[视频生成] 任务已提交 → {stage.split(':',1)[1]}")
|
||||
if pbar: pbar.update_absolute(5, 100)
|
||||
elif stage == "downloading":
|
||||
print("[视频生成] 下载视频...")
|
||||
if pbar: pbar.update_absolute(99, 100)
|
||||
elif stage == "done":
|
||||
print("[视频生成] 完成")
|
||||
if pbar: pbar.update_absolute(100, 100)
|
||||
|
||||
def on_progress(pct: int):
|
||||
mapped = 5 + int(pct * 0.94)
|
||||
if pbar: pbar.update_absolute(mapped, 100)
|
||||
|
||||
try:
|
||||
result_path = await client.generate_async(
|
||||
endpoint_type=endpoint_type,
|
||||
body=body,
|
||||
save_path=save_path,
|
||||
on_stage=on_stage,
|
||||
on_progress=on_progress,
|
||||
)
|
||||
return (InputImpl.VideoFromFile(result_path),)
|
||||
finally:
|
||||
# 查询余额
|
||||
try:
|
||||
_balance_client = GeminiAPIClient()
|
||||
balance_data = _balance_client.query_balance_sync()
|
||||
balance_info = _balance_client.format_balance_info(balance_data)
|
||||
print(f"自研视频模型: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class KlingFirstLastFrame:
|
||||
"""Kling 3.0 首尾帧到视频节点"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"首帧": ("IMAGE",),
|
||||
"尾帧": ("IMAGE",),
|
||||
"提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"时长": ([5, 10, 15],),
|
||||
"生成音频": (["打开", "关闭"], {"default": "打开"}),
|
||||
"模型": (["v3"],),
|
||||
"分辨率": (["1080p", "720p"],),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("视频",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/Kling"
|
||||
|
||||
async def generate(self, **kwargs):
|
||||
first_frame = kwargs["首帧"]
|
||||
end_frame = kwargs["尾帧"]
|
||||
prompt = kwargs["提示词"]
|
||||
duration = kwargs["时长"]
|
||||
generate_audio = kwargs["生成音频"]
|
||||
model_base = kwargs["模型"]
|
||||
model_base = "kling-" + model_base # v3 → kling-v3（后端值还原）
|
||||
resolution = kwargs["分辨率"]
|
||||
seed = kwargs.get("seed", 0) # noqa: F841 — 触发 ComfyUI 缓存刷新
|
||||
|
||||
_validate_prompt(prompt, required=True)
|
||||
|
||||
# 时长校验
|
||||
if duration not in (5, 10, 15):
|
||||
raise ValueError(f"时长仅支持 5、10、15 秒，当前值为 {duration}，请重新选择。")
|
||||
|
||||
# 拼接模型名：kling-v3-{mode}-{dur}s-{voice}
|
||||
mode = "pro" if resolution == "1080p" else "std"
|
||||
voice = "voice" if generate_audio == "打开" else "novoice"
|
||||
model_name = f"{model_base}-{mode}-{duration}s-{voice}"
|
||||
|
||||
# 图片校验 & 转 base64
|
||||
_validate_image(first_frame, "首帧")
|
||||
_validate_image(end_frame, "尾帧")
|
||||
image_b64 = _tensor_to_base64(first_frame)
|
||||
image_tail_b64 = _tensor_to_base64(end_frame)
|
||||
|
||||
# ── 按规范编码 prompt 和 sound ──────────────────────────
|
||||
import json, base64
|
||||
sound = "on" if generate_audio == "打开" else "off"
|
||||
|
||||
body = {
|
||||
"model": model_name,
|
||||
"image": image_b64,
|
||||
"mode": mode,
|
||||
"duration": duration,
|
||||
"metadata": {
|
||||
"image_tail": image_tail_b64,
|
||||
},
|
||||
}
|
||||
|
||||
if sound == "on":
|
||||
ms_payload = {
|
||||
"prompt": prompt,
|
||||
"sound": "on",
|
||||
}
|
||||
encoded = base64.b64encode(
|
||||
json.dumps(ms_payload, ensure_ascii=False).encode("utf-8")
|
||||
).decode("utf-8")
|
||||
body["prompt"] = f"__MS__:{encoded}"
|
||||
else:
|
||||
body["prompt"] = prompt
|
||||
|
||||
# 保存路径
|
||||
video_dir = _get_video_output_dir()
|
||||
counter = _get_next_counter(video_dir, "kling")
|
||||
save_path = os.path.join(video_dir, f"kling_{counter:05d}.mp4")
|
||||
|
||||
client = KlingClient()
|
||||
|
||||
# 进度条：0~100 步
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
pbar = ProgressBar(100)
|
||||
except Exception:
|
||||
pbar = None
|
||||
|
||||
def on_stage(stage: str):
|
||||
if stage == "submitting":
|
||||
print("[视频生成] 提交中...")
|
||||
if pbar:
|
||||
pbar.update_absolute(0, 100)
|
||||
elif stage.startswith("submitted:"):
|
||||
print(f"[视频生成] 任务已提交 → {stage.split(':',1)[1]}")
|
||||
if pbar:
|
||||
pbar.update_absolute(5, 100)
|
||||
elif stage == "downloading":
|
||||
print("[视频生成] 下载视频...")
|
||||
if pbar:
|
||||
pbar.update_absolute(99, 100)
|
||||
elif stage == "done":
|
||||
print("[视频生成] 完成")
|
||||
if pbar:
|
||||
pbar.update_absolute(100, 100)
|
||||
|
||||
def on_progress(pct: int):
|
||||
# pct 来自 API progress 字段，如 50 表示 50%
|
||||
# 生成阶段占 5~99 区间
|
||||
mapped = 5 + int(pct * 0.94)
|
||||
if pbar:
|
||||
pbar.update_absolute(mapped, 100)
|
||||
|
||||
try:
|
||||
result_path = await client.generate_async(
|
||||
endpoint_type="image2video",
|
||||
body=body,
|
||||
save_path=save_path,
|
||||
on_stage=on_stage,
|
||||
on_progress=on_progress,
|
||||
)
|
||||
return (InputImpl.VideoFromFile(result_path),)
|
||||
finally:
|
||||
# 查询余额
|
||||
try:
|
||||
_balance_client = GeminiAPIClient()
|
||||
balance_data = _balance_client.query_balance_sync()
|
||||
balance_info = _balance_client.format_balance_info(balance_data)
|
||||
print(f"自研视频模型: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class KlingMotionControlTest:
|
||||
"""Kling 动作控制（测试）节点 —— reference_video 接受 VIDEO 类型输入"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"提示词": ("STRING", {"multiline": True, "default": ""}),
|
||||
"参考图片": ("IMAGE",),
|
||||
"参考视频": ("VIDEO",),
|
||||
},
|
||||
"optional": {
|
||||
"保留原声": ("BOOLEAN", {"default": True}),
|
||||
"人物朝向": (["video", "image"],),
|
||||
"画质模式": (["专家", "标准"],),
|
||||
"模型版本": (["v3"],),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("视频",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "comfyui_o1key/Kling"
|
||||
|
||||
async def generate(self, **kwargs):
|
||||
"""动作控制（测试）：VIDEO 类型参考视频 + 图片人物动作迁移"""
|
||||
import base64
|
||||
|
||||
prompt = kwargs["提示词"]
|
||||
reference_image = kwargs["参考图片"]
|
||||
reference_video = kwargs["参考视频"]
|
||||
keep_original_sound = kwargs.get("保留原声", True)
|
||||
character_orientation = kwargs.get("人物朝向", "video")
|
||||
mode = kwargs.get("画质模式", "专家")
|
||||
mode = "pro" if mode == "专家" else "std" # 映射为 API 参数值
|
||||
model = kwargs.get("模型版本", "v3")
|
||||
model = "kling-" + model # v3 → kling-v3（后端值还原）
|
||||
seed = kwargs.get("seed", 0) # noqa: F841 — 触发 ComfyUI 缓存刷新
|
||||
|
||||
# ── 校验提示词 ────────────────────────────────────────────────
|
||||
_validate_prompt(prompt, required=True)
|
||||
|
||||
# ── 校验参考图片 ──────────────────────────────────────────────
|
||||
_validate_image(reference_image, "参考图片")
|
||||
image_b64 = _tensor_to_base64(reference_image)
|
||||
|
||||
# ── 从 VIDEO 对象获取本地文件路径并读取 ───────────────────────
|
||||
# ComfyUI VIDEO 对象有 .source_path 或通过 VideoFromFile 构造
|
||||
video_path = None
|
||||
if hasattr(reference_video, "source_path"):
|
||||
video_path = reference_video.source_path
|
||||
elif hasattr(reference_video, "path"):
|
||||
video_path = reference_video.path
|
||||
elif isinstance(reference_video, str):
|
||||
video_path = reference_video.strip()
|
||||
|
||||
if not video_path or not os.path.isfile(video_path):
|
||||
raise ValueError(
|
||||
f"无法获取参考视频文件路径，请确保连接的是本地视频文件。"
|
||||
f"（当前路径：{video_path}）"
|
||||
)
|
||||
|
||||
# ── 校验视频时长约束 ──────────────────────────────────────────
|
||||
# 人物朝向="video" → 3~30 秒；人物朝向="image" → 3~10 秒
|
||||
try:
|
||||
import subprocess, json as _json
|
||||
ffprobe_cmd = [
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_format",
|
||||
video_path,
|
||||
]
|
||||
result_proc = subprocess.run(ffprobe_cmd, capture_output=True, text=True, timeout=30)
|
||||
if result_proc.returncode == 0:
|
||||
info = _json.loads(result_proc.stdout)
|
||||
duration_sec = float(info.get("format", {}).get("duration", 0))
|
||||
if character_orientation == "video":
|
||||
if not (3 <= duration_sec <= 30):
|
||||
raise ValueError(
|
||||
f"当人物朝向为 'video' 时，"
|
||||
f"参考视频时长须在 3~30 秒之间，当前为 {duration_sec:.1f}s。"
|
||||
)
|
||||
else: # "image"
|
||||
if not (3 <= duration_sec <= 10):
|
||||
raise ValueError(
|
||||
f"当人物朝向为 'image' 时，"
|
||||
f"参考视频时长须在 3~10 秒之间，当前为 {duration_sec:.1f}s。"
|
||||
)
|
||||
except FileNotFoundError:
|
||||
# ffprobe 不可用时跳过时长校验，但打印提示
|
||||
print("[动作控制] 警告：ffprobe 未找到，跳过视频时长校验。")
|
||||
except ValueError:
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"[动作控制] 时长校验异常（已跳过）：{e}")
|
||||
|
||||
# ── 视频转 base64 ─────────────────────────────────────────────
|
||||
with open(video_path, "rb") as f:
|
||||
video_b64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
# ── 构建请求体 ────────────────────────────────────────────────
|
||||
body = {
|
||||
"prompt": prompt,
|
||||
"character_orientation": character_orientation,
|
||||
"mode": mode,
|
||||
"model": model,
|
||||
"keep_original_sound": "yes" if keep_original_sound else "no",
|
||||
"image": image_b64,
|
||||
"video": video_b64,
|
||||
}
|
||||
|
||||
# ── 保存路径 ──────────────────────────────────────────────────
|
||||
video_dir = _get_video_output_dir()
|
||||
counter = _get_next_counter(video_dir, "kling_motion_test")
|
||||
save_path = os.path.join(video_dir, f"kling_motion_test_{counter:05d}.mp4")
|
||||
|
||||
client = KlingClient()
|
||||
|
||||
# ── 进度条 ────────────────────────────────────────────────────
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
pbar = ProgressBar(100)
|
||||
except Exception:
|
||||
pbar = None
|
||||
|
||||
def on_stage(stage: str):
|
||||
if stage == "submitting":
|
||||
print("[动作控制] 提交中...")
|
||||
if pbar: pbar.update_absolute(0, 100)
|
||||
elif stage.startswith("submitted:"):
|
||||
print(f"[动作控制] 任务已提交 → {stage.split(':',1)[1]}")
|
||||
if pbar: pbar.update_absolute(5, 100)
|
||||
elif stage == "downloading":
|
||||
print("[动作控制] 下载视频...")
|
||||
if pbar: pbar.update_absolute(99, 100)
|
||||
elif stage == "done":
|
||||
print("[动作控制] 完成")
|
||||
if pbar: pbar.update_absolute(100, 100)
|
||||
|
||||
def on_progress(pct: int):
|
||||
mapped = 5 + int(pct * 0.94)
|
||||
if pbar: pbar.update_absolute(mapped, 100)
|
||||
|
||||
try:
|
||||
result_path = await client.generate_async(
|
||||
endpoint_type="motion_control",
|
||||
body=body,
|
||||
save_path=save_path,
|
||||
on_stage=on_stage,
|
||||
on_progress=on_progress,
|
||||
)
|
||||
return (InputImpl.VideoFromFile(result_path),)
|
||||
finally:
|
||||
# 查询余额
|
||||
try:
|
||||
_balance_client = GeminiAPIClient()
|
||||
balance_data = _balance_client.query_balance_sync()
|
||||
balance_info = _balance_client.format_balance_info(balance_data)
|
||||
print(f"自研视频模型: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class AspectRatioPreset:
|
||||
"""图片宽高比预设节点"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"图像": ("IMAGE",),
|
||||
"宽高比": (["智能", "16:9", "9:16", "4:3", "3:4", "1:1"], {"default": "智能"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("图像",)
|
||||
FUNCTION = "resize"
|
||||
CATEGORY = "comfyui_o1key/Utils"
|
||||
|
||||
def resize(self, 图像, 宽高比):
|
||||
import torch
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
pil_images = tensor_to_pil(图像)
|
||||
img = pil_images[0]
|
||||
w, h = img.size
|
||||
img_ratio = w / h
|
||||
|
||||
# 确定原图所属的宽高比家族
|
||||
ratios = {"16:9": 16/9, "9:16": 9/16, "4:3": 4/3, "3:4": 3/4, "1:1": 1.0}
|
||||
closest_ratio = min(ratios.keys(), key=lambda k: abs(ratios[k] - img_ratio))
|
||||
|
||||
# 智能模式：使用最接近的比例
|
||||
if 宽高比 == "智能":
|
||||
宽高比 = closest_ratio
|
||||
|
||||
# 解析目标比例
|
||||
target_w, target_h = map(int, 宽高比.split(":"))
|
||||
target_ratio = target_w / target_h
|
||||
|
||||
# 确定分辨率级别（1K/2K）
|
||||
max_dim = max(w, h)
|
||||
if max_dim <= 1080:
|
||||
base = 1080
|
||||
elif max_dim <= 2160:
|
||||
base = 2160
|
||||
else:
|
||||
base = 2160
|
||||
|
||||
# 计算目标尺寸
|
||||
if target_ratio >= 1:
|
||||
target_width = base
|
||||
target_height = int(base / target_ratio)
|
||||
else:
|
||||
target_height = base
|
||||
target_width = int(base * target_ratio)
|
||||
|
||||
# 判断是否同家族（横向家族：16:9, 4:3；纵向家族：9:16, 3:4；正方形：1:1）
|
||||
horizontal_family = ["16:9", "4:3"]
|
||||
vertical_family = ["9:16", "3:4"]
|
||||
|
||||
same_family = False
|
||||
if closest_ratio in horizontal_family and 宽高比 in horizontal_family:
|
||||
same_family = True
|
||||
elif closest_ratio in vertical_family and 宽高比 in vertical_family:
|
||||
same_family = True
|
||||
elif closest_ratio == "1:1" and 宽高比 == "1:1":
|
||||
same_family = True
|
||||
|
||||
# 同家族：直接缩放或裁剪（无白底）
|
||||
if same_family:
|
||||
if img_ratio > target_ratio:
|
||||
# 图像更宽，以高度为准缩放后裁剪
|
||||
scale = target_height / h
|
||||
scaled_w = int(w * scale)
|
||||
scaled_h = target_height
|
||||
scaled = img.resize((scaled_w, scaled_h), Image.LANCZOS)
|
||||
left = (scaled_w - target_width) // 2
|
||||
result = scaled.crop((left, 0, left + target_width, target_height))
|
||||
else:
|
||||
# 图像更高，以宽度为准缩放后裁剪
|
||||
scale = target_width / w
|
||||
scaled_w = target_width
|
||||
scaled_h = int(h * scale)
|
||||
scaled = img.resize((scaled_w, scaled_h), Image.LANCZOS)
|
||||
top = (scaled_h - target_height) // 2
|
||||
result = scaled.crop((0, top, target_width, top + target_height))
|
||||
|
||||
# 不同家族：保持宽高比 + 白底填充
|
||||
else:
|
||||
if img_ratio > target_ratio:
|
||||
scaled_w = target_width
|
||||
scaled_h = int(target_width / img_ratio)
|
||||
else:
|
||||
scaled_h = target_height
|
||||
scaled_w = int(target_height * img_ratio)
|
||||
|
||||
scaled = img.resize((scaled_w, scaled_h), Image.LANCZOS)
|
||||
canvas = Image.new("RGB", (target_width, target_height), (255, 255, 255))
|
||||
paste_x = (target_width - scaled_w) // 2
|
||||
paste_y = (target_height - scaled_h) // 2
|
||||
canvas.paste(scaled, (paste_x, paste_y))
|
||||
result = canvas
|
||||
|
||||
# 转回 tensor
|
||||
arr = np.array(result).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(arr).unsqueeze(0)
|
||||
|
||||
return (tensor,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"KlingVideo": KlingVideo,
|
||||
"KlingFirstLastFrame": KlingFirstLastFrame,
|
||||
"KlingMotionControlTest": KlingMotionControlTest,
|
||||
"AspectRatioPreset": AspectRatioPreset,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"KlingVideo": "自研模型 3.0 视频",
|
||||
"KlingFirstLastFrame": "自研模型 3.0 首尾帧到视频",
|
||||
"KlingMotionControlTest": "自研模型 动作控制（测试）",
|
||||
"AspectRatioPreset": "图片宽高比预设",
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
"""
|
||||
LoadFile 节点
|
||||
ComfyUI 自定义节点,用于加载文件并转换为 FILE 类型数据
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Tuple
|
||||
|
||||
from ..utils.file_types import FileData, DOCUMENT_MIME_TYPES, FILE_SIZE_LIMITS
|
||||
|
||||
|
||||
class LoadFile:
|
||||
"""
|
||||
LoadFile 节点
|
||||
|
||||
功能：
|
||||
- 从文件系统加载文件
|
||||
- 支持 PDF 和 TXT 文件
|
||||
- 转换为 FILE 类型数据（包含 base64 编码内容）
|
||||
- 验证文件大小和格式
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"文件路径": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
# 返回值类型
|
||||
RETURN_TYPES = ("FILE", "STRING")
|
||||
RETURN_NAMES = ("文件", "文件信息")
|
||||
|
||||
# 执行函数名
|
||||
FUNCTION = "load_file"
|
||||
|
||||
# 节点分类
|
||||
CATEGORY = "file/input"
|
||||
|
||||
def load_file(self, 文件路径: str) -> Tuple[FileData, str]:
|
||||
"""
|
||||
加载文件并转换为 FILE 类型
|
||||
|
||||
Args:
|
||||
文件路径: 文件的完整路径（支持绝对路径和相对路径）
|
||||
|
||||
Returns:
|
||||
(FileData, 文件信息预览)
|
||||
|
||||
Raises:
|
||||
ValueError: 文件不存在、不支持的文件类型或文件过大
|
||||
"""
|
||||
try:
|
||||
# 清理路径（去除空格和引号）
|
||||
file_path = 文件路径.strip().strip('"').strip("'")
|
||||
|
||||
if not file_path:
|
||||
raise ValueError("文件路径不能为空")
|
||||
|
||||
# 转换为 Path 对象
|
||||
path = Path(file_path)
|
||||
|
||||
# 如果是相对路径，转换为绝对路径
|
||||
if not path.is_absolute():
|
||||
# 相对于当前工作目录
|
||||
path = Path.cwd() / path
|
||||
|
||||
# 验证文件是否存在
|
||||
if not path.exists():
|
||||
raise ValueError(f"文件不存在: {file_path}")
|
||||
|
||||
if not path.is_file():
|
||||
raise ValueError(f"路径不是文件: {file_path}")
|
||||
|
||||
# 获取文件信息
|
||||
extension = path.suffix.lower()
|
||||
filename = path.stem
|
||||
file_size = path.stat().st_size
|
||||
|
||||
# 验证文件类型
|
||||
if extension not in DOCUMENT_MIME_TYPES:
|
||||
supported_types = ", ".join(DOCUMENT_MIME_TYPES.keys())
|
||||
raise ValueError(
|
||||
f"不支持的文件类型: {extension}\n"
|
||||
f"支持的类型: {supported_types}"
|
||||
)
|
||||
|
||||
# 获取 MIME 类型
|
||||
mime_type = DOCUMENT_MIME_TYPES[extension]
|
||||
|
||||
# 验证文件大小
|
||||
size_limit = FILE_SIZE_LIMITS.get(extension, 20 * 1024 * 1024)
|
||||
if file_size > size_limit:
|
||||
raise ValueError(
|
||||
f"文件过大 ({file_size / 1024 / 1024:.2f}MB)，"
|
||||
f"最大支持 {size_limit / 1024 / 1024:.0f}MB"
|
||||
)
|
||||
|
||||
# 读取文件并转换为 base64
|
||||
print(f"LoadFile: 正在加载文件 {filename}{extension}")
|
||||
print(f"LoadFile: 文件大小 = {file_size / 1024:.2f}KB")
|
||||
|
||||
with open(path, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
# Base64 编码
|
||||
b64_str = base64.b64encode(file_bytes).decode("utf-8")
|
||||
|
||||
# 创建 FileData 对象
|
||||
file_data = FileData(
|
||||
path=str(path),
|
||||
filename=filename,
|
||||
extension=extension,
|
||||
mime_type=mime_type,
|
||||
data=b64_str,
|
||||
size=file_size
|
||||
)
|
||||
|
||||
# 生成文件信息预览
|
||||
file_info = (
|
||||
f"文件名: {filename}{extension}\n"
|
||||
f"类型: {mime_type}\n"
|
||||
f"大小: {file_size / 1024:.2f}KB\n"
|
||||
f"路径: {path}"
|
||||
)
|
||||
|
||||
print(f"LoadFile: 加载成功")
|
||||
|
||||
return (file_data, file_info)
|
||||
|
||||
except ValueError as e:
|
||||
print(f"LoadFile: 输入错误 - {str(e)}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
print(f"LoadFile: 未知错误 - {str(e)}")
|
||||
raise
|
||||
@@ -0,0 +1,165 @@
|
||||
"""
|
||||
多分辨率图像预览节点
|
||||
ComfyUI 自定义节点，支持同时预览多张不同分辨率的图像
|
||||
|
||||
背景：
|
||||
ComfyUI 原生「预览图像」节点要求 batch 内所有图片分辨率相同（因为它们被
|
||||
stack 成一个 [B, H, W, C] tensor）。当 API 返回多张不同尺寸的图片时
|
||||
（例如 nano-banana-2 同时返回 1K + 2K），原生节点会报错。
|
||||
|
||||
解决方案：
|
||||
声明 INPUT_IS_LIST = True，ComfyUI 会将连入的所有图像作为
|
||||
Python list[Tensor] 传入，而不是强行 stack 成单个 tensor。
|
||||
节点逐张单独保存为临时 PNG，再通过 ui.images 列表返回给前端并列展示，
|
||||
完全不受分辨率一致性的限制。
|
||||
"""
|
||||
|
||||
import os
|
||||
import uuid
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
|
||||
def _get_temp_dir() -> str:
|
||||
"""获取 ComfyUI temp 目录，不可用时回退到系统临时目录"""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
return folder_paths.get_temp_directory()
|
||||
import tempfile
|
||||
return tempfile.gettempdir()
|
||||
|
||||
|
||||
def _tensor_to_pil(tensor) -> list:
|
||||
"""
|
||||
将单个 IMAGE tensor 转换为 PIL Image 列表。
|
||||
|
||||
ComfyUI IMAGE tensor 格式：[B, H, W, C]，float32，值域 [0, 1]
|
||||
支持：
|
||||
- 单张图 tensor: shape [H, W, C] 或 [1, H, W, C]
|
||||
- batch tensor: shape [B, H, W, C]（B 张相同尺寸图）
|
||||
"""
|
||||
import torch
|
||||
if not isinstance(tensor, torch.Tensor):
|
||||
return []
|
||||
|
||||
if tensor.ndim == 3:
|
||||
tensor = tensor.unsqueeze(0)
|
||||
|
||||
results = []
|
||||
for i in range(tensor.shape[0]):
|
||||
img_np = tensor[i].cpu().numpy()
|
||||
img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)
|
||||
results.append(Image.fromarray(img_np))
|
||||
return results
|
||||
|
||||
|
||||
class MultiResPreview:
|
||||
"""
|
||||
多分辨率图像预览节点
|
||||
|
||||
功能：
|
||||
- 单个「图像」输入端口，支持接入批次图像
|
||||
- INPUT_IS_LIST = True：ComfyUI 将每张图作为独立 tensor 传入，
|
||||
不强制要求尺寸相同，彻底解决不同分辨率无法共存的问题
|
||||
- 每张图像独立保存为临时 PNG，在节点上并列展示所有图像
|
||||
|
||||
用法：
|
||||
将 Nano Banana 节点的输出直接连入「图像」端口即可，
|
||||
无论返回几张、分辨率是否相同，都能正确展示。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"图像": ("IMAGE",),
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
},
|
||||
}
|
||||
|
||||
# 关键：告知 ComfyUI 以 list[Tensor] 而非 stacked Tensor 传入图像
|
||||
# 这样不同分辨率的图片可以共存于同一个输入中
|
||||
INPUT_IS_LIST = True
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "preview"
|
||||
CATEGORY = "image"
|
||||
|
||||
DESCRIPTION = (
|
||||
"多分辨率图像预览节点。\n"
|
||||
"单个图像输入端口，支持任意数量、任意分辨率的批次图像。\n"
|
||||
"解决了原生「预览图像」节点要求 batch 内图片尺寸相同的限制。\n"
|
||||
"常用场景：nano-banana-2 同时返回 1K + 2K 图时，直接连入本节点即可。"
|
||||
)
|
||||
|
||||
def preview(self, 图像, prompt=None, extra_pnginfo=None) -> dict:
|
||||
"""
|
||||
逐张将图像保存到 temp 目录，返回 ui.images 供前端展示。
|
||||
|
||||
Args:
|
||||
图像: list[Tensor]，每个元素是一张或一批图（INPUT_IS_LIST）
|
||||
prompt: ComfyUI 注入的 prompt 元数据（可选）
|
||||
extra_pnginfo: ComfyUI 注入的额外 PNG 信息（可选）
|
||||
|
||||
Returns:
|
||||
{"ui": {"images": [...]}} 格式，每项对应一张图
|
||||
"""
|
||||
temp_dir = _get_temp_dir()
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
# 构建 PNG 元数据（与原生预览节点行为一致）
|
||||
metadata = PngInfo()
|
||||
# INPUT_IS_LIST 时 hidden 值也会被包装成 list，取第一个元素
|
||||
_prompt = prompt[0] if isinstance(prompt, list) else prompt
|
||||
_extra = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo
|
||||
if _prompt is not None:
|
||||
try:
|
||||
metadata.add_text("prompt", json.dumps(_prompt))
|
||||
except Exception:
|
||||
pass
|
||||
if _extra is not None:
|
||||
try:
|
||||
for k, v in _extra.items():
|
||||
metadata.add_text(k, json.dumps(v))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
saved = []
|
||||
total_input = 0
|
||||
total_saved = 0
|
||||
|
||||
# 图像 是 list[Tensor]，逐个处理（每个 Tensor 可能自身是个 batch）
|
||||
for tensor in 图像:
|
||||
pil_images = _tensor_to_pil(tensor)
|
||||
total_input += len(pil_images)
|
||||
|
||||
for pil_img in pil_images:
|
||||
try:
|
||||
filename = f"multi_res_preview_{uuid.uuid4().hex[:12]}.png"
|
||||
filepath = os.path.join(temp_dir, filename)
|
||||
pil_img.save(filepath, pnginfo=metadata, compress_level=1)
|
||||
|
||||
saved.append({
|
||||
"filename": filename,
|
||||
"subfolder": "",
|
||||
"type": "temp",
|
||||
})
|
||||
total_saved += 1
|
||||
except Exception as e:
|
||||
print(f"多分辨率预览: ⚠️ 保存图像失败 - {e}")
|
||||
|
||||
if total_input == 0:
|
||||
print("多分辨率预览: ⚠️ 没有接收到任何图像")
|
||||
|
||||
return {"ui": {"images": saved}}
|
||||
@@ -0,0 +1,903 @@
|
||||
"""
|
||||
Nano Banana Pro 节点
|
||||
ComfyUI 自定义节点，用于调用 Gemini 模型生成图像
|
||||
"""
|
||||
|
||||
import time
|
||||
import math
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
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 ..utils.file_utils import ImageInfo, generate_timestamp_filename, save_image
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
from ..models_config import (
|
||||
get_enabled_models, get_model_description,
|
||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
||||
get_model_supported_resolutions, get_all_supported_resolutions
|
||||
)
|
||||
|
||||
# 检查 folder_paths 是否可用
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
# 导入 ComfyUI 原生进度条
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ NanoBananaPro: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")
|
||||
|
||||
# 内存监控（可选）
|
||||
try:
|
||||
import psutil
|
||||
MEMORY_MONITOR_AVAILABLE = True
|
||||
except ImportError:
|
||||
MEMORY_MONITOR_AVAILABLE = False
|
||||
print("⚠️ NanoBananaPro: psutil 不可用，内存监控功能禁用")
|
||||
|
||||
# ============================================================================
|
||||
# 调试日志配置
|
||||
# ============================================================================
|
||||
# 是否启用调试日志（打印完整的 API 响应内容）
|
||||
# 设置为 True 以启用调试日志，False 以禁用
|
||||
DEBUG_LOG_ENABLED = False
|
||||
# 是否启用请求体日志（打印发送给 API 的请求体，base64 图片数据将自动截断）
|
||||
# 设置为 True 以启用请求体日志，False 以禁用
|
||||
REQUEST_LOG_ENABLED = False
|
||||
# ============================================================================
|
||||
|
||||
_NODE = "Nano Banana Pro"
|
||||
|
||||
|
||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
||||
"""
|
||||
将 PIL Image 列表转换为 ComfyUI tensor，安全处理多张不同尺寸的情况。
|
||||
|
||||
ComfyUI 的 IMAGE tensor 格式为 [B, H, W, C]，要求 batch 内所有图尺寸相同。
|
||||
当 API 返回多张不同分辨率的图时（主图 + 附图），直接 stack 会崩溃。
|
||||
|
||||
策略：
|
||||
- 所有图均已按原始分辨率保存到磁盘（调用此函数前已完成）
|
||||
- 以第一张图的尺寸为基准，只将尺寸相同的图纳入 tensor 输出
|
||||
- 尺寸不同的图跳过（不 resize、不丢弃磁盘文件），并打印日志提示
|
||||
- 若没有任何图与第一张尺寸相同（极罕见），则只输出第一张
|
||||
"""
|
||||
if not images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
return pil_to_tensor([placeholder])
|
||||
|
||||
base_size = images[0].size # PIL size = (W, H)
|
||||
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]}×{img.size[1]}" for img in skipped)
|
||||
print(
|
||||
f"{node_label}: API 额外返回了 {len(skipped)} 张不同尺寸的图 ({sizes_str})，"
|
||||
f"已按原始分辨率保存到磁盘，tensor 输出仅包含与主图尺寸相同的 {len(matched)} 张 "
|
||||
f"({base_size[0]}×{base_size[1]})"
|
||||
)
|
||||
|
||||
return pil_to_tensor(matched if matched else [images[0]])
|
||||
|
||||
|
||||
class NanoBananaPro:
|
||||
"""
|
||||
Nano Banana Pro 节点
|
||||
|
||||
功能：
|
||||
- 文生图：基于提示词生成图像
|
||||
- 图生图：基于输入图像和提示词生成新图像
|
||||
- 批量生成：支持并发生成多张图像
|
||||
|
||||
注意：
|
||||
- 支持的模型列表从 models_config.py 动态加载
|
||||
- 要添加/禁用模型，请编辑 models_config.py 文件
|
||||
"""
|
||||
|
||||
# 支持的模型列表（从配置文件动态加载）
|
||||
MODELS = None # 将在 INPUT_TYPES 中动态获取
|
||||
|
||||
# 支持的宽高比列表（全量：所有启用模型的并集，动态加载）
|
||||
# 实际渲染时通过 get_all_supported_aspect_ratios() 获取
|
||||
ASPECT_RATIOS = [
|
||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
||||
"1:4", "4:1", "1:8", "8:1"
|
||||
]
|
||||
|
||||
# 支持的分辨率列表（全量兜底，实际由 get_all_supported_resolutions() 动态生成）
|
||||
RESOLUTIONS = ["512", "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 中启用至少一个模型"]
|
||||
|
||||
# 动态获取所有启用模型支持的宽高比（去重合并）
|
||||
all_aspect_ratios = get_all_supported_aspect_ratios()
|
||||
if not all_aspect_ratios:
|
||||
all_aspect_ratios = cls.ASPECT_RATIOS
|
||||
|
||||
# 动态获取所有启用模型支持的分辨率（去重合并）
|
||||
all_resolutions = get_all_supported_resolutions()
|
||||
if not all_resolutions:
|
||||
all_resolutions = cls.RESOLUTIONS
|
||||
|
||||
# 创建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]
|
||||
}),
|
||||
"宽高比": (all_aspect_ratios, {
|
||||
"default": "1:1"
|
||||
}),
|
||||
"分辨率": (all_resolutions, {
|
||||
"default": "2K"
|
||||
}),
|
||||
"生图数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 1000,
|
||||
"step": 1
|
||||
}),
|
||||
"像素缩放": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"label_on": "打开",
|
||||
"label_off": "关闭"
|
||||
}),
|
||||
"分辨率像素": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.1,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"display": "number"
|
||||
}),
|
||||
"谷歌搜索（联网）": (["关闭", "打开"], {
|
||||
"default": "关闭"
|
||||
}),
|
||||
"图片搜索（联网）": (["关闭", "打开"], {
|
||||
"default": "关闭"
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff
|
||||
})
|
||||
},
|
||||
"optional": optional_inputs
|
||||
}
|
||||
|
||||
# 返回值类型
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("输出图像",)
|
||||
|
||||
# 导入 ComfyUI 的文件夹路径管理
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
# 执行函数名
|
||||
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]"
|
||||
)
|
||||
|
||||
async def _generate_single_task(
|
||||
self,
|
||||
session: aiohttp.ClientSession,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: List[Image.Image],
|
||||
output_folder: str,
|
||||
global_task_index: int,
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
) -> dict:
|
||||
"""执行单个生成任务，生成后立即保存到磁盘"""
|
||||
result = {
|
||||
"global_task_index": global_task_index,
|
||||
"prompt": prompt,
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"saved_files": [],
|
||||
"error": None
|
||||
}
|
||||
|
||||
try:
|
||||
gen_result = await self.client.generate_single_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images if images else None,
|
||||
session=session,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
if gen_result:
|
||||
images_list, _ = gen_result
|
||||
for gen_img in images_list:
|
||||
output_path = generate_timestamp_filename(
|
||||
output_folder=output_folder,
|
||||
extension=".png"
|
||||
)
|
||||
save_image(gen_img, output_path)
|
||||
result["saved_files"].append(output_path)
|
||||
gen_img = None # 释放内存
|
||||
|
||||
result["success"] = True
|
||||
result["generated_count"] = len(images_list)
|
||||
except Exception as e:
|
||||
result["error"] = str(e)
|
||||
|
||||
return result
|
||||
|
||||
async def _process_batch_async(
|
||||
self,
|
||||
prompts: List[str],
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images_per_prompt: int,
|
||||
input_images: List[Image.Image],
|
||||
output_folder: str,
|
||||
pbar=None,
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
) -> List[dict]:
|
||||
"""异步批量处理：每个提示词独立调用 API，生成后立即写磁盘"""
|
||||
# 构建任务列表：(prompt, sub_index) 用于 images_per_prompt > 1 的情况
|
||||
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)
|
||||
num_prompts = len(prompts)
|
||||
print(f"Nano Banana Pro: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务")
|
||||
|
||||
max_concurrent = 10
|
||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
||||
|
||||
all_results = []
|
||||
completed = 0
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
for batch_idx in range(num_batches):
|
||||
start_idx = batch_idx * max_concurrent
|
||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||||
|
||||
tasks = []
|
||||
for i in range(start_idx, end_idx):
|
||||
_, _, prompt = tasks_def[i]
|
||||
task = asyncio.create_task(
|
||||
self._generate_single_task(
|
||||
session=session,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_images,
|
||||
output_folder=output_folder,
|
||||
global_task_index=i,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
batch_results = []
|
||||
for coro in asyncio.as_completed(tasks):
|
||||
result_data = None
|
||||
try:
|
||||
result = await coro
|
||||
if isinstance(result, Exception):
|
||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": [], "prompt": ""}
|
||||
else:
|
||||
result_data = result
|
||||
batch_results.append(result_data)
|
||||
except Exception as e:
|
||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": [], "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 Pro: [{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 Pro: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
all_results.extend(batch_results)
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
return all_results
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str,
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
生图数量: int,
|
||||
像素缩放: bool,
|
||||
分辨率像素: float,
|
||||
seed: int,
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""
|
||||
生成图像
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
模型: 模型名称
|
||||
宽高比: 宽高比
|
||||
分辨率: 分辨率
|
||||
生图数量: 批次大小
|
||||
像素缩放: 是否启用像素缩放
|
||||
分辨率像素: 目标像素数（百万像素）
|
||||
seed: 随机种子
|
||||
**kwargs: 搜索开关（谷歌搜索（联网）/ 图片搜索（联网））及动态参考图输入 (参考图1-9)
|
||||
注：两个搜索参数名含全角括号，不能作为 Python 形参，从 kwargs 中提取
|
||||
|
||||
注意：
|
||||
调试日志功能已移至文件顶部配置，通过修改 DEBUG_LOG_ENABLED 常量控制
|
||||
|
||||
Returns:
|
||||
生成的图像张量 (IMAGE,)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# 从 kwargs 提取搜索参数（界面显示为「关闭/打开」，转为 bool 供调用）
|
||||
enable_grounding: bool = (kwargs.pop("谷歌搜索（联网）", "关闭") == "打开")
|
||||
enable_image_search: bool = (kwargs.pop("图片搜索（联网）", "关闭") == "打开")
|
||||
|
||||
# 创建 ComfyUI 原生进度条
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(生图数量)
|
||||
|
||||
try:
|
||||
# 设置随机种子（用于本地随机操作）
|
||||
random.seed(seed)
|
||||
np.random.seed(seed % (2**32))
|
||||
|
||||
# 内存监控初始化
|
||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
||||
import psutil
|
||||
process = psutil.Process()
|
||||
initial_memory = process.memory_info().rss / 1024 / 1024
|
||||
print(f"Nano Banana Pro: 初始内存使用: {initial_memory:.1f} MB")
|
||||
|
||||
# 初始化 API 客户端
|
||||
if self.client is None:
|
||||
try:
|
||||
self.client = GeminiAPIClient()
|
||||
except ValueError as e:
|
||||
raise ValueError(f"初始化失败: {str(e)}")
|
||||
|
||||
# 校验分辨率与模型的兼容性
|
||||
supported_resolutions = get_model_supported_resolutions(模型)
|
||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
||||
raise ValueError(
|
||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的分辨率：{', '.join(supported_resolutions)}"
|
||||
)
|
||||
|
||||
# 校验宽高比与模型的兼容性
|
||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
||||
if supported_ratios and 宽高比 not in supported_ratios:
|
||||
raise ValueError(
|
||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的宽高比：{', '.join(supported_ratios)}"
|
||||
)
|
||||
|
||||
# 校验图片搜索（联网）与模型的兼容性
|
||||
# 仅 nano-banana-2-限时特价 和 gemini-3.1-flash-image-preview 支持图片搜索
|
||||
IMAGE_SEARCH_UNSUPPORTED_MODELS = ["nano-banana-pro-限时特价", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview"]
|
||||
if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS:
|
||||
raise ValueError(
|
||||
f"模型 \"{模型}\" 不支持【图片搜索（联网）】功能！"
|
||||
f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用"
|
||||
)
|
||||
|
||||
# 收集独立输入的参考图
|
||||
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
|
||||
|
||||
# 解析批量提示词
|
||||
batch_prompts = parse_batch_prompts(prompt)
|
||||
|
||||
# 打印首行概览
|
||||
# 图片搜索（联网）开启时隐含谷歌搜索接地，与客户端请求逻辑保持一致
|
||||
grounding_str = ""
|
||||
if enable_image_search:
|
||||
grounding_str = " | 谷歌图片搜索接地"
|
||||
elif enable_grounding:
|
||||
grounding_str = " | 谷歌搜索接地"
|
||||
|
||||
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 Pro: {mode_str} | {分辨率} {宽高比} | 共{total_images}张{grounding_str}")
|
||||
|
||||
# 大批量警告
|
||||
if total_images > 100:
|
||||
print(f"⚠️ Nano Banana Pro: 警告！批量生成 {total_images} 张图片，内存占用可能较高")
|
||||
print(f"⚠️ 建议：分批执行或减少生图数量")
|
||||
else:
|
||||
# 单提示词模式
|
||||
mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式"
|
||||
print(f"Nano Banana Pro: {mode_str} | {分辨率} {宽高比} | {生图数量}张{grounding_str}")
|
||||
|
||||
# 大批量警告
|
||||
if 生图数量 > 100:
|
||||
print(f"⚠️ Nano Banana Pro: 警告！批量生成 {生图数量} 张图片，内存占用可能较高")
|
||||
print(f"⚠️ 建议：分批执行或减少生图数量")
|
||||
|
||||
# 统计变量
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
# 进度回调 - 打印错误信息并更新进度条，添加内存监控
|
||||
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} 成功 ✓")
|
||||
else:
|
||||
fail_count += 1
|
||||
# 打印完整的错误信息（用于排查问题）
|
||||
if error_msg:
|
||||
print(f"Nano Banana Pro: 任务 {current}/{total} 失败 ✗")
|
||||
print(f"原始错误详情:\n{error_msg}")
|
||||
else:
|
||||
print(f"Nano Banana Pro: 任务 {current}/{total} 失败 ✗")
|
||||
|
||||
# 更新 ComfyUI 原生进度条
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
# 内存监控（每完成10个任务检查一次）
|
||||
if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0:
|
||||
import gc
|
||||
gc.collect() # 强制垃圾回收
|
||||
current_memory = process.memory_info().rss / 1024 / 1024
|
||||
memory_increase = current_memory - initial_memory
|
||||
print(f"Nano Banana Pro: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)")
|
||||
|
||||
# 内存警告阈值（2GB）
|
||||
if current_memory > 2000:
|
||||
print(f"⚠️ Nano Banana Pro: 内存使用过高！建议减少生图数量或分批执行")
|
||||
|
||||
# 根据是否有批量提示词选择生成模式
|
||||
if batch_prompts:
|
||||
num_prompts = len(batch_prompts)
|
||||
total_images = num_prompts * 生图数量
|
||||
|
||||
# ===== 批量提示词模式：异步并发+磁盘保存 =====
|
||||
if pbar is not None:
|
||||
pbar = ProgressBar(total_images)
|
||||
|
||||
# 确定保存路径
|
||||
output_folder = ""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
output_folder = folder_paths.get_output_directory()
|
||||
print(f"Nano Banana Pro: 磁盘保存模式 → {output_folder}")
|
||||
else:
|
||||
raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用")
|
||||
|
||||
import os
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
|
||||
def run_async_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
self._process_batch_async(
|
||||
prompts=batch_prompts,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images_per_prompt=生图数量,
|
||||
input_images=input_images,
|
||||
output_folder=output_folder,
|
||||
pbar=pbar,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async_in_thread)
|
||||
try:
|
||||
results = future.result(timeout=3600)
|
||||
except TimeoutError:
|
||||
raise RuntimeError("任务执行超时（1小时），请减少提示词数量或检查网络连接")
|
||||
|
||||
# 统计结果
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||
all_saved_files = []
|
||||
for r in results:
|
||||
all_saved_files.extend(r.get("saved_files", []))
|
||||
|
||||
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_images} | 失败: {fail_count}")
|
||||
|
||||
# 失败详情
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
if failed_results:
|
||||
for fr in failed_results:
|
||||
idx = fr.get("global_task_index", -1) + 1
|
||||
prompt_snippet = (fr.get("prompt", "") or "")[:30]
|
||||
error_msg = fr.get("error", "未知错误")
|
||||
print(f" 失败 #{idx}: {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → {error_msg}")
|
||||
|
||||
# 从磁盘加载最后 10 张图片
|
||||
output_images = []
|
||||
max_output_images = 10
|
||||
recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):]
|
||||
for file_path in recent_files:
|
||||
try:
|
||||
img = Image.open(file_path)
|
||||
output_images.append(img)
|
||||
except Exception as e:
|
||||
print(f"Nano Banana Pro: 无法加载 {file_path} - {e}")
|
||||
|
||||
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)
|
||||
print(f"Nano Banana Pro: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张")
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
else:
|
||||
# 单提示词模式
|
||||
if 生图数量 == 1:
|
||||
# 单张：同步生成 + 保存到磁盘 + 输出 tensor
|
||||
generated_images = self.client.generate_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
batch_size=1,
|
||||
images=input_images,
|
||||
progress_callback=progress_callback,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
# 单张：保存到磁盘
|
||||
import os
|
||||
output_folder = ""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
output_folder = folder_paths.get_output_directory()
|
||||
print(f"Nano Banana Pro: 磁盘保存模式 → {output_folder}")
|
||||
else:
|
||||
raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用")
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
for gen_img in generated_images:
|
||||
output_path = generate_timestamp_filename(output_folder=output_folder)
|
||||
save_image(gen_img, output_path)
|
||||
else:
|
||||
# 多张：异步并发 + 磁盘保存（与批量提示词逻辑一致）
|
||||
print(f"Nano Banana Pro: 单提示词×{生图数量}张 → 异步并发模式")
|
||||
|
||||
if pbar is not None:
|
||||
pbar = ProgressBar(生图数量)
|
||||
|
||||
output_folder = ""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
output_folder = folder_paths.get_output_directory()
|
||||
print(f"Nano Banana Pro: 磁盘保存模式 → {output_folder}")
|
||||
else:
|
||||
raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用")
|
||||
|
||||
import os
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
|
||||
def run_async_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
self._process_batch_async(
|
||||
prompts=[prompt],
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images_per_prompt=生图数量,
|
||||
input_images=input_images,
|
||||
output_folder=output_folder,
|
||||
pbar=pbar,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async_in_thread)
|
||||
try:
|
||||
results = future.result(timeout=3600)
|
||||
except TimeoutError:
|
||||
raise RuntimeError("任务执行超时（1小时），请减少生图数量或检查网络连接")
|
||||
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||
all_saved_files = []
|
||||
for r in results:
|
||||
all_saved_files.extend(r.get("saved_files", []))
|
||||
|
||||
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}/{生图数量} | 失败: {fail_count}")
|
||||
|
||||
# 失败详情
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
if failed_results:
|
||||
for fr in failed_results:
|
||||
idx = fr.get("global_task_index", -1) + 1
|
||||
error_msg = fr.get("error", "未知错误")
|
||||
print(f" 失败 #{idx}: {prompt[:30]}{'...' if len(prompt) >= 30 else ''} → {error_msg}")
|
||||
|
||||
# 从磁盘加载最后 10 张图片
|
||||
output_images = []
|
||||
max_output_images = 10
|
||||
recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):]
|
||||
for file_path in recent_files:
|
||||
try:
|
||||
img = Image.open(file_path)
|
||||
output_images.append(img)
|
||||
except Exception as e:
|
||||
print(f"Nano Banana Pro: 无法加载 {file_path} - {e}")
|
||||
|
||||
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)
|
||||
print(f"Nano Banana Pro: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张")
|
||||
# 不生成 prompts_map.txt（单提示词无需映射）
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
|
||||
|
||||
# 优化：限制输出图片数量，避免内存爆炸
|
||||
max_output_images = 20 # 最多输出20张图片到ComfyUI
|
||||
|
||||
if len(generated_images) > max_output_images:
|
||||
print(f"Nano Banana Pro: 生成 {len(generated_images)} 张图片，限制输出前 {max_output_images} 张到ComfyUI")
|
||||
output_images = generated_images[:max_output_images]
|
||||
else:
|
||||
output_images = generated_images
|
||||
|
||||
# 转换输出图像
|
||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||
|
||||
# 计算耗时并打印最终统计
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed < 1:
|
||||
time_str = f"{elapsed:.3f}s"
|
||||
else:
|
||||
time_str = f"{elapsed:.2f}s"
|
||||
|
||||
# 打印最终汇总
|
||||
if fail_count > 0:
|
||||
print(f"[4/4] 完成！总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张")
|
||||
else:
|
||||
print(f"[4/4] 完成！总耗时 {time_str} | 成功 {len(generated_images)}张")
|
||||
|
||||
# 最终内存清理
|
||||
import gc
|
||||
gc.collect()
|
||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
||||
final_memory = process.memory_info().rss / 1024 / 1024
|
||||
print(f"Nano Banana Pro: 最终内存使用: {final_memory:.1f} MB")
|
||||
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as e:
|
||||
# 检测是否为授权错误
|
||||
if str(e) == "未授权！":
|
||||
print("请联系作者授权后方可使用！")
|
||||
raise ValueError("未授权！") from None
|
||||
else:
|
||||
# 用户输入错误 - 打印完整错误信息
|
||||
error_msg = str(e)
|
||||
print(f"Nano Banana Pro: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
# 打印完整错误信息
|
||||
error_full = str(e)
|
||||
print(f"Nano Banana Pro: ❌ {error_full}")
|
||||
raise RuntimeError(error_full) from None
|
||||
|
||||
except Exception as e:
|
||||
# 其他未知错误 - 打印完整错误信息
|
||||
error_msg = str(e)
|
||||
print(f"Nano Banana Pro: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
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:
|
||||
pass
|
||||
|
||||
# 最终内存清理
|
||||
import gc
|
||||
gc.collect()
|
||||
print(f"Nano Banana Pro: 最终内存清理完成")
|
||||
@@ -0,0 +1,656 @@
|
||||
"""
|
||||
Nano Banana v2 节点
|
||||
NanoBananaPro 的完全复刻，唯一改动：
|
||||
|
||||
将原来 9 个独立「参考图1~9」输入端
|
||||
改为 1 个「参考图」输入端（可选），配合「加载图像（批量）」节点使用。
|
||||
|
||||
「加载图像（批量）」输出 is_output_list=True（list[Tensor]），
|
||||
本节点声明 INPUT_IS_LIST = True 来整体接收该列表，
|
||||
然后在 generate() 开头对所有参数统一解包，其余业务逻辑与原节点完全一致。
|
||||
"""
|
||||
|
||||
import os
|
||||
import gc
|
||||
import time
|
||||
import math
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
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 ..utils.file_utils import ImageInfo, generate_timestamp_filename, save_image
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
from ..models_config import (
|
||||
get_enabled_models, get_model_description,
|
||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
||||
get_model_supported_resolutions, get_all_supported_resolutions
|
||||
)
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
|
||||
try:
|
||||
import psutil
|
||||
MEMORY_MONITOR_AVAILABLE = True
|
||||
except ImportError:
|
||||
MEMORY_MONITOR_AVAILABLE = False
|
||||
|
||||
DEBUG_LOG_ENABLED = False
|
||||
REQUEST_LOG_ENABLED = False
|
||||
|
||||
_NODE = "Nano Banana v2"
|
||||
|
||||
|
||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
||||
"""
|
||||
将 PIL Image 列表转换为 ComfyUI tensor，安全处理多张不同尺寸的情况。
|
||||
|
||||
ComfyUI 的 IMAGE tensor 格式为 [B, H, W, C]，要求 batch 内所有图尺寸相同。
|
||||
当 API 返回多张不同分辨率的图时（主图 + 附图），直接 stack 会崩溃。
|
||||
|
||||
策略：
|
||||
- 所有图均已按原始分辨率保存到磁盘（调用此函数前已完成）
|
||||
- 以第一张图的尺寸为基准，只将尺寸相同的图纳入 tensor 输出
|
||||
- 尺寸不同的图跳过（不 resize、不丢弃磁盘文件），并打印日志提示
|
||||
- 若没有任何图与第一张尺寸相同（极罕见），则只输出第一张
|
||||
"""
|
||||
if not images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
return pil_to_tensor([placeholder])
|
||||
|
||||
base_size = images[0].size # PIL size = (W, H)
|
||||
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]}×{img.size[1]}" for img in skipped)
|
||||
print(
|
||||
f"{node_label}: API 额外返回了 {len(skipped)} 张不同尺寸的图 ({sizes_str})，"
|
||||
f"已按原始分辨率保存到磁盘，tensor 输出仅包含与主图尺寸相同的 {len(matched)} 张 "
|
||||
f"({base_size[0]}×{base_size[1]})"
|
||||
)
|
||||
|
||||
return pil_to_tensor(matched if matched else [images[0]])
|
||||
|
||||
|
||||
class NanaBananaV2:
|
||||
"""
|
||||
Nano Banana v2
|
||||
|
||||
与 NanoBananaPro 完全一致，参考图输入方式不同：
|
||||
- 原版：9 个独立可选端口（参考图1~9）
|
||||
- v2：1 个可选端口「参考图」，配合「加载图像（批量）」可传入任意数量图片
|
||||
"""
|
||||
|
||||
ASPECT_RATIOS = [
|
||||
"1:1", "4:3", "3:4", "16:9", "9:16",
|
||||
"2:3", "3:2", "4:5", "5:4", "21:9",
|
||||
"1:4", "4:1", "1:8", "8:1"
|
||||
]
|
||||
RESOLUTIONS = ["512", "1K", "2K", "4K"]
|
||||
|
||||
def __init__(self):
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
enabled_models = get_enabled_models()
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
||||
|
||||
all_aspect_ratios = get_all_supported_aspect_ratios() or cls.ASPECT_RATIOS
|
||||
all_resolutions = get_all_supported_resolutions() or cls.RESOLUTIONS
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "一个中国女子的OOTD",
|
||||
"multiline": True
|
||||
}),
|
||||
"模型": (enabled_models, {"default": enabled_models[0]}),
|
||||
"宽高比": (all_aspect_ratios, {"default": "1:1"}),
|
||||
"分辨率": (all_resolutions, {"default": "2K"}),
|
||||
"生图数量": ("INT", {"default": 1, "min": 1, "max": 1000, "step": 1}),
|
||||
"像素缩放": ("BOOLEAN", {"default": True, "label_on": "打开", "label_off": "关闭"}),
|
||||
"分辨率像素": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.1, "display": "number"}),
|
||||
"谷歌搜索（联网）": (["关闭", "打开"], {"default": "关闭"}),
|
||||
"图片搜索（联网）": (["关闭", "打开"], {"default": "关闭"}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
},
|
||||
"optional": {
|
||||
# 单个参考图端口，接受普通 IMAGE 或「加载图像（批量）」输出的列表
|
||||
"参考图": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("输出图像",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "image/generation"
|
||||
|
||||
# 声明 INPUT_IS_LIST，使 ComfyUI 将「加载图像（批量）」的 list[Tensor]
|
||||
# 整体传入而非逐张迭代执行，同时其余所有参数也会被包进 list，需解包。
|
||||
INPUT_IS_LIST = True
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 以下方法与 NanoBananaPro 完全相同，仅 generate() 开头增加了解包逻辑
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def resize_to_megapixels(self, image: Image.Image, target_megapixels: float) -> Image.Image:
|
||||
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 = max(1, int(image.width * scale))
|
||||
new_height = max(1, int(image.height * scale))
|
||||
return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
async def _generate_single_task(
|
||||
self,
|
||||
session: aiohttp.ClientSession,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: List[Image.Image],
|
||||
output_folder: str,
|
||||
global_task_index: int,
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
) -> dict:
|
||||
result = {
|
||||
"global_task_index": global_task_index,
|
||||
"prompt": prompt,
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"saved_files": [],
|
||||
"error": None
|
||||
}
|
||||
try:
|
||||
gen_result = await self.client.generate_single_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images if images else None,
|
||||
session=session,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
if gen_result:
|
||||
images_list, _ = gen_result
|
||||
for gen_img in images_list:
|
||||
output_path = generate_timestamp_filename(
|
||||
output_folder=output_folder,
|
||||
extension=".png"
|
||||
)
|
||||
save_image(gen_img, output_path)
|
||||
result["saved_files"].append(output_path)
|
||||
gen_img = None
|
||||
result["success"] = True
|
||||
result["generated_count"] = len(images_list)
|
||||
except Exception as e:
|
||||
result["error"] = str(e)
|
||||
return result
|
||||
|
||||
async def _process_batch_async(
|
||||
self,
|
||||
prompts: List[str],
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images_per_prompt: int,
|
||||
input_images: List[Image.Image],
|
||||
output_folder: str,
|
||||
pbar=None,
|
||||
enable_grounding: bool = False,
|
||||
enable_image_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)
|
||||
num_prompts = len(prompts)
|
||||
print(f"{_NODE}: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务")
|
||||
|
||||
max_concurrent = 10
|
||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
||||
all_results = []
|
||||
completed = 0
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
for batch_idx in range(num_batches):
|
||||
start_idx = batch_idx * max_concurrent
|
||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||||
tasks = []
|
||||
for i in range(start_idx, end_idx):
|
||||
_, _, prompt = tasks_def[i]
|
||||
task = asyncio.create_task(
|
||||
self._generate_single_task(
|
||||
session=session,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_images,
|
||||
output_folder=output_folder,
|
||||
global_task_index=i,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
batch_results = []
|
||||
for coro in asyncio.as_completed(tasks):
|
||||
result_data = None
|
||||
try:
|
||||
result = await coro
|
||||
if isinstance(result, Exception):
|
||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": [], "prompt": ""}
|
||||
else:
|
||||
result_data = result
|
||||
batch_results.append(result_data)
|
||||
except Exception as e:
|
||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": [], "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"{_NODE}: [{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"{_NODE}: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
all_results.extend(batch_results)
|
||||
gc.collect()
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
return all_results
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt,
|
||||
模型,
|
||||
宽高比,
|
||||
分辨率,
|
||||
生图数量,
|
||||
像素缩放,
|
||||
分辨率像素,
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
# ----------------------------------------------------------------
|
||||
# INPUT_IS_LIST=True 时，所有参数均为 list，先统一解包为标量
|
||||
# ----------------------------------------------------------------
|
||||
prompt = prompt[0] if isinstance(prompt, list) else prompt
|
||||
模型 = 模型[0] if isinstance(模型, list) else 模型
|
||||
宽高比 = 宽高比[0] if isinstance(宽高比, list) else 宽高比
|
||||
分辨率 = 分辨率[0] if isinstance(分辨率, list) else 分辨率
|
||||
生图数量 = 生图数量[0] if isinstance(生图数量, list) else 生图数量
|
||||
像素缩放 = 像素缩放[0] if isinstance(像素缩放, list) else 像素缩放
|
||||
分辨率像素 = 分辨率像素[0] if isinstance(分辨率像素, list) else 分辨率像素
|
||||
|
||||
# seed 也在 kwargs 里（含全角括号的参数名无法作为形参）
|
||||
seed_raw = kwargs.pop("seed", [0])
|
||||
seed: int = seed_raw[0] if isinstance(seed_raw, list) else seed_raw
|
||||
|
||||
# 搜索开关同理
|
||||
grounding_raw = kwargs.pop("谷歌搜索（联网）", ["关闭"])
|
||||
image_search_raw = kwargs.pop("图片搜索（联网）", ["关闭"])
|
||||
enable_grounding: bool = (grounding_raw[0] if isinstance(grounding_raw, list) else grounding_raw) == "打开"
|
||||
enable_image_search: bool = (image_search_raw[0] if isinstance(image_search_raw, list) else image_search_raw) == "打开"
|
||||
|
||||
# ----------------------------------------------------------------
|
||||
# 收集参考图：兼容两种来源
|
||||
# 1. 「加载图像（批量）」→ is_output_list=True → list[Tensor]
|
||||
# INPUT_IS_LIST 下传入的是 list[list[Tensor]] 或 list[Tensor]，需展平
|
||||
# 2. 普通 IMAGE 端口（单 tensor 或 batch tensor）→ list 中只有 1 个元素
|
||||
# ----------------------------------------------------------------
|
||||
ref_raw = kwargs.pop("参考图", None)
|
||||
input_images: List[Image.Image] = []
|
||||
|
||||
if ref_raw is not None:
|
||||
# INPUT_IS_LIST 下，可选端口若连接则为 list；元素可能是 Tensor 或 list[Tensor]
|
||||
items = ref_raw if isinstance(ref_raw, list) else [ref_raw]
|
||||
for item in items:
|
||||
if item is None:
|
||||
continue
|
||||
if isinstance(item, list):
|
||||
# 来自 is_output_list 的嵌套 list，继续展平
|
||||
for sub in item:
|
||||
if sub is not None and isinstance(sub, torch.Tensor):
|
||||
input_images.extend(tensor_to_pil(sub))
|
||||
elif isinstance(item, torch.Tensor):
|
||||
input_images.extend(tensor_to_pil(item))
|
||||
|
||||
# ----------------------------------------------------------------
|
||||
# 以下逻辑与 NanoBananaPro.generate() 完全一致
|
||||
# ----------------------------------------------------------------
|
||||
start_time = time.time()
|
||||
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(生图数量)
|
||||
|
||||
try:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed % (2 ** 32))
|
||||
|
||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
||||
process = psutil.Process()
|
||||
initial_memory = process.memory_info().rss / 1024 / 1024
|
||||
print(f"{_NODE}: 初始内存使用: {initial_memory:.1f} MB")
|
||||
|
||||
if self.client is None:
|
||||
try:
|
||||
self.client = GeminiAPIClient()
|
||||
except ValueError as e:
|
||||
raise ValueError(f"初始化失败: {str(e)}")
|
||||
|
||||
# 校验分辨率
|
||||
supported_resolutions = get_model_supported_resolutions(模型)
|
||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
||||
raise ValueError(
|
||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的分辨率：{', '.join(supported_resolutions)}"
|
||||
)
|
||||
|
||||
# 校验宽高比
|
||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
||||
if supported_ratios and 宽高比 not in supported_ratios:
|
||||
raise ValueError(
|
||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的宽高比：{', '.join(supported_ratios)}"
|
||||
)
|
||||
|
||||
# 校验图片搜索与模型兼容性
|
||||
IMAGE_SEARCH_UNSUPPORTED_MODELS = [
|
||||
"nano-banana-pro-限时特价", "nano-banana-pro-官方计费", "gemini-3-pro-image-preview"
|
||||
]
|
||||
if enable_image_search and 模型 in IMAGE_SEARCH_UNSUPPORTED_MODELS:
|
||||
raise ValueError(
|
||||
f"模型 \"{模型}\" 不支持【图片搜索（联网）】功能！"
|
||||
f"请切换到 nano-banana-2-限时特价 或 gemini-3.1-flash-image-preview 后再使用"
|
||||
)
|
||||
|
||||
# 验证输入图像数量上限
|
||||
if len(input_images) > 14:
|
||||
raise ValueError(
|
||||
f"输入图像数量 {len(input_images)} 超过限制 14 张，请减少输入图像数量"
|
||||
)
|
||||
|
||||
# 像素缩放
|
||||
if input_images and 像素缩放:
|
||||
input_images = [self.resize_to_megapixels(img, 分辨率像素) for img in input_images]
|
||||
|
||||
# 解析批量提示词
|
||||
batch_prompts = parse_batch_prompts(prompt)
|
||||
|
||||
# 打印概览
|
||||
grounding_str = ""
|
||||
if enable_image_search:
|
||||
grounding_str = " | 谷歌图片搜索接地"
|
||||
elif enable_grounding:
|
||||
grounding_str = " | 谷歌搜索接地"
|
||||
|
||||
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"{_NODE}: {mode_str} | {分辨率} {宽高比} | 共{total_images}张{grounding_str}")
|
||||
if total_images > 100:
|
||||
print(f"⚠️ {_NODE}: 警告！批量生成 {total_images} 张图片，内存占用可能较高")
|
||||
print(f"⚠️ 建议：分批执行或减少生图数量")
|
||||
else:
|
||||
mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式"
|
||||
print(f"{_NODE}: {mode_str} | {分辨率} {宽高比} | {生图数量}张{grounding_str}")
|
||||
if 生图数量 > 100:
|
||||
print(f"⚠️ {_NODE}: 警告！批量生成 {生图数量} 张图片，内存占用可能较高")
|
||||
print(f"⚠️ 建议：分批执行或减少生图数量")
|
||||
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
def progress_callback(current, total, success, error_msg=None):
|
||||
nonlocal success_count, fail_count
|
||||
if success:
|
||||
success_count += 1
|
||||
print(f"{_NODE}: 任务 {current}/{total} 成功 ✓")
|
||||
else:
|
||||
fail_count += 1
|
||||
if error_msg:
|
||||
print(f"{_NODE}: 任务 {current}/{total} 失败 ✗")
|
||||
print(f"原始错误详情:\n{error_msg}")
|
||||
else:
|
||||
print(f"{_NODE}: 任务 {current}/{total} 失败 ✗")
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0:
|
||||
gc.collect()
|
||||
current_memory = process.memory_info().rss / 1024 / 1024
|
||||
memory_increase = current_memory - initial_memory
|
||||
print(f"{_NODE}: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)")
|
||||
if current_memory > 2000:
|
||||
print(f"⚠️ {_NODE}: 内存使用过高！建议减少生图数量或分批执行")
|
||||
|
||||
def _get_output_folder():
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
folder = folder_paths.get_output_directory()
|
||||
return folder
|
||||
raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用")
|
||||
|
||||
def run_async_in_thread(coro_fn):
|
||||
def _run():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(coro_fn())
|
||||
finally:
|
||||
loop.close()
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(_run)
|
||||
try:
|
||||
return future.result(timeout=3600)
|
||||
except TimeoutError:
|
||||
raise RuntimeError("任务执行超时（1小时），请减少数量或检查网络连接")
|
||||
|
||||
# ── 批量提示词模式 ──────────────────────────────────────
|
||||
if batch_prompts:
|
||||
num_prompts = len(batch_prompts)
|
||||
total_images = num_prompts * 生图数量
|
||||
if pbar is not None:
|
||||
pbar = ProgressBar(total_images)
|
||||
|
||||
output_folder = _get_output_folder()
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
|
||||
results = run_async_in_thread(lambda: self._process_batch_async(
|
||||
prompts=batch_prompts,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images_per_prompt=生图数量,
|
||||
input_images=input_images,
|
||||
output_folder=output_folder,
|
||||
pbar=pbar,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
))
|
||||
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
all_saved_files = [f for r in results for f in r.get("saved_files", [])]
|
||||
|
||||
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_images} | 失败: {fail_count}")
|
||||
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
for fr in failed_results:
|
||||
idx = fr.get("global_task_index", -1) + 1
|
||||
snippet = (fr.get("prompt", "") or "")[:30]
|
||||
print(f" 失败 #{idx}: {snippet}{'...' if len(snippet) >= 30 else ''} → {fr.get('error', '未知错误')}")
|
||||
|
||||
output_images = []
|
||||
for fp in all_saved_files[-min(10, len(all_saved_files)):]:
|
||||
try:
|
||||
output_images.append(Image.open(fp))
|
||||
except Exception as e:
|
||||
print(f"{_NODE}: 无法加载 {fp} - {e}")
|
||||
|
||||
if not output_images:
|
||||
output_images = [Image.new('RGB', (512, 512), color=(128, 128, 128))]
|
||||
|
||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||
print(f"{_NODE}: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张")
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
|
||||
# ── 单提示词模式 ────────────────────────────────────────
|
||||
if 生图数量 == 1:
|
||||
generated_images = self.client.generate_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
batch_size=1,
|
||||
images=input_images,
|
||||
progress_callback=progress_callback,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
)
|
||||
output_folder = _get_output_folder()
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
for gen_img in generated_images:
|
||||
output_path = generate_timestamp_filename(output_folder=output_folder)
|
||||
save_image(gen_img, output_path)
|
||||
else:
|
||||
print(f"{_NODE}: 单提示词×{生图数量}张 → 异步并发模式")
|
||||
if pbar is not None:
|
||||
pbar = ProgressBar(生图数量)
|
||||
|
||||
output_folder = _get_output_folder()
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
|
||||
results = run_async_in_thread(lambda: self._process_batch_async(
|
||||
prompts=[prompt],
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images_per_prompt=生图数量,
|
||||
input_images=input_images,
|
||||
output_folder=output_folder,
|
||||
pbar=pbar,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
))
|
||||
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
all_saved_files = [f for r in results for f in r.get("saved_files", [])]
|
||||
|
||||
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}/{生图数量} | 失败: {fail_count}")
|
||||
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
for fr in failed_results:
|
||||
idx = fr.get("global_task_index", -1) + 1
|
||||
print(f" 失败 #{idx}: {prompt[:30]}{'...' if len(prompt) >= 30 else ''} → {fr.get('error', '未知错误')}")
|
||||
|
||||
output_images = []
|
||||
for fp in all_saved_files[-min(10, len(all_saved_files)):]:
|
||||
try:
|
||||
output_images.append(Image.open(fp))
|
||||
except Exception as e:
|
||||
print(f"{_NODE}: 无法加载 {fp} - {e}")
|
||||
|
||||
if not output_images:
|
||||
output_images = [Image.new('RGB', (512, 512), color=(128, 128, 128))]
|
||||
|
||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||
print(f"{_NODE}: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张")
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
|
||||
# 单张同步模式的输出路径（生图数量==1 走到这里）
|
||||
max_output_images = 20
|
||||
if len(generated_images) > max_output_images:
|
||||
print(f"{_NODE}: 生成 {len(generated_images)} 张图片，限制输出前 {max_output_images} 张到ComfyUI")
|
||||
output_images = generated_images[:max_output_images]
|
||||
else:
|
||||
output_images = generated_images
|
||||
|
||||
output_tensor = _images_to_tensor_safe(output_images, _NODE)
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||||
if fail_count > 0:
|
||||
print(f"[4/4] 完成！总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张")
|
||||
else:
|
||||
print(f"[4/4] 完成！总耗时 {time_str} | 成功 {len(generated_images)}张")
|
||||
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as e:
|
||||
if str(e) == "未授权！":
|
||||
print("请联系作者授权后方可使用！")
|
||||
raise ValueError("未授权！") from None
|
||||
error_msg = str(e)
|
||||
print(f"{_NODE}: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
error_full = str(e)
|
||||
print(f"{_NODE}: ❌ {error_full}")
|
||||
raise RuntimeError(error_full) from None
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"{_NODE}: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
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"{_NODE}: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
gc.collect()
|
||||
@@ -0,0 +1,826 @@
|
||||
"""
|
||||
全能生图 节点
|
||||
ComfyUI 自定义节点，用于调用 Gemini 模型生成图像
|
||||
"""
|
||||
|
||||
import time
|
||||
import math
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
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 ..utils.file_utils import ImageInfo, generate_timestamp_filename, save_image
|
||||
from ..clients.openai_client import OpenAIAPIClient
|
||||
from ..models_config import (
|
||||
get_enabled_models, get_model_description,
|
||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
||||
get_model_supported_resolutions, get_all_supported_resolutions
|
||||
)
|
||||
|
||||
# 检查 folder_paths 是否可用
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
# 导入 ComfyUI 原生进度条
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ 全能生图: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")
|
||||
|
||||
# 内存监控（可选）
|
||||
try:
|
||||
import psutil
|
||||
MEMORY_MONITOR_AVAILABLE = True
|
||||
except ImportError:
|
||||
MEMORY_MONITOR_AVAILABLE = False
|
||||
print("⚠️ 全能生图: psutil 不可用，内存监控功能禁用")
|
||||
|
||||
# ============================================================================
|
||||
# 调试日志配置
|
||||
# ============================================================================
|
||||
# 是否启用调试日志（打印完整的 API 响应内容）
|
||||
# 设置为 True 以启用调试日志，False 以禁用
|
||||
DEBUG_LOG_ENABLED = False
|
||||
# 是否启用请求体日志（打印发送给 API 的请求体，base64 图片数据将自动截断）
|
||||
# 设置为 True 以启用请求体日志，False 以禁用
|
||||
REQUEST_LOG_ENABLED = False
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class QuanNengShengTu:
|
||||
"""
|
||||
全能生图 节点
|
||||
|
||||
功能：
|
||||
- 文生图：基于提示词生成图像
|
||||
- 图生图：基于输入图像和提示词生成新图像
|
||||
- 批量生成：支持并发生成多张图像
|
||||
|
||||
注意：
|
||||
- 支持的模型列表从 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",
|
||||
"1:4", "4:1", "1:8", "8:1"
|
||||
]
|
||||
|
||||
# 支持的分辨率列表（全量兜底，实际由 get_all_supported_resolutions() 动态生成）
|
||||
RESOLUTIONS = ["512", "1K", "2K", "4K"]
|
||||
|
||||
def __init__(self):
|
||||
"""初始化节点"""
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
|
||||
ComfyUI 节点规范：
|
||||
- required: 必选参数
|
||||
- optional: 可选参数
|
||||
"""
|
||||
# 从配置文件动态获取启用的模型列表
|
||||
enabled_models = get_enabled_models()
|
||||
|
||||
# 过滤掉包含"限时特价"的模型
|
||||
enabled_models = [m for m in enabled_models if "限时特价" not in m]
|
||||
|
||||
# 如果没有启用的模型，使用空列表（会导致节点不可用，提示用户配置）
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个模型"]
|
||||
|
||||
# 动态获取所有启用模型支持的宽高比（去重合并）
|
||||
all_aspect_ratios = get_all_supported_aspect_ratios()
|
||||
if not all_aspect_ratios:
|
||||
all_aspect_ratios = cls.ASPECT_RATIOS
|
||||
|
||||
# 动态获取所有启用模型支持的分辨率（去重合并）
|
||||
all_resolutions = get_all_supported_resolutions()
|
||||
if not all_resolutions:
|
||||
all_resolutions = cls.RESOLUTIONS
|
||||
|
||||
# 创建9个独立的图像输入
|
||||
optional_inputs = {}
|
||||
for i in range(1, 10): # 1-9
|
||||
optional_inputs[f"参考图{i}"] = ("IMAGE",)
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"提示词": ("STRING", {
|
||||
"default": "一个中国女子的OOTD",
|
||||
"multiline": True
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0]
|
||||
}),
|
||||
"宽高比": (all_aspect_ratios, {
|
||||
"default": "1:1"
|
||||
}),
|
||||
"分辨率": (all_resolutions, {
|
||||
"default": "2K"
|
||||
}),
|
||||
"生图数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 1000,
|
||||
"step": 1
|
||||
}),
|
||||
"像素缩放": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"label_on": "打开",
|
||||
"label_off": "关闭"
|
||||
}),
|
||||
"分辨率像素": ("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 = ("输出图像",)
|
||||
|
||||
# 导入 ComfyUI 的文件夹路径管理
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
# 执行函数名
|
||||
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
|
||||
"""
|
||||
# 计算当前像素数
|
||||
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]"
|
||||
)
|
||||
|
||||
async def _generate_single_task(
|
||||
self,
|
||||
session: aiohttp.ClientSession,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: List[Image.Image],
|
||||
output_folder: str,
|
||||
global_task_index: int,
|
||||
) -> dict:
|
||||
"""执行单个生成任务，生成后立即保存到磁盘"""
|
||||
result = {
|
||||
"global_task_index": global_task_index,
|
||||
"prompt": prompt,
|
||||
"success": False,
|
||||
"generated_count": 0,
|
||||
"saved_files": [],
|
||||
"error": None
|
||||
}
|
||||
|
||||
try:
|
||||
gen_result = await self.client.generate_single_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images if images else None,
|
||||
session=session,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=False,
|
||||
enable_image_search=False
|
||||
)
|
||||
if gen_result:
|
||||
images_list, _ = gen_result
|
||||
for gen_img in images_list:
|
||||
output_path = generate_timestamp_filename(
|
||||
output_folder=output_folder,
|
||||
extension=".png"
|
||||
)
|
||||
save_image(gen_img, output_path)
|
||||
result["saved_files"].append(output_path)
|
||||
gen_img = None # 释放内存
|
||||
|
||||
result["success"] = True
|
||||
result["generated_count"] = len(images_list)
|
||||
except Exception as e:
|
||||
result["error"] = str(e)
|
||||
|
||||
return result
|
||||
|
||||
async def _process_batch_async(
|
||||
self,
|
||||
prompts: List[str],
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images_per_prompt: int,
|
||||
input_images: List[Image.Image],
|
||||
output_folder: str,
|
||||
pbar=None,
|
||||
) -> List[dict]:
|
||||
"""异步批量处理：每个提示词独立调用 API，生成后立即写磁盘"""
|
||||
# 构建任务列表：(prompt, sub_index) 用于 images_per_prompt > 1 的情况
|
||||
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)
|
||||
num_prompts = len(prompts)
|
||||
print(f"全能生图: 批量提示词模式 | {num_prompts}个提示词 × {images_per_prompt}张/提示词 | 共{total_tasks}任务")
|
||||
|
||||
max_concurrent = 10
|
||||
num_batches = math.ceil(total_tasks / max_concurrent)
|
||||
|
||||
all_results = []
|
||||
completed = 0
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
for batch_idx in range(num_batches):
|
||||
start_idx = batch_idx * max_concurrent
|
||||
end_idx = min(start_idx + max_concurrent, total_tasks)
|
||||
|
||||
tasks = []
|
||||
for i in range(start_idx, end_idx):
|
||||
_, _, prompt = tasks_def[i]
|
||||
task = asyncio.create_task(
|
||||
self._generate_single_task(
|
||||
session=session,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_images,
|
||||
output_folder=output_folder,
|
||||
global_task_index=i,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
batch_results = []
|
||||
for coro in asyncio.as_completed(tasks):
|
||||
result_data = None
|
||||
try:
|
||||
result = await coro
|
||||
if isinstance(result, Exception):
|
||||
result_data = {"success": False, "error": str(result), "generated_count": 0, "saved_files": [], "prompt": ""}
|
||||
else:
|
||||
result_data = result
|
||||
batch_results.append(result_data)
|
||||
except Exception as e:
|
||||
result_data = {"success": False, "error": str(e), "generated_count": 0, "saved_files": [], "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"全能生图: [{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"全能生图: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
all_results.extend(batch_results)
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
return all_results
|
||||
|
||||
def generate(
|
||||
self,
|
||||
提示词: str,
|
||||
模型: str,
|
||||
宽高比: str,
|
||||
分辨率: str,
|
||||
生图数量: int,
|
||||
像素缩放: bool,
|
||||
分辨率像素: float,
|
||||
seed: int,
|
||||
**kwargs
|
||||
) -> Tuple[torch.Tensor]:
|
||||
"""
|
||||
生成图像
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
模型: 模型名称
|
||||
宽高比: 宽高比
|
||||
分辨率: 分辨率
|
||||
生图数量: 批次大小
|
||||
像素缩放: 是否启用像素缩放
|
||||
分辨率像素: 目标像素数（百万像素）
|
||||
seed: 随机种子
|
||||
**kwargs: 动态参考图输入 (参考图1-9)
|
||||
|
||||
注意：
|
||||
调试日志功能已移至文件顶部配置，通过修改 DEBUG_LOG_ENABLED 常量控制
|
||||
|
||||
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))
|
||||
|
||||
# 内存监控初始化
|
||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
||||
import psutil
|
||||
process = psutil.Process()
|
||||
initial_memory = process.memory_info().rss / 1024 / 1024
|
||||
print(f"全能生图: 初始内存使用: {initial_memory:.1f} MB")
|
||||
|
||||
# 初始化 API 客户端
|
||||
if self.client is None:
|
||||
try:
|
||||
self.client = OpenAIAPIClient()
|
||||
except ValueError as e:
|
||||
raise ValueError(f"初始化失败: {str(e)}")
|
||||
|
||||
# 校验分辨率与模型的兼容性
|
||||
supported_resolutions = get_model_supported_resolutions(模型)
|
||||
if supported_resolutions and 分辨率 not in supported_resolutions:
|
||||
raise ValueError(
|
||||
f"分辨率 \"{分辨率}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的分辨率：{', '.join(supported_resolutions)}"
|
||||
)
|
||||
|
||||
# 校验宽高比与模型的兼容性
|
||||
supported_ratios = get_model_supported_aspect_ratios(模型)
|
||||
if supported_ratios and 宽高比 not in supported_ratios:
|
||||
raise ValueError(
|
||||
f"宽高比 \"{宽高比}\" 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的宽高比：{', '.join(supported_ratios)}"
|
||||
)
|
||||
|
||||
# 收集独立输入的参考图
|
||||
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
|
||||
|
||||
# 解析批量提示词
|
||||
batch_prompts = parse_batch_prompts(提示词)
|
||||
|
||||
# 打印首行概览
|
||||
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"全能生图: {mode_str} | {分辨率} {宽高比} | 共{total_images}张")
|
||||
|
||||
# 大批量警告
|
||||
if total_images > 100:
|
||||
print(f"⚠️ 全能生图: 警告！批量生成 {total_images} 张图片，内存占用可能较高")
|
||||
print(f"⚠️ 建议：分批执行或减少生图数量")
|
||||
else:
|
||||
# 单提示词模式
|
||||
mode_str = f"图生图模式 (输入{len(input_images)}张)" if input_images else "文生图模式"
|
||||
print(f"全能生图: {mode_str} | {分辨率} {宽高比} | {生图数量}张")
|
||||
|
||||
# 大批量警告
|
||||
if 生图数量 > 100:
|
||||
print(f"⚠️ 全能生图: 警告！批量生成 {生图数量} 张图片，内存占用可能较高")
|
||||
print(f"⚠️ 建议：分批执行或减少生图数量")
|
||||
|
||||
# 统计变量
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
# 进度回调 - 打印错误信息并更新进度条，添加内存监控
|
||||
def progress_callback(current, total, success, error_msg=None):
|
||||
nonlocal success_count, fail_count
|
||||
if success:
|
||||
success_count += 1
|
||||
print(f"全能生图: 任务 {current}/{total} 成功 ✓")
|
||||
else:
|
||||
fail_count += 1
|
||||
if error_msg:
|
||||
print(f"全能生图: 任务 {current}/{total} 失败 ✗")
|
||||
print(f"原始错误详情:\n{error_msg}")
|
||||
else:
|
||||
print(f"全能生图: 任务 {current}/{total} 失败 ✗")
|
||||
|
||||
# 更新 ComfyUI 原生进度条
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
# 内存监控（每完成10个任务检查一次）
|
||||
if MEMORY_MONITOR_AVAILABLE and total > 50 and current % 10 == 0:
|
||||
import gc
|
||||
gc.collect() # 强制垃圾回收
|
||||
current_memory = process.memory_info().rss / 1024 / 1024
|
||||
memory_increase = current_memory - initial_memory
|
||||
print(f"全能生图: 内存使用: {current_memory:.1f} MB (+{memory_increase:.1f} MB)")
|
||||
|
||||
# 内存警告阈值（2GB）
|
||||
if current_memory > 2000:
|
||||
print(f"⚠️ 全能生图: 内存使用过高！建议减少生图数量或分批执行")
|
||||
|
||||
# 根据是否有批量提示词选择生成模式
|
||||
if batch_prompts:
|
||||
num_prompts = len(batch_prompts)
|
||||
total_images = num_prompts * 生图数量
|
||||
|
||||
# ===== 批量提示词模式：异步并发+磁盘保存 =====
|
||||
if pbar is not None:
|
||||
pbar = ProgressBar(total_images)
|
||||
|
||||
# 确定保存路径
|
||||
output_folder = ""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
output_folder = folder_paths.get_output_directory()
|
||||
print(f"全能生图: 磁盘保存模式 → {output_folder}")
|
||||
else:
|
||||
raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用")
|
||||
|
||||
import os
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
|
||||
def run_async_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
self._process_batch_async(
|
||||
prompts=batch_prompts,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images_per_prompt=生图数量,
|
||||
input_images=input_images,
|
||||
output_folder=output_folder,
|
||||
pbar=pbar,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async_in_thread)
|
||||
try:
|
||||
results = future.result(timeout=3600)
|
||||
except TimeoutError:
|
||||
raise RuntimeError("任务执行超时（1小时），请减少提示词数量或检查网络连接")
|
||||
|
||||
# 统计结果
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||
all_saved_files = []
|
||||
for r in results:
|
||||
all_saved_files.extend(r.get("saved_files", []))
|
||||
|
||||
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_images} | 失败: {fail_count}")
|
||||
|
||||
# 失败详情
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
if failed_results:
|
||||
for fr in failed_results:
|
||||
idx = fr.get("global_task_index", -1) + 1
|
||||
prompt_snippet = (fr.get("prompt", "") or "")[:30]
|
||||
error_msg = fr.get("error", "未知错误")
|
||||
print(f" 失败 #{idx}: {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → {error_msg}")
|
||||
|
||||
# 从磁盘加载最后 10 张图片
|
||||
output_images = []
|
||||
max_output_images = 10
|
||||
recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):]
|
||||
for file_path in recent_files:
|
||||
try:
|
||||
img = Image.open(file_path)
|
||||
output_images.append(img)
|
||||
except Exception as e:
|
||||
print(f"全能生图: 无法加载 {file_path} - {e}")
|
||||
|
||||
if not output_images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
output_images = [placeholder]
|
||||
|
||||
output_tensor = pil_to_tensor(output_images)
|
||||
print(f"全能生图: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张")
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
else:
|
||||
# 单提示词模式
|
||||
if 生图数量 == 1:
|
||||
# 单张：同步生成 + 保存到磁盘 + 输出 tensor
|
||||
generated_images = self.client.generate_sync(
|
||||
prompt=提示词,
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
batch_size=1,
|
||||
images=input_images,
|
||||
progress_callback=progress_callback,
|
||||
debug=DEBUG_LOG_ENABLED,
|
||||
debug_request=REQUEST_LOG_ENABLED,
|
||||
enable_grounding=False,
|
||||
enable_image_search=False
|
||||
)
|
||||
# 单张：保存到磁盘
|
||||
import os
|
||||
output_folder = ""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
output_folder = folder_paths.get_output_directory()
|
||||
print(f"全能生图: 磁盘保存模式 → {output_folder}")
|
||||
else:
|
||||
raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用")
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
for gen_img in generated_images:
|
||||
output_path = generate_timestamp_filename(output_folder=output_folder)
|
||||
save_image(gen_img, output_path)
|
||||
else:
|
||||
# 多张：异步并发 + 磁盘保存（与批量提示词逻辑一致）
|
||||
print(f"全能生图: 单提示词×{生图数量}张 → 异步并发模式")
|
||||
|
||||
if pbar is not None:
|
||||
pbar = ProgressBar(生图数量)
|
||||
|
||||
output_folder = ""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
output_folder = folder_paths.get_output_directory()
|
||||
print(f"全能生图: 磁盘保存模式 → {output_folder}")
|
||||
else:
|
||||
raise ValueError("无法获取 ComfyUI output 目录，请检查 folder_paths 是否可用")
|
||||
|
||||
import os
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
|
||||
def run_async_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
self._process_batch_async(
|
||||
prompts=[提示词],
|
||||
model=模型,
|
||||
resolution=分辨率,
|
||||
aspect_ratio=宽高比,
|
||||
images_per_prompt=生图数量,
|
||||
input_images=input_images,
|
||||
output_folder=output_folder,
|
||||
pbar=pbar,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_async_in_thread)
|
||||
try:
|
||||
results = future.result(timeout=3600)
|
||||
except TimeoutError:
|
||||
raise RuntimeError("任务执行超时（1小时），请减少生图数量或检查网络连接")
|
||||
|
||||
success_count = sum(1 for r in results if r.get("success", False))
|
||||
fail_count = len(results) - success_count
|
||||
total_generated = sum(r.get("generated_count", 0) for r in results)
|
||||
all_saved_files = []
|
||||
for r in results:
|
||||
all_saved_files.extend(r.get("saved_files", []))
|
||||
|
||||
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}/{生图数量} | 失败: {fail_count}")
|
||||
|
||||
# 失败详情
|
||||
failed_results = [r for r in results if not r.get("success", False)]
|
||||
if failed_results:
|
||||
for fr in failed_results:
|
||||
idx = fr.get("global_task_index", -1) + 1
|
||||
error_msg = fr.get("error", "未知错误")
|
||||
print(f" 失败 #{idx}: {提示词[:30]}{'...' if len(提示词) >= 30 else ''} → {error_msg}")
|
||||
|
||||
# 从磁盘加载最后 10 张图片
|
||||
output_images = []
|
||||
max_output_images = 10
|
||||
recent_files = all_saved_files[-min(max_output_images, len(all_saved_files)):]
|
||||
for file_path in recent_files:
|
||||
try:
|
||||
img = Image.open(file_path)
|
||||
output_images.append(img)
|
||||
except Exception as e:
|
||||
print(f"全能生图: 无法加载 {file_path} - {e}")
|
||||
|
||||
if not output_images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
output_images = [placeholder]
|
||||
|
||||
output_tensor = pil_to_tensor(output_images)
|
||||
print(f"全能生图: 共保存 {len(all_saved_files)} 张图片到磁盘，节点输出最后 {len(output_images)} 张")
|
||||
|
||||
import gc
|
||||
gc.collect()
|
||||
return (output_tensor,)
|
||||
|
||||
|
||||
# 优化：限制输出图片数量，避免内存爆炸
|
||||
max_output_images = 20 # 最多输出20张图片到ComfyUI
|
||||
|
||||
if len(generated_images) > max_output_images:
|
||||
print(f"全能生图: 生成 {len(generated_images)} 张图片，限制输出前 {max_output_images} 张到ComfyUI")
|
||||
output_images = generated_images[:max_output_images]
|
||||
else:
|
||||
output_images = generated_images
|
||||
|
||||
# 转换输出图像
|
||||
output_tensor = pil_to_tensor(output_images)
|
||||
|
||||
# 计算耗时并打印最终统计
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed < 1:
|
||||
time_str = f"{elapsed:.3f}s"
|
||||
else:
|
||||
time_str = f"{elapsed:.2f}s"
|
||||
|
||||
# 打印最终汇总
|
||||
if fail_count > 0:
|
||||
print(f"[4/4] 完成！总耗时 {time_str} | 成功 {success_count}张 | 失败 {fail_count}张")
|
||||
else:
|
||||
print(f"[4/4] 完成！总耗时 {time_str} | 成功 {len(generated_images)}张")
|
||||
|
||||
# 最终内存清理
|
||||
import gc
|
||||
gc.collect()
|
||||
if MEMORY_MONITOR_AVAILABLE and 生图数量 > 50:
|
||||
final_memory = process.memory_info().rss / 1024 / 1024
|
||||
print(f"全能生图: 最终内存使用: {final_memory:.1f} MB")
|
||||
|
||||
return (output_tensor,)
|
||||
|
||||
except ValueError as e:
|
||||
# 检测是否为授权错误
|
||||
if str(e) == "未授权！":
|
||||
print("请联系作者授权后方可使用！")
|
||||
raise ValueError("未授权！") from None
|
||||
else:
|
||||
error_msg = str(e)
|
||||
print(f"全能生图: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
error_full = str(e)
|
||||
print(f"全能生图: ❌ {error_full}")
|
||||
raise RuntimeError(error_full) from None
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"全能生图: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
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"全能生图: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 最终内存清理
|
||||
import gc
|
||||
gc.collect()
|
||||
print(f"全能生图: 最终内存清理完成")
|
||||
@@ -0,0 +1,371 @@
|
||||
"""
|
||||
图像元数据去除节点
|
||||
替代 ComfyUI 原生"保存图像"节点，保存时不写入提示词、工作流等 AI 元数据
|
||||
|
||||
提供两种节点：
|
||||
1. SaveCleanImage - 接收 IMAGE 张量，去除元数据后直接保存到 output 目录
|
||||
2. BatchCleanMetadata - 指定文件夹路径，批量去除已有图片中的元数据
|
||||
"""
|
||||
|
||||
import os
|
||||
from datetime import datetime
|
||||
import random
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..utils.file_utils import _get_port_suffix
|
||||
|
||||
# 尝试导入 ComfyUI 的 folder_paths
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
# 支持的图片格式
|
||||
SUPPORTED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.webp', '.bmp', '.tiff', '.tif'}
|
||||
|
||||
|
||||
def _get_output_dir() -> str:
|
||||
"""
|
||||
获取 ComfyUI output 目录
|
||||
|
||||
Returns:
|
||||
output 目录的绝对路径
|
||||
"""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
return folder_paths.get_output_directory()
|
||||
# fallback: 相对于插件目录推断
|
||||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
return os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output")
|
||||
|
||||
|
||||
def _get_next_counter(directory: str, prefix: str) -> int:
|
||||
"""
|
||||
扫描目录，获取下一个可用的文件计数器
|
||||
|
||||
Args:
|
||||
directory: 目标目录
|
||||
prefix: 文件名前缀
|
||||
|
||||
Returns:
|
||||
下一个计数器值
|
||||
"""
|
||||
if not os.path.exists(directory):
|
||||
return 1
|
||||
|
||||
if prefix:
|
||||
pattern = re.compile(rf'^{re.escape(prefix)}_(\d+)')
|
||||
else:
|
||||
pattern = re.compile(rf'^(\d+)\.')
|
||||
max_counter = 0
|
||||
|
||||
for f in os.listdir(directory):
|
||||
m = pattern.match(f)
|
||||
if m:
|
||||
counter = int(m.group(1))
|
||||
max_counter = max(max_counter, counter)
|
||||
|
||||
return max_counter + 1
|
||||
|
||||
|
||||
def _save_image_clean(image: Image.Image, path: str, fmt: str = None, quality: int = 95) -> None:
|
||||
"""
|
||||
保存图像，不包含任何元数据
|
||||
|
||||
通过提取纯像素数据并重建全新的 Image 对象，确保没有任何元数据残留。
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
path: 保存路径
|
||||
fmt: 图像格式（PNG/JPEG/WEBP），为 None 时根据扩展名推断
|
||||
quality: JPEG/WEBP 质量（1-100）
|
||||
"""
|
||||
# 确保 RGB 模式
|
||||
if image.mode != 'RGB':
|
||||
image = image.convert('RGB')
|
||||
|
||||
# 提取纯像素数据，重建全新的 Image 对象
|
||||
# 使用 tobytes() + frombytes() 确保只保留像素数据，彻底断开与原图像的关联
|
||||
pixel_data = image.tobytes()
|
||||
clean = Image.frombytes('RGB', image.size, pixel_data)
|
||||
|
||||
# 显式清空 info 字典，确保不会有任何残留元数据
|
||||
clean.info = {}
|
||||
|
||||
# 推断格式
|
||||
if fmt is None:
|
||||
ext = os.path.splitext(path)[1].lower()
|
||||
format_map = {
|
||||
'.png': 'PNG',
|
||||
'.jpg': 'JPEG',
|
||||
'.jpeg': 'JPEG',
|
||||
'.webp': 'WEBP',
|
||||
'.bmp': 'BMP',
|
||||
'.tiff': 'TIFF',
|
||||
'.tif': 'TIFF',
|
||||
}
|
||||
fmt = format_map.get(ext, 'PNG')
|
||||
|
||||
# 构建保存参数（确保不写入任何元数据）
|
||||
save_kwargs = {}
|
||||
if fmt == 'PNG':
|
||||
save_kwargs['pnginfo'] = PngInfo() # 空的 PngInfo，不包含任何文本块
|
||||
elif fmt == 'JPEG':
|
||||
save_kwargs['quality'] = quality
|
||||
# 不传 exif 参数，自然不会写入 EXIF 数据
|
||||
elif fmt == 'WEBP':
|
||||
save_kwargs['quality'] = quality
|
||||
save_kwargs['exif'] = b"" # 显式清空 EXIF
|
||||
|
||||
clean.save(path, format=fmt, **save_kwargs)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 节点 1：保存干净图像
|
||||
# ============================================================================
|
||||
|
||||
class SaveCleanImage:
|
||||
"""
|
||||
保存干净图像节点（不含元数据）
|
||||
|
||||
功能：
|
||||
- 接收 IMAGE 张量（支持单图和批次）
|
||||
- 去除所有元数据后保存到 ComfyUI/output 目录
|
||||
- 文件名自动添加 nometa 标识，方便辨认
|
||||
- 支持 PNG/JPEG/WEBP 格式
|
||||
- 作为终端节点，替代 ComfyUI 原生"保存图像"节点
|
||||
|
||||
使用场景：
|
||||
- 生图完成后，直接保存不含 AI 元数据的干净图像
|
||||
- 分享图像时不暴露提示词和工作流
|
||||
"""
|
||||
|
||||
SAVE_FORMATS = ["PNG", "JPEG", "WEBP"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
|
||||
Returns:
|
||||
输入参数配置字典
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"图像": ("IMAGE",),
|
||||
"文件名前缀": ("STRING", {"default": "ComfyUI_nometa"}),
|
||||
"保存格式": (cls.SAVE_FORMATS, {"default": "PNG"}),
|
||||
},
|
||||
"optional": {
|
||||
"JPEG/WEBP质量": ("INT", {
|
||||
"default": 95,
|
||||
"min": 1,
|
||||
"max": 100,
|
||||
"step": 1
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "save_clean"
|
||||
CATEGORY = "image"
|
||||
|
||||
DESCRIPTION = (
|
||||
"保存干净图像（不含元数据）。\n"
|
||||
"替代 ComfyUI 原生'保存图像'节点，保存时不写入提示词、工作流等 AI 元数据。\n"
|
||||
"文件保存到 ComfyUI/output 目录。"
|
||||
)
|
||||
|
||||
def save_clean(
|
||||
self,
|
||||
图像: torch.Tensor,
|
||||
文件名前缀: str = "ComfyUI_nometa",
|
||||
保存格式: str = "PNG",
|
||||
**kwargs
|
||||
) -> dict:
|
||||
"""
|
||||
去除元数据并保存图像
|
||||
|
||||
Args:
|
||||
图像: ComfyUI 图像张量 [B, H, W, C]
|
||||
文件名前缀: 保存文件名前缀
|
||||
保存格式: 图像格式（PNG/JPEG/WEBP）
|
||||
**kwargs: 可选参数（JPEG/WEBP质量）
|
||||
|
||||
Returns:
|
||||
UI 结果字典，包含保存的图像信息用于前端预览
|
||||
"""
|
||||
quality = kwargs.get("JPEG/WEBP质量", 95)
|
||||
|
||||
output_dir = _get_output_dir()
|
||||
port_suffix = _get_port_suffix()
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# 格式与扩展名映射
|
||||
ext_map = {"PNG": ".png", "JPEG": ".jpg", "WEBP": ".webp"}
|
||||
ext = ext_map.get(保存格式, ".png")
|
||||
|
||||
# 转换为 PIL 图像
|
||||
pil_images = tensor_to_pil(图像)
|
||||
|
||||
results = []
|
||||
saved_paths = []
|
||||
for img in pil_images:
|
||||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
ms = random.randint(0, 999)
|
||||
|
||||
while True:
|
||||
if 文件名前缀:
|
||||
filename = f"{文件名前缀}_{ts}_{ms:03d}{port_suffix}{ext}"
|
||||
else:
|
||||
filename = f"{ts}_{ms:03d}{port_suffix}{ext}"
|
||||
filepath = os.path.join(output_dir, filename)
|
||||
if not os.path.exists(filepath):
|
||||
break
|
||||
ms = (ms + 1) % 1000
|
||||
|
||||
_save_image_clean(img, filepath, fmt=保存格式, quality=quality)
|
||||
|
||||
results.append({
|
||||
"filename": filename,
|
||||
"subfolder": "",
|
||||
"type": "output"
|
||||
})
|
||||
saved_paths.append(filepath)
|
||||
|
||||
# 打印详细日志，方便用户定位保存的文件
|
||||
print(f"保存干净图像: 已保存 {len(pil_images)} 张无元数据图像 (格式: {保存格式})")
|
||||
for p in saved_paths:
|
||||
print(f" → {p}")
|
||||
|
||||
return {"ui": {"images": results}}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 节点 2：批量去除元数据
|
||||
# ============================================================================
|
||||
|
||||
class BatchCleanMetadata:
|
||||
"""
|
||||
批量去除文件夹中图片元数据的节点
|
||||
|
||||
功能：
|
||||
- 指定文件夹路径，批量处理其中所有图片
|
||||
- 去除 EXIF、PNG tEXt 块、ComfyUI 工作流等所有元数据
|
||||
- 支持保存到原目录（添加 _nometa 后缀）或覆盖原文件
|
||||
- 支持 PNG/JPG/JPEG/WEBP/BMP/TIFF 格式
|
||||
|
||||
使用场景：
|
||||
- 已经保存了一批含有 AI 元数据的图片，需要批量清理
|
||||
- 批量处理指定文件夹中的所有图片
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
定义输入参数
|
||||
|
||||
Returns:
|
||||
输入参数配置字典
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"文件夹路径": ("STRING", {"default": ""}),
|
||||
"覆盖原文件": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("处理结果",)
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "batch_clean"
|
||||
CATEGORY = "image"
|
||||
|
||||
DESCRIPTION = (
|
||||
"批量去除文件夹中图片的元数据。\n"
|
||||
"支持 PNG/JPG/JPEG/WEBP/BMP/TIFF 格式。\n"
|
||||
"默认在原文件名后添加 _nometa 后缀保存，也可选择覆盖原文件。"
|
||||
)
|
||||
|
||||
def batch_clean(
|
||||
self,
|
||||
文件夹路径: str,
|
||||
覆盖原文件: bool = False,
|
||||
) -> tuple:
|
||||
"""
|
||||
批量去除文件夹中图片的元数据
|
||||
|
||||
Args:
|
||||
文件夹路径: 待处理图片所在的文件夹路径
|
||||
覆盖原文件: 是否覆盖原文件（False 则添加 _nometa 后缀）
|
||||
|
||||
Returns:
|
||||
处理结果字符串
|
||||
|
||||
Raises:
|
||||
ValueError: 文件夹路径无效
|
||||
"""
|
||||
if not 文件夹路径 or not 文件夹路径.strip():
|
||||
raise ValueError("请输入文件夹路径")
|
||||
|
||||
folder = 文件夹路径.strip()
|
||||
|
||||
if not os.path.isdir(folder):
|
||||
raise ValueError(f"文件夹路径无效或不存在: {folder}")
|
||||
|
||||
# 扫描支持的图片文件
|
||||
files = []
|
||||
for f in sorted(os.listdir(folder)):
|
||||
ext = os.path.splitext(f)[1].lower()
|
||||
if ext in SUPPORTED_EXTENSIONS:
|
||||
files.append(f)
|
||||
|
||||
if not files:
|
||||
msg = f"文件夹中未找到支持的图片文件 ({', '.join(SUPPORTED_EXTENSIONS)})"
|
||||
print(f"批量去除元数据: {msg}")
|
||||
return (msg,)
|
||||
|
||||
print(f"批量去除元数据: 找到 {len(files)} 张图片，开始处理...")
|
||||
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
for f in files:
|
||||
try:
|
||||
src_path = os.path.join(folder, f)
|
||||
img = Image.open(src_path)
|
||||
|
||||
if 覆盖原文件:
|
||||
dst_path = src_path
|
||||
else:
|
||||
name, ext = os.path.splitext(f)
|
||||
dst_path = os.path.join(folder, f"{name}_nometa{ext}")
|
||||
|
||||
_save_image_clean(img, dst_path)
|
||||
success_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f"批量去除元数据: 处理 {f} 失败 - {str(e)}")
|
||||
fail_count += 1
|
||||
|
||||
# 构建结果消息
|
||||
if fail_count > 0:
|
||||
msg = f"处理完成: 成功 {success_count} 张, 失败 {fail_count} 张"
|
||||
else:
|
||||
msg = f"处理完成: 全部 {success_count} 张成功"
|
||||
|
||||
if not 覆盖原文件:
|
||||
msg += " (已添加 _nometa 后缀)"
|
||||
else:
|
||||
msg += " (已覆盖原文件)"
|
||||
|
||||
print(f"批量去除元数据: {msg}")
|
||||
|
||||
return (msg,)
|
||||
@@ -0,0 +1,526 @@
|
||||
"""
|
||||
Sora 视频生成节点
|
||||
ComfyUI 自定义节点，调用 Sora API 生成视频
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from math import gcd
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..clients.sora_client import SoraClient
|
||||
from ..models_config import (
|
||||
get_enabled_sora_models,
|
||||
get_all_sora_seconds,
|
||||
get_all_sora_sizes,
|
||||
get_sora_supported_seconds,
|
||||
get_sora_supported_sizes,
|
||||
get_sora_seconds_with_labels,
|
||||
get_sora_sizes_with_labels,
|
||||
SORA_MODELS,
|
||||
)
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ SoraVideo: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")
|
||||
|
||||
|
||||
def _size_to_display(size: str) -> str:
|
||||
"""
|
||||
将 'WxH' 格式的分辨率转换为友好显示名。
|
||||
|
||||
例如：
|
||||
"720x1280" → "720P 9:16"
|
||||
"1280x720" → "720P 16:9"
|
||||
"1024x1792" → "1K 4:7"
|
||||
"1792x1024" → "1K 7:4"
|
||||
|
||||
Args:
|
||||
size: 分辨率字符串，格式 "WxH"
|
||||
|
||||
Returns:
|
||||
友好显示名字符串
|
||||
"""
|
||||
parts = size.lower().split("x")
|
||||
w, h = int(parts[0]), int(parts[1])
|
||||
short_side = min(w, h)
|
||||
if short_side >= 3840:
|
||||
res = "4K"
|
||||
elif short_side >= 1920:
|
||||
res = "2K"
|
||||
elif short_side >= 1080:
|
||||
res = "1K"
|
||||
elif short_side >= 720:
|
||||
res = "720P"
|
||||
elif short_side >= 480:
|
||||
res = "480P"
|
||||
else:
|
||||
res = f"{short_side}P"
|
||||
g = gcd(w, h)
|
||||
ratio = f"{w // g}:{h // g}"
|
||||
return f"{res} {ratio} ({size})"
|
||||
|
||||
|
||||
def _build_size_display_map(sizes: list) -> dict:
|
||||
"""
|
||||
构建 显示名 → 实际值 映射字典。
|
||||
|
||||
Args:
|
||||
sizes: 实际分辨率列表，如 ["720x1280", "1280x720"]
|
||||
|
||||
Returns:
|
||||
字典，key 为显示名，value 为实际分辨率字符串
|
||||
"""
|
||||
mapping = {}
|
||||
for size in sizes:
|
||||
display = _size_to_display(size)
|
||||
if display in mapping:
|
||||
# 极少数情况下防止重名
|
||||
display = f"{display} ({size})"
|
||||
mapping[display] = size
|
||||
return mapping
|
||||
|
||||
|
||||
def _get_video_output_dir() -> str:
|
||||
"""获取视频输出目录: ComfyUI/output/video"""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
base = folder_paths.get_output_directory()
|
||||
else:
|
||||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
base = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output")
|
||||
video_dir = os.path.join(base, "video")
|
||||
os.makedirs(video_dir, exist_ok=True)
|
||||
return video_dir
|
||||
|
||||
|
||||
def _get_next_counter(directory: str, prefix: str) -> int:
|
||||
"""扫描目录，获取下一个可用的文件计数器"""
|
||||
if not os.path.exists(directory):
|
||||
return 1
|
||||
pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)")
|
||||
max_counter = 0
|
||||
for f in os.listdir(directory):
|
||||
m = pattern.match(f)
|
||||
if m:
|
||||
max_counter = max(max_counter, int(m.group(1)))
|
||||
return max_counter + 1
|
||||
|
||||
|
||||
def _fit_image_to_target(image, target_size: str):
|
||||
"""
|
||||
将参考图片按 "等比缩放覆盖 + 居中裁剪" 策略适配到目标分辨率。
|
||||
|
||||
策略 (Cover Crop)：
|
||||
1. 比较图片宽高比和目标宽高比
|
||||
2. 等比缩放，使图片最短边刚好覆盖目标对应边（图片完全覆盖目标区域）
|
||||
3. 居中裁剪多余部分，得到精确目标尺寸
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
target_size: 目标分辨率字符串，格式 "WxH"（如 "720x1280"）
|
||||
|
||||
Returns:
|
||||
适配后的 PIL Image 对象
|
||||
"""
|
||||
from PIL import Image as PILImage
|
||||
|
||||
# 解析目标尺寸
|
||||
parts = target_size.lower().split("x")
|
||||
target_w, target_h = int(parts[0]), int(parts[1])
|
||||
|
||||
src_w, src_h = image.size
|
||||
src_ratio = src_w / src_h
|
||||
target_ratio = target_w / target_h
|
||||
|
||||
# 宽高比一致且尺寸不超过目标，无需处理
|
||||
if abs(src_ratio - target_ratio) < 0.01 and src_w <= target_w and src_h <= target_h:
|
||||
return image
|
||||
|
||||
print(f"Sora: 参考图片 {src_w}x{src_h} (比例 {src_ratio:.2f}) → 目标 {target_w}x{target_h} (比例 {target_ratio:.2f})")
|
||||
|
||||
# 获取高质量重采样滤波器
|
||||
resample = PILImage.Resampling.LANCZOS if hasattr(PILImage, "Resampling") else PILImage.LANCZOS
|
||||
|
||||
# Cover Crop: 缩放使图片完全覆盖目标区域，然后居中裁剪
|
||||
if src_ratio > target_ratio:
|
||||
# 图片更宽：以高度为基准缩放，裁左右
|
||||
scale = target_h / src_h
|
||||
new_w = round(src_w * scale)
|
||||
new_h = target_h
|
||||
image = image.resize((new_w, new_h), resample=resample)
|
||||
# 居中裁剪宽度
|
||||
left = (new_w - target_w) // 2
|
||||
image = image.crop((left, 0, left + target_w, target_h))
|
||||
else:
|
||||
# 图片更高（或一样）：以宽度为基准缩放，裁上下
|
||||
scale = target_w / src_w
|
||||
new_w = target_w
|
||||
new_h = round(src_h * scale)
|
||||
image = image.resize((new_w, new_h), resample=resample)
|
||||
# 居中裁剪高度
|
||||
top = (new_h - target_h) // 2
|
||||
image = image.crop((0, top, target_w, top + target_h))
|
||||
|
||||
print(f"Sora: 参考图片已适配为 {image.size[0]}x{image.size[1]}")
|
||||
return image
|
||||
|
||||
|
||||
def _compress_image_for_upload(
|
||||
image,
|
||||
target_size: Optional[str] = None,
|
||||
) -> bytes:
|
||||
"""
|
||||
将 PIL Image 适配目标分辨率并编码为 PNG 字节，用于上传。
|
||||
|
||||
============================================================
|
||||
⚠️ 已验证可用的标准做法，请勿随意修改以下编码逻辑！
|
||||
============================================================
|
||||
经过多轮调试（2026-02-28），以下参数组合为唯一验证成功的方案：
|
||||
|
||||
1. 图片格式：PNG（format="PNG"）
|
||||
- 不可改为 JPEG —— API 会校验 Content-Type，抓包确认服务端使用 image/png
|
||||
- 不可使用 base64 字符串 —— 会报 "expected a file, got a string"
|
||||
- 不可使用 data URI —— 服务端不识别，返回 500
|
||||
|
||||
2. 图片尺寸：必须与视频分辨率完全一致（target_size）
|
||||
- 不可缩放降采样 —— 会报 "Inpaint image must match the requested width and height"
|
||||
- 尺寸由 _fit_image_to_target() 保证（等比缩放 + 居中裁剪）
|
||||
|
||||
3. 上传方式：由调用方（sora_client.py）以 multipart/form-data 文件字段上传
|
||||
- filename="reference.png", content_type="image/png"
|
||||
- 不可改回 application/json —— 服务端校验 input_reference 必须为 file 类型
|
||||
============================================================
|
||||
|
||||
Args:
|
||||
image: PIL Image 对象
|
||||
target_size: 目标分辨率字符串 "WxH"（如 "720x1280"）
|
||||
|
||||
Returns:
|
||||
PNG 格式的二进制字节
|
||||
"""
|
||||
from io import BytesIO
|
||||
|
||||
# 统一转换为 RGB（去除透明通道及其他模式）
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
|
||||
# 适配到目标分辨率（等比缩放 + 居中裁剪）
|
||||
# ⚠️ 必须保持此尺寸不变，API 强制要求参考图片与视频分辨率完全一致
|
||||
if target_size:
|
||||
image = _fit_image_to_target(image, target_size)
|
||||
|
||||
# ⚠️ 必须使用 PNG 格式，不可改为 JPEG 或其他格式
|
||||
buffered = BytesIO()
|
||||
image.save(buffered, format="PNG")
|
||||
size_kb = buffered.tell() / 1024
|
||||
print(f"Sora: 参考图片编码为 PNG，{size_kb:.0f} KB ({image.size[0]}x{image.size[1]})")
|
||||
return buffered.getvalue()
|
||||
|
||||
|
||||
class SoraVideo:
|
||||
"""
|
||||
Sora 视频生成节点
|
||||
|
||||
功能：
|
||||
- 文生视频：基于提示词生成视频
|
||||
- 图生视频：基于参考图片和提示词生成视频
|
||||
- 异步轮询：自动等待生成完成并下载
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
from ..models_config import SECONDS_DISPLAY_MAP, RESOLUTION_DISPLAY_MAP
|
||||
|
||||
enabled_models = get_enabled_sora_models()
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用至少一个 Sora 模型"]
|
||||
|
||||
# 构建秒数选项列表（按数字顺序排序）
|
||||
# 格式: ["4", "8", "10", "12", "15", "25(pro)"]
|
||||
all_seconds_display = []
|
||||
seen_seconds = set()
|
||||
for model_id in enabled_models:
|
||||
supported = get_sora_supported_seconds(model_id)
|
||||
for s in supported:
|
||||
if s not in seen_seconds:
|
||||
seen_seconds.add(s)
|
||||
display = SECONDS_DISPLAY_MAP.get(s, str(s))
|
||||
all_seconds_display.append((s, display))
|
||||
# 按秒数数值排序
|
||||
all_seconds_display = sorted(all_seconds_display, key=lambda x: x[0])
|
||||
seconds_options = [d for _, d in all_seconds_display] if all_seconds_display else ["4", "8", "12"]
|
||||
|
||||
# 构建分辨率选项列表（去重）
|
||||
# 格式: ["720P", "1080P"]
|
||||
seen_resolutions = set()
|
||||
for model_id in enabled_models:
|
||||
supported = get_sora_supported_sizes(model_id)
|
||||
for size in supported:
|
||||
if size in RESOLUTION_DISPLAY_MAP:
|
||||
res_name, _ = RESOLUTION_DISPLAY_MAP[size]
|
||||
seen_resolutions.add(res_name)
|
||||
resolution_options = sorted(list(seen_resolutions)) if seen_resolutions else ["720P"]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "A calico cat playing a piano on stage",
|
||||
"multiline": True,
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0],
|
||||
}),
|
||||
"分辨率": (resolution_options, {
|
||||
"default": resolution_options[0] if resolution_options else "720P",
|
||||
}),
|
||||
"宽高比": (["竖屏", "横屏"], {
|
||||
"default": "竖屏",
|
||||
}),
|
||||
"视频时长": (seconds_options, {
|
||||
"default": seconds_options[0] if seconds_options else "4",
|
||||
}),
|
||||
"生成数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"参考图片": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("预览视频",)
|
||||
FUNCTION = "generate_video"
|
||||
CATEGORY = "video/generation"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Sora 视频生成节点。\n"
|
||||
"支持文生视频和图生视频，自动轮询任务状态并下载视频。\n"
|
||||
"视频保存到 ComfyUI/output/video/ 目录。\n\n"
|
||||
"【模型说明】\n"
|
||||
"• sora-2：官方模型，支持 4/8/12秒、720P 分辨率\n"
|
||||
"• sora-2-pro：增强模型，支持全时长（含25秒）、1080P 分辨率\n\n"
|
||||
"【时长说明】\n"
|
||||
"• 25(pro)：仅 sora-2-pro 支持的25秒时长\n\n"
|
||||
"【分辨率说明】\n"
|
||||
"• 720P：sora-2 和 sora-2-pro 均支持\n"
|
||||
"• 1080P：仅 sora-2-pro 支持的高清分辨率"
|
||||
)
|
||||
|
||||
def generate_video(
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str,
|
||||
**kwargs,
|
||||
) -> Tuple[str]:
|
||||
from ..models_config import SECONDS_DISPLAY_MAP, RESOLUTION_DISPLAY_MAP
|
||||
|
||||
视频时长_display = kwargs.pop("视频时长", "4")
|
||||
分辨率_display = kwargs.pop("分辨率", "720P")
|
||||
宽高比 = kwargs.pop("宽高比", "竖屏")
|
||||
生成数量 = kwargs.pop("生成数量", 1)
|
||||
seed = kwargs.pop("seed", 0)
|
||||
start_time = time.time()
|
||||
|
||||
# 解析秒数显示值（如 "25(pro)" → 25）
|
||||
seconds = 4 # 默认
|
||||
for actual, display in SECONDS_DISPLAY_MAP.items():
|
||||
if display == 视频时长_display:
|
||||
seconds = actual
|
||||
break
|
||||
# 如果找不到映射，尝试直接解析数字
|
||||
if seconds == 4 and 视频时长_display != "4":
|
||||
try:
|
||||
seconds = int(视频时长_display.replace("(pro)", ""))
|
||||
except ValueError:
|
||||
seconds = 4
|
||||
|
||||
# 根据分辨率和宽高比确定实际分辨率值
|
||||
分辨率 = "720x1280" # 默认
|
||||
for actual, (res_name, orientation) in RESOLUTION_DISPLAY_MAP.items():
|
||||
if res_name == 分辨率_display and orientation == 宽高比:
|
||||
分辨率 = actual
|
||||
break
|
||||
|
||||
# 检查参考图片
|
||||
ref_image = kwargs.get("参考图片")
|
||||
ref_image_bytes = None
|
||||
if ref_image is not None:
|
||||
pil_images = tensor_to_pil(ref_image)
|
||||
if pil_images:
|
||||
ref_image_bytes = _compress_image_for_upload(pil_images[0], target_size=分辨率)
|
||||
|
||||
mode_str = "图生视频 (含参考图)" if ref_image_bytes else "文生视频"
|
||||
# 获取用户友好的显示值用于日志
|
||||
seconds_display = SECONDS_DISPLAY_MAP.get(seconds, str(seconds))
|
||||
res_display = f"{分辨率_display} {宽高比}"
|
||||
if 生成数量 > 1:
|
||||
print(f"Sora: {mode_str} | 并发{生成数量}个 | {模型} | {seconds_display} | {res_display}")
|
||||
else:
|
||||
print(f"Sora: {mode_str} | {模型} | {seconds_display} | {res_display}")
|
||||
|
||||
# 校验参数兼容性
|
||||
supported_seconds = get_sora_supported_seconds(模型)
|
||||
if supported_seconds and seconds not in supported_seconds:
|
||||
# 构建带标签的支持时长列表
|
||||
supported_labels = []
|
||||
for s in supported_seconds:
|
||||
display = SECONDS_DISPLAY_MAP.get(s, str(s))
|
||||
supported_labels.append(display)
|
||||
raise ValueError(
|
||||
f"时长 {SECONDS_DISPLAY_MAP.get(seconds, str(seconds))} 与模型 \"{模型}\" 不兼容！\n"
|
||||
f"该模型支持的时长: {', '.join(supported_labels)}"
|
||||
)
|
||||
|
||||
supported_sizes = get_sora_supported_sizes(模型)
|
||||
if supported_sizes and 分辨率 not in supported_sizes:
|
||||
# 检查该分辨率是否为Pro独占
|
||||
pro_only_sizes = ["1024x1792", "1792x1024"]
|
||||
_, orientation = RESOLUTION_DISPLAY_MAP.get(分辨率, (分辨率, ""))
|
||||
extra_hint = f"\n提示：1080P {orientation} 为 sora-2-pro 独占，请切换模型或选择720P。" if 分辨率 in pro_only_sizes else ""
|
||||
raise ValueError(
|
||||
f"分辨率 \"{分辨率_display} {宽高比}\" 与模型 \"{模型}\" 不兼容！"
|
||||
f"支持的分辨率: {', '.join(supported_sizes)}" + extra_hint
|
||||
)
|
||||
|
||||
# 准备保存路径
|
||||
video_dir = _get_video_output_dir()
|
||||
counter = _get_next_counter(video_dir, "sora")
|
||||
|
||||
# ProgressBar
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(生成数量 if 生成数量 > 1 else 100)
|
||||
|
||||
try:
|
||||
if self.client is None:
|
||||
self.client = SoraClient()
|
||||
|
||||
if 生成数量 == 1:
|
||||
# ── 单个视频：保留详细进度（提交→轮询→下载）
|
||||
save_path = os.path.join(video_dir, f"sora_{counter:05d}.mp4")
|
||||
last_progress = [0]
|
||||
|
||||
def progress_callback(progress_pct: int):
|
||||
print(
|
||||
f"\rSora: 生成中... 进度: {progress_pct}%",
|
||||
end="", flush=True
|
||||
)
|
||||
if pbar is not None and progress_pct > last_progress[0]:
|
||||
pbar.update(progress_pct - last_progress[0])
|
||||
last_progress[0] = progress_pct
|
||||
|
||||
def on_stage(stage: str):
|
||||
if stage == "submitting":
|
||||
print("Sora: 正在提交视频生成任务...")
|
||||
elif stage.startswith("submitted:"):
|
||||
vid = stage.split(":", 1)[1]
|
||||
print(f"Sora: 视频任务已提交，ID: {vid}")
|
||||
elif stage == "polling":
|
||||
print("Sora: 等待视频生成...")
|
||||
elif stage == "downloading":
|
||||
print("") # 换行（结束 \r 行）
|
||||
print("Sora: 视频生成完成，正在下载...")
|
||||
|
||||
result_path = self.client.generate_video_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
seconds=seconds,
|
||||
size=分辨率,
|
||||
save_path=save_path,
|
||||
input_reference_bytes=ref_image_bytes,
|
||||
seed=seed,
|
||||
progress_callback=progress_callback,
|
||||
on_stage=on_stage,
|
||||
)
|
||||
result_paths = [result_path]
|
||||
|
||||
else:
|
||||
# ── 批量并发：同时提交多个任务
|
||||
save_paths = [
|
||||
os.path.join(video_dir, f"sora_{counter + i:05d}.mp4")
|
||||
for i in range(生成数量)
|
||||
]
|
||||
success_count = [0]
|
||||
fail_count = [0]
|
||||
|
||||
def batch_progress_callback(current: int, total: int, success: bool, error_msg):
|
||||
if success:
|
||||
success_count[0] += 1
|
||||
print(f"Sora: 第 {current}/{total} 个视频完成 ✓")
|
||||
else:
|
||||
fail_count[0] += 1
|
||||
print(f"Sora: 第 {current}/{total} 个视频失败 ✗")
|
||||
if error_msg:
|
||||
print(f"原始错误详情:\n{error_msg}")
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
print(f"Sora: 正在并发提交 {生成数量} 个视频任务，请耐心等待...")
|
||||
result_paths = self.client.generate_batch_videos_sync(
|
||||
prompt=prompt,
|
||||
model=模型,
|
||||
seconds=seconds,
|
||||
size=分辨率,
|
||||
save_paths=save_paths,
|
||||
input_reference_bytes=ref_image_bytes,
|
||||
seed=seed,
|
||||
progress_callback=batch_progress_callback,
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.2f}s" if elapsed >= 1 else f"{elapsed:.3f}s"
|
||||
print(f"Sora: 完成！总耗时 {time_str} | 已生成 {len(result_paths)} 个视频")
|
||||
for p in result_paths:
|
||||
print(f" → {p}")
|
||||
|
||||
output_path = "\n".join(result_paths)
|
||||
return (output_path,)
|
||||
|
||||
except ValueError as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nSora: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nSora: ❌ {error_msg}")
|
||||
raise RuntimeError(error_msg) from None
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nSora: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
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"Sora: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,304 @@
|
||||
"""
|
||||
全能LLM对话助手节点
|
||||
ComfyUI 自定义节点，通过 OpenAI 兼容协议调用市面上主流的 AI 对话大模型
|
||||
支持多模态（图片输入），单轮对话，非流式输出
|
||||
|
||||
API 密钥和地址通过插件统一配置（环境变量或 .config 文件），与 Google Gemini 节点一致
|
||||
"""
|
||||
|
||||
import time
|
||||
import base64
|
||||
import json
|
||||
from io import BytesIO
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
|
||||
# ============================================================================
|
||||
# 模型配置
|
||||
# ============================================================================
|
||||
|
||||
SUPPORTED_MODELS = [
|
||||
"gpt-5.4",
|
||||
"gemini-3-flash-preview",
|
||||
"gemini-3.1-flash-lite-preview",
|
||||
"gemini-3.1-pro-preview",
|
||||
"deepseek-v3.2",
|
||||
"kimi-k2.5",
|
||||
"doubao-seed-2-0-pro-260215",
|
||||
"qwen3.5-plus-2026-02-15",
|
||||
"qwen3.5-plus",
|
||||
]
|
||||
|
||||
# 图片缩放最大尺寸
|
||||
MAX_IMAGE_DIMENSION = 1568
|
||||
|
||||
# 图片最大文件大小（20MB）
|
||||
MAX_IMAGE_SIZE = 20 * 1024 * 1024
|
||||
|
||||
|
||||
class UniversalLLMChat:
|
||||
"""
|
||||
全能LLM对话助手
|
||||
|
||||
功能：
|
||||
- 通过 OpenAI 兼容协议调用主流大模型
|
||||
- 支持多模态（图片输入）
|
||||
- 单轮对话，非流式输出
|
||||
- API 密钥和地址继承插件统一配置
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._api_key = None
|
||||
self._base_url = None
|
||||
|
||||
def _ensure_config(self):
|
||||
"""延迟加载配置，首次调用时初始化"""
|
||||
if self._api_key is None:
|
||||
self._api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
self._base_url = get_api_base_url()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"模型": (SUPPORTED_MODELS, {
|
||||
"default": SUPPORTED_MODELS[0]
|
||||
}),
|
||||
"提示词": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"图片": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("回复",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "text/generation"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def _resize_image(self, img: Image.Image) -> Image.Image:
|
||||
"""如果图片过长边超过限制，等比缩放"""
|
||||
w, h = img.size
|
||||
max_dim = max(w, h)
|
||||
if max_dim > MAX_IMAGE_DIMENSION:
|
||||
scale = MAX_IMAGE_DIMENSION / max_dim
|
||||
new_w, new_h = int(w * scale), int(h * scale)
|
||||
print(f"全能LLM: 图片缩放 {w}x{h} -> {new_w}x{new_h}")
|
||||
return img.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
||||
return img
|
||||
|
||||
def _image_to_data_url(self, img: Image.Image) -> str:
|
||||
"""将 PIL Image 转为 data URL（JPEG base64）"""
|
||||
img = self._resize_image(img)
|
||||
if img.mode in ('RGBA', 'P'):
|
||||
img = img.convert('RGB')
|
||||
|
||||
for quality in [92, 82, 72, 60, 45]:
|
||||
buf = BytesIO()
|
||||
img.save(buf, format='JPEG', quality=quality, optimize=True)
|
||||
data = buf.getvalue()
|
||||
if len(data) <= MAX_IMAGE_SIZE:
|
||||
b64 = base64.b64encode(data).decode('utf-8')
|
||||
return f"data:image/jpeg;base64,{b64}"
|
||||
|
||||
b64 = base64.b64encode(data).decode('utf-8')
|
||||
return f"data:image/jpeg;base64,{b64}"
|
||||
|
||||
def _build_messages(
|
||||
self,
|
||||
prompt: str,
|
||||
images: Optional[torch.Tensor] = None,
|
||||
) -> list:
|
||||
"""构建 OpenAI 格式的 messages 数组"""
|
||||
image_data_urls = []
|
||||
pil_images_cache = [] # 保留 PIL Image 用于总体积重新编码
|
||||
|
||||
if images is not None:
|
||||
pil_images = tensor_to_pil(images)
|
||||
for img in pil_images:
|
||||
img_resized = self._resize_image(img)
|
||||
if img_resized.mode in ('RGBA', 'P'):
|
||||
img_resized = img_resized.convert('RGB')
|
||||
pil_images_cache.append(img_resized)
|
||||
image_data_urls.append(self._image_to_data_url(img_resized))
|
||||
|
||||
# 多图总体积控制
|
||||
if pil_images_cache and len(pil_images_cache) > 1:
|
||||
total_bytes = sum(
|
||||
len(base64.b64decode(url.split(',', 1)[1])) for url in image_data_urls
|
||||
)
|
||||
if total_bytes > MAX_IMAGE_SIZE:
|
||||
print(f"全能LLM: 图片总体积 {total_bytes / 1024 / 1024:.2f}MB 超过 {MAX_IMAGE_SIZE // 1024 // 1024}MB 限制，正在压缩...")
|
||||
|
||||
# 降质量
|
||||
compressed = False
|
||||
for quality in [80, 70, 60, 50, 40, 30, 20]:
|
||||
new_urls = []
|
||||
for img in pil_images_cache:
|
||||
buf = BytesIO()
|
||||
img.save(buf, format='JPEG', quality=quality, optimize=True)
|
||||
b64 = base64.b64encode(buf.getvalue()).decode('utf-8')
|
||||
new_urls.append(f"data:image/jpeg;base64,{b64}")
|
||||
total_bytes = sum(len(base64.b64decode(u.split(',', 1)[1])) for u in new_urls)
|
||||
if total_bytes <= MAX_IMAGE_SIZE:
|
||||
image_data_urls = new_urls
|
||||
print(f"全能LLM: 图片压缩完成，总体积 {total_bytes / 1024 / 1024:.2f}MB ({len(pil_images_cache)}张图片，质量{quality})")
|
||||
compressed = True
|
||||
break
|
||||
|
||||
# 降分辨率
|
||||
if not compressed:
|
||||
for scale in [0.75, 0.5, 0.35]:
|
||||
new_urls = []
|
||||
for img in pil_images_cache:
|
||||
w, h = img.size
|
||||
resized = img.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
|
||||
buf = BytesIO()
|
||||
resized.save(buf, format='JPEG', quality=20, optimize=True)
|
||||
b64 = base64.b64encode(buf.getvalue()).decode('utf-8')
|
||||
new_urls.append(f"data:image/jpeg;base64,{b64}")
|
||||
total_bytes = sum(len(base64.b64decode(u.split(',', 1)[1])) for u in new_urls)
|
||||
if total_bytes <= MAX_IMAGE_SIZE:
|
||||
image_data_urls = new_urls
|
||||
print(f"全能LLM: 图片压缩完成，总体积 {total_bytes / 1024 / 1024:.2f}MB ({len(pil_images_cache)}张图片，缩放{int(scale*100)}%)")
|
||||
compressed = True
|
||||
break
|
||||
|
||||
if not compressed:
|
||||
print(f"全能LLM: 无法将 {len(pil_images_cache)} 张图片压缩到 {MAX_IMAGE_SIZE // 1024 // 1024}MB 以内，请减少图片数量或降低分辨率")
|
||||
raise ValueError(f"图片总体积 {total_bytes / 1024 / 1024:.2f}MB 超过限制，无法压缩到 {MAX_IMAGE_SIZE // 1024 // 1024}MB 以内")
|
||||
|
||||
if not image_data_urls:
|
||||
return [{"role": "user", "content": prompt}]
|
||||
|
||||
content_parts = []
|
||||
for url in image_data_urls:
|
||||
content_parts.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": url}
|
||||
})
|
||||
content_parts.append({
|
||||
"type": "text",
|
||||
"text": prompt
|
||||
})
|
||||
|
||||
return [{"role": "user", "content": content_parts}]
|
||||
|
||||
def generate(
|
||||
self,
|
||||
模型: str,
|
||||
提示词: str,
|
||||
图片: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[str]:
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
self._ensure_config()
|
||||
|
||||
# 构建 messages
|
||||
messages = self._build_messages(提示词, 图片)
|
||||
|
||||
img_count = len(tensor_to_pil(图片)) if 图片 is not None else 0
|
||||
input_desc = "文本" + (f" + {img_count}张图片" if img_count > 0 else "")
|
||||
|
||||
print(f"全能LLM: 模型 = {模型}")
|
||||
print(f"全能LLM: 输入 = {input_desc}")
|
||||
|
||||
# 构建请求体
|
||||
request_body = {
|
||||
"model": 模型,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
# 发送请求（在独立线程中运行异步请求，避免与 ComfyUI 事件循环冲突）
|
||||
import aiohttp
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
async def _do_request():
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {self._api_key}",
|
||||
}
|
||||
url = f"{self._base_url}/v1/chat/completions"
|
||||
timeout = aiohttp.ClientTimeout(total=120)
|
||||
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.post(url, headers=headers, json=request_body) as resp:
|
||||
status = resp.status
|
||||
body = await resp.text()
|
||||
|
||||
if status != 200:
|
||||
try:
|
||||
err_data = json.loads(body)
|
||||
err_msg = err_data.get("error", {}).get("message", body[:200])
|
||||
except Exception:
|
||||
err_msg = body[:200]
|
||||
|
||||
if status == 401:
|
||||
raise ValueError(f"认证失败：API Key 无效或已过期")
|
||||
elif status == 403:
|
||||
raise ValueError(f"无权访问模型 {模型}")
|
||||
elif status == 429:
|
||||
raise ValueError(f"请求频率超限，请稍后重试")
|
||||
elif status == 404:
|
||||
raise ValueError(f"模型 {模型} 不存在或 API 地址错误")
|
||||
else:
|
||||
raise RuntimeError(f"API 错误 ({status}): {err_msg}")
|
||||
|
||||
return json.loads(body)
|
||||
|
||||
def _run_in_thread():
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
return loop.run_until_complete(_do_request())
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
response_data = pool.submit(_run_in_thread).result()
|
||||
|
||||
# 解析响应
|
||||
choices = response_data.get("choices", [])
|
||||
if not choices:
|
||||
raise RuntimeError("API 返回了空响应（无 choices）")
|
||||
|
||||
reply = choices[0].get("message", {}).get("content", "")
|
||||
|
||||
# Token 用量
|
||||
usage = response_data.get("usage", {})
|
||||
prompt_tokens = usage.get("prompt_tokens", 0)
|
||||
completion_tokens = usage.get("completion_tokens", 0)
|
||||
total_tokens = usage.get("total_tokens", 0)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"全能LLM: 生成完成 (耗时: {elapsed:.2f}s)")
|
||||
print(f"全能LLM: Token 用量 — 输入: {prompt_tokens}, 输出: {completion_tokens}, 合计: {total_tokens}")
|
||||
if reply:
|
||||
preview = reply[:100] + "..." if len(reply) > 100 else reply
|
||||
print(f"全能LLM: 回复预览: {preview}")
|
||||
|
||||
return (reply,)
|
||||
|
||||
except ValueError as e:
|
||||
if str(e) == "未授权！":
|
||||
print("全能LLM: 请联系作者授权后方可使用！")
|
||||
raise ValueError("未授权！") from None
|
||||
error_msg = str(e).split('\n')[0]
|
||||
print(f"全能LLM: ❌ {error_msg}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e).split('\n')[0]
|
||||
print(f"全能LLM: ❌ {error_msg}")
|
||||
raise RuntimeError(error_msg) from None
|
||||
@@ -0,0 +1,422 @@
|
||||
"""
|
||||
Google Veo 视频生成节点
|
||||
ComfyUI 自定义节点，调用 Veo API 生成视频
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..clients.veo_client import VeoClient
|
||||
from ..models_config import (
|
||||
get_enabled_veo_models,
|
||||
VEO_MODELS,
|
||||
VEO_RESOLUTION_MAP,
|
||||
)
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
except ImportError:
|
||||
PROGRESS_BAR_AVAILABLE = False
|
||||
print("⚠️ GoogleVeo: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")
|
||||
|
||||
|
||||
def _get_video_output_dir() -> str:
|
||||
"""获取视频输出目录: ComfyUI/output/video"""
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
base = folder_paths.get_output_directory()
|
||||
else:
|
||||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
base = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output")
|
||||
video_dir = os.path.join(base, "video")
|
||||
os.makedirs(video_dir, exist_ok=True)
|
||||
return video_dir
|
||||
|
||||
|
||||
def _get_next_counter(directory: str, prefix: str) -> int:
|
||||
"""扫描目录，获取下一个可用的文件计数器"""
|
||||
if not os.path.exists(directory):
|
||||
return 1
|
||||
pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)")
|
||||
max_counter = 0
|
||||
for f in os.listdir(directory):
|
||||
m = pattern.match(f)
|
||||
if m:
|
||||
max_counter = max(max_counter, int(m.group(1)))
|
||||
return max_counter + 1
|
||||
|
||||
|
||||
def _fit_image_to_target(image, target_size: str):
|
||||
"""
|
||||
将参考图片按 "等比缩放覆盖 + 居中裁剪" 策略适配到目标分辨率。
|
||||
"""
|
||||
from PIL import Image as PILImage
|
||||
|
||||
parts = target_size.lower().split("x")
|
||||
target_w, target_h = int(parts[0]), int(parts[1])
|
||||
|
||||
src_w, src_h = image.size
|
||||
src_ratio = src_w / src_h
|
||||
target_ratio = target_w / target_h
|
||||
|
||||
if abs(src_ratio - target_ratio) < 0.01 and src_w <= target_w and src_h <= target_h:
|
||||
return image
|
||||
|
||||
print(f"Veo: 参考图片 {src_w}x{src_h} (比例 {src_ratio:.2f}) → 目标 {target_w}x{target_h} (比例 {target_ratio:.2f})")
|
||||
|
||||
resample = PILImage.Resampling.LANCZOS if hasattr(PILImage, "Resampling") else PILImage.LANCZOS
|
||||
|
||||
if src_ratio > target_ratio:
|
||||
scale = target_h / src_h
|
||||
new_w = round(src_w * scale)
|
||||
new_h = target_h
|
||||
image = image.resize((new_w, new_h), resample=resample)
|
||||
left = (new_w - target_w) // 2
|
||||
image = image.crop((left, 0, left + target_w, target_h))
|
||||
else:
|
||||
scale = target_w / src_w
|
||||
new_w = target_w
|
||||
new_h = round(src_h * scale)
|
||||
image = image.resize((new_w, new_h), resample=resample)
|
||||
top = (new_h - target_h) // 2
|
||||
image = image.crop((0, top, target_w, top + target_h))
|
||||
|
||||
print(f"Veo: 参考图片已适配为 {image.size[0]}x{image.size[1]}")
|
||||
return image
|
||||
|
||||
|
||||
def _compress_image_to_bytes(image, target_size: Optional[str] = None) -> bytes:
|
||||
"""
|
||||
将 PIL Image 适配目标分辨率并编码为 PNG 字节
|
||||
"""
|
||||
from io import BytesIO
|
||||
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
|
||||
if target_size:
|
||||
image = _fit_image_to_target(image, target_size)
|
||||
|
||||
buffered = BytesIO()
|
||||
image.save(buffered, format="PNG")
|
||||
size_kb = buffered.tell() / 1024
|
||||
print(f"Veo: 参考图片编码为 PNG，{size_kb:.0f} KB ({image.size[0]}x{image.size[1]})")
|
||||
return buffered.getvalue()
|
||||
|
||||
|
||||
class GoogleVeo:
|
||||
"""
|
||||
Google Veo 视频生成节点
|
||||
|
||||
功能：
|
||||
- 文生视频：基于提示词生成视频
|
||||
- 图生视频：基于首帧/尾帧/参考图生成视频
|
||||
- 异步轮询：自动等待生成完成并下载
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.client = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
enabled_models = get_enabled_veo_models()
|
||||
if not enabled_models:
|
||||
enabled_models = ["请在 models_config.py 中启用 Veo 模型"]
|
||||
|
||||
# 分辨率选项
|
||||
resolution_options = ["720p", "1080p", "4K"]
|
||||
|
||||
# 宽高比选项
|
||||
aspect_ratio_options = ["16:9", "9:16"]
|
||||
|
||||
# 视频秒数选项
|
||||
seconds_options = ["4", "6", "8"]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
"default": "A calico cat playing a piano on stage",
|
||||
"multiline": True,
|
||||
}),
|
||||
"模型": (enabled_models, {
|
||||
"default": enabled_models[0] if enabled_models else "Veo3.1",
|
||||
}),
|
||||
"分辨率": (resolution_options, {
|
||||
"default": "720p",
|
||||
}),
|
||||
"宽高比": (aspect_ratio_options, {
|
||||
"default": "9:16",
|
||||
}),
|
||||
"视频时长": (seconds_options, {
|
||||
"default": "8",
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff,
|
||||
}),
|
||||
"生成数量": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"首帧": ("IMAGE",),
|
||||
"尾帧": ("IMAGE",),
|
||||
"参考图": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("预览视频",)
|
||||
FUNCTION = "generate_video"
|
||||
CATEGORY = "video/generation"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Google Veo 视频生成节点。\n"
|
||||
"支持文生视频和图生视频（图生视频支持首帧、尾帧、参考图）。\n"
|
||||
"视频保存到 ComfyUI/output/video/ 目录。\n\n"
|
||||
"【模型说明】\n"
|
||||
"• Veo3.1：Google 最新视频生成模型\n\n"
|
||||
"【分辨率说明】\n"
|
||||
"• 720p：标清\n"
|
||||
"• 1080p：高清\n"
|
||||
"• 4K：超高清\n\n"
|
||||
"【时长说明】\n"
|
||||
"• 4秒：短视频\n"
|
||||
"• 6秒：标准\n"
|
||||
"• 8秒：长视频（默认）\n\n"
|
||||
"【图生视频说明】\n"
|
||||
"• 首帧：视频开始的第一帧图像\n"
|
||||
"• 尾帧：视频结束时的最后一帧图像\n"
|
||||
"• 参考图：参考图像（与首帧/尾帧配合使用）\n"
|
||||
"• 至少需要提供首帧或参考图之一"
|
||||
)
|
||||
|
||||
def generate_video(
|
||||
self,
|
||||
prompt: str,
|
||||
模型: str,
|
||||
**kwargs,
|
||||
) -> Tuple[str]:
|
||||
分辨率 = kwargs.pop("分辨率", "720p")
|
||||
宽高比 = kwargs.pop("宽高比", "9:16")
|
||||
视频时长 = kwargs.pop("视频时长", "8")
|
||||
seed = kwargs.pop("seed", 0)
|
||||
生成数量 = kwargs.pop("生成数量", 1)
|
||||
start_time = time.time()
|
||||
|
||||
# 解析视频时长
|
||||
seconds = int(视频时长)
|
||||
|
||||
# 解析分辨率和宽高比，映射到模型名称
|
||||
size_key = f"{分辨率}_{宽高比}"
|
||||
actual_size = VEO_RESOLUTION_MAP.get(size_key)
|
||||
if not actual_size:
|
||||
# 默认值
|
||||
actual_size = "720x1280" # 720p 9:16
|
||||
|
||||
# 检查是否有参考图输入
|
||||
首帧 = kwargs.get("首帧")
|
||||
尾帧 = kwargs.get("尾帧")
|
||||
参考图 = kwargs.get("参考图")
|
||||
|
||||
has_image = 首帧 is not None or 尾帧 is not None or 参考图 is not None
|
||||
|
||||
# 根据是否有图片选择模型前缀
|
||||
if has_image:
|
||||
model_prefix = "veo3.1"
|
||||
else:
|
||||
model_prefix = "veo3.1"
|
||||
|
||||
# 构建完整模型名称
|
||||
# 格式: veo3.1-portrait / veo3.1-landscape / veo3.1-portrait-fl / veo3.1-landscape-fl 等
|
||||
if 分辨率 == "720p":
|
||||
res_suffix = ""
|
||||
if 宽高比 == "9:16":
|
||||
orientation = "portrait"
|
||||
else:
|
||||
orientation = "landscape"
|
||||
elif 分辨率 == "1080p":
|
||||
res_suffix = "-hd"
|
||||
if 宽高比 == "9:16":
|
||||
orientation = "portrait"
|
||||
else:
|
||||
orientation = "landscape"
|
||||
else: # 4K
|
||||
res_suffix = "-4k"
|
||||
if 宽高比 == "9:16":
|
||||
orientation = "portrait"
|
||||
else:
|
||||
orientation = "landscape"
|
||||
|
||||
# 图生视频添加 -fl 后缀
|
||||
if has_image:
|
||||
model_suffix = f"-{orientation}-fl{res_suffix}"
|
||||
else:
|
||||
model_suffix = f"-{orientation}{res_suffix}"
|
||||
|
||||
model = f"{model_prefix}{model_suffix}"
|
||||
|
||||
# 准备图片字节
|
||||
first_frame_bytes = None
|
||||
last_frame_bytes = None
|
||||
reference_bytes = None
|
||||
|
||||
if 首帧 is not None:
|
||||
pil_images = tensor_to_pil(首帧)
|
||||
if pil_images:
|
||||
first_frame_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size)
|
||||
|
||||
if 尾帧 is not None:
|
||||
pil_images = tensor_to_pil(尾帧)
|
||||
if pil_images:
|
||||
last_frame_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size)
|
||||
|
||||
if 参考图 is not None:
|
||||
pil_images = tensor_to_pil(参考图)
|
||||
if pil_images:
|
||||
reference_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size)
|
||||
|
||||
mode_str = "图生视频" if has_image else "文生视频"
|
||||
print(f"Veo: {mode_str} | 并发{生成数量}个 | 模型: {model} | {seconds}秒 | {分辨率} {宽高比}")
|
||||
|
||||
# 准备保存路径
|
||||
video_dir = _get_video_output_dir()
|
||||
counter = _get_next_counter(video_dir, "veo")
|
||||
|
||||
# ProgressBar
|
||||
pbar = None
|
||||
if PROGRESS_BAR_AVAILABLE:
|
||||
pbar = ProgressBar(生成数量 if 生成数量 > 1 else 100)
|
||||
|
||||
try:
|
||||
if self.client is None:
|
||||
self.client = VeoClient()
|
||||
|
||||
if 生成数量 == 1:
|
||||
save_path = os.path.join(video_dir, f"veo_{counter:05d}.mp4")
|
||||
last_progress = [0]
|
||||
|
||||
def progress_callback(progress_pct: int):
|
||||
print(
|
||||
f"\rVeo: 生成中... 进度: {progress_pct}%",
|
||||
end="", flush=True
|
||||
)
|
||||
if pbar is not None and progress_pct > last_progress[0]:
|
||||
pbar.update(progress_pct - last_progress[0])
|
||||
last_progress[0] = progress_pct
|
||||
|
||||
def on_stage(stage: str):
|
||||
if stage == "submitting":
|
||||
print("Veo: 正在提交视频生成任务...")
|
||||
elif stage.startswith("submitted:"):
|
||||
vid = stage.split(":", 1)[1]
|
||||
print(f"Veo: 视频任务已提交，ID: {vid}")
|
||||
elif stage == "polling":
|
||||
print("Veo: 等待视频生成...")
|
||||
elif stage == "downloading":
|
||||
print("")
|
||||
print("Veo: 视频生成完成，正在下载...")
|
||||
|
||||
result_path = self.client.generate_video_sync(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=actual_size,
|
||||
save_path=save_path,
|
||||
first_frame_bytes=first_frame_bytes,
|
||||
last_frame_bytes=last_frame_bytes,
|
||||
reference_bytes=reference_bytes,
|
||||
seed=seed,
|
||||
progress_callback=progress_callback,
|
||||
on_stage=on_stage,
|
||||
)
|
||||
result_paths = [result_path]
|
||||
|
||||
else:
|
||||
save_paths = [
|
||||
os.path.join(video_dir, f"veo_{counter + i:05d}.mp4")
|
||||
for i in range(生成数量)
|
||||
]
|
||||
success_count = [0]
|
||||
|
||||
def batch_progress_callback(current: int, total: int, success: bool, error_msg):
|
||||
if success:
|
||||
success_count[0] += 1
|
||||
print(f"Veo: 第 {current}/{total} 个视频完成 ✓")
|
||||
else:
|
||||
print(f"Veo: 第 {current}/{total} 个视频失败 ✗")
|
||||
if error_msg:
|
||||
print(f"原始错误详情:\n{error_msg}")
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
print(f"Veo: 正在并发提交 {生成数量} 个视频任务，请耐心等待...")
|
||||
result_paths = self.client.generate_batch_videos_sync(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=actual_size,
|
||||
save_paths=save_paths,
|
||||
first_frame_bytes=first_frame_bytes,
|
||||
last_frame_bytes=last_frame_bytes,
|
||||
reference_bytes=reference_bytes,
|
||||
seed=seed,
|
||||
progress_callback=batch_progress_callback,
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.2f}s" if elapsed >= 1 else f"{elapsed:.3f}s"
|
||||
print(f"Veo: 完成！总耗时 {time_str} | 已生成 {len(result_paths)} 个视频")
|
||||
for p in result_paths:
|
||||
print(f" → {p}")
|
||||
|
||||
output_path = "\n".join(result_paths)
|
||||
return (output_path,)
|
||||
|
||||
except ValueError as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nVeo: ❌ {error_msg}")
|
||||
raise ValueError(error_msg) from None
|
||||
|
||||
except RuntimeError as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nVeo: ❌ {error_msg}")
|
||||
raise RuntimeError(error_msg) from None
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
print(f"\nVeo: ❌ {error_msg}")
|
||||
raise type(e)(error_msg) from None
|
||||
|
||||
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"Veo: {balance_info}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"GoogleVeo": GoogleVeo,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"GoogleVeo": "Google Veo - ab",
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
"""
|
||||
通用视频预览节点
|
||||
ComfyUI 自定义节点，接收视频文件路径并在前端展示预览
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
SUPPORTED_VIDEO_EXTENSIONS = {".mp4", ".webm", ".mov", ".avi", ".mkv", ".flv", ".wmv", ".3gp"}
|
||||
|
||||
|
||||
def _get_output_dir() -> str:
|
||||
if FOLDER_PATHS_AVAILABLE:
|
||||
return folder_paths.get_output_directory()
|
||||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
return os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output")
|
||||
|
||||
|
||||
class VideoPreview:
|
||||
"""
|
||||
通用视频预览节点
|
||||
|
||||
功能：
|
||||
- 接收视频文件路径（STRING）
|
||||
- 在 ComfyUI 前端节点上内嵌 <video> 播放器进行预览
|
||||
- 支持 mp4, webm, mov, avi, mkv 等主流格式
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"预览视频": ("STRING", {"forceInput": True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "preview"
|
||||
CATEGORY = "video"
|
||||
|
||||
DESCRIPTION = (
|
||||
"通用视频预览节点。\n"
|
||||
"接收视频文件路径，在节点上显示视频播放器。\n"
|
||||
"支持 mp4, webm, mov, avi, mkv 等主流视频格式。"
|
||||
)
|
||||
|
||||
def preview(self, **kwargs) -> dict:
|
||||
video_path = kwargs.get("预览视频", "")
|
||||
if not video_path or not video_path.strip():
|
||||
raise ValueError("视频路径为空")
|
||||
|
||||
video_path = video_path.strip()
|
||||
|
||||
if not os.path.isfile(video_path):
|
||||
raise ValueError(f"视频文件不存在: {video_path}")
|
||||
|
||||
ext = os.path.splitext(video_path)[1].lower()
|
||||
if ext not in SUPPORTED_VIDEO_EXTENSIONS:
|
||||
raise ValueError(
|
||||
f"不支持的视频格式 '{ext}'，"
|
||||
f"支持: {', '.join(sorted(SUPPORTED_VIDEO_EXTENSIONS))}"
|
||||
)
|
||||
|
||||
output_dir = _get_output_dir()
|
||||
abs_video = os.path.abspath(video_path)
|
||||
abs_output = os.path.abspath(output_dir)
|
||||
|
||||
if abs_video.startswith(abs_output):
|
||||
rel_path = os.path.relpath(abs_video, abs_output)
|
||||
subfolder = os.path.dirname(rel_path).replace("\\", "/")
|
||||
filename = os.path.basename(rel_path)
|
||||
file_type = "output"
|
||||
else:
|
||||
filename = os.path.basename(abs_video)
|
||||
subfolder = ""
|
||||
file_type = "output"
|
||||
|
||||
# 如果文件不在 output 目录下，复制一份到 output/video/
|
||||
target_dir = os.path.join(output_dir, "video")
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
target_path = os.path.join(target_dir, filename)
|
||||
|
||||
if not os.path.exists(target_path) or abs_video != os.path.abspath(target_path):
|
||||
import shutil
|
||||
shutil.copy2(abs_video, target_path)
|
||||
|
||||
subfolder = "video"
|
||||
|
||||
file_size = os.path.getsize(abs_video)
|
||||
size_mb = file_size / (1024 * 1024)
|
||||
print(f"视频预览: {filename} ({size_mb:.1f}MB)")
|
||||
|
||||
return {
|
||||
"ui": {
|
||||
"videos": [{
|
||||
"filename": filename,
|
||||
"subfolder": subfolder,
|
||||
"type": file_type,
|
||||
}],
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user