commit edc57248fad2bd4bc4cf0c719e5dcf5b862f01dd Author: kulefanle <82266556+kulefanle@users.noreply.github.com> Date: Sun May 10 18:08:58 2026 +0800 Add __init__.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..03fac67 --- /dev/null +++ b/__init__.py @@ -0,0 +1,440 @@ +""" +ComfyUI custom node: Proxy GPT Image Generator + +This node is designed for OpenAI-compatible proxy services that expose either +`/v1/images/generations` or `/v1/chat/completions` for image generation. + +Risk notes: +- API keys are accepted as node inputs, but they are never printed, returned, or + written to disk by this node. Do not share workflow JSON files that contain + real keys, because ComfyUI stores widget values inside exported workflows. +- The node downloads image URLs returned by the proxy. Only use a proxy/base URL + you trust, because remote image URLs can reveal your server IP and consume + bandwidth. +- If no valid image is found, the node raises an error instead of returning a + blank placeholder. This makes authorization, parsing, or proxy errors visible. +""" + +from __future__ import annotations + +import base64 +import io +import json +import re +from typing import Any, Iterable +from urllib.parse import urljoin + +import numpy as np +import requests +import torch +from PIL import Image + + +IMAGE_URL_RE = re.compile( + r"https?://[^\s\"'<>]+?\.(?:png|jpg|jpeg|webp)(?:\?[^\s\"'<>]*)?", + re.IGNORECASE, +) +DATA_URL_RE = re.compile(r"data:image/[^;]+;base64,([A-Za-z0-9+/=\s]+)") +BASE64_RE = re.compile(r"^[A-Za-z0-9+/=\s]{200,}$") + + +def _join_endpoint(base_url: str, route: str) -> str: + """Build a stable endpoint without silently producing duplicate slashes.""" + clean_base = base_url.strip().rstrip("/") + clean_route = route.strip() + if not clean_base: + raise ValueError("Base URL is empty.") + if clean_route.startswith("http://") or clean_route.startswith("https://"): + return clean_route + return urljoin(clean_base + "/", clean_route.lstrip("/")) + + +def _image_to_data_url(image: torch.Tensor, fmt: str = "PNG") -> str: + """ + Convert a ComfyUI IMAGE tensor to a data URL for chat/completions proxies. + + ComfyUI images are float tensors in [B, H, W, C]. The proxy usually accepts + a standard data URL in an OpenAI-style `image_url` content part. + """ + if image is None: + raise ValueError("Reference image is missing.") + + image_np = image[0].detach().cpu().numpy() + image_np = np.clip(image_np * 255.0, 0, 255).astype(np.uint8) + pil_image = Image.fromarray(image_np).convert("RGB") + + buffer = io.BytesIO() + pil_image.save(buffer, format=fmt) + encoded = base64.b64encode(buffer.getvalue()).decode("ascii") + return f"data:image/{fmt.lower()};base64,{encoded}" + + +def _bytes_to_comfy_image(image_bytes: bytes) -> torch.Tensor: + """ + Decode image bytes into a ComfyUI IMAGE tensor. + + RGBA images are composited onto a white background. This avoids the common + "transparent PNG looks like a white/blank image" trap while still returning + a normal RGB tensor to PreviewImage. + """ + pil_image = Image.open(io.BytesIO(image_bytes)) + if pil_image.mode == "RGBA": + background = Image.new("RGBA", pil_image.size, (255, 255, 255, 255)) + pil_image = Image.alpha_composite(background, pil_image).convert("RGB") + else: + pil_image = pil_image.convert("RGB") + + image_np = np.asarray(pil_image).astype(np.float32) / 255.0 + return torch.from_numpy(image_np)[None,] + + +def _decode_data_url(value: str) -> bytes | None: + """Decode a `data:image/...;base64,...` string if the value contains one.""" + match = DATA_URL_RE.search(value) + if not match: + return None + return base64.b64decode(match.group(1).replace("\n", "").replace(" ", "")) + + +def _decode_probable_base64(value: str) -> bytes | None: + """ + Decode long raw base64 strings. + + The length and alphabet checks reduce the risk of treating normal text such + as `auto` or an error message as an image. + """ + compact = value.strip().replace("\n", "").replace(" ", "") + if len(compact) < 200 or not BASE64_RE.match(compact): + return None + try: + decoded = base64.b64decode(compact, validate=True) + except Exception: + return None + if decoded[:8] == b"\x89PNG\r\n\x1a\n" or decoded[:3] == b"\xff\xd8\xff" or decoded[:4] == b"RIFF": + return decoded + return None + + +def _iter_strings(value: Any) -> Iterable[str]: + """Walk a JSON-like object and yield candidate strings from common fields.""" + if isinstance(value, str): + yield value + return + + if isinstance(value, list): + for item in value: + yield from _iter_strings(item) + return + + if isinstance(value, dict): + priority_keys = ( + "b64_json", + "url", + "image_url", + "image", + "file_url", + "content", + "text", + "output", + "data", + "images", + ) + for key in priority_keys: + if key in value: + item = value[key] + if isinstance(item, dict) and "url" in item: + yield from _iter_strings(item["url"]) + else: + yield from _iter_strings(item) + for key, item in value.items(): + if key not in priority_keys: + yield from _iter_strings(item) + + +def _extract_first_image_candidate(payload: Any) -> tuple[str, bytes | None]: + """ + Find the first usable image from a proxy response. + + Supported forms: + - OpenAI Images API: data[0].b64_json or data[0].url + - Chat proxies: message content containing a data URL, raw base64, markdown + image URL, or direct HTTP image URL + - Custom proxies: top-level `images`, `image`, `url`, or `file_url` + """ + for text in _iter_strings(payload): + data_url_bytes = _decode_data_url(text) + if data_url_bytes: + return ("data_url", data_url_bytes) + + raw_base64_bytes = _decode_probable_base64(text) + if raw_base64_bytes: + return ("base64", raw_base64_bytes) + + url_match = IMAGE_URL_RE.search(text) + if url_match: + return (url_match.group(0), None) + + stripped = text.strip() + if stripped.startswith("http://") or stripped.startswith("https://"): + return (stripped, None) + + return ("", None) + + +def _summarize_payload(payload: Any, limit: int = 1600) -> str: + """ + Return a compact response summary for PreviewAny. + + This intentionally redacts large base64-like strings so the UI stays usable + and exported workflows do not become huge because of response text. + """ + try: + text = json.dumps(payload, ensure_ascii=False, indent=2) + except TypeError: + text = str(payload) + text = re.sub(r"([A-Za-z0-9+/=]{180,})", "[base64 omitted]", text) + return text[:limit] + + +class ProxyGPTImageGenerator: + """ComfyUI node that calls an OpenAI-compatible image proxy.""" + + @classmethod + def INPUT_TYPES(cls) -> dict[str, Any]: + return { + "required": { + "api_key": ( + "STRING", + { + "default": "", + "multiline": False, + "password": True, + }, + ), + "base_url": ( + "STRING", + { + "default": "https://greenapi.ink", + "multiline": False, + }, + ), + "api_route": ( + [ + "/v1/images/generations", + "/v1/chat/completions", + ], + { + "default": "/v1/images/generations", + }, + ), + "model": ( + "STRING", + { + "default": "gpt-image-2", + "multiline": False, + }, + ), + "prompt": ( + "STRING", + { + "default": "a red apple on a wooden table, realistic photo", + "multiline": True, + }, + ), + "size": ( + [ + "1024x1024", + "1024x1536", + "1536x1024", + "auto", + ], + { + "default": "1024x1024", + }, + ), + "quality": ( + [ + "low", + "medium", + "high", + "auto", + ], + { + "default": "medium", + }, + ), + "background": ( + [ + "opaque", + "transparent", + "auto", + ], + { + "default": "opaque", + }, + ), + "output_format": ( + [ + "png", + "jpeg", + "webp", + ], + { + "default": "png", + }, + ), + "timeout_seconds": ( + "INT", + { + "default": 180, + "min": 10, + "max": 900, + "step": 10, + }, + ), + }, + "optional": { + "reference_image_1": ("IMAGE",), + "reference_image_2": ("IMAGE",), + "reference_image_3": ("IMAGE",), + "reference_image_4": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING", "STRING") + RETURN_NAMES = ("image", "image_source", "response_summary") + FUNCTION = "generate" + CATEGORY = "api/proxy" + + def generate( + self, + api_key: str, + base_url: str, + api_route: str, + model: str, + prompt: str, + size: str, + quality: str, + background: str, + output_format: str, + timeout_seconds: int, + reference_image_1: torch.Tensor | None = None, + reference_image_2: torch.Tensor | None = None, + reference_image_3: torch.Tensor | None = None, + reference_image_4: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, str, str]: + clean_key = api_key.strip() + clean_prompt = prompt.strip() + + if not clean_key: + raise ValueError("API key is empty. Fill a valid key from the same proxy service as Base URL.") + if not clean_prompt: + raise ValueError("Prompt is empty. The proxy may return a blank result when prompt is empty.") + + endpoint = _join_endpoint(base_url, api_route) + headers = { + "Authorization": f"Bearer {clean_key}", + "Content-Type": "application/json", + "Accept": "application/json, image/*", + } + + references = [ + image + for image in (reference_image_1, reference_image_2, reference_image_3, reference_image_4) + if image is not None + ] + + if api_route.endswith("/chat/completions"): + # Many proxy services map image generation models onto the chat API. + # The content-array format keeps text and optional references explicit. + content: list[dict[str, Any]] = [{"type": "text", "text": clean_prompt}] + for image in references: + content.append( + { + "type": "image_url", + "image_url": {"url": _image_to_data_url(image)}, + } + ) + + payload: dict[str, Any] = { + "model": model.strip(), + "messages": [{"role": "user", "content": content}], + "size": size, + "quality": quality, + "background": background, + "output_format": output_format, + } + else: + # OpenAI Images-style JSON body. Some proxies also accept `images` + # as data URLs for image-to-image; unsupported proxies should return + # a clear 4xx error rather than a hidden blank image. + payload = { + "model": model.strip(), + "prompt": clean_prompt, + "n": 1, + "size": size, + "quality": quality, + "background": background, + "output_format": output_format, + } + if references: + payload["images"] = [_image_to_data_url(image) for image in references] + + try: + response = requests.post( + endpoint, + headers=headers, + json=payload, + timeout=int(timeout_seconds), + ) + except requests.RequestException as exc: + raise RuntimeError(f"Failed to connect to proxy endpoint: {endpoint}. Detail: {exc}") from exc + + content_type = response.headers.get("content-type", "").lower() + + if not response.ok: + # Do not include headers or API key in this error. The response body + # is useful for diagnosing 401/403/429/model-not-found issues. + body = response.text[:1200] + raise RuntimeError(f"Proxy returned HTTP {response.status_code} for {endpoint}: {body}") + + if content_type.startswith("image/"): + image_tensor = _bytes_to_comfy_image(response.content) + return (image_tensor, endpoint, "Proxy returned image bytes directly.") + + try: + payload_response: Any = response.json() + except ValueError as exc: + body = response.text[:1200] + raise RuntimeError(f"Proxy did not return JSON or image bytes. Body: {body}") from exc + + image_source, inline_bytes = _extract_first_image_candidate(payload_response) + summary = _summarize_payload(payload_response) + + if inline_bytes: + image_tensor = _bytes_to_comfy_image(inline_bytes) + return (image_tensor, image_source, summary) + + if image_source.startswith("http://") or image_source.startswith("https://"): + try: + image_response = requests.get(image_source, timeout=int(timeout_seconds)) + image_response.raise_for_status() + except requests.RequestException as exc: + raise RuntimeError(f"Image URL was returned but could not be downloaded: {image_source}") from exc + + image_tensor = _bytes_to_comfy_image(image_response.content) + return (image_tensor, image_source, summary) + + raise RuntimeError( + "Proxy response did not contain a usable image URL, data URL, or b64_json. " + f"Response summary: {summary}" + ) + + +NODE_CLASS_MAPPINGS = { + "ProxyGPTImageGenerator": ProxyGPTImageGenerator, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ProxyGPTImageGenerator": "Proxy GPT Image Generator", +}