325 lines
9.9 KiB
Python
325 lines
9.9 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from urllib.parse import urlparse
|
|
|
|
import numpy as np
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
DEFAULT_TIMEOUT_SECONDS = 600
|
|
CONNECT_TIMEOUT_SECONDS = 30
|
|
IMAGE_B64_KEYS = ("b64_json", "image_b64", "b64", "base64", "image_base64")
|
|
IMAGE_URL_KEYS = ("url", "image_url", "output_url", "result_url", "download_url")
|
|
NESTED_IMAGE_KEYS = ("image", "output", "result")
|
|
|
|
|
|
def _parse_json_object(value, field_name):
|
|
value = (value or "").strip()
|
|
if not value:
|
|
return {}
|
|
parsed = json.loads(value)
|
|
if not isinstance(parsed, dict):
|
|
raise ValueError(f"{field_name} must be a JSON object.")
|
|
return parsed
|
|
|
|
|
|
def _api_key_from_input(api_key):
|
|
api_key = api_key.strip() or os.environ.get("OPENAI_API_KEY", "").strip()
|
|
if not api_key:
|
|
raise ValueError("Missing API key. Set OPENAI_API_KEY or fill the api_key input.")
|
|
return api_key
|
|
|
|
|
|
def _endpoint_url(api_base, endpoint_path):
|
|
api_base = api_base.rstrip("/")
|
|
endpoint_path = endpoint_path if endpoint_path.startswith("/") else f"/{endpoint_path}"
|
|
return f"{api_base}{endpoint_path}"
|
|
|
|
|
|
def _request_timeout(timeout_seconds):
|
|
read_timeout = max(int(timeout_seconds), CONNECT_TIMEOUT_SECONDS)
|
|
return (CONNECT_TIMEOUT_SECONDS, read_timeout)
|
|
|
|
|
|
def _format_timeout_message(operation, timeout_seconds):
|
|
return (
|
|
f"{operation} timed out after {timeout_seconds}s. "
|
|
"Image APIs and proxy endpoints can take several minutes. "
|
|
"Increase timeout_seconds, reduce n, or check whether the provider requires async polling."
|
|
)
|
|
|
|
|
|
def _post_json(url, headers, body, timeout_seconds, operation):
|
|
try:
|
|
return requests.post(
|
|
url,
|
|
headers=headers,
|
|
json=body,
|
|
timeout=_request_timeout(timeout_seconds),
|
|
)
|
|
except requests.exceptions.Timeout as exc:
|
|
raise TimeoutError(_format_timeout_message(operation, timeout_seconds)) from exc
|
|
except requests.exceptions.RequestException as exc:
|
|
raise RuntimeError(f"{operation} request failed before receiving a response: {exc}") from exc
|
|
|
|
|
|
def _get_image_url(image_url, timeout_seconds):
|
|
try:
|
|
response = requests.get(
|
|
image_url,
|
|
headers={"Connection": "close"},
|
|
timeout=_request_timeout(timeout_seconds),
|
|
)
|
|
response.raise_for_status()
|
|
return response
|
|
except requests.exceptions.Timeout as exc:
|
|
raise TimeoutError(_format_timeout_message("Image URL download", timeout_seconds)) from exc
|
|
except requests.exceptions.RequestException as exc:
|
|
raise RuntimeError(f"Image URL download failed: {exc}") from exc
|
|
|
|
|
|
def _looks_like_url(value):
|
|
if not isinstance(value, str):
|
|
return False
|
|
parsed = urlparse(value)
|
|
return parsed.scheme in {"http", "https"} and bool(parsed.netloc)
|
|
|
|
|
|
def _normalize_data_items(payload):
|
|
if isinstance(payload, list):
|
|
return payload
|
|
if not isinstance(payload, dict):
|
|
return [payload]
|
|
|
|
data = payload.get("data")
|
|
if isinstance(data, list):
|
|
return data
|
|
if isinstance(data, dict):
|
|
return [data]
|
|
|
|
for key in ("images", "outputs", "results"):
|
|
value = payload.get(key)
|
|
if isinstance(value, list):
|
|
return value
|
|
if isinstance(value, dict):
|
|
return [value]
|
|
|
|
return [payload]
|
|
|
|
|
|
def _get_first_value(mapping, keys):
|
|
for key in keys:
|
|
value = mapping.get(key)
|
|
if value:
|
|
return value
|
|
return None
|
|
|
|
|
|
def _extract_image_value(item):
|
|
if isinstance(item, str):
|
|
if _looks_like_url(item):
|
|
return "url", item
|
|
return "b64", item
|
|
|
|
if not isinstance(item, dict):
|
|
return None, None
|
|
|
|
b64_image = _get_first_value(item, IMAGE_B64_KEYS)
|
|
if b64_image:
|
|
return "b64", b64_image
|
|
|
|
image_url = _get_first_value(item, IMAGE_URL_KEYS)
|
|
if image_url:
|
|
return "url", image_url
|
|
|
|
for key in NESTED_IMAGE_KEYS:
|
|
value = item.get(key)
|
|
if isinstance(value, dict):
|
|
image_type, image_value = _extract_image_value(value)
|
|
if image_value:
|
|
return image_type, image_value
|
|
if isinstance(value, str):
|
|
if _looks_like_url(value):
|
|
return "url", value
|
|
return "b64", value
|
|
|
|
return None, None
|
|
|
|
|
|
def _b64_to_tensor(b64_image):
|
|
if "," in b64_image and b64_image.strip().startswith("data:"):
|
|
b64_image = b64_image.split(",", 1)[1]
|
|
image_bytes = base64.b64decode(b64_image)
|
|
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
|
array = np.asarray(image).astype(np.float32) / 255.0
|
|
return torch.from_numpy(array)
|
|
|
|
|
|
def _response_to_tensor_batch(payload, timeout_seconds):
|
|
items = _normalize_data_items(payload)
|
|
if not items:
|
|
raise RuntimeError(f"Image API response has no image items: {payload}")
|
|
|
|
images = []
|
|
for item in items:
|
|
image_type, image_value = _extract_image_value(item)
|
|
if image_type == "b64":
|
|
images.append(_b64_to_tensor(image_value))
|
|
continue
|
|
|
|
if image_type == "url":
|
|
response = _get_image_url(image_value, timeout_seconds)
|
|
image = Image.open(io.BytesIO(response.content)).convert("RGB")
|
|
array = np.asarray(image).astype(np.float32) / 255.0
|
|
images.append(torch.from_numpy(array))
|
|
continue
|
|
|
|
raise RuntimeError(
|
|
"Image item has no recognized image field. "
|
|
f"Supported base64 keys: {IMAGE_B64_KEYS}; supported URL keys: {IMAGE_URL_KEYS}. "
|
|
f"Full item: {item}"
|
|
)
|
|
|
|
return torch.stack(images, dim=0)
|
|
|
|
|
|
def _response_json_text(payload):
|
|
return json.dumps(payload, ensure_ascii=False, indent=2)
|
|
|
|
|
|
def _maybe_save_debug_response(enabled, operation, payload):
|
|
if not enabled:
|
|
return
|
|
safe_operation = operation.lower().replace(" ", "_")
|
|
filename = f"comfyui_openai_compatible_{safe_operation}_{int(time.time())}.json"
|
|
path = os.path.join(tempfile.gettempdir(), filename)
|
|
with open(path, "w", encoding="utf-8") as file:
|
|
file.write(_response_json_text(payload))
|
|
print(f"[OpenAI Compatible Image] Saved raw {operation} response to: {path}")
|
|
|
|
|
|
class OpenAICompatibleImageGenerate:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": "A clean product photo of a translucent glass teapot on a walnut table, soft daylight, realistic.",
|
|
},
|
|
),
|
|
"api_base": (
|
|
"STRING",
|
|
{"default": "https://api.openai.com/v1"},
|
|
),
|
|
"model": (
|
|
"STRING",
|
|
{"default": "gpt-image-2"},
|
|
),
|
|
"size": (
|
|
"STRING",
|
|
{"default": "1024x1024"},
|
|
),
|
|
"quality": (
|
|
["auto", "low", "medium", "high"],
|
|
{"default": "auto"},
|
|
),
|
|
"n": (
|
|
"INT",
|
|
{"default": 1, "min": 1, "max": 10, "step": 1},
|
|
),
|
|
"timeout_seconds": (
|
|
"INT",
|
|
{"default": DEFAULT_TIMEOUT_SECONDS, "min": 30, "max": 3600, "step": 30},
|
|
),
|
|
},
|
|
"optional": {
|
|
"api_key": (
|
|
"STRING",
|
|
{"default": "", "multiline": False},
|
|
),
|
|
"endpoint_path": (
|
|
"STRING",
|
|
{"default": "/images/generations"},
|
|
),
|
|
"extra_headers_json": (
|
|
"STRING",
|
|
{"multiline": True, "default": "{}"},
|
|
),
|
|
"extra_body_json": (
|
|
"STRING",
|
|
{"multiline": True, "default": "{}"},
|
|
),
|
|
"save_raw_response": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "response_json")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "OpenAI Compatible/Image"
|
|
|
|
def generate(
|
|
self,
|
|
prompt,
|
|
api_base,
|
|
model,
|
|
size,
|
|
quality,
|
|
n,
|
|
timeout_seconds,
|
|
api_key="",
|
|
endpoint_path="/images/generations",
|
|
extra_headers_json="{}",
|
|
extra_body_json="{}",
|
|
save_raw_response=False,
|
|
):
|
|
api_key = _api_key_from_input(api_key)
|
|
url = _endpoint_url(api_base, endpoint_path)
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
"Connection": "close",
|
|
}
|
|
headers.update(_parse_json_object(extra_headers_json, "extra_headers_json"))
|
|
|
|
body = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"n": n,
|
|
}
|
|
if size != "auto":
|
|
body["size"] = size
|
|
if quality != "auto":
|
|
body["quality"] = quality
|
|
body.update(_parse_json_object(extra_body_json, "extra_body_json"))
|
|
|
|
response = _post_json(url, headers, body, timeout_seconds, "Image generation")
|
|
if response.status_code >= 400:
|
|
raise RuntimeError(f"Image API request failed ({response.status_code}): {response.text}")
|
|
|
|
payload = response.json()
|
|
_maybe_save_debug_response(save_raw_response, "Image generation", payload)
|
|
return (_response_to_tensor_batch(payload, timeout_seconds), _response_json_text(payload))
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"OpenAICompatibleImageGenerate": OpenAICompatibleImageGenerate,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"OpenAICompatibleImageGenerate": "OpenAI Compatible Image Generate",
|
|
}
|