256 lines
9.2 KiB
Python
256 lines
9.2 KiB
Python
"""
|
|
PD: OpenAI GPT Image 2 Node (ComfyUI Auth Token)
|
|
"""
|
|
|
|
import time
|
|
|
|
import requests
|
|
import torch
|
|
|
|
from ._image_gpt_common import (
|
|
GPT_IMAGE_2_PRESET_SIZES,
|
|
build_openai_edit_files_from_sources,
|
|
calculate_exact_cost_gpt_image_2_from_usage,
|
|
calculate_token_price_2,
|
|
decode_response_images,
|
|
empty_image,
|
|
estimate_price_range_2,
|
|
format_credits_value,
|
|
format_rmb_range,
|
|
format_rmb_value,
|
|
post_with_retries,
|
|
resolve_gpt_image_2_size,
|
|
validate_gpt_image_2_size,
|
|
)
|
|
|
|
|
|
class PDOpenAIGPTImage2AuthToken:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"auth_token": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"placeholder": "Paste your Auth Token here"
|
|
}),
|
|
"prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Describe what you want to generate or edit"
|
|
}),
|
|
"model": (["gpt-image-2"], {
|
|
"default": "gpt-image-2",
|
|
"tooltip": "Locked GPT Image 2 model"
|
|
}),
|
|
"quality": (["auto", "low", "medium", "high"], {
|
|
"default": "auto",
|
|
"tooltip": "GPT Image 2 quality"
|
|
}),
|
|
"background": (["auto", "opaque"], {
|
|
"default": "auto",
|
|
"tooltip": "GPT Image 2 does not support transparent background"
|
|
}),
|
|
"size": (GPT_IMAGE_2_PRESET_SIZES, {
|
|
"default": "1024x1024",
|
|
"tooltip": "Preset output size, or Custom with custom_width/custom_height"
|
|
}),
|
|
"use_custom_size": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Enable custom_width/custom_height size override."
|
|
}),
|
|
"custom_width": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 3840,
|
|
"step": 16,
|
|
"tooltip": "Custom width. Must be paired with custom_height."
|
|
}),
|
|
"custom_height": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 3840,
|
|
"step": 16,
|
|
"tooltip": "Custom height. Must be paired with custom_width."
|
|
}),
|
|
"num_images": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 8,
|
|
"step": 1,
|
|
"tooltip": "How many images to generate"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xFFFFFFFFFFFFFFFF
|
|
}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE", {"default": None, "tooltip": "Input image for editing"}),
|
|
"mask": ("MASK", {"default": None, "tooltip": "Mask for inpainting"}),
|
|
"files": ("GEMINI_INPUT_FILES", {"default": None, "tooltip": "Packed multi-image files input"}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "STRING")
|
|
RETURN_NAMES = ("image", "info", "price")
|
|
FUNCTION = "generate_image"
|
|
CATEGORY = "PD_Tools/Image_Generation"
|
|
|
|
def generate_image(
|
|
self,
|
|
auth_token,
|
|
prompt,
|
|
model,
|
|
quality,
|
|
background,
|
|
size,
|
|
use_custom_size,
|
|
custom_width,
|
|
custom_height,
|
|
num_images,
|
|
seed,
|
|
unique_id,
|
|
image=None,
|
|
mask=None,
|
|
files=None,
|
|
):
|
|
if not auth_token or not auth_token.strip():
|
|
error_msg = "Error: Auth Token is required"
|
|
print(f"[IMAGE_GPT2_AUTHTOKEN] {error_msg}")
|
|
return (empty_image(), error_msg, "")
|
|
|
|
try:
|
|
resolved_size = validate_gpt_image_2_size(
|
|
size,
|
|
custom_width,
|
|
custom_height,
|
|
use_custom_size=use_custom_size,
|
|
)
|
|
except Exception as e:
|
|
return (empty_image(), f"Error: {str(e)}", "")
|
|
|
|
auth_token = auth_token.strip()
|
|
if auth_token.lower().startswith("bearer "):
|
|
auth_token = auth_token[7:].strip()
|
|
|
|
base_url = "https://api.comfy.org"
|
|
is_edit_mode = image is not None or files is not None
|
|
start_time = time.time()
|
|
|
|
try:
|
|
if is_edit_mode:
|
|
url = f"{base_url}/proxy/openai/images/edits"
|
|
request_files = build_openai_edit_files_from_sources(image=image, mask=mask, input_files=files)
|
|
|
|
data = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"n": str(num_images),
|
|
"size": resolved_size,
|
|
"quality": quality,
|
|
"background": background,
|
|
"moderation": "low"
|
|
}
|
|
headers = {
|
|
"Authorization": f"Bearer {auth_token}",
|
|
"Connection": "close",
|
|
"User-Agent": "ComfyUI-PD-Node/2.1"
|
|
}
|
|
response = post_with_retries(url, headers, data=data, files=request_files, timeout=300)
|
|
else:
|
|
url = f"{base_url}/proxy/openai/images/generations"
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"n": num_images,
|
|
"size": resolved_size,
|
|
"quality": quality,
|
|
"background": background,
|
|
"moderation": "low"
|
|
}
|
|
headers = {
|
|
"Authorization": f"Bearer {auth_token}",
|
|
"Content-Type": "application/json",
|
|
"Connection": "close",
|
|
"User-Agent": "ComfyUI-PD-Node/2.1"
|
|
}
|
|
response = post_with_retries(url, headers, json_payload=payload, timeout=300)
|
|
|
|
if response.status_code != 200:
|
|
error_msg = f"API Error {response.status_code}: {response.text[:200]}"
|
|
print(f"[IMAGE_GPT2_AUTHTOKEN] {error_msg}")
|
|
return (empty_image(), error_msg, "")
|
|
|
|
result = response.json()
|
|
image_tensors = decode_response_images(result)
|
|
if not image_tensors:
|
|
error_msg = "Error: Failed to process output images"
|
|
print(f"[IMAGE_GPT2_AUTHTOKEN] {error_msg}")
|
|
return (empty_image(), error_msg, "")
|
|
|
|
final_image = torch.cat(image_tensors, dim=0)
|
|
duration = time.time() - start_time
|
|
min_cost, max_cost = estimate_price_range_2(quality, num_images)
|
|
token_cost = calculate_token_price_2(num_images)
|
|
size_note = resolve_gpt_image_2_size(
|
|
size,
|
|
custom_width,
|
|
custom_height,
|
|
use_custom_size=use_custom_size,
|
|
)
|
|
exact_cost = calculate_exact_cost_gpt_image_2_from_usage(result)
|
|
|
|
mode_label = "编辑" if is_edit_mode else "生成"
|
|
info_lines = [
|
|
f"模型: {model}",
|
|
"状态: 成功",
|
|
f"模式: {mode_label}",
|
|
f"耗时: {duration:.2f} 秒",
|
|
f"输出数量: {len(image_tensors)}",
|
|
f"尺寸预设: {size}",
|
|
f"实际尺寸: {size_note}",
|
|
f"质量: {quality}",
|
|
f"背景: {background}",
|
|
]
|
|
|
|
if exact_cost is not None:
|
|
price_lines = [
|
|
f"Tokens: 图像:{exact_cost['image_tokens']} | 文本:{exact_cost['text_tokens']} | 输出:{exact_cost['output_tokens']}",
|
|
f"实际价格(USD): ${exact_cost['usd']:.6f}",
|
|
f"实际价格(RMB): {format_rmb_value(exact_cost['usd'])}",
|
|
f"实际价格(Credits): {format_credits_value(exact_cost['usd'])}",
|
|
]
|
|
else:
|
|
price_lines = [
|
|
"Tokens: 未返回精确 usage，以下为估算",
|
|
f"估算价格(USD): ~${min_cost:.3f}-${max_cost:.3f}",
|
|
f"估算价格(RMB): {format_rmb_range(min_cost, max_cost)}",
|
|
f"Token 估算(USD): ${token_cost:.4f}",
|
|
f"Token 估算(RMB): {format_rmb_value(token_cost)}",
|
|
]
|
|
|
|
return (final_image, "\n".join(info_lines), "\n".join(price_lines))
|
|
|
|
except requests.exceptions.Timeout:
|
|
return (empty_image(), "Error: Request timeout (300s)", "")
|
|
except requests.exceptions.ConnectionError as e:
|
|
return (empty_image(), f"Error: Connection failed - {str(e)}", "")
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
return (empty_image(), f"Error: {str(e)}", "")
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PDOpenAIGPTImage2AuthToken": PDOpenAIGPTImage2AuthToken
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PDOpenAIGPTImage2AuthToken": "PD: GPT Image 2 (comfyui_AuthToken)"
|
|
}
|