"""GPT Image V3 批量跑图节点。""" import asyncio import gc import os import time from concurrent.futures import ThreadPoolExecutor from typing import List, Optional, Tuple from PIL import Image from comfy_api.latest import io from ..clients.gpt_image_client import ( GPT_IMAGE_MODEL_OPTIONS, GPT_IMAGE_ROUTE_OPTIONS, GptImageClient, resolve_gpt_image_model, ) from .gpt_image import GPT_IMAGE_25_QUALITY_OPTIONS, resolve_gpt_image_quality from ..utils.config import get_base_url_by_route from ..utils.file_utils import ( ImageInfo, load_images_from_folder, pair_images_by_name, pair_images_cartesian, pair_images_indexed, save_image, ) from ..utils.image_utils import parse_batch_prompts, pil_to_tensor, tensor_to_pil from ..utils.o1key_image_catalog import ( GPT_IMAGE_BACKGROUND_OPTIONS, GPT_IMAGE_OUTPUT_FORMAT_OPTIONS, ) try: from comfy.model_management import InterruptProcessingException except ImportError: InterruptProcessingException = RuntimeError try: from comfy.utils import ProgressBar _PROGRESS_AVAILABLE = True except ImportError: _PROGRESS_AVAILABLE = False try: import folder_paths _FOLDER_PATHS_AVAILABLE = True except ImportError: _FOLDER_PATHS_AVAILABLE = False _MAX_PATHS = 5 _MAX_REFERENCES = 9 _PAIRING_MODES = ["不配对", "相同文件名", "同序号", "全匹配"] _RESOLUTION_OPTIONS = [ "智能", "1024x1024（1K 正方形 1:1）", "1536x1024（1K 横版 3:2）", "1024x1536（1K 竖版 2:3）", "1360x1024（1K 横版 4:3）", "1024x1360（1K 竖版 3:4）", "1824x1024（1K 横版 16:9）", "1024x1824（1K 竖版 9:16）", "2048x2048（2K 正方形 1:1）", "3072x2048（2K 横版 3:2）", "2048x3072（2K 竖版 2:3）", "2736x2048（2K 横版 4:3）", "2048x2736（2K 竖版 3:4）", "3648x2048（2K 横版 16:9）", "2048x3648（2K 竖版 9:16）", "2880x2880（4K 正方形 1:1）", "3504x2336（4K 横版 3:2）", "2336x3504（4K 竖版 2:3）", "3264x2448（4K 横版 4:3）", "2448x3264（4K 竖版 3:4）", "3840x2160（4K 横版 16:9）", "2160x3840（4K 竖版 9:16）", ] def _path_name(index: int) -> str: return "参考图1（主图）" if index == 1 else f"参考图{index}" def _path_count(value) -> int: try: return max(1, min(_MAX_PATHS, int(str(value).split("个", 1)[0]))) except (TypeError, ValueError): return 1 def _path_option(count: int): inputs = [ io.String.Input( _path_name(index), default="", placeholder="填写图片文件夹路径", tooltip="主图文件夹路径" if index == 1 else f"第 {index} 个参考图文件夹路径", ) for index in range(1, count + 1) ] if count >= 2: inputs.append(io.Combo.Input( "图片配对模式", options=_PAIRING_MODES, default="不配对", tooltip="支持不配对、相同文件名、同序号和全部组合。", )) return io.DynamicCombo.Option(f"{count}个路径", inputs) def _collect_group(value) -> list: if isinstance(value, dict): return [item for item in value.values() if item is not None] return [] if value is None else [value] def _resolve_size(value: str) -> str: value = (value or "").strip() if not value or value == "智能": return "auto" return value.split("（", 1)[0].strip().lower().replace("×", "x").replace("*", "x") class O1keyGPTImageBatch(io.ComfyNode): """动态文件夹、Autogrow 参考图和并发异步任务版 GPT Image 批量节点。""" @classmethod def define_schema(cls): references = io.Autogrow.Input( "参考图组", template=io.Autogrow.TemplateNames( input=io.Image.Input("参考图"), names=[f"参考图{i}" for i in range(1, _MAX_REFERENCES + 1)], min=0, ), tooltip="固定追加到每个批量任务末尾；端口序号接在图片路径数量之后，连接后自动增加，最多 9 张。", ) return io.Schema( node_id="O1keyGPTImageBatch", display_name="GPT Image 批量跑图", category="o1key/image", inputs=[ io.String.Input("prompt", default="", multiline=True, tooltip="可用独占一行的 --- 分隔多条提示词。"), io.Combo.Input( "模型", options=GPT_IMAGE_MODEL_OPTIONS, default="gpt-image-2.5-sunburst", ), io.Combo.Input("模型线路", options=GPT_IMAGE_ROUTE_OPTIONS, default="畅速"), io.Combo.Input("分辨率", options=_RESOLUTION_OPTIONS, default="智能"), io.Int.Input("生图数量", default=1, min=1, max=8, step=1), io.Combo.Input( "质量", options=GPT_IMAGE_25_QUALITY_OPTIONS, default="自动", tooltip="GPT Image 2.5 另支持超高=xhigh、最高=max。", ), io.DynamicCombo.Input( "图片路径数量", options=[_path_option(count) for count in range(1, _MAX_PATHS + 1)], tooltip="按需显示 1～5 个图片文件夹路径。", ), io.Mask.Input("遮罩", optional=True, tooltip="应用到每个任务的第一张参考图。"), references, io.Combo.Input("图片输出格式", options=["原始", "JPEG", "PNG", "WebP"], default="原始"), io.Combo.Input( "背景", options=list(GPT_IMAGE_BACKGROUND_OPTIONS), default="auto", tooltip="透明背景仅支持 PNG 或 WebP 输出格式。", ), io.Combo.Input("图片保存命名规则", options=["和原始图片名保持一致", "自然数字"], default="和原始图片名保持一致"), io.String.Input("图片保存路径", default="", placeholder="留空时保存到 ComfyUI output 目录"), io.Combo.Input( "缩放图片", options=["不缩放", "智能缩放"], default="智能缩放", tooltip="请求体超过 18 MiB 时，智能缩放会等比缩小占用最大的参考图。", ), io.Int.Input( "seed", default=0, min=0, max=2**31 - 1, step=1, control_after_generate=io.ControlAfterGenerate.randomize, ), ], outputs=[io.Image.Output(display_name="输出图像")], accept_all_inputs=True, ) @staticmethod def _normalize_pairing(value: str) -> str: return { "按相同图片命名": "相同文件名", "1*N": "全匹配", }.get(value, value if value in _PAIRING_MODES else "不配对") @staticmethod def _manual_images(values: list) -> List[ImageInfo]: images = [] for input_index, tensor in enumerate(values, 1): for frame_index, image in enumerate(tensor_to_pil(tensor)): images.append(ImageInfo(image, f"manual_{input_index}_{frame_index}", ".png", "")) return images @classmethod def _create_pairs( cls, image_lists: List[List[ImageInfo]], pairing_mode: str, manual_images: List[ImageInfo], ) -> List[Tuple[ImageInfo, ...]]: pairing_mode = cls._normalize_pairing(pairing_mode) if pairing_mode == "不配对": if len(image_lists) > 1: raise ValueError("「不配对」模式只支持一个配对路径。") if image_lists: base_pairs = [(item,) for item in image_lists[0]] else: return [] elif not image_lists: return [] elif len(image_lists) == 1: base_pairs = [(item,) for item in image_lists[0]] elif pairing_mode == "相同文件名": base_pairs = list(pair_images_by_name(*image_lists)) elif pairing_mode == "同序号": base_pairs = list(pair_images_indexed(*image_lists)) else: base_pairs = list(pair_images_cartesian(*image_lists)) manual_tuple = tuple(manual_images) return [pair + manual_tuple for pair in base_pairs] @staticmethod def _output_folder(path: str) -> str: folder = (path or "").strip() if not folder and _FOLDER_PATHS_AVAILABLE: folder = folder_paths.get_output_directory() if not folder: raise ValueError("未设置保存路径，且无法获取 ComfyUI output 目录") os.makedirs(folder, exist_ok=True) return folder @staticmethod def _save_images( images: List[Image.Image], folder: str, image_format: str, naming_rule: str, task_index: int, base_filename: Optional[str], ) -> List[str]: saved = [] for image_index, image in enumerate(images, 1): if image_format == "原始": fmt = str(getattr(image, "format", None) or "PNG").upper() fmt = "JPEG" if fmt in ("JPG", "JPEG") else "WEBP" if fmt == "WEBP" else "PNG" else: fmt = image_format.upper() ext = {"JPEG": ".jpg", "WEBP": ".webp"}.get(fmt, ".png") if naming_rule == "自然数字": stem = str(task_index + 1) if image_index > 1: stem += f"_{image_index}" else: stem = base_filename or f"task_{task_index + 1}" if image_index > 1: stem += f"+{image_index - 1}" path = os.path.join(folder, f"{stem}{ext}") collision = 1 while os.path.exists(path): path = os.path.join(folder, f"{stem}+{collision}{ext}") collision += 1 if fmt == "JPEG": image.convert("RGB").save(path, format="JPEG", quality=100, subsampling=0) elif fmt == "WEBP": image.save(path, format="WEBP", lossless=True, quality=100) else: save_image(image, path) saved.append(path) return saved @staticmethod def _progress_callback(progress_values, index, pbar): if pbar is None: return None def update(value): try: progress_values[index] = max(0, min(100, int(value))) except (TypeError, ValueError): return pbar.update_absolute(sum(progress_values), len(progress_values) * 100) return update @classmethod async def _run_task( cls, client, pair, task_prompt, task_index, total_tasks, model, quality, size, image_count, seed, mask, output_format, background, resize_mode, folder, naming_rule, save_lock, progress_callback, ) -> dict: try: task_started = time.time() images = await client.generate_image_async( prompt=task_prompt, model=model, quality=quality, size=size, n=image_count, seed=seed, image_tensor=[pil_to_tensor([item.image]) for item in pair], mask_tensor=mask, output_format=output_format, background=background, progress_callback=progress_callback, special_price_parallel=True, log_downloads=True, log_prefix=f"[GPT Image Batch] [{task_index + 1}/{total_tasks}]", task_submitted_callback=lambda task_id, status, elapsed: print( f"[GPT Image Batch] [{task_index + 1}/{total_tasks}] 已提交 | " f"task_id={task_id} | 状态={status} | 耗时={elapsed:.1f}s" ), log_request_start=False, resize_mode=resize_mode, ) async with save_lock: files = cls._save_images( images, folder, output_format if output_format != "png" else "PNG", naming_rule, task_index, pair[0].filename if pair else None, ) print( f"[GPT Image Batch] [{task_index + 1}/{total_tasks}] 完成 ✓ | " f"生成={len(images)} 张 | 耗时={time.time() - task_started:.1f}s" ) return {"task_index": task_index, "success": True, "generated_count": len(images), "saved_files": files, "error": None} except (InterruptProcessingException, asyncio.CancelledError): raise except Exception as error: message = str(error).splitlines()[0] print(f"[GPT Image Batch] [{task_index + 1}/{total_tasks}] ❌ {message}") return {"task_index": task_index, "success": False, "generated_count": 0, "saved_files": [], "error": message} @classmethod async def _process_async( cls, client, task_defs, model, quality, size, image_count, seed, mask, output_format, background, resize_mode, folder, naming_rule, pbar, ) -> List[dict]: total = len(task_defs) progress_values = [0] * total save_lock = asyncio.Lock() tasks = [ asyncio.create_task(cls._run_task( client, pair, prompt, index, total, model, quality, size, image_count, seed, mask, output_format, background, resize_mode, folder, naming_rule, save_lock, cls._progress_callback(progress_values, index, pbar), )) for index, pair, prompt in task_defs ] batch = await asyncio.gather(*tasks, return_exceptions=True) results = [] for index, item in enumerate(batch): if isinstance(item, (InterruptProcessingException, asyncio.CancelledError)): raise item if isinstance(item, BaseException): item = {"task_index": index, "success": False, "generated_count": 0, "saved_files": [], "error": str(item)} results.append(item) gc.collect() print(f"[GPT Image Batch] 进度 {total}/{total}") return results @classmethod def execute( cls, prompt, 模型="gpt-image-2.5-sunburst", 模型线路="畅速", 分辨率="智能", 生图数量=1, 质量="自动", 图片路径数量=None, 遮罩=None, seed=0, 图片输出格式="原始", 图片保存命名规则="和原始图片名保持一致", 图片保存路径="", 缩放图片="智能缩放", 背景="auto", **kwargs, ) -> io.NodeOutput: if not prompt or not str(prompt).strip(): raise ValueError("提示词不能为空") if 缩放图片 not in {"不缩放", "智能缩放"}: raise ValueError("缩放图片参数无效") if 背景 not in GPT_IMAGE_BACKGROUND_OPTIONS: raise ValueError("GPT Image 背景参数无效") actual_model = resolve_gpt_image_model(模型, 模型线路) # 兼容旧节点通过命名参数调用时使用的字段名。 if 图片输出格式 == "原始" and kwargs.get("图片格式"): 图片输出格式 = kwargs["图片格式"] if not str(图片保存路径 or "").strip() and kwargs.get("保存路径"): 图片保存路径 = kwargs["保存路径"] legacy_group = kwargs.get("图片文件夹数量") nested = 图片路径数量 if isinstance(图片路径数量, dict) else ( legacy_group if isinstance(legacy_group, dict) else None ) values = nested or kwargs selected_count = _path_count(values.get( "图片路径数量", values.get("图片文件夹数量", 图片路径数量 or legacy_group) )) paths = [ values.get(_path_name(i), values.get(f"图片路径{i}", kwargs.get(f"文件夹{i}", ""))) for i in range(1, _MAX_PATHS + 1) ] if nested is not None: paths = paths[:selected_count] + [""] * (_MAX_PATHS - selected_count) if not any(str(path).strip() for path in paths if path is not None): raise ValueError("请至少填写一个图片文件夹路径") pairing = cls._normalize_pairing(values.get("图片配对模式", kwargs.get("图片配对模式", "不配对"))) image_lists = [] for index, path in enumerate(paths, 1): if path and str(path).strip(): image_lists.append(load_images_from_folder(str(path).strip())) reference_values = _collect_group(kwargs.get("参考图组")) if not reference_values: reference_values = [kwargs[f"参考图{i}"] for i in range(1, _MAX_REFERENCES + 1) if kwargs.get(f"参考图{i}") is not None] pairs = cls._create_pairs( image_lists, pairing, cls._manual_images(reference_values) ) if not pairs: raise ValueError("图片配对结果为空，请检查路径和配对模式") batch_prompts = parse_batch_prompts(prompt) prompts = batch_prompts or [prompt] task_defs = [] for pair in pairs: for task_prompt in prompts: task_defs.append((len(task_defs), pair, task_prompt)) folder = cls._output_folder(图片保存路径) quality = resolve_gpt_image_quality(模型, 质量) output_format = {"原始": "png", "JPEG": "jpeg", "PNG": "png", "WebP": "webp"}.get(图片输出格式, "png") if output_format not in GPT_IMAGE_OUTPUT_FORMAT_OPTIONS: raise ValueError("GPT Image 输出格式无效") if 背景 == "transparent" and output_format == "jpeg": raise ValueError("GPT Image 透明背景仅支持 PNG 或 WebP 输出格式") client = GptImageClient() client.base_url = get_base_url_by_route() client.response_log_enabled = False client.poll_log_enabled = False print( f"[GPT Image Batch] 开始 | 任务={len(task_defs)} | 全并发 | " f"模型={actual_model} | 每任务={生图数量} 张" ) pbar = ProgressBar(len(task_defs) * 100) if _PROGRESS_AVAILABLE else None start_time = time.time() def run_async(): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: coro = cls._process_async( client, task_defs, actual_model, quality, _resolve_size(分辨率), 生图数量, seed, 遮罩, output_format, 背景, 缩放图片, folder, 图片保存命名规则, pbar, ) return loop.run_until_complete(client._run_with_interrupt(coro)) finally: loop.close() try: with ThreadPoolExecutor(max_workers=1, thread_name_prefix="gpt-image-batch") as executor: results = executor.submit(run_async).result() finally: try: print(f"[GPT Image Batch] {client.format_balance_info(client.query_balance_sync())}") except Exception: pass successful = [item for item in results if item.get("success")] if not successful: raise RuntimeError("所有 GPT Image 批量任务均生成失败") saved_files = [path for item in successful for path in item["saved_files"]] preview = [] for path in saved_files[-10:]: try: image = Image.open(path) image.load() preview.append(image.copy()) except Exception as error: print(f"[GPT Image Batch] 预览加载失败 {path}: {error}") output = GptImageClient._pil_list_to_tensor(preview) generated = sum(item["generated_count"] for item in successful) print( f"[GPT Image Batch] 完成 | 成功={len(successful)}/{len(results)} " f"| 生成={generated} | 耗时={time.time() - start_time:.1f}s | 保存={folder}" ) return io.NodeOutput(output)