Publish current ComfyUI O1Key code baseline
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
This commit is contained in:
Jony
@@ -0,0 +1,479 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user