Add Comfy API image and chat nodes
This commit is contained in:
pandy
@@ -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` 发送。
|
||||
|
||||
@@ -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)"}
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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)",
|
||||
}
|
||||
@@ -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]])
|
||||
@@ -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]])
|
||||
Reference in New Issue
Block a user