""" 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)