diff --git a/README.md b/README.md index aad0c55..1781818 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,25 @@ pip install -r comfyui-PD_comfy-api-node/requirements.txt ## 📝 更新日志 +### 2026-10-11 - Grok 文生图 +- 新增 `PDGrokImageAPIKey` 文生图节点，复用现有 Grok 模型列表；API Key 由节点输入提供。同步更新 Grok 编辑节点的取消检查与失败状态。 + +### 2026-10-11 - OpenAI ChatGPT / GPT Image 2.5 +- 新增 `PD: OpenAI ChatGPT Latest (ComfyUI Key)`（`PDOpenAIChatLatestAPIKey`），在 `PD_Tools/Text_Generation`。模型列表复用本地官方节点的 18 个选项，支持批量图片、`OPENAI_INPUT_FILES` 文件输入及 `OPENAI_CHAT_CONFIG` 高级设置；可连接官方 OpenAI ChatGPT Input Files / Advanced Options 节点。输出文字。 +- 新增 `PD: OpenAI GPT Image 2.5 (ComfyUI Key)`（`PDOpenAIGPTImage25APIKey`），在 `PD_Tools/Image_Generation`。支持 `gpt-image-2.5-flare` / `gpt-image-2.5-sunburst` / `gpt-image-2` / `gpt-image-1.5` / `gpt-image-1`；动态尺寸、质量、背景、最多 16 张参考图、单图遮罩和 1–8 张输出。保留 RGBA 透明度，自动对齐同批输出尺寸。seed 沿用官方控件，当前后端未实现。 +- 两者手动填写 Comfy.org API Key，复用本地官方价格估算角标；完成状态显示服务返回的 usage，估算角标不代表实际扣费。旧 ChatGPT/GPT Image 节点保留，需包含上述官方模型节点的新版 ComfyUI。 +- ChatGPT 每次独立请求；官方 `persist_context` 已废弃且无效，因此新节点省略该开关，不在本地保存历史。Responses 格式参考 [OpenAI Docs](https://developers.openai.com/api/docs/guides/conversation-state)，Comfy.org 代理契约以本地 `comfy_api_nodes/nodes_openai.py` 为准。 +- 请求/轮询/下载期间每秒显示耗时，每 0.1 秒检查取消，整次执行限时 600 秒；仅一次生成 POST，不自动重试。取消不保证退还远端已产生的费用。 +- 离线测试：`python custom_nodes/comfyui-PD_comfy-api-node/tests/test_openai_latest.py`。覆盖全部模型、文件/高级设置、异步轮询、RGBA/尺寸对齐、遮罩、慢请求取消、超时不重试及请求前校验；本次未提交真实付费请求。 + +### 2026-10-11 - Nano Banana 2.1 +- 新增 `PD: Nano Banana 2.1 (ComfyUI Key)`（`PDGeminiNanoBanana21`），位于 `PD_Tools/Image_Generation`；旧 Nano Banana 2 节点保留。 +- 复用当前官方节点的动态模型参数：Nano Banana 2.1、Nano Banana 2、Nano Banana 2 Lite；Lite 仅支持 1K，2.1 支持 MINIMAL/MEDIUM/HIGH。支持最多 14 张参考图及 `GEMINI_INPUT_FILES`，输出成品图、文字、第一张思考图；无思考图时与官方节点一样输出空白占位图。 +- 手动填写 Comfy.org API Key；无需通过官方节点的登录入口。需包含 `GeminiNanoBanana2V2` 的新版 ComfyUI，本机已验证版本为 0.39.0。费用角标沿用官方估算，实际扣费以账户账单为准。 +- 按官方请求契约，2.1 不发送 temperature/top_p；所有三个模型当前都不发送 seed。控件提示已注明，固定 seed 不保证复现。 +- 等待期间每秒显示耗时，取消检查间隔 0.1 秒；单次执行只提交一次生成请求，不自动重试。取消后不继续轮询、下载或返回部分结果；已受理的远端请求不保证取消扣费。 +- 离线模拟验证：`python custom_nodes/comfyui-PD_comfy-api-node/tests/test_nano_banana21.py`，覆盖三个模型请求、思考图分离、慢生成/下载取消、超时不重试及请求前校验；不产生 API 费用。 + ### v2.0.1 - 2026.0424 - 修复 `PD_comfyplus_image (Multi-Size Pack)` -> GPT Image 2 `files` 输入不兼容的问题。 - 在发送 OpenAI 兼容的 multipart 请求前，对打包图片的 MIME type 做归一化处理，确保 `GeminiMimeType.image_png` / `image_png` 会以 `image/png` 发送。 diff --git a/py/IMAGE_GeminiNanoBanana21.py b/py/IMAGE_GeminiNanoBanana21.py new file mode 100644 index 0000000..fb63ead --- /dev/null +++ b/py/IMAGE_GeminiNanoBanana21.py @@ -0,0 +1,165 @@ +import base64 +import io +import time +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout + +import numpy as np +import requests +import torch +from PIL import Image + +import comfy.model_management +from comfy_api.latest import io as IO +from comfy_api_nodes.nodes_gemini import GeminiNanoBanana2V2 +from server import PromptServer + +from .IMAGE_GeminiNanoBanana2 import tensor_to_base64 + + +MODELS = { + "Gemini Nano Banana 2.1": "gemini-nano-banana-2.1", + "Nano Banana 2 (Gemini 3.1 Flash Image)": "gemini-3.1-flash-image", + "Nano Banana 2 Lite": "gemini-3.1-flash-lite-image", +} + + +class PDGeminiNanoBanana21(IO.ComfyNode): + @classmethod + def define_schema(cls): + schema = GeminiNanoBanana2V2.define_schema() + schema.node_id = "PDGeminiNanoBanana21" + schema.display_name = "PD: Nano Banana 2.1 (ComfyUI Key)" + schema.category = "PD_Tools/Image_Generation" + schema.inputs.insert(0, IO.String.Input("comfy_api_key", default="", tooltip="Comfy.org API Key")) + for item in schema.inputs: + if item.id == "seed": + item.tooltip = "Matches the official node. The current Comfy.org endpoint does not receive seed; a fixed seed does not guarantee repeatable output." + elif item.id in ("temperature", "top_p"): + item.tooltip += " Ignored by Nano Banana 2.1, matching the official endpoint." + schema.hidden = [IO.Hidden.unique_id] + schema.is_api_node = False + return schema + + @classmethod + def execute(cls, comfy_api_key, prompt, model, seed, response_modalities, + system_prompt="", temperature=1.0, top_p=0.95): + if not comfy_api_key.strip(): + raise ValueError("请填写 ComfyUI API Key。") + if not prompt.strip(): + raise ValueError("请填写提示词。") + choice = model["model"] + model_id = MODELS[choice] + if choice == "Nano Banana 2 Lite" and model["resolution"] != "1K": + raise ValueError("Nano Banana 2 Lite 仅支持 1K。") + if choice != "Gemini Nano Banana 2.1" and model["thinking_level"] == "MEDIUM": + raise ValueError("MEDIUM 仅适用于 Nano Banana 2.1。") + images = list((model.get("images") or {}).values()) + if sum(image.shape[0] for image in images if image is not None) > 14: + raise ValueError("最多支持 14 张参考图。") + parts = [{"text": prompt}] + for image in images: + if image is not None: + for item in image: + comfy.model_management.throw_exception_if_processing_interrupted() + parts.append({"inlineData": {"mimeType": "image/png", "data": tensor_to_base64(item)}}) + for part in model.get("files") or []: + parts.append(part.model_dump(exclude_none=True) if hasattr(part, "model_dump") else part) + image_config = {"imageSize": model["resolution"]} + if model["aspect_ratio"] != "auto": + image_config["aspectRatio"] = model["aspect_ratio"] + config = { + "responseModalities": ["IMAGE"] if response_modalities == "IMAGE" else ["TEXT", "IMAGE"], + "imageConfig": image_config, + "thinkingConfig": {"thinkingLevel": model["thinking_level"]}, + } + # The official 2.1 endpoint omits sampling parameters and seed. + if choice != "Gemini Nano Banana 2.1": + config.update(temperature=temperature, topP=top_p) + payload = {"contents": [{"role": "user", "parts": parts}], "generationConfig": config} + if system_prompt: + payload["systemInstruction"] = {"parts": [{"text": system_prompt}]} + + started = time.monotonic() + session = requests.Session() + executor = ThreadPoolExecutor(max_workers=1) + future = None + + def status(text): + PromptServer.instance.send_sync("progress_text", { + "node": cls.hidden.unique_id, + "text": f"{text} | {time.monotonic() - started:.1f}s", + }) + + def request(method, url, **kwargs): + nonlocal future + comfy.model_management.throw_exception_if_processing_interrupted() + request_started = time.monotonic() + future = executor.submit(session.request, method, url, timeout=(15, 180), **kwargs) + last_status = 0.0 + while True: + comfy.model_management.throw_exception_if_processing_interrupted() + if time.monotonic() - request_started > 195: + raise requests.Timeout("Request exceeded 195 seconds") + if time.monotonic() - last_status >= 1: + status("Running") + last_status = time.monotonic() + try: + return future.result(timeout=0.1) + except FutureTimeout: + if future.done(): + raise + + def decode(data): + comfy.model_management.throw_exception_if_processing_interrupted() + with Image.open(io.BytesIO(data)) as image: + return torch.from_numpy(np.array(image.convert("RGB")).astype(np.float32) / 255.0).unsqueeze(0) + + try: + response = request("POST", f"https://api.comfy.org/proxy/vertexai/gemini/{model_id}", + headers={"X-API-KEY": comfy_api_key.strip()}, json=payload) + response.raise_for_status() + result = response.json() + generated, thoughts, texts = [], [], [] + for candidate in result.get("candidates", []): + for part in candidate.get("content", {}).get("parts", []): + comfy.model_management.throw_exception_if_processing_interrupted() + if part.get("text"): + texts.append(part["text"]) + inline = part.get("inlineData") or part.get("inline_data") + file_data = part.get("fileData") or part.get("file_data") + if inline and inline.get("mimeType", inline.get("mime_type", "")).startswith("image/"): + image = decode(base64.b64decode(inline["data"])) + elif file_data and file_data.get("mimeType", file_data.get("mime_type", "")).startswith("image/"): + downloaded = request("GET", file_data.get("fileUri") or file_data["file_uri"]) + downloaded.raise_for_status() + image = decode(downloaded.content) + else: + continue + (thoughts if part.get("thought") else generated).append(image) + if not generated: + raise ValueError("Gemini 未返回成品图：" + ("\n".join(texts) or str(result.get("promptFeedback", {})))) + output = torch.cat(generated) + thought = thoughts[0] if thoughts else torch.zeros((1, 1024, 1024, 3)) + comfy.model_management.throw_exception_if_processing_interrupted() + usage = result.get("usageMetadata", {}) + status(f"Completed | Tokens: {usage.get('totalTokenCount', 'N/A')}") + return IO.NodeOutput(output, "\n".join(texts), thought) + except comfy.model_management.InterruptProcessingException: + session.close() + if future is not None: + future.cancel() + status("Cancelled") + raise + except (requests.Timeout, FutureTimeout): + status("Timed out") + raise + except Exception: + status("Failed") + raise + finally: + session.close() + executor.shutdown(wait=False, cancel_futures=True) + + +NODE_CLASS_MAPPINGS = {"PDGeminiNanoBanana21": PDGeminiNanoBanana21} +NODE_DISPLAY_NAME_MAPPINGS = {"PDGeminiNanoBanana21": "PD: Nano Banana 2.1 (ComfyUI Key)"} diff --git a/py/IMAGE_GrokEdit_apikey.py b/py/IMAGE_GrokEdit_apikey.py index e6dac6d..81c5478 100644 --- a/py/IMAGE_GrokEdit_apikey.py +++ b/py/IMAGE_GrokEdit_apikey.py @@ -79,6 +79,7 @@ def send_status(node_id, text): def request_with_interrupt(session, method, url, node_id, status, start_time, **kwargs): + mm.throw_exception_if_processing_interrupted() executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) future = executor.submit(method, url, **kwargs) last_elapsed = -1 @@ -260,6 +261,7 @@ class PDGrokImageEditAPIKey: timeout=600, ) if response.status_code != 200: + send_status(unique_id, f"Failed\nTime elapsed: {int(time.monotonic() - start_time)}s") return (empty_image(), f"API Error {response.status_code}: {response.text[:500]}") result = response.json() @@ -281,6 +283,7 @@ class PDGrokImageEditAPIKey: output_images.append(load_image_to_tensor(base64.b64decode(item["b64_json"]))) if not output_images: + send_status(unique_id, f"Failed\nTime elapsed: {int(time.monotonic() - start_time)}s") return (empty_image(), f"Error: response contained no images: {str(result)[:500]}") image = torch.cat(output_images, dim=0) @@ -289,6 +292,8 @@ class PDGrokImageEditAPIKey: send_status(unique_id, f"Completed\nTime elapsed: {int(elapsed)}s") print(f"[PD_GROK_EDIT] Completed {len(output_images)} image(s) in {elapsed:.2f}s") return (image, price) + except mm.InterruptProcessingException: + raise except requests.exceptions.Timeout: send_status(unique_id, f"Timed out\nTime elapsed: {int(time.monotonic() - start_time)}s") return (empty_image(), "Error: request timeout (600s)") @@ -307,5 +312,5 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { - "PDGrokImageEditAPIKey": "grok edit (comfyapikey)", + "PDGrokImageEditAPIKey": "PD-grok-image-edit", } diff --git a/py/IMAGE_Grok_apikey.py b/py/IMAGE_Grok_apikey.py new file mode 100644 index 0000000..5e89182 --- /dev/null +++ b/py/IMAGE_Grok_apikey.py @@ -0,0 +1,204 @@ +import base64 +import time + +import comfy.model_management as mm +import requests +import torch + +from .IMAGE_GrokEdit_apikey import ( + ASPECT_RATIOS, + BASE_URL, + MODELS, + QUALITY_MODELS, + request_with_interrupt, + send_status, +) +from ._image_gpt_common import ( + empty_image, + format_credits_value, + format_rmb_value, + load_image_to_tensor, + usd_to_comfy_credits, +) + + +GENERATION_ENDPOINT = "/proxy/xai/v1/images/generations" + + +def estimate_price(model, resolution, quality, number_of_images): + is_1k = resolution.lower() == "1k" + if model == "grok-imagine-image-2.0": + rate = (0.04 if is_1k else 0.06) if quality == "low" else (0.06 if is_1k else 0.08) + elif model == "grok-imagine-image-quality": + rate = 0.05 if is_1k else 0.07 + elif model == "grok-imagine-image-pro": + rate = 0.07 + else: + rate = 0.02 + return rate * number_of_images + + +def format_price(result, response, model, resolution, quality, number_of_images): + usage = result.get("usage") or {} + ticks = usage.get("cost_in_usd_ticks") + credits_header = response.headers.get("X-Comfy-Credits-Used") + + credits = None + try: + credits = float(credits_header) if credits_header is not None else None + except ValueError: + pass + + usd = ticks / 10_000_000_000 if isinstance(ticks, (int, float)) else None + if usd is None and credits is not None: + usd = credits / usd_to_comfy_credits(1.0) + if credits is None and usd is not None: + credits = usd_to_comfy_credits(usd) + + if usd is None: + usd = estimate_price(model, resolution, quality, number_of_images) + return "\n".join([ + f"Estimated price (USD): ${usd:.6f}", + f"Estimated price (RMB): {format_rmb_value(usd)}", + f"Estimated price (Credits): {format_credits_value(usd)}", + ]) + + lines = [ + f"Actual price (USD): ${usd:.6f}", + f"Actual price (RMB): {format_rmb_value(usd)}", + ] + if credits is not None: + lines.append(f"Credits used: {credits:.2f} credits") + return "\n".join(lines) + + +class PDGrokImageAPIKey: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "api_key": ("STRING", { + "default": "", + "multiline": False, + "placeholder": "Your ComfyUI API Key", + }), + "prompt": ("STRING", {"default": "", "multiline": True}), + "model": (MODELS, {"default": "grok-imagine-image-2.0"}), + "aspect_ratio": (ASPECT_RATIOS[1:], {"default": "1:1"}), + "resolution": (["1K", "2K"], {"default": "1K"}), + "quality": (["medium", "low"], {"default": "medium"}), + "number_of_images": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1}), + "seed": ("INT", {"default": 0, "min": 0, "max": 2147483647, "step": 1}), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + }, + } + + RETURN_TYPES = ("IMAGE", "STRING") + RETURN_NAMES = ("image", "show_price") + FUNCTION = "generate_image" + CATEGORY = "PD_Tools/Image_Generation" + + def generate_image( + self, + api_key, + prompt, + model, + aspect_ratio, + resolution, + quality, + number_of_images, + seed, + unique_id, + ): + if not api_key or not api_key.strip(): + return (empty_image(), "Error: ComfyUI API Key is required") + if not prompt or not prompt.strip(): + return (empty_image(), "Error: prompt is required") + + payload = { + "model": model, + "prompt": prompt.strip(), + "aspect_ratio": aspect_ratio, + "n": number_of_images, + "seed": seed, + "response_format": "url", + "resolution": resolution.lower(), + "quality": quality if model in QUALITY_MODELS else None, + } + headers = { + "X-API-KEY": api_key.strip(), + "Accept": "application/json", + "Content-Type": "application/json", + "User-Agent": "ComfyUI-PD-Node/2.1", + } + + start_time = time.monotonic() + session = requests.Session() + try: + response = request_with_interrupt( + session, + session.post, + f"{BASE_URL}{GENERATION_ENDPOINT}", + unique_id, + "Waiting for Grok server", + start_time, + headers=headers, + json=payload, + timeout=600, + ) + if response.status_code != 200: + send_status(unique_id, f"Failed\nTime elapsed: {int(time.monotonic() - start_time)}s") + return (empty_image(), f"API Error {response.status_code}: {response.text[:500]}") + + result = response.json() + output_images = [] + for item in result.get("data") or []: + if item.get("url"): + download = request_with_interrupt( + session, + session.get, + item["url"], + unique_id, + "Downloading result", + start_time, + timeout=120, + ) + download.raise_for_status() + output_images.append(load_image_to_tensor(download.content)) + elif item.get("b64_json"): + output_images.append(load_image_to_tensor(base64.b64decode(item["b64_json"]))) + + if not output_images: + send_status(unique_id, f"Failed\nTime elapsed: {int(time.monotonic() - start_time)}s") + return (empty_image(), f"Error: response contained no images: {str(result)[:500]}") + + image = torch.cat(output_images, dim=0) + price = format_price(result, response, model, resolution, quality, number_of_images) + elapsed = time.monotonic() - start_time + send_status(unique_id, f"Completed\nTime elapsed: {int(elapsed)}s") + print(f"[PD_GROK_IMAGE] {model} completed {len(output_images)} image(s) in {elapsed:.2f}s") + return (image, price) + except mm.InterruptProcessingException: + raise + except requests.exceptions.Timeout: + send_status(unique_id, f"Timed out\nTime elapsed: {int(time.monotonic() - start_time)}s") + return (empty_image(), "Error: request timeout (600s)") + except requests.exceptions.ConnectionError as error: + send_status(unique_id, f"Connection failed\nTime elapsed: {int(time.monotonic() - start_time)}s") + return (empty_image(), f"Error: connection failed - {error}") + except Exception as error: + send_status(unique_id, f"Failed\nTime elapsed: {int(time.monotonic() - start_time)}s") + return (empty_image(), f"Error: {error}") + finally: + session.close() + + +NODE_CLASS_MAPPINGS = { + "PDGrokImageAPIKey": PDGrokImageAPIKey, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PDGrokImageAPIKey": "PD-grok-text to image", +} diff --git a/py/OPENAI_latest_apikey.py b/py/OPENAI_latest_apikey.py new file mode 100644 index 0000000..b1c31ed --- /dev/null +++ b/py/OPENAI_latest_apikey.py @@ -0,0 +1,222 @@ +import base64 +import io +import time +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout + +import numpy as np +import requests +import torch +from PIL import Image + +import comfy.model_management +from comfy.utils import common_upscale +from comfy_api.latest import IO +from comfy_api_nodes.nodes_openai import ( + OpenAIChatNode, OpenAIGPTImageNodeV2, GPT_IMAGE_MODELS, + SUPPORTED_REASONING_EFFORTS, SupportedOpenAIModel, downscale_image_tensor, +) +from server import PromptServer + +from ._image_gpt_common import tensor_to_bytes, validate_gpt_image_2_size + + +BASE_URL = "https://api.comfy.org" + + +class _Request: + def __init__(self, node_id, api_key): + if not api_key.strip(): + raise ValueError("请填写 Comfy.org API Key。") + self.node_id = node_id + self.started = time.monotonic() + self.session = requests.Session() + self.headers = {"X-API-KEY": api_key.strip(), "Accept": "application/json"} + self.executor = ThreadPoolExecutor(max_workers=1) + self.future = None + self.response_id = "" + self.usage = {} + self.last_status = 0 + + def status(self, text): + PromptServer.instance.send_sync("progress_text", { + "node": self.node_id, + "text": f"{text} | {time.monotonic() - self.started:.1f}s {self.response_id}", + }) + + def check(self): + comfy.model_management.throw_exception_if_processing_interrupted() + if time.monotonic() - self.started > 600: + raise requests.Timeout("Operation exceeded 600 seconds") + if time.monotonic() - self.last_status >= 1: + self.status("Running") + self.last_status = time.monotonic() + + def request(self, method, url, **kwargs): + self.check() + headers = self.headers if url.startswith(BASE_URL + "/") else {} + self.future = self.executor.submit(self.session.request, method, url, + headers=headers, timeout=(15, 180), **kwargs) + while True: + self.check() + try: + response = self.future.result(timeout=0.1) + response.raise_for_status() + return response + except FutureTimeout: + if self.future.done(): + raise + + def sleep(self, seconds): + until = time.monotonic() + seconds + while time.monotonic() < until: + self.check() + time.sleep(min(0.1, max(0, until - time.monotonic()))) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback): + self.session.close() + if self.future is not None: + self.future.cancel() + self.executor.shutdown(wait=False, cancel_futures=True) + if exc_type is not None and issubclass(exc_type, comfy.model_management.InterruptProcessingException): + self.status("Cancelled") + elif exc_type is not None and issubclass(exc_type, (requests.Timeout, FutureTimeout)): + self.status("Timed out") + elif exc_type is not None: + self.status("Failed") + else: + self.status(f"Completed | Usage: {self.usage}") + + +def _schema(official, node_id, name, category): + schema = official.define_schema() + schema.node_id = node_id + schema.display_name = name + schema.category = category + schema.inputs.insert(0, IO.String.Input("comfy_api_key", default="", tooltip="Comfy.org API Key")) + schema.hidden = [IO.Hidden.unique_id] + schema.is_api_node = False + return schema + + +class PDOpenAIChatLatestAPIKey(IO.ComfyNode): + @classmethod + def define_schema(cls): + schema = _schema(OpenAIChatNode, "PDOpenAIChatLatestAPIKey", "PD: OpenAI ChatGPT Latest (ComfyUI Key)", "PD_Tools/Text_Generation") + # The official persist_context widget is deprecated and does nothing. + schema.inputs = [item for item in schema.inputs if item.id != "persist_context"] + return schema + + @classmethod + def execute(cls, comfy_api_key, prompt, model, images=None, files=None, advanced_options=None): + if model not in [item.value for item in SupportedOpenAIModel]: + raise ValueError(f"Unknown model: {model}") + if not prompt.strip(): + raise ValueError("请填写提示词。") + options = advanced_options.model_dump(exclude_none=True) if advanced_options else {} + effort = (options.get("reasoning") or {}).get("effort") + if effort is not None and effort not in SUPPORTED_REASONING_EFFORTS[model]: + raise ValueError(f"{model} 不支持 reasoning_effort={effort}，请选择 default 或该模型支持的等级。") + content = OpenAIChatNode.create_input_message_contents(prompt, images, files) + payload = {**options, "model": model, "input": [{"role": "user", "content": [ + item.model_dump(exclude_none=True) for item in content + ]}], "store": True, "stream": False} + with _Request(cls.hidden.unique_id, comfy_api_key) as api: + result = api.request("POST", BASE_URL + "/proxy/openai/v1/responses", json=payload).json() + api.response_id = result.get("id", "") + while result.get("status") in ("queued", "in_progress"): + if not api.response_id: + raise ValueError("OpenAI response missing ID") + api.sleep(1) + result = api.request("GET", BASE_URL + "/proxy/openai/v1/responses/" + api.response_id).json() + api.check() + if result.get("error") or result.get("status") not in ("completed", "incomplete"): + raise ValueError(f"OpenAI response {result.get('status')}: {result.get('error')}") + texts = [item["text"] for output in result.get("output", []) if output.get("type") == "message" + for item in output.get("content", []) if item.get("type") == "output_text" and item.get("text")] + if not texts: + raise ValueError("OpenAI 未返回文字。") + api.usage = result.get("usage", {}) + return IO.NodeOutput("\n".join(texts)) + + +class PDOpenAIGPTImage25APIKey(IO.ComfyNode): + @classmethod + def define_schema(cls): + return _schema(OpenAIGPTImageNodeV2, "PDOpenAIGPTImage25APIKey", "PD: OpenAI GPT Image 2.5 (ComfyUI Key)", "PD_Tools/Image_Generation") + + @classmethod + def execute(cls, comfy_api_key, prompt, model, n, seed): + model_id = model["model"] + if model_id not in GPT_IMAGE_MODELS: + raise ValueError(f"Unknown model: {model_id}") + if not prompt.strip(): + raise ValueError("请填写提示词。") + if not 1 <= n <= 8: + raise ValueError("图片数量须为 1–8。") + size = model["size"] + if size == "Custom": + size = validate_gpt_image_2_size(size, model["custom_width"], model["custom_height"]) + payload = {"model": model_id, "prompt": prompt, "n": n, "size": size, + "quality": model["quality"], "background": model["background"], "moderation": "low"} + if model_id == "gpt-image-2" and payload["background"] == "transparent": + raise ValueError("gpt-image-2 不支持透明背景。") + images = [item for image in (model.get("images") or {}).values() if image is not None for item in image] + if len(images) > 16: + raise ValueError("最多支持 16 张参考图。") + mask = model.get("mask") + if mask is not None and (len(images) != 1 or mask.shape != (1, *images[0].shape[:2])): + raise ValueError("遮罩要求恰好一张参考图，且尺寸相同。") + with _Request(cls.hidden.unique_id, comfy_api_key) as api: + uploads = [] + for index, image in enumerate(images): + api.check() + uploads.append(("image" if len(images) == 1 else "image[]", + (f"image_{index}.png", tensor_to_bytes(image).getvalue(), "image/png"))) + if mask is not None: + rgba = torch.zeros((*images[0].shape[:2], 4)) + rgba[:, :, 3] = 1 - mask[0].cpu() + scaled = downscale_image_tensor(rgba.unsqueeze(0), total_pixels=2048 * 2048).squeeze(0) + data = io.BytesIO() + Image.fromarray((scaled.numpy() * 255).astype(np.uint8)).save(data, format="PNG") + uploads.append(("mask", ("mask.png", data.getvalue(), "image/png"))) + if uploads: + result = api.request("POST", BASE_URL + "/proxy/openai/images/edits", data=payload, files=uploads).json() + else: + result = api.request("POST", BASE_URL + "/proxy/openai/images/generations", json=payload).json() + tensors = [] + for item in result.get("data", []): + api.check() + if item.get("b64_json"): + data = base64.b64decode(item["b64_json"]) + elif item.get("url"): + url = item["url"] + if url.startswith("/proxy/"): + url = BASE_URL + url + data = api.request("GET", url).content + else: + raise ValueError("图片响应缺少 URL 或 base64 数据。") + api.check() + with Image.open(io.BytesIO(data)) as image: + tensors.append(torch.from_numpy(np.array(image.convert("RGBA")).astype(np.float32) / 255.0)) + if not tensors: + raise ValueError("OpenAI 未返回图片。") + height, width = tensors[0].shape[:2] + for index, image in enumerate(tensors): + if image.shape[:2] != (height, width): + tensors[index] = common_upscale(image.movedim(-1, 0).unsqueeze(0), width, height, "bilinear", "center")[0].movedim(0, -1) + api.check() + api.usage = result.get("usage", {}) + return IO.NodeOutput(torch.stack(tensors)) + + +NODE_CLASS_MAPPINGS = { + "PDOpenAIChatLatestAPIKey": PDOpenAIChatLatestAPIKey, + "PDOpenAIGPTImage25APIKey": PDOpenAIGPTImage25APIKey, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "PDOpenAIChatLatestAPIKey": "PD: OpenAI ChatGPT Latest (ComfyUI Key)", + "PDOpenAIGPTImage25APIKey": "PD: OpenAI GPT Image 2.5 (ComfyUI Key)", +} diff --git a/tests/test_nano_banana21.py b/tests/test_nano_banana21.py new file mode 100644 index 0000000..5edc040 --- /dev/null +++ b/tests/test_nano_banana21.py @@ -0,0 +1,160 @@ +import base64 +import importlib +import io +import sys +import threading +import time +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +import requests +import torch +from PIL import Image + + +ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(ROOT)) +sys.argv = [sys.argv[0], "--cpu"] +package = types.ModuleType("pd_api_test") +package.__path__ = [str(Path(__file__).resolve().parents[1])] +sys.modules[package.__name__] = package +module = importlib.import_module("pd_api_test.py.IMAGE_GeminiNanoBanana21") +Node = module.PDGeminiNanoBanana21 +Node.hidden = types.SimpleNamespace(unique_id="test-node") + + +def image_part(color, thought=False): + data = io.BytesIO() + Image.new("RGB", (8, 8), color).save(data, format="PNG") + return {"thought": thought, "inlineData": {"mimeType": "image/png", "data": base64.b64encode(data.getvalue()).decode()}} + + +class Response: + def raise_for_status(self): + pass + + def json(self): + return {"candidates": [{"content": {"parts": [ + image_part("red", True), image_part("green"), {"text": "generated"}, + ]}}], "usageMetadata": {"totalTokenCount": 123}} + + +class Session: + def __init__(self, slow=False, timeout=False, download=False): + self.slow = slow + self.timeout = timeout + self.download = download + self.calls = [] + self.closed = threading.Event() + + def request(self, method, url, **kwargs): + self.calls.append((method, url, kwargs)) + if self.timeout: + raise requests.Timeout("simulated") + if self.slow and (not self.download or method == "GET"): + self.closed.wait(5) + response = Response() + if self.download and method == "POST": + response.json = lambda: {"candidates": [{"content": {"parts": [ + {"fileData": {"mimeType": "image/png", "fileUri": "https://example.test/image.png"}}, + ]}}]} + return response + + def close(self): + self.closed.set() + + +class NanoBananaTests(unittest.TestCase): + def setUp(self): + self.statuses = [] + self.cancel = threading.Event() + self.patches = [ + patch.object(module.PromptServer, "instance", types.SimpleNamespace(send_sync=self.status), create=True), + patch.object(module.comfy.model_management, "throw_exception_if_processing_interrupted", self.interrupted), + ] + for item in self.patches: + item.start() + + def tearDown(self): + for item in reversed(self.patches): + item.stop() + + def status(self, event, data): + self.assertEqual(event, "progress_text") + self.statuses.append((time.monotonic(), data["text"])) + + def interrupted(self): + if self.cancel.is_set(): + raise module.comfy.model_management.InterruptProcessingException() + + def run_node(self, choice="Gemini Nano Banana 2.1", **model_args): + model = {"model": choice, "resolution": "1K", "aspect_ratio": "1:1", "thinking_level": "MINIMAL"} + model.update(model_args) + return Node.execute("test-key", "test prompt", model, 42, "IMAGE+TEXT") + + def test_models_payload_and_separate_thought_image(self): + for choice, endpoint in module.MODELS.items(): + session = Session() + with patch.object(module.requests, "Session", return_value=session): + result = self.run_node(choice) + self.assertEqual(len(session.calls), 1) + self.assertTrue(session.calls[0][1].endswith(endpoint)) + config = session.calls[0][2]["json"]["generationConfig"] + self.assertEqual("temperature" in config, choice != "Gemini Nano Banana 2.1") + self.assertNotIn("seed", config) + self.assertEqual(result.result[1], "generated") + self.assertTrue(torch.equal(result.result[0][0, 0, 0], torch.tensor([0., 128 / 255, 0.]))) + self.assertTrue(torch.equal(result.result[2][0, 0, 0], torch.tensor([1., 0., 0.]))) + + def test_cancel_slow_generation_and_download(self): + for download in (False, True): + self.cancel.clear() + self.statuses.clear() + session = Session(slow=True, download=download) + errors = [] + def execute(): + try: + self.run_node() + except module.comfy.model_management.InterruptProcessingException as exc: + errors.append(exc) + with patch.object(module.requests, "Session", return_value=session): + worker = threading.Thread(target=execute) + worker.start() + time.sleep(1.25) + started = time.monotonic() + self.cancel.set() + worker.join(0.5) + self.assertFalse(worker.is_alive()) + self.assertLess(time.monotonic() - started, 0.5) + self.assertEqual(len(errors), 1) + self.assertTrue(session.closed.is_set()) + self.assertEqual(sum(call[0] == "POST" for call in session.calls), 1) + running = [item for item in self.statuses if item[1].startswith("Running")] + self.assertGreaterEqual(len(running), 2) + self.assertTrue(all("s" in item[1] for item in running)) + self.assertTrue(self.statuses[-1][1].startswith("Cancelled")) + + def test_timeout_does_not_retry(self): + session = Session(timeout=True) + with patch.object(module.requests, "Session", return_value=session), self.assertRaises(requests.Timeout): + self.run_node() + self.assertEqual(len(session.calls), 1) + self.assertTrue(self.statuses[-1][1].startswith("Timed out")) + + def test_validation_and_cancel_before_submission(self): + session = Session() + with patch.object(module.requests, "Session", return_value=session): + with self.assertRaises(ValueError): + self.run_node("Nano Banana 2 Lite", resolution="4K") + with self.assertRaises(ValueError): + self.run_node(images={"image_1": torch.zeros((15, 8, 8, 3))}) + self.cancel.set() + with self.assertRaises(module.comfy.model_management.InterruptProcessingException): + self.run_node() + self.assertEqual(session.calls, []) + + +if __name__ == "__main__": + unittest.main(argv=[sys.argv[0]]) diff --git a/tests/test_openai_latest.py b/tests/test_openai_latest.py new file mode 100644 index 0000000..be88af3 --- /dev/null +++ b/tests/test_openai_latest.py @@ -0,0 +1,188 @@ +import base64 +import importlib +import io +import sys +import threading +import time +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +import requests +import torch +from PIL import Image + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) +sys.argv = [sys.argv[0], "--cpu"] +package = types.ModuleType("pd_openai_test") +package.__path__ = [str(Path(__file__).resolve().parents[1])] +sys.modules[package.__name__] = package +module = importlib.import_module("pd_openai_test.py.OPENAI_latest_apikey") +from comfy_api_nodes.apis.openai import InputFileContent, ModelResponseProperties, Reasoning + +Chat = module.PDOpenAIChatLatestAPIKey +ImageNode = module.PDOpenAIGPTImage25APIKey +for node in (Chat, ImageNode): + node.hidden = types.SimpleNamespace(unique_id="test-node") + + +def png(size=(8, 8)): + data = io.BytesIO() + Image.new("RGBA", size, (255, 0, 0, 80)).save(data, format="PNG") + return data.getvalue() + + +class Response: + def __init__(self, data): + self.data = data + self.content = png() + + def json(self): + return self.data + + def raise_for_status(self): + pass + + +class Session: + def __init__(self, slow=None, timeout=False, queued=False, url_image=False): + self.calls = [] + self.closed = threading.Event() + self.slow = slow + self.timeout = timeout + self.queued = queued + self.url_image = url_image + + def request(self, method, url, **kwargs): + self.calls.append((method, url, kwargs)) + if self.timeout: + raise requests.Timeout("simulated") + if method == self.slow: + self.closed.wait(5) + if "/responses" in url: + if method == "POST" and self.queued: + return Response({"id": "resp_test", "status": "queued"}) + return Response({"id": "resp_test", "status": "completed", "usage": {"total_tokens": 12}, + "output": [{"type": "message", "content": [ + {"type": "output_text", "text": "one"}, {"type": "output_text", "text": "two"}, + ]}]}) + if method == "GET": + return Response({}) + if self.url_image: + return Response({"data": [{"url": "https://example.test/image.png"}]}) + return Response({"data": [{"b64_json": base64.b64encode(png()).decode()}, + {"b64_json": base64.b64encode(png((10, 12))).decode()}]}) + + def close(self): + self.closed.set() + + +class OpenAITests(unittest.TestCase): + def setUp(self): + self.cancel = threading.Event() + self.events = [] + self.patches = [ + patch.object(module.PromptServer, "instance", types.SimpleNamespace(send_sync=lambda event, data: self.events.append(data)), create=True), + patch.object(module.comfy.model_management, "throw_exception_if_processing_interrupted", self.check), + ] + for item in self.patches: + item.start() + + def tearDown(self): + for item in reversed(self.patches): + item.stop() + + def check(self): + if self.cancel.is_set(): + raise module.comfy.model_management.InterruptProcessingException() + + def chat(self, **kwargs): + return Chat.execute("test-key", "test prompt", kwargs.pop("model", "gpt-6-astra"), **kwargs) + + def image(self, **kwargs): + model = {"model": "gpt-image-2.5-flare", "size": "1024x1024", "quality": "low", "background": "transparent"} + model.update(kwargs) + return ImageNode.execute("test-key", "test prompt", model, 1, 0) + + def test_chat_models_files_options_and_poll(self): + self.assertEqual(len(list(module.SupportedOpenAIModel)), 18) + for model in module.SupportedOpenAIModel: + session = Session(queued=True) + with patch.object(module.requests, "Session", return_value=session), patch.object(module._Request, "sleep"): + result = self.chat(model=model.value, images=torch.zeros((1, 8, 8, 3)), + files=[InputFileContent(type="input_file", filename="test.pdf", file_data="data:application/pdf;base64,AA==")], + advanced_options=ModelResponseProperties(instructions="system", max_output_tokens=32)) + self.assertEqual(result.result, ("one\ntwo",)) + self.assertEqual([item[0] for item in session.calls], ["POST", "GET"]) + body = session.calls[0][2]["json"] + self.assertEqual(body["model"], model.value) + self.assertEqual(body["instructions"], "system") + self.assertEqual([item["type"] for item in body["input"][0]["content"]], ["input_text", "input_image", "input_file"]) + self.assertNotIn("previous_response_id", body) + + def test_image_models_generation_edit_mask_and_alpha(self): + for model in module.GPT_IMAGE_MODELS: + session = Session() + with patch.object(module.requests, "Session", return_value=session): + result = self.image(model=model, background="auto") + self.assertEqual(result.result[0].shape, (2, 8, 8, 4)) + self.assertAlmostEqual(result.result[0][0, 0, 0, 3].item(), 80 / 255) + self.assertTrue(session.calls[0][1].endswith("/images/generations")) + self.assertNotIn("seed", session.calls[0][2]["json"]) + session = Session() + with patch.object(module.requests, "Session", return_value=session): + self.image(size="Custom", custom_width=1024, custom_height=1024, + images={"image_1": torch.zeros((1, 8, 8, 3))}, mask=torch.ones((1, 8, 8))) + self.assertTrue(session.calls[0][1].endswith("/images/edits")) + self.assertEqual([item[0] for item in session.calls[0][2]["files"]], ["image", "mask"]) + with Image.open(io.BytesIO(session.calls[0][2]["files"][1][1][1])) as mask: + self.assertEqual(mask.getpixel((0, 0))[3], 0) + + def test_cancellation_generation_poll_and_download(self): + cases = [(self.chat, Session(slow="POST")), (self.chat, Session(slow="GET", queued=True)), + (self.image, Session(slow="POST")), (self.image, Session(slow="GET", url_image=True))] + for action, session in cases: + self.cancel.clear() + self.events.clear() + errors = [] + def run(): + try: + action() + except module.comfy.model_management.InterruptProcessingException as exc: + errors.append(exc) + with patch.object(module.requests, "Session", return_value=session): + worker = threading.Thread(target=run) + worker.start() + time.sleep(1.25) + started = time.monotonic() + self.cancel.set() + worker.join(0.5) + self.assertFalse(worker.is_alive()) + self.assertLess(time.monotonic() - started, 0.5) + self.assertEqual(len(errors), 1) + self.assertEqual(sum(item[0] == "POST" for item in session.calls), 1) + self.assertTrue(session.closed.is_set()) + self.assertTrue(self.events[-1]["text"].startswith("Cancelled")) + self.assertGreaterEqual(sum(item["text"].startswith("Running") for item in self.events), 2) + + def test_timeout_and_invalid_inputs_never_retry(self): + session = Session(timeout=True) + with patch.object(module.requests, "Session", return_value=session), self.assertRaises(requests.Timeout): + self.chat() + self.assertEqual(len(session.calls), 1) + session = Session() + with patch.object(module.requests, "Session", return_value=session): + for action in ( + lambda: self.chat(model="gpt-4.1", advanced_options=ModelResponseProperties(reasoning=Reasoning(effort="high"))), + lambda: self.image(images={"image_1": torch.zeros((17, 8, 8, 3))}), + lambda: self.image(mask=torch.zeros((1, 8, 8))), + lambda: self.image(size="Custom", custom_width=1025, custom_height=1024), + ): + with self.assertRaises(ValueError): + action() + self.assertEqual(session.calls, []) + + +if __name__ == "__main__": + unittest.main(argv=[sys.argv[0]])