418 lines
12 KiB
Python
418 lines
12 KiB
Python
import base64
|
|
import io
|
|
import time
|
|
|
|
import numpy as np
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
GPT_IMAGE_EDIT_MAX_PIXELS = 2048 * 2048
|
|
GPT_IMAGE_2_PRESET_SIZES = [
|
|
"1024x1024",
|
|
"1536x1024",
|
|
"1024x1536",
|
|
"2048x2048",
|
|
"2048x1152",
|
|
"3840x2160",
|
|
"2160x3840",
|
|
"Custom",
|
|
"auto",
|
|
]
|
|
|
|
SUPPORTED_OPENAI_IMAGE_MIME_TYPES = {
|
|
"image/jpeg": ("image/jpeg", "jpg"),
|
|
"image/jpg": ("image/jpeg", "jpg"),
|
|
"image/png": ("image/png", "png"),
|
|
"image/webp": ("image/webp", "webp"),
|
|
"geminimimetype.image_jpeg": ("image/jpeg", "jpg"),
|
|
"geminimimetype.image_png": ("image/png", "png"),
|
|
"geminimimetype.image_webp": ("image/webp", "webp"),
|
|
"image_jpeg": ("image/jpeg", "jpg"),
|
|
"image_png": ("image/png", "png"),
|
|
"image_webp": ("image/webp", "webp"),
|
|
}
|
|
|
|
# USD/CNY reference rate used for RMB estimation in node info output.
|
|
# Based on the most recent publicly indexed April 2026 daily spot pages we could verify:
|
|
# 2026-04-03: 1 USD ≈ 6.88154 CNY
|
|
USD_TO_CNY_RATE = 6.8815
|
|
USD_TO_COMFY_CREDITS = 211.0
|
|
|
|
|
|
def empty_image():
|
|
return torch.zeros((1, 512, 512, 3))
|
|
|
|
|
|
def tensor_to_bytes(tensor, max_pixels=GPT_IMAGE_EDIT_MAX_PIXELS):
|
|
if tensor.ndim == 4:
|
|
tensor = tensor[0]
|
|
|
|
height, width = tensor.shape[:2]
|
|
current_pixels = height * width
|
|
if current_pixels > max_pixels:
|
|
scale = (max_pixels / current_pixels) ** 0.5
|
|
new_height = max(1, int(height * scale))
|
|
new_width = max(1, int(width * scale))
|
|
array = np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8)
|
|
image = Image.fromarray(array)
|
|
image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
else:
|
|
array = np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8)
|
|
image = Image.fromarray(array)
|
|
|
|
buffered = io.BytesIO()
|
|
image.save(buffered, format="PNG")
|
|
buffered.seek(0)
|
|
return buffered
|
|
|
|
|
|
def load_image_to_tensor(image_data):
|
|
try:
|
|
if isinstance(image_data, str) and image_data.startswith("http"):
|
|
response = requests.get(image_data, timeout=30)
|
|
response.raise_for_status()
|
|
img = Image.open(io.BytesIO(response.content))
|
|
else:
|
|
img = Image.open(io.BytesIO(image_data))
|
|
|
|
img = img.convert("RGB")
|
|
image_np = np.array(img).astype(np.float32) / 255.0
|
|
return torch.from_numpy(image_np).unsqueeze(0)
|
|
except Exception:
|
|
return empty_image()
|
|
|
|
|
|
def mask_to_png_bytes(mask):
|
|
mask_tensor = mask
|
|
if mask_tensor.ndim == 3:
|
|
mask_tensor = mask_tensor[0]
|
|
|
|
mask_h, mask_w = mask_tensor.shape
|
|
rgba = np.zeros((mask_h, mask_w, 4), dtype=np.uint8)
|
|
mask_np = mask_tensor.cpu().numpy()
|
|
alpha = ((1.0 - mask_np) * 255).astype(np.uint8)
|
|
rgba[..., 3] = alpha
|
|
|
|
mask_pil = Image.fromarray(rgba)
|
|
mask_buffer = io.BytesIO()
|
|
mask_pil.save(mask_buffer, format="PNG")
|
|
mask_buffer.seek(0)
|
|
return mask_buffer
|
|
|
|
|
|
def build_openai_edit_files(image, mask=None):
|
|
files = []
|
|
batch_size = image.shape[0] if image.ndim == 4 else 1
|
|
|
|
if image.ndim == 4:
|
|
for i in range(batch_size):
|
|
image_bytes = tensor_to_bytes(image[i : i + 1]).getvalue()
|
|
field_name = "image" if batch_size == 1 else "image[]"
|
|
files.append((field_name, (f"image_{i}.png", image_bytes, "image/png")))
|
|
else:
|
|
image_bytes = tensor_to_bytes(image).getvalue()
|
|
files.append(("image", ("image_0.png", image_bytes, "image/png")))
|
|
|
|
if mask is not None:
|
|
if image.ndim == 4 and image.shape[0] != 1:
|
|
raise ValueError("Mask editing only supports a single input image")
|
|
mask_bytes = mask_to_png_bytes(mask).getvalue()
|
|
files.append(("mask", ("mask.png", mask_bytes, "image/png")))
|
|
|
|
return files
|
|
|
|
|
|
def normalize_openai_image_mime_type(mime_type):
|
|
enum_value = getattr(mime_type, "value", mime_type)
|
|
if enum_value is None:
|
|
return ("image/png", "png")
|
|
|
|
mime_key = str(enum_value).strip().lower()
|
|
if not mime_key:
|
|
return ("image/png", "png")
|
|
|
|
return SUPPORTED_OPENAI_IMAGE_MIME_TYPES.get(mime_key, (None, None))
|
|
|
|
|
|
def append_gemini_input_files_to_openai_files(files, input_files):
|
|
if not input_files:
|
|
return files
|
|
|
|
image_index = sum(1 for field_name, _ in files if field_name in ("image", "image[]"))
|
|
for file_part in input_files:
|
|
if not isinstance(file_part, dict):
|
|
continue
|
|
|
|
inline_data = file_part.get("inlineData") or file_part.get("inline_data")
|
|
if not isinstance(inline_data, dict):
|
|
continue
|
|
|
|
mime_type, extension = normalize_openai_image_mime_type(
|
|
inline_data.get("mimeType") or inline_data.get("mime_type")
|
|
)
|
|
if not mime_type:
|
|
continue
|
|
|
|
data_b64 = inline_data.get("data")
|
|
if not data_b64:
|
|
continue
|
|
|
|
image_bytes = base64.b64decode(data_b64)
|
|
field_name = "image" if image_index == 0 and len(files) == 0 else "image[]"
|
|
files.append((field_name, (f"image_{image_index}.{extension}", image_bytes, mime_type)))
|
|
image_index += 1
|
|
|
|
return files
|
|
|
|
|
|
def normalize_mime_type(mime_type):
|
|
if hasattr(mime_type, "value"):
|
|
return mime_type.value
|
|
if isinstance(mime_type, str) and "/" not in mime_type and "." in mime_type:
|
|
return mime_type.rsplit(".", 1)[-1].replace("_", "/")
|
|
return mime_type
|
|
|
|
|
|
def build_openai_edit_files_from_sources(image=None, mask=None, input_files=None):
|
|
files = []
|
|
|
|
if image is not None:
|
|
files.extend(build_openai_edit_files(image, mask))
|
|
elif mask is not None:
|
|
raise ValueError("Mask editing requires an input image")
|
|
|
|
append_gemini_input_files_to_openai_files(files, input_files)
|
|
|
|
image_count = sum(1 for field_name, _ in files if field_name in ("image", "image[]"))
|
|
if image_count == 0:
|
|
raise ValueError("No valid images were provided")
|
|
|
|
if image_count > 1:
|
|
normalized_files = []
|
|
for field_name, file_value in files:
|
|
if field_name in ("image", "image[]"):
|
|
normalized_files.append(("image[]", file_value))
|
|
else:
|
|
normalized_files.append((field_name, file_value))
|
|
files = normalized_files
|
|
|
|
if mask is not None and image is None:
|
|
raise ValueError("Mask editing does not support files-only input")
|
|
|
|
if mask is not None and image_count != 1:
|
|
raise ValueError("Mask editing only supports a single input image")
|
|
|
|
return files
|
|
|
|
|
|
def post_with_retries(url, headers, *, json_payload=None, data=None, files=None, timeout=300, retries=3, retry_delay=1.0):
|
|
last_error = None
|
|
for attempt in range(retries):
|
|
try:
|
|
response = requests.post(
|
|
url,
|
|
headers=headers,
|
|
json=json_payload,
|
|
data=data,
|
|
files=files,
|
|
timeout=timeout,
|
|
)
|
|
return response
|
|
except (requests.exceptions.SSLError, requests.exceptions.ConnectionError) as exc:
|
|
last_error = exc
|
|
if attempt == retries - 1:
|
|
raise
|
|
time.sleep(retry_delay * (attempt + 1))
|
|
raise last_error
|
|
|
|
|
|
def decode_response_images(result):
|
|
data_list = result.get("data", [])
|
|
image_tensors = []
|
|
for item in data_list:
|
|
if "url" in item:
|
|
img_tensor = load_image_to_tensor(item["url"])
|
|
image_tensors.append(img_tensor)
|
|
elif "b64_json" in item:
|
|
img_data = base64.b64decode(item["b64_json"])
|
|
img_tensor = load_image_to_tensor(img_data)
|
|
image_tensors.append(img_tensor)
|
|
return image_tensors
|
|
|
|
|
|
def estimate_price_range_1_or_1_5(model, quality, n):
|
|
ranges = {
|
|
"gpt-image-1": {
|
|
"low": (0.011, 0.016),
|
|
"medium": (0.042, 0.063),
|
|
"high": (0.167, 0.250),
|
|
},
|
|
"gpt-image-1.5": {
|
|
"low": (0.009, 0.013),
|
|
"medium": (0.034, 0.050),
|
|
"high": (0.133, 0.200),
|
|
},
|
|
}
|
|
model_ranges = ranges.get(model, ranges["gpt-image-1"])
|
|
min_cost, max_cost = model_ranges.get(quality, model_ranges["medium"])
|
|
return (min_cost * n, max_cost * n)
|
|
|
|
|
|
def calculate_token_price_1_or_1_5(model, n):
|
|
if model == "gpt-image-1":
|
|
cost = (1000 * 10 + 2000 * 40) / 1_000_000
|
|
else:
|
|
cost = (1000 * 8 + 2000 * 32) / 1_000_000
|
|
return cost * n
|
|
|
|
|
|
def estimate_price_range_2(quality, n):
|
|
ranges = {
|
|
"auto": (0.005, 0.250),
|
|
"low": (0.005, 0.030),
|
|
"medium": (0.041, 0.120),
|
|
"high": (0.165, 0.250),
|
|
}
|
|
min_cost, max_cost = ranges.get(quality, ranges["auto"])
|
|
return (min_cost * n, max_cost * n)
|
|
|
|
|
|
def calculate_token_price_2(n):
|
|
cost = (1000 * 8 + 2000 * 30) / 1_000_000
|
|
return cost * n
|
|
|
|
|
|
def _custom_size_enabled(use_custom_size):
|
|
if isinstance(use_custom_size, str):
|
|
return use_custom_size.strip().lower() in ("true", "1", "yes", "on")
|
|
return bool(use_custom_size)
|
|
|
|
|
|
def resolve_gpt_image_2_size(
|
|
size,
|
|
custom_width,
|
|
custom_height,
|
|
size_input=None,
|
|
width_input=None,
|
|
height_input=None,
|
|
use_custom_size=True,
|
|
):
|
|
if not _custom_size_enabled(use_custom_size):
|
|
return size
|
|
|
|
custom_requested = str(size).strip().lower() == "custom"
|
|
|
|
if size_input is not None:
|
|
size_input = str(size_input).strip()
|
|
if size_input:
|
|
return size_input
|
|
|
|
if width_input is not None and height_input is not None:
|
|
try:
|
|
width_input = int(width_input)
|
|
height_input = int(height_input)
|
|
except (TypeError, ValueError):
|
|
width_input = 0
|
|
height_input = 0
|
|
|
|
if width_input > 0 and height_input > 0:
|
|
return f"{width_input}x{height_input}"
|
|
|
|
if custom_requested or custom_width > 0 or custom_height > 0:
|
|
if custom_width <= 0 or custom_height <= 0:
|
|
raise ValueError("Custom size requires both custom_width and custom_height to be greater than 0")
|
|
return f"{custom_width}x{custom_height}"
|
|
return size
|
|
|
|
|
|
def usd_to_cny(usd_amount):
|
|
return usd_amount * USD_TO_CNY_RATE
|
|
|
|
|
|
def format_rmb_range(min_usd, max_usd):
|
|
return f"约￥{usd_to_cny(min_usd):.2f}-￥{usd_to_cny(max_usd):.2f}"
|
|
|
|
|
|
def format_rmb_value(usd_amount):
|
|
return f"约￥{usd_to_cny(usd_amount):.2f}"
|
|
|
|
|
|
def usd_to_comfy_credits(usd_amount):
|
|
return usd_amount * USD_TO_COMFY_CREDITS
|
|
|
|
|
|
def format_credits_value(usd_amount):
|
|
return f"约 {usd_to_comfy_credits(usd_amount):.2f} credits"
|
|
|
|
|
|
def calculate_exact_cost_gpt_image_2_from_usage(result):
|
|
usage = result.get("usage") or {}
|
|
input_details = usage.get("input_tokens_details") or {}
|
|
|
|
image_tokens = input_details.get("image_tokens") or 0
|
|
text_tokens = input_details.get("text_tokens") or 0
|
|
output_tokens = usage.get("output_tokens") or 0
|
|
|
|
if image_tokens == 0 and text_tokens == 0 and output_tokens == 0:
|
|
return None
|
|
|
|
usd = (
|
|
(image_tokens * 8.0) +
|
|
(text_tokens * 5.0) +
|
|
(output_tokens * 30.0)
|
|
) / 1_000_000.0
|
|
|
|
return {
|
|
"usd": usd,
|
|
"image_tokens": image_tokens,
|
|
"text_tokens": text_tokens,
|
|
"output_tokens": output_tokens,
|
|
}
|
|
|
|
|
|
def validate_gpt_image_2_size(
|
|
size,
|
|
custom_width=0,
|
|
custom_height=0,
|
|
size_input=None,
|
|
width_input=None,
|
|
height_input=None,
|
|
use_custom_size=True,
|
|
):
|
|
resolved_size = resolve_gpt_image_2_size(
|
|
size,
|
|
custom_width,
|
|
custom_height,
|
|
size_input=size_input,
|
|
width_input=width_input,
|
|
height_input=height_input,
|
|
use_custom_size=use_custom_size,
|
|
)
|
|
if resolved_size == "auto":
|
|
return resolved_size
|
|
|
|
try:
|
|
width_str, height_str = resolved_size.lower().split("x")
|
|
width = int(width_str)
|
|
height = int(height_str)
|
|
except Exception as exc:
|
|
raise ValueError(f"Invalid size: {resolved_size}") from exc
|
|
|
|
if width > 3840 or height > 3840:
|
|
raise ValueError("Width and height must not exceed 3840")
|
|
if width % 16 != 0 or height % 16 != 0:
|
|
raise ValueError("Width and height must be multiples of 16")
|
|
|
|
shorter_edge = min(width, height)
|
|
longer_edge = max(width, height)
|
|
if shorter_edge == 0 or (longer_edge / shorter_edge) > 3:
|
|
raise ValueError("Aspect ratio must not exceed 3:1")
|
|
|
|
total_pixels = width * height
|
|
if total_pixels < 655360 or total_pixels > 8294400:
|
|
raise ValueError("Total pixels must be between 655360 and 8294400")
|
|
|
|
return resolved_size
|