- GPT Image 模型下拉选项映射为实际 API 参数(次卡→gpt-image-2-c,按量→gpt-image-2) - 全局:含 "high load" 关键词的错误统一展示为"模型过载,请稍后重试!" - GPT 独立:500 错误展示为"触发内容风控,或服务器繁忙!" Co-Authored-By: Claude Opus 4.7 <[email protected]>
273 lines
12 KiB
Python
273 lines
12 KiB
Python
"""
|
||
o1key GPT Image 节点
|
||
支持 gpt-image-1 / gpt-image-1.5 模型的文生图、图生图、图像编辑(带蒙版)
|
||
"""
|
||
|
||
import time
|
||
from ..clients.gpt_image_client import GptImageClient
|
||
from ..utils.image_utils import parse_batch_prompts
|
||
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||
|
||
try:
|
||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||
_INTERRUPT_AVAILABLE = True
|
||
except ImportError:
|
||
_INTERRUPT_AVAILABLE = False
|
||
processing_interrupted = lambda: False
|
||
InterruptProcessingException = RuntimeError
|
||
|
||
|
||
class O1keyGPTImage:
|
||
"""
|
||
o1key GPT Image 节点
|
||
|
||
功能:
|
||
- 文生图:仅提供 prompt
|
||
- 图生图:提供 prompt + 图片(无遮罩)
|
||
- 图像编辑:提供 prompt + 图片 + 遮罩(白色区域将被替换)
|
||
- 批量模式:prompt 中用单独一行 --- 分隔多条提示词
|
||
|
||
参数:
|
||
- prompt : 文本提示词(多行;用 --- 独占一行分隔批量提示词)
|
||
- 模型 : 模型选择
|
||
- 分辨率 : 图像尺寸(auto 让 API 自动决定)
|
||
- 生图数量 : 每条提示词生成数量 1-8
|
||
- 质量 : 生成质量
|
||
- seed : 随机种子(0 表示不指定)
|
||
- 图片 : 可选参考图(用于图生图或编辑)
|
||
- 遮罩 : 可选蒙版(白色区域将被替换)
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
# 创建9个独立的参考图输入
|
||
optional_inputs = {}
|
||
for i in range(1, 10):
|
||
optional_inputs[f"参考图{i}"] = ("IMAGE", {
|
||
"tooltip": f"Optional reference image {i} for image editing.",
|
||
})
|
||
|
||
optional_inputs["模型"] = ([
|
||
"gpt-image-2-按量",
|
||
"gpt-image-2-次卡",
|
||
], {
|
||
"default": "gpt-image-2-次卡",
|
||
})
|
||
optional_inputs["网络"] = (NETWORK_ROUTE_OPTIONS, {
|
||
"default": "全球加速",
|
||
})
|
||
optional_inputs["分辨率"] = ([
|
||
"智能",
|
||
# ── 1K ──
|
||
"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)",
|
||
# ── 2K ──
|
||
"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)",
|
||
# ── 4K ──
|
||
"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)",
|
||
], {
|
||
"default": "智能",
|
||
"tooltip": "Image size (智能 = API decides)",
|
||
})
|
||
optional_inputs["生图数量"] = ("INT", {
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 8,
|
||
"step": 1,
|
||
"display": "number",
|
||
"tooltip": "How many images to generate per prompt",
|
||
})
|
||
optional_inputs["质量"] = (["高", "中", "低", "自动"], {
|
||
"default": "自动",
|
||
"tooltip": "Image quality: 高=high, 中=medium, 低=low, 自动=auto",
|
||
})
|
||
optional_inputs["seed"] = ("INT", {
|
||
"default": 0,
|
||
"min": 0,
|
||
"max": 2**31 - 1,
|
||
"step": 1,
|
||
"display": "number",
|
||
"control_after_generate": True,
|
||
"tooltip": "Random seed (0 = not specified)",
|
||
})
|
||
optional_inputs["遮罩"] = ("MASK", {
|
||
"tooltip": "Optional mask for inpainting (white areas will be replaced)",
|
||
})
|
||
|
||
return {
|
||
"required": {
|
||
"prompt": ("STRING", {
|
||
"default": "",
|
||
"multiline": True,
|
||
"tooltip": "Text prompt for GPT Image. Use --- on its own line to separate batch prompts.",
|
||
}),
|
||
},
|
||
"optional": optional_inputs,
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("IMAGE",)
|
||
FUNCTION = "generate"
|
||
CATEGORY = "o1key/image"
|
||
OUTPUT_NODE = False
|
||
|
||
def generate(
|
||
self,
|
||
prompt: str,
|
||
模型: str = "gpt-image-2-次卡",
|
||
网络: str = "全球加速",
|
||
分辨率: str = "auto",
|
||
质量: str = "自动",
|
||
生图数量: int = 1,
|
||
seed: int = 0,
|
||
遮罩=None,
|
||
**kwargs,
|
||
):
|
||
"""
|
||
生成图像(文生图 / 图生图 / 图像编辑 / 批量提示词)
|
||
|
||
路由逻辑:
|
||
- 无图片 → generations 接口(文生图)
|
||
- 有图片,无遮罩 → edits 接口(图生图)
|
||
- 有图片,有遮罩 → edits 接口(图像编辑 + 蒙版)
|
||
- prompt 含 --- → 批量模式,逐条调用上述接口
|
||
"""
|
||
start_time = time.time()
|
||
|
||
# ── 0. 收集多参考图输入 ────────────────────────────────────────────────
|
||
reference_tensors = []
|
||
for i in range(1, 10):
|
||
key = f"参考图{i}"
|
||
if key in kwargs and kwargs[key] is not None:
|
||
reference_tensors.append(kwargs[key])
|
||
|
||
图片 = reference_tensors if reference_tensors else None
|
||
|
||
# ── 1. 参数校验 ───────────────────────────────────────────────────────
|
||
if 遮罩 is not None and 图片 is None:
|
||
raise ValueError("提供了遮罩但未提供图片,请同时提供图片和遮罩")
|
||
|
||
# ── 2. 解析分辨率显示值 → API 参数值 ──────────────────────────────────
|
||
size = "auto" if 分辨率 == "智能" else 分辨率.split("(")[0].strip()
|
||
|
||
# ── 2b. 解析模型显示值 → API 参数值 ───────────────────────────────────
|
||
_model_map = {"gpt-image-2-次卡": "gpt-image-2-c", "gpt-image-2-按量": "gpt-image-2"}
|
||
model = _model_map.get(模型, 模型)
|
||
|
||
# ── 2c. 解析质量显示值 → API 参数值 ───────────────────────────────────
|
||
_quality_map = {"高": "high", "中": "medium", "低": "low", "自动": "auto"}
|
||
quality = _quality_map.get(质量, "auto")
|
||
|
||
# ── 3. 创建客户端 ─────────────────────────────────────────────────────
|
||
try:
|
||
client = GptImageClient()
|
||
client.base_url = get_base_url_by_route(网络)
|
||
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 = []
|
||
|
||
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.run_sync(
|
||
prompt=p,
|
||
model=model,
|
||
quality=quality,
|
||
size=size,
|
||
n=生图数量,
|
||
seed=seed,
|
||
image_tensor=图片,
|
||
mask_tensor=遮罩,
|
||
)
|
||
all_pil_images.extend(pil_images)
|
||
snippet = p[:30] + ("..." if len(p) >= 30 else "")
|
||
print(f"[o1key GPT Image] [{idx}/{total}] ✓ {snippet}")
|
||
except InterruptProcessingException:
|
||
raise
|
||
except Exception as e:
|
||
error_msg = str(e).split('\n')[0]
|
||
snippet = p[:30] + ("..." if len(p) >= 30 else "")
|
||
print(f"[o1key GPT Image] [{idx}/{total}] ❌ {snippet} → {error_msg}")
|
||
else:
|
||
# 单提示词模式
|
||
if not prompt or not prompt.strip():
|
||
raise ValueError("提示词不能为空")
|
||
try:
|
||
pil_images = client.run_sync(
|
||
prompt=prompt,
|
||
model=model,
|
||
quality=quality,
|
||
size=size,
|
||
n=生图数量,
|
||
seed=seed,
|
||
image_tensor=图片,
|
||
mask_tensor=遮罩,
|
||
)
|
||
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:
|
||
self._print_balance(client)
|
||
|
||
def _print_balance(self, 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
|