OpenAI 画像生成 API を直接呼び出す ComfyUI カスタムノード。 - OpenAIImage: image 入力の有無で t2i / i2i を自動切替 - 未接続: /v1/images/generations (JSON) - 接続あり: /v1/images/edits (multipart) - gpt-image / dall-e のパラメータ違いを自動調整 - API key はノードパラメータで渡す - moderation / auth / quota などのエラーを分類して ComfyUI に表示 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
259 lines
9.5 KiB
Python
259 lines
9.5 KiB
Python
"""
|
|
ComfyUI custom node that calls the OpenAI Images API directly with a
|
|
user-supplied API key.
|
|
|
|
If `image` is connected -> POST /v1/images/edits (multipart, image-to-image)
|
|
If `image` is NOT connected -> POST /v1/images/generations (JSON, text-to-image)
|
|
|
|
Docs: https://developers.openai.com/api/docs/guides/image-generation
|
|
"""
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
import logging
|
|
|
|
import numpy as np
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
|
|
logger = logging.getLogger("ComfyUI-OpenAI-Image-API")
|
|
|
|
API_BASE = "https://api.openai.com/v1"
|
|
DEFAULT_TIMEOUT = 180 # gpt-image can take up to ~2 min
|
|
|
|
MODERATION_HINTS = (
|
|
"moderation", "moderated", "content policy", "safety", "blocked",
|
|
"violates", "violation", "not allowed", "disallowed", "prohibited",
|
|
"nsfw", "explicit", "harmful", "rejected",
|
|
)
|
|
|
|
|
|
class OpenAIAPIError(RuntimeError):
|
|
def __init__(self, message: str, *, status: int | None = None, code: str | None = None,
|
|
raw: object | None = None, kind: str = "api"):
|
|
super().__init__(message)
|
|
self.status = status
|
|
self.code = code
|
|
self.raw = raw
|
|
self.kind = kind
|
|
|
|
|
|
def _classify_error(status: int | None, code: str | None, message: str) -> str:
|
|
text = f"{code or ''} {message or ''}".lower()
|
|
if any(h in text for h in MODERATION_HINTS):
|
|
return "moderation"
|
|
if status in (401, 403):
|
|
return "auth"
|
|
if status == 429 or (code and "quota" in code.lower()):
|
|
return "rate_limit"
|
|
if status is not None and 400 <= status < 500:
|
|
return "bad_request"
|
|
return "api"
|
|
|
|
|
|
def _parse_error_body(resp: requests.Response) -> tuple[str, str | None, object]:
|
|
raw: object = resp.text
|
|
try:
|
|
data = resp.json()
|
|
raw = data
|
|
except (json.JSONDecodeError, ValueError):
|
|
return (resp.text or f"HTTP {resp.status_code}"), None, raw
|
|
|
|
if isinstance(data, dict):
|
|
err = data.get("error")
|
|
if isinstance(err, dict):
|
|
msg = err.get("message") or err.get("error") or json.dumps(err)
|
|
code = err.get("code") or err.get("type")
|
|
param = err.get("param")
|
|
if param:
|
|
msg = f"{msg} (param: {param})"
|
|
return msg, code, raw
|
|
if isinstance(err, str):
|
|
return err, data.get("code"), raw
|
|
if data.get("message"):
|
|
return data["message"], data.get("code"), raw
|
|
return json.dumps(data)[:500], None, raw
|
|
|
|
|
|
def _format_user_error(node_label: str, message: str, kind: str, status: int | None, code: str | None) -> str:
|
|
prefix = {
|
|
"moderation": f"[{node_label}] OpenAI Content Moderation",
|
|
"auth": f"[{node_label}] OpenAI Auth Error",
|
|
"rate_limit": f"[{node_label}] OpenAI Rate Limit / Quota",
|
|
"bad_request": f"[{node_label}] OpenAI Bad Request",
|
|
"api": f"[{node_label}] OpenAI API Error",
|
|
}.get(kind, f"[{node_label}] OpenAI Error")
|
|
bits = []
|
|
if status is not None:
|
|
bits.append(f"HTTP {status}")
|
|
if code:
|
|
bits.append(f"code={code}")
|
|
suffix = f" ({', '.join(bits)})" if bits else ""
|
|
return f"{prefix}{suffix}: {message}"
|
|
|
|
|
|
def _tensor_to_png_bytes(image_tensor: torch.Tensor) -> bytes:
|
|
"""ComfyUI IMAGE [B,H,W,C] in [0,1] -> PNG bytes (first frame)."""
|
|
arr = image_tensor[0].detach().cpu().numpy()
|
|
arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
|
|
img = Image.fromarray(arr)
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
def _bytes_to_image_tensor(data: bytes) -> torch.Tensor:
|
|
img = Image.open(io.BytesIO(data)).convert("RGB")
|
|
arr = np.array(img).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr).unsqueeze(0)
|
|
|
|
|
|
def _decode_response_image(item: dict) -> torch.Tensor:
|
|
if item.get("b64_json"):
|
|
return _bytes_to_image_tensor(base64.b64decode(item["b64_json"]))
|
|
if item.get("url"):
|
|
r = requests.get(item["url"], timeout=DEFAULT_TIMEOUT)
|
|
r.raise_for_status()
|
|
return _bytes_to_image_tensor(r.content)
|
|
raise RuntimeError("OpenAI API response item has neither url nor b64_json")
|
|
|
|
|
|
def _do_request(method: str, path: str, api_key: str, *,
|
|
json_body: dict | None = None,
|
|
form_data: dict | None = None,
|
|
files: list | None = None) -> dict:
|
|
url = f"{API_BASE}{path}"
|
|
headers = {"Authorization": f"Bearer {api_key}"}
|
|
try:
|
|
resp = requests.request(
|
|
method, url, headers=headers,
|
|
json=json_body, data=form_data, files=files,
|
|
timeout=DEFAULT_TIMEOUT,
|
|
)
|
|
except requests.Timeout:
|
|
raise OpenAIAPIError(f"OpenAI {path} timed out after {DEFAULT_TIMEOUT}s", kind="api")
|
|
except requests.RequestException as e:
|
|
raise OpenAIAPIError(f"OpenAI {path} network error: {e}", kind="api")
|
|
|
|
if resp.ok:
|
|
try:
|
|
return resp.json()
|
|
except ValueError:
|
|
raise OpenAIAPIError(f"OpenAI {path}: invalid JSON in 200 response",
|
|
status=resp.status_code, kind="api")
|
|
|
|
msg, code, raw = _parse_error_body(resp)
|
|
kind = _classify_error(resp.status_code, code, msg)
|
|
logger.warning("OpenAI %s %s -> %s [%s] %s", method, path, resp.status_code, code, msg)
|
|
raise OpenAIAPIError(msg, status=resp.status_code, code=code, raw=raw, kind=kind)
|
|
|
|
|
|
class OpenAIImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"api_key": ("STRING", {"default": "", "multiline": False, "password": True}),
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"model": (
|
|
[
|
|
"gpt-image-1",
|
|
"gpt-image-1-mini",
|
|
"gpt-image-2",
|
|
"gpt-image-1.5",
|
|
"dall-e-3",
|
|
"dall-e-2",
|
|
],
|
|
{"default": "gpt-image-1"},
|
|
),
|
|
"size": (
|
|
[
|
|
"auto",
|
|
"1024x1024", "1536x1024", "1024x1536", "2048x2048",
|
|
"1792x1024", "1024x1792",
|
|
"512x512", "256x256",
|
|
],
|
|
{"default": "auto"},
|
|
),
|
|
"quality": (
|
|
["auto", "low", "medium", "high", "standard", "hd"],
|
|
{"default": "auto"},
|
|
),
|
|
"background": (
|
|
["auto", "opaque", "transparent"],
|
|
{"default": "auto"},
|
|
),
|
|
"n": ("INT", {"default": 1, "min": 1, "max": 10}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "revised_prompt")
|
|
FUNCTION = "run"
|
|
CATEGORY = "OpenAI"
|
|
|
|
def run(self, api_key, prompt, model, size, quality, background, n, image=None):
|
|
if not api_key:
|
|
raise ValueError("[OpenAI Image] api_key is required")
|
|
if not prompt or not prompt.strip():
|
|
raise ValueError("[OpenAI Image] prompt is required")
|
|
|
|
is_gpt_image = model.startswith("gpt-image")
|
|
is_i2i = image is not None
|
|
|
|
try:
|
|
if is_i2i:
|
|
data = self._call_edits(api_key, prompt, model, size, quality, background, n, image, is_gpt_image)
|
|
label = "OpenAI Image i2i"
|
|
else:
|
|
data = self._call_generations(api_key, prompt, model, size, quality, background, n, is_gpt_image)
|
|
label = "OpenAI Image t2i"
|
|
except OpenAIAPIError as e:
|
|
label = "OpenAI Image i2i" if is_i2i else "OpenAI Image t2i"
|
|
raise RuntimeError(_format_user_error(label, str(e), e.kind, e.status, e.code)) from None
|
|
|
|
items = data.get("data") or []
|
|
if not items:
|
|
raise RuntimeError(f"[{label}] OpenAI returned empty data: {data}")
|
|
|
|
try:
|
|
tensors = [_decode_response_image(it) for it in items]
|
|
except requests.RequestException as e:
|
|
raise RuntimeError(f"[{label}] failed to download result image: {e}") from None
|
|
out = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0)
|
|
revised = next((it.get("revised_prompt", "") for it in items if it.get("revised_prompt")), "")
|
|
return (out, revised)
|
|
|
|
@staticmethod
|
|
def _call_generations(api_key, prompt, model, size, quality, background, n, is_gpt_image) -> dict:
|
|
body: dict = {"model": model, "prompt": prompt, "n": int(n)}
|
|
if size != "auto":
|
|
body["size"] = size
|
|
if quality != "auto":
|
|
body["quality"] = quality
|
|
if is_gpt_image and background != "auto":
|
|
body["background"] = background
|
|
if not is_gpt_image:
|
|
body["response_format"] = "b64_json"
|
|
return _do_request("POST", "/images/generations", api_key, json_body=body)
|
|
|
|
@staticmethod
|
|
def _call_edits(api_key, prompt, model, size, quality, background, n, image, is_gpt_image) -> dict:
|
|
png_bytes = _tensor_to_png_bytes(image)
|
|
files = [("image", ("image.png", png_bytes, "image/png"))]
|
|
form: dict = {"model": model, "prompt": prompt, "n": str(int(n))}
|
|
if size != "auto":
|
|
form["size"] = size
|
|
if quality != "auto":
|
|
form["quality"] = quality
|
|
if is_gpt_image and background != "auto":
|
|
form["background"] = background
|
|
if not is_gpt_image:
|
|
form["response_format"] = "b64_json"
|
|
return _do_request("POST", "/images/edits", api_key, form_data=form, files=files)
|