Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
885 lines
35 KiB
Python
885 lines
35 KiB
Python
"""
|
||
o1key GPT Image 节点
|
||
支持 GPT Image 2 / 2.5 系列的文生图、图生图和带蒙版图像编辑
|
||
"""
|
||
|
||
import os
|
||
import time
|
||
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 ..utils.image_utils import parse_batch_prompts, pil_to_tensor, tensor_to_pil
|
||
from ..utils.config import get_base_url_by_route
|
||
from ..utils.o1key_image_catalog import (
|
||
GPT_IMAGE_BACKGROUND_OPTIONS,
|
||
GPT_IMAGE_EXACT_SIZE_OPTIONS,
|
||
GPT_IMAGE_OUTPUT_FORMAT_OPTIONS,
|
||
GPT_IMAGE_25_QUALITY_OPTIONS,
|
||
resolve_gpt_image_quality,
|
||
resolve_gpt_image_size,
|
||
)
|
||
from ..utils.file_utils import (
|
||
ImageInfo,
|
||
generate_timestamp_filename,
|
||
load_images_from_folder,
|
||
pair_images_by_name,
|
||
pair_images_cartesian,
|
||
save_image,
|
||
)
|
||
|
||
try:
|
||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||
_INTERRUPT_AVAILABLE = True
|
||
except ImportError:
|
||
_INTERRUPT_AVAILABLE = False
|
||
processing_interrupted = lambda: False
|
||
InterruptProcessingException = RuntimeError
|
||
|
||
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
|
||
|
||
|
||
def _make_node_progress_callback(progress_bar, task_index: int, total_tasks: int):
|
||
if progress_bar is None:
|
||
return None
|
||
|
||
total_units = max(1, total_tasks) * 100
|
||
base_units = max(0, task_index - 1) * 100
|
||
last_pct = {"value": -1}
|
||
|
||
def _callback(pct: int):
|
||
try:
|
||
pct_value = int(round(float(pct)))
|
||
except (TypeError, ValueError):
|
||
return
|
||
pct_value = max(0, min(100, pct_value))
|
||
if pct_value < last_pct["value"]:
|
||
return
|
||
last_pct["value"] = pct_value
|
||
progress_bar.update_absolute(
|
||
min(total_units, base_units + pct_value),
|
||
total_units,
|
||
)
|
||
|
||
return _callback
|
||
|
||
|
||
def _resolve_async_size(value: str) -> str:
|
||
return resolve_gpt_image_size(value)
|
||
|
||
|
||
MAX_GPT_IMAGE_REFERENCES = 9
|
||
def _collect_autogrow_inputs(value) -> list:
|
||
"""收集已连接的 Autogrow 输入，并兼容单个旧值。"""
|
||
if value is None:
|
||
return []
|
||
if isinstance(value, dict):
|
||
return [item for item in value.values() if item is not None]
|
||
return [value]
|
||
|
||
|
||
class O1keyGPTImage(io.ComfyNode):
|
||
"""
|
||
o1key GPT Image 节点
|
||
|
||
功能：
|
||
- 文生图：仅提供 prompt
|
||
- 图生图：提供 prompt + 图片（无遮罩）
|
||
- 图像编辑：提供 prompt + 图片 + 遮罩（白色区域将被替换）
|
||
- 批量模式：prompt 中用单独一行 --- 分隔多条提示词
|
||
|
||
参数：
|
||
- prompt : 文本提示词（多行；用 --- 独占一行分隔批量提示词）
|
||
- 模型 : GPT Image 主模型
|
||
- 模型线路 : 畅速、直连或专线
|
||
- 分辨率 : 图像尺寸（auto 让 API 自动决定）
|
||
- 生图数量 : 每条提示词生成数量 1-8
|
||
- 质量 : 生成质量
|
||
- seed : 随机种子（0 表示不指定）
|
||
- 图片 : 可选参考图（用于图生图或编辑）
|
||
- 遮罩 : 可选蒙版（白色区域将被替换）
|
||
"""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
reference_images = io.Autogrow.Input(
|
||
"参考图组",
|
||
template=io.Autogrow.TemplateNames(
|
||
input=io.Image.Input("参考图"),
|
||
names=[f"参考图{i}" for i in range(1, MAX_GPT_IMAGE_REFERENCES + 1)],
|
||
min=0,
|
||
),
|
||
tooltip=f"连接后自动增加输入端口，合计最多 {MAX_GPT_IMAGE_REFERENCES} 张参考图。",
|
||
)
|
||
return io.Schema(
|
||
node_id="O1keyGPTImage",
|
||
display_name="gpt image",
|
||
category="o1key/image",
|
||
inputs=[
|
||
io.String.Input(
|
||
"prompt",
|
||
default="",
|
||
multiline=True,
|
||
tooltip="Text prompt for GPT Image. Use --- on its own line to separate batch prompts.",
|
||
),
|
||
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=GPT_IMAGE_EXACT_SIZE_OPTIONS,
|
||
default="智能",
|
||
tooltip="Image size (智能 = API decides)",
|
||
),
|
||
io.Int.Input(
|
||
"生图数量",
|
||
default=1,
|
||
min=1,
|
||
max=8,
|
||
step=1,
|
||
display_mode=io.NumberDisplay.number,
|
||
tooltip="How many images to generate per prompt",
|
||
),
|
||
io.Combo.Input(
|
||
"质量",
|
||
options=GPT_IMAGE_25_QUALITY_OPTIONS,
|
||
default="自动",
|
||
tooltip="GPT Image 2 支持高/中/低/自动；GPT Image 2.5 另支持超高=xhigh、最高=max。",
|
||
),
|
||
io.Combo.Input(
|
||
"输出格式",
|
||
options=["png", "jpeg", "webp"],
|
||
default="png",
|
||
tooltip="Generated image output format",
|
||
),
|
||
io.Combo.Input(
|
||
"背景",
|
||
options=list(GPT_IMAGE_BACKGROUND_OPTIONS),
|
||
default="auto",
|
||
tooltip="透明背景仅支持 PNG 或 WebP 输出格式。",
|
||
),
|
||
io.Mask.Input(
|
||
"遮罩",
|
||
optional=True,
|
||
tooltip="Optional mask for inpainting (white areas will be replaced)",
|
||
),
|
||
reference_images,
|
||
io.Combo.Input(
|
||
"缩放图片",
|
||
options=["不缩放", "智能缩放"],
|
||
default="智能缩放",
|
||
tooltip="请求体超过 18 MiB 时，智能缩放会等比缩小占用最大的参考图。",
|
||
),
|
||
io.Int.Input(
|
||
"seed",
|
||
default=0,
|
||
min=0,
|
||
max=2**31 - 1,
|
||
step=1,
|
||
display_mode=io.NumberDisplay.number,
|
||
control_after_generate=io.ControlAfterGenerate.randomize,
|
||
tooltip="Random seed (0 = not specified)",
|
||
),
|
||
],
|
||
outputs=[io.Image.Output(display_name="IMAGE")],
|
||
# 兼容 Autogrow 改造前保存的参考图1～参考图9端口。
|
||
accept_all_inputs=True,
|
||
)
|
||
|
||
@classmethod
|
||
def execute(
|
||
cls,
|
||
prompt: str,
|
||
模型: str = "gpt-image-2.5-sunburst",
|
||
模型线路: str = "畅速",
|
||
分辨率: str = "智能",
|
||
质量: str = "自动",
|
||
输出格式: str = "png",
|
||
生图数量: int = 1,
|
||
seed: int = 0,
|
||
遮罩=None,
|
||
缩放图片: str = "智能缩放",
|
||
背景: str = "auto",
|
||
**kwargs,
|
||
) -> io.NodeOutput:
|
||
result = cls.generate(
|
||
prompt=prompt,
|
||
模型=模型,
|
||
模型线路=模型线路,
|
||
分辨率=分辨率,
|
||
质量=质量,
|
||
输出格式=输出格式,
|
||
生图数量=生图数量,
|
||
seed=seed,
|
||
遮罩=遮罩,
|
||
缩放图片=缩放图片,
|
||
背景=背景,
|
||
**kwargs,
|
||
)
|
||
return io.NodeOutput(*result)
|
||
|
||
@classmethod
|
||
def generate(
|
||
cls,
|
||
prompt: str,
|
||
模型: str = "gpt-image-2.5-sunburst",
|
||
模型线路: str = "畅速",
|
||
分辨率: str = "智能",
|
||
质量: str = "自动",
|
||
输出格式: str = "png",
|
||
生图数量: int = 1,
|
||
seed: int = 0,
|
||
遮罩=None,
|
||
缩放图片: str = "智能缩放",
|
||
背景: str = "auto",
|
||
**kwargs,
|
||
):
|
||
"""
|
||
生成图像（文生图 / 图生图 / 图像编辑 / 批量提示词）
|
||
|
||
路由逻辑：
|
||
- 无图片 → generations 接口（文生图）
|
||
- 有图片，无遮罩 → edits 接口（图生图）
|
||
- 有图片，有遮罩 → edits 接口（图像编辑 + 蒙版）
|
||
- prompt 含 --- → 批量模式，逐条调用上述接口
|
||
"""
|
||
start_time = time.time()
|
||
# ── 0. 收集多参考图输入 ────────────────────────────────────────────────
|
||
reference_tensors = _collect_autogrow_inputs(kwargs.get("参考图组"))
|
||
if not reference_tensors:
|
||
# 兼容 Autogrow 改造前保存的固定参考图端口。
|
||
reference_tensors = [
|
||
kwargs[f"参考图{i}"]
|
||
for i in range(1, MAX_GPT_IMAGE_REFERENCES + 1)
|
||
if kwargs.get(f"参考图{i}") is not None
|
||
]
|
||
|
||
图片 = reference_tensors if reference_tensors else None
|
||
|
||
# ── 1. 参数校验 ───────────────────────────────────────────────────────
|
||
if 遮罩 is not None and 图片 is None:
|
||
raise ValueError("提供了遮罩但未提供图片，请同时提供图片和遮罩")
|
||
if 缩放图片 not in {"不缩放", "智能缩放"}:
|
||
raise ValueError("缩放图片参数无效")
|
||
if 输出格式 not in GPT_IMAGE_OUTPUT_FORMAT_OPTIONS:
|
||
raise ValueError("GPT Image 输出格式无效")
|
||
if 背景 not in GPT_IMAGE_BACKGROUND_OPTIONS:
|
||
raise ValueError("GPT Image 背景参数无效")
|
||
if 背景 == "transparent" and 输出格式 == "jpeg":
|
||
raise ValueError("GPT Image 透明背景仅支持 PNG 或 WebP 输出格式")
|
||
|
||
# ── 2. 解析分辨率显示值 → API 参数值 ──────────────────────────────────
|
||
size = _resolve_async_size(分辨率)
|
||
|
||
# ── 2b. 主模型与线路共同解析为 API 模型名 ───────────────────────────
|
||
model = resolve_gpt_image_model(模型, 模型线路)
|
||
|
||
# ── 2c. 解析质量显示值 → API 参数值 ───────────────────────────────────
|
||
quality = resolve_gpt_image_quality(模型, 质量)
|
||
|
||
# ── 3. 创建客户端 ─────────────────────────────────────────────────────
|
||
try:
|
||
client = GptImageClient()
|
||
client.base_url = get_base_url_by_route()
|
||
client.response_log_enabled = False
|
||
client.poll_log_enabled = False
|
||
except ValueError as e:
|
||
if str(e) == "未授权！":
|
||
print("[o1key GPT Image] 请联系作者授权后方可使用！")
|
||
raise ValueError("未授权！") from None
|
||
raise
|
||
|
||
try:
|
||
# ── 4. 解析批量提示词 ─────────────────────────────────────────────
|
||
batch_prompts = parse_batch_prompts(prompt)
|
||
|
||
# ── 5. 调用 API ───────────────────────────────────────────────────
|
||
all_pil_images = []
|
||
def _submitted(task_id, status, elapsed):
|
||
print(f"[o1key GPT Image] 已提交 | task_id={task_id} | 状态={status} | 耗时={elapsed:.1f}s")
|
||
|
||
def _completed(task_id, image_count, elapsed, urls):
|
||
del urls
|
||
print(f"[o1key GPT Image] 完成 ✓ | task_id={task_id} | 生成={image_count} 张 | 耗时={elapsed:.1f}s")
|
||
progress_total = len(batch_prompts) if batch_prompts else 1
|
||
progress_bar = ProgressBar(progress_total * 100) if _PROGRESS_BAR_AVAILABLE else None
|
||
|
||
if batch_prompts:
|
||
# 批量模式：逐条提示词调用
|
||
total = len(batch_prompts)
|
||
print(f"[o1key GPT Image] 批量模式 | {total} 条提示词 | 每条生成 {生图数量} 张")
|
||
for idx, p in enumerate(batch_prompts, 1):
|
||
if _INTERRUPT_AVAILABLE and processing_interrupted():
|
||
print("[o1key GPT Image] 用户取消，已中断批量生成")
|
||
raise InterruptProcessingException()
|
||
try:
|
||
pil_images = client.generate_image_async_sync(
|
||
prompt=p,
|
||
model=model,
|
||
quality=quality,
|
||
size=size,
|
||
n=生图数量,
|
||
seed=seed,
|
||
image_tensor=图片,
|
||
mask_tensor=遮罩,
|
||
output_format=输出格式,
|
||
background=背景,
|
||
progress_callback=_make_node_progress_callback(progress_bar, idx, total),
|
||
special_price_parallel=True,
|
||
task_submitted_callback=_submitted,
|
||
task_completed_callback=_completed,
|
||
log_request_start=False,
|
||
log_downloads=True,
|
||
log_prefix=f"[o1key GPT Image] [{idx}/{total}]",
|
||
resize_mode=缩放图片,
|
||
)
|
||
all_pil_images.extend(pil_images)
|
||
except InterruptProcessingException:
|
||
raise
|
||
except Exception as e:
|
||
error_msg = str(e).split('\n')[0]
|
||
print(f"[o1key GPT Image] [{idx}/{total}] ❌ {error_msg}")
|
||
if progress_bar is not None:
|
||
progress_bar.update_absolute(idx * 100, total * 100)
|
||
else:
|
||
# 单提示词模式
|
||
if not prompt or not prompt.strip():
|
||
raise ValueError("提示词不能为空")
|
||
try:
|
||
pil_images = client.generate_image_async_sync(
|
||
prompt=prompt,
|
||
model=model,
|
||
quality=quality,
|
||
size=size,
|
||
n=生图数量,
|
||
seed=seed,
|
||
image_tensor=图片,
|
||
mask_tensor=遮罩,
|
||
output_format=输出格式,
|
||
background=背景,
|
||
progress_callback=_make_node_progress_callback(progress_bar, 1, 1),
|
||
special_price_parallel=True,
|
||
task_submitted_callback=_submitted,
|
||
task_completed_callback=_completed,
|
||
log_request_start=False,
|
||
log_downloads=True,
|
||
resize_mode=缩放图片,
|
||
)
|
||
all_pil_images.extend(pil_images)
|
||
except InterruptProcessingException:
|
||
raise
|
||
except Exception as e:
|
||
error_msg = str(e).split('\n')[0]
|
||
print(f"[o1key GPT Image] ❌ {error_msg}")
|
||
raise RuntimeError(error_msg) from None
|
||
|
||
# ── 6. 检查是否有可用图像 ─────────────────────────────────────────
|
||
if not all_pil_images:
|
||
raise RuntimeError("所有提示词均生成失败，无可用图像输出")
|
||
|
||
# ── 7. PIL → tensor ───────────────────────────────────────────────
|
||
output_tensor = GptImageClient._pil_list_to_tensor(all_pil_images)
|
||
|
||
# ── 8. 完成日志 ───────────────────────────────────────────────────
|
||
elapsed = time.time() - start_time
|
||
print(
|
||
f"[o1key GPT Image] 完成！耗时 {elapsed:.1f}s，"
|
||
f"输出 {output_tensor.shape[0]} 张 "
|
||
f"{output_tensor.shape[2]}×{output_tensor.shape[1]}"
|
||
)
|
||
|
||
return (output_tensor,)
|
||
|
||
finally:
|
||
cls._print_balance(client)
|
||
|
||
@staticmethod
|
||
def _print_balance(client):
|
||
try:
|
||
balance_data = client.query_balance_sync()
|
||
balance_info = client.format_balance_info(balance_data)
|
||
print(f"[o1key GPT Image] {balance_info}")
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
class _LegacyO1keyGPTImageBatch:
|
||
"""
|
||
o1key GPT Image 批量节点
|
||
|
||
复用 BatchNanoBananaPro 的批量思路：
|
||
- 从文件夹批量加载图片
|
||
- 按文件名同名 / 1*N / 不配对 三种模式创建任务
|
||
- 可追加节点手动输入参考图
|
||
- prompt 支持用独占一行 --- 展开为多提示词任务
|
||
- 每个任务调用 GPT Image 客户端并保存到磁盘
|
||
"""
|
||
|
||
PAIRING_MODES = ["按相同图片命名", "1*N", "不配对"]
|
||
IMAGE_FORMATS = ["原始", "JPEG", "PNG", "WebP"]
|
||
MODEL_OPTIONS = GPT_IMAGE_ROUTE_OPTIONS
|
||
QUALITY_OPTIONS = ["高", "中", "低", "自动"]
|
||
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）",
|
||
]
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
optional_inputs = {}
|
||
for image_index in range(1, 10):
|
||
optional_inputs[f"参考图{image_index}"] = ("IMAGE", {
|
||
"tooltip": "追加到每个批量任务末尾的固定参考图。",
|
||
})
|
||
|
||
optional_inputs["遮罩"] = ("MASK", {
|
||
"tooltip": "可选蒙版，会应用到每个任务的第一张参考图；请确保尺寸一致。",
|
||
})
|
||
optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, {
|
||
"default": "不配对",
|
||
"tooltip": "文件夹图片的组合方式；手动参考图只追加，不参与配对。",
|
||
})
|
||
|
||
return {
|
||
"required": {
|
||
"prompt": ("STRING", {
|
||
"default": "",
|
||
"multiline": True,
|
||
"tooltip": "提示词；可用独占一行的 --- 分隔多条批量提示词。",
|
||
}),
|
||
"模型线路": (cls.MODEL_OPTIONS, {
|
||
"default": "畅速",
|
||
}),
|
||
"分辨率": (cls.RESOLUTION_OPTIONS, {
|
||
"default": "智能",
|
||
}),
|
||
"生图数量": ("INT", {
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 8,
|
||
"step": 1,
|
||
"display": "number",
|
||
}),
|
||
"质量": (cls.QUALITY_OPTIONS, {
|
||
"default": "自动",
|
||
}),
|
||
"seed": ("INT", {
|
||
"default": 0,
|
||
"min": 0,
|
||
"max": 2**31 - 1,
|
||
"step": 1,
|
||
"display": "number",
|
||
"control_after_generate": True,
|
||
}),
|
||
"图片格式": (cls.IMAGE_FORMATS, {
|
||
"default": "原始",
|
||
}),
|
||
"文件夹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,
|
||
}),
|
||
"保存路径": ("STRING", {
|
||
"default": "",
|
||
"multiline": False,
|
||
"tooltip": "为空时优先使用 ComfyUI 默认 output 目录。",
|
||
}),
|
||
},
|
||
"optional": optional_inputs,
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("IMAGE",)
|
||
FUNCTION = "process_batch"
|
||
CATEGORY = "o1key/image"
|
||
OUTPUT_NODE = False
|
||
|
||
def _load_folders(self, folders: List[str]) -> List[List[ImageInfo]]:
|
||
image_lists = []
|
||
for folder_index, folder in enumerate(folders, 1):
|
||
if not folder or not folder.strip():
|
||
continue
|
||
try:
|
||
loaded_images = load_images_from_folder(folder)
|
||
if loaded_images:
|
||
image_lists.append(loaded_images)
|
||
except ValueError as error:
|
||
print(f"[o1key GPT Image Batch] 文件夹{folder_index} 加载失败 - {error}")
|
||
return image_lists
|
||
|
||
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 [
|
||
(folder_image,) + tuple(manual_images)
|
||
for folder_image in image_lists[0]
|
||
]
|
||
if image_lists:
|
||
return [(folder_image,) for folder_image in image_lists[0]]
|
||
return []
|
||
|
||
if not image_lists:
|
||
return []
|
||
|
||
if len(image_lists) == 1:
|
||
base_pairs = [(folder_image,) for folder_image 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
|
||
|
||
def _collect_manual_images(self, kwargs) -> List[ImageInfo]:
|
||
manual_images = []
|
||
for image_index in range(1, 10):
|
||
key = f"参考图{image_index}"
|
||
if key not in kwargs or kwargs[key] is None:
|
||
continue
|
||
for tensor_index, image in enumerate(tensor_to_pil(kwargs[key])):
|
||
manual_images.append(ImageInfo(
|
||
image=image,
|
||
filename=f"manual_{image_index}_{tensor_index}",
|
||
extension=".png",
|
||
source_path="",
|
||
))
|
||
return manual_images
|
||
|
||
@staticmethod
|
||
def _pair_to_tensors(pair: Tuple[ImageInfo, ...]) -> List:
|
||
return [pil_to_tensor([image_info.image]) for image_info in pair]
|
||
|
||
@staticmethod
|
||
def _resolve_size(分辨率: str) -> str:
|
||
return _resolve_async_size(分辨率)
|
||
|
||
@staticmethod
|
||
def _resolve_model(模型线路: str) -> str:
|
||
# 直接返回，客户端会映射到 API 值
|
||
return 模型线路
|
||
|
||
@staticmethod
|
||
def _resolve_quality(质量: str) -> str:
|
||
quality_map = {"高": "high", "中": "medium", "低": "low", "自动": "auto"}
|
||
return quality_map.get(质量, "auto")
|
||
|
||
@staticmethod
|
||
def _resolve_output_format(图片格式: str) -> str:
|
||
output_format_map = {
|
||
"JPEG": "jpeg",
|
||
"PNG": "png",
|
||
"WebP": "webp",
|
||
}
|
||
return output_format_map.get(图片格式, "png")
|
||
|
||
@staticmethod
|
||
def _ensure_output_folder(保存路径: str) -> str:
|
||
output_folder = (保存路径 or "").strip()
|
||
if not output_folder and _FOLDER_PATHS_AVAILABLE:
|
||
output_folder = folder_paths.get_output_directory()
|
||
print(f"[o1key GPT Image Batch] 未设置保存路径，使用 ComfyUI 默认 output 目录: {output_folder}")
|
||
|
||
if not output_folder:
|
||
raise ValueError("未设置保存路径，且当前环境无法获取 ComfyUI 默认 output 目录")
|
||
|
||
os.makedirs(output_folder, exist_ok=True)
|
||
test_path = os.path.join(output_folder, ".write_test")
|
||
with open(test_path, "w", encoding="utf-8") as test_file:
|
||
test_file.write("test")
|
||
os.remove(test_path)
|
||
return output_folder
|
||
|
||
@staticmethod
|
||
def _save_images(
|
||
images: List[Image.Image],
|
||
output_folder: str,
|
||
image_format: str,
|
||
base_filename: Optional[str] = None,
|
||
) -> List[str]:
|
||
format_ext_map = {"JPEG": ".jpg", "PNG": ".png", "WebP": ".webp"}
|
||
save_ext = format_ext_map.get(image_format, ".png")
|
||
saved_files = []
|
||
|
||
for image in images:
|
||
if base_filename:
|
||
counter = 0
|
||
while True:
|
||
suffix = "" if counter == 0 else f"+{counter}"
|
||
filename = f"{base_filename}{suffix}{save_ext}"
|
||
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=save_ext,
|
||
)
|
||
|
||
if image_format == "JPEG":
|
||
if image.mode != "RGB":
|
||
image = image.convert("RGB")
|
||
image.save(output_path, quality=100)
|
||
elif image_format == "WebP":
|
||
image.save(output_path, lossless=True)
|
||
else:
|
||
save_image(image, output_path)
|
||
|
||
saved_files.append(output_path)
|
||
|
||
return saved_files
|
||
|
||
def process_batch(
|
||
self,
|
||
prompt: str,
|
||
模型线路: str,
|
||
分辨率: str,
|
||
生图数量: int,
|
||
质量: str,
|
||
seed: int,
|
||
图片格式: str,
|
||
文件夹1: str,
|
||
文件夹2: str,
|
||
文件夹3: str,
|
||
文件夹4: str,
|
||
文件夹5: str,
|
||
保存路径: str = "",
|
||
图片配对模式: str = "不配对",
|
||
遮罩=None,
|
||
**kwargs,
|
||
):
|
||
start_time = time.time()
|
||
client = None
|
||
|
||
try:
|
||
if not prompt or not prompt.strip():
|
||
raise ValueError("提示词不能为空")
|
||
|
||
folders = [文件夹1, 文件夹2, 文件夹3, 文件夹4, 文件夹5]
|
||
if not any(folder and folder.strip() for folder in folders):
|
||
raise ValueError("请至少填写一个文件夹路径，该节点专为批量文件夹处理设计")
|
||
|
||
image_lists = self._load_folders(folders)
|
||
total_folder_images = sum(len(image_list) for image_list in image_lists)
|
||
if total_folder_images == 0:
|
||
raise ValueError("文件夹中未找到任何图片，请检查文件夹路径是否正确")
|
||
|
||
manual_images = self._collect_manual_images(kwargs)
|
||
pairs = self._create_pairs(
|
||
image_lists=image_lists,
|
||
pairing_mode=图片配对模式,
|
||
manual_images=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 batch_prompt in batch_prompts:
|
||
expanded_pairs.append(pair)
|
||
expanded_prompts.append(batch_prompt)
|
||
pairs = expanded_pairs
|
||
prompts_per_task = expanded_prompts
|
||
|
||
total_tasks = len(pairs)
|
||
if batch_prompts:
|
||
print(
|
||
f"[o1key GPT Image Batch] 批量任务 | {图片配对模式} × "
|
||
f"{len(batch_prompts)} 个提示词 | 共 {total_tasks} 任务"
|
||
)
|
||
else:
|
||
print(f"[o1key GPT Image Batch] 批量任务 | {图片配对模式} | 共 {total_tasks} 任务")
|
||
|
||
output_folder = self._ensure_output_folder(保存路径)
|
||
size = self._resolve_size(分辨率)
|
||
model = self._resolve_model(模型线路)
|
||
quality = self._resolve_quality(质量)
|
||
output_format = self._resolve_output_format(图片格式)
|
||
|
||
client = GptImageClient()
|
||
client.base_url = get_base_url_by_route()
|
||
|
||
progress_bar = ProgressBar(total_tasks * 100) if _PROGRESS_BAR_AVAILABLE else None
|
||
results = []
|
||
all_saved_files = []
|
||
|
||
for task_index, pair in enumerate(pairs, 1):
|
||
if _INTERRUPT_AVAILABLE and processing_interrupted():
|
||
print("[o1key GPT Image Batch] 用户取消，已中断批量生成")
|
||
raise InterruptProcessingException()
|
||
|
||
task_prompt = prompts_per_task[task_index - 1] if prompts_per_task else prompt
|
||
base_filename = pair[0].filename if pair else None
|
||
result = {
|
||
"task_index": task_index,
|
||
"success": False,
|
||
"generated_count": 0,
|
||
"saved_files": [],
|
||
"error": None,
|
||
}
|
||
|
||
try:
|
||
pil_images = client.generate_image_async_sync(
|
||
prompt=task_prompt,
|
||
model=model,
|
||
quality=quality,
|
||
size=size,
|
||
n=生图数量,
|
||
seed=seed,
|
||
image_tensor=self._pair_to_tensors(pair),
|
||
mask_tensor=遮罩,
|
||
output_format=output_format,
|
||
progress_callback=_make_node_progress_callback(progress_bar, task_index, total_tasks),
|
||
)
|
||
saved_files = self._save_images(
|
||
images=pil_images,
|
||
output_folder=output_folder,
|
||
image_format=图片格式,
|
||
base_filename=base_filename,
|
||
)
|
||
result["success"] = bool(pil_images)
|
||
result["generated_count"] = len(pil_images)
|
||
result["saved_files"] = saved_files
|
||
all_saved_files.extend(saved_files)
|
||
print(f"[o1key GPT Image Batch] [{task_index}/{total_tasks}] ✓ {base_filename or 'task'}")
|
||
except InterruptProcessingException:
|
||
raise
|
||
except Exception as error:
|
||
error_msg = str(error).split("\n")[0]
|
||
result["error"] = error_msg
|
||
print(f"[o1key GPT Image Batch] [{task_index}/{total_tasks}] ❌ {base_filename or 'task'} → {error_msg}")
|
||
|
||
results.append(result)
|
||
if progress_bar is not None:
|
||
progress_bar.update_absolute(task_index * 100, total_tasks * 100)
|
||
|
||
success_count = sum(1 for result in results if result.get("success", False))
|
||
total_generated = sum(result.get("generated_count", 0) for result in results)
|
||
if success_count == 0:
|
||
raise RuntimeError("所有批量任务均生成失败，无可用图像输出")
|
||
|
||
output_images = []
|
||
for file_path in all_saved_files[-10:]:
|
||
try:
|
||
loaded_image = Image.open(file_path)
|
||
loaded_image.load()
|
||
output_images.append(loaded_image)
|
||
except Exception as error:
|
||
print(f"[o1key GPT Image Batch] 无法加载输出图片 {file_path} - {error}")
|
||
|
||
if not output_images:
|
||
output_images = [Image.new("RGBA", (512, 512), (128, 128, 128, 255))]
|
||
|
||
output_tensor = GptImageClient._pil_list_to_tensor(output_images)
|
||
elapsed = time.time() - start_time
|
||
print("=" * 60)
|
||
print(
|
||
f"[o1key GPT Image Batch] 完成！耗时 {elapsed:.1f}s | "
|
||
f"成功 {success_count}/{total_tasks} | 生成 {total_generated} 张"
|
||
)
|
||
print(f"[o1key GPT Image Batch] 保存路径: {output_folder}")
|
||
if all_saved_files:
|
||
print(f"[o1key GPT Image Batch] 最新保存文件: {all_saved_files[-1]}")
|
||
|
||
failed_results = [result for result in results if not result.get("success", False)]
|
||
if failed_results:
|
||
print(f"[o1key GPT Image Batch] 失败任务: {len(failed_results)} 个")
|
||
for failed_result in failed_results[:3]:
|
||
print(
|
||
f" - #{failed_result.get('task_index')}: "
|
||
f"{failed_result.get('error', '未知错误')}"
|
||
)
|
||
|
||
return (output_tensor,)
|
||
|
||
except ValueError as error:
|
||
if str(error) == "未授权！":
|
||
print("[o1key GPT Image Batch] 请联系作者授权后方可使用！")
|
||
raise ValueError("未授权！") from None
|
||
raise ValueError(str(error)) from None
|
||
except RuntimeError as error:
|
||
raise RuntimeError(str(error)) from None
|
||
finally:
|
||
if client is not None:
|
||
try:
|
||
balance_data = client.query_balance_sync()
|
||
balance_info = client.format_balance_info(balance_data)
|
||
print(f"[o1key GPT Image Batch] {balance_info}")
|
||
except Exception:
|
||
pass
|