Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
345 lines
13 KiB
Python
345 lines
13 KiB
Python
"""Offline regression tests for the standalone GPT Image nodes."""
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from io import BytesIO
|
|
from types import ModuleType
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
PLUGIN_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
CUSTOM_NODES_DIR = os.path.dirname(PLUGIN_DIR)
|
|
COMFY_ROOT = os.path.dirname(CUSTOM_NODES_DIR)
|
|
for child in ("nodes", "utils", "clients"):
|
|
package = ModuleType(f"comfyui_o1key.{child}")
|
|
package.__path__ = [os.path.join(PLUGIN_DIR, child)]
|
|
sys.modules[package.__name__] = package
|
|
plugin_package = ModuleType("comfyui_o1key")
|
|
plugin_package.__path__ = [PLUGIN_DIR]
|
|
sys.modules[plugin_package.__name__] = plugin_package
|
|
model_management = ModuleType("comfy.model_management")
|
|
model_management.processing_interrupted = lambda: False
|
|
model_management.InterruptProcessingException = RuntimeError
|
|
sys.modules[model_management.__name__] = model_management
|
|
sys.path.insert(0, COMFY_ROOT)
|
|
sys.path.insert(0, CUSTOM_NODES_DIR)
|
|
|
|
from comfyui_o1key.clients.gpt_image_client import ( # noqa: E402
|
|
GptImageClient,
|
|
resolve_gpt_image_model,
|
|
)
|
|
from comfyui_o1key.nodes.gpt_image import ( # noqa: E402
|
|
O1keyGPTImage,
|
|
resolve_gpt_image_quality,
|
|
)
|
|
from comfyui_o1key.nodes.gpt_image_batch import ( # noqa: E402
|
|
O1keyGPTImageBatch,
|
|
_path_option,
|
|
)
|
|
from comfyui_o1key.utils.file_utils import ImageInfo # noqa: E402
|
|
|
|
|
|
class GPTImageNodeSchemaTests(unittest.TestCase):
|
|
def assert_combo_values_are_strings(self, schema):
|
|
for item in schema.inputs:
|
|
if item.io_type != "COMBO":
|
|
continue
|
|
self.assertTrue(
|
|
all(isinstance(option, str) for option in item.options),
|
|
item.id,
|
|
)
|
|
if item.default is not None:
|
|
self.assertIsInstance(item.default, str, item.id)
|
|
|
|
def test_single_node_keeps_request_controls_without_color_correction(self):
|
|
schema = O1keyGPTImage.define_schema()
|
|
schema.validate()
|
|
self.assert_combo_values_are_strings(schema)
|
|
inputs = {item.id: item for item in schema.inputs}
|
|
|
|
self.assertEqual(inputs["缩放图片"].options, ["不缩放", "智能缩放"])
|
|
self.assertEqual(inputs["缩放图片"].default, "智能缩放")
|
|
self.assertEqual(inputs["模型"].default, "gpt-image-2.5-sunburst")
|
|
self.assertNotIn("色彩纠正", inputs)
|
|
self.assertEqual(inputs["背景"].options, ["auto", "transparent", "opaque"])
|
|
self.assertNotIn("内容审查强度", inputs)
|
|
self.assertEqual(
|
|
[item.id for item in schema.inputs],
|
|
[
|
|
"prompt", "模型", "模型线路", "分辨率", "生图数量", "质量",
|
|
"输出格式", "背景", "遮罩", "参考图组", "缩放图片", "seed",
|
|
],
|
|
)
|
|
self.assertEqual(
|
|
inputs["质量"].options,
|
|
["高", "中", "低", "自动", "超高", "最高"],
|
|
)
|
|
|
|
def test_25_quality_values_are_model_specific(self):
|
|
self.assertEqual(resolve_gpt_image_quality("gpt-image-2.5-sunburst", "超高"), "xhigh")
|
|
self.assertEqual(resolve_gpt_image_quality("gpt-image-2.5-flare", "最高"), "max")
|
|
self.assertEqual(resolve_gpt_image_quality("gpt-image-2", "高"), "high")
|
|
with self.assertRaisesRegex(ValueError, "仅支持 GPT Image 2.5"):
|
|
resolve_gpt_image_quality("gpt-image-2", "超高")
|
|
|
|
def test_gpt_image_25_models_resolve_all_route_ids(self):
|
|
expected = {
|
|
("gpt-image-2.5-sunburst", "畅速"): "gpt-image-2.5-sunburst-sp",
|
|
("gpt-image-2.5-sunburst", "直连"): "gpt-image-2.5-sunburst-sd",
|
|
("gpt-image-2.5-sunburst", "专线"): "gpt-image-2.5-sunburst",
|
|
("gpt-image-2.5-flare", "畅速"): "gpt-image-2.5-flare-sp",
|
|
("gpt-image-2.5-flare", "直连"): "gpt-image-2.5-flare-sd",
|
|
("gpt-image-2.5-flare", "专线"): "gpt-image-2.5-flare",
|
|
}
|
|
for (model, route), actual_model in expected.items():
|
|
with self.subTest(model=model, route=route):
|
|
self.assertEqual(resolve_gpt_image_model(model, route), actual_model)
|
|
|
|
def test_batch_node_keeps_operational_controls_outside_advanced_inputs(self):
|
|
schema = O1keyGPTImageBatch.define_schema()
|
|
schema.validate()
|
|
self.assert_combo_values_are_strings(schema)
|
|
inputs = {item.id: item for item in schema.inputs}
|
|
|
|
self.assertEqual(inputs["模型"].default, "gpt-image-2.5-sunburst")
|
|
self.assertEqual(inputs["缩放图片"].default, "智能缩放")
|
|
self.assertEqual(inputs["质量"].options, ["高", "中", "低", "自动", "超高", "最高"])
|
|
self.assertNotIn("色彩纠正", inputs)
|
|
self.assertNotIn("内容审查强度", inputs)
|
|
self.assertEqual(
|
|
[item.id for item in schema.inputs],
|
|
[
|
|
"prompt", "模型", "模型线路", "分辨率", "生图数量", "质量",
|
|
"图片路径数量", "遮罩", "参考图组", "图片输出格式",
|
|
"背景", "图片保存命名规则", "图片保存路径", "缩放图片", "seed",
|
|
],
|
|
)
|
|
self.assertNotIn("并发数", inputs)
|
|
self.assertEqual(
|
|
[item.id for item in _path_option(1).inputs],
|
|
["参考图1（主图）"],
|
|
)
|
|
self.assertEqual(
|
|
[item.id for item in _path_option(2).inputs],
|
|
["参考图1（主图）", "参考图2", "图片配对模式"],
|
|
)
|
|
|
|
for input_name in (
|
|
"seed",
|
|
"图片输出格式",
|
|
"图片保存命名规则",
|
|
"图片保存路径",
|
|
"缩放图片",
|
|
"背景",
|
|
):
|
|
self.assertFalse(inputs[input_name].advanced, input_name)
|
|
|
|
class GPTImageNodeExecutionTests(unittest.TestCase):
|
|
def test_single_node_forwards_request_controls_without_postprocess(self):
|
|
calls = []
|
|
generated = torch.ones((1, 4, 6, 3), dtype=torch.float32)
|
|
|
|
class FakeClient:
|
|
def __init__(self):
|
|
self.base_url = ""
|
|
self.response_log_enabled = True
|
|
self.poll_log_enabled = True
|
|
|
|
def generate_image_async_sync(self, **kwargs):
|
|
calls.append(kwargs)
|
|
return [Image.new("RGB", (6, 4), "green")]
|
|
|
|
@staticmethod
|
|
def _pil_list_to_tensor(_images):
|
|
return generated
|
|
|
|
with (
|
|
patch("comfyui_o1key.nodes.gpt_image.GptImageClient", FakeClient),
|
|
patch(
|
|
"comfyui_o1key.nodes.gpt_image.get_base_url_by_route",
|
|
return_value="https://example.invalid",
|
|
),
|
|
):
|
|
result = O1keyGPTImage.generate(
|
|
prompt="测试",
|
|
输出格式="webp",
|
|
缩放图片="智能缩放",
|
|
背景="transparent",
|
|
)
|
|
|
|
self.assertIs(result[0], generated)
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertEqual(calls[0]["model"], "gpt-image-2.5-sunburst-sp")
|
|
self.assertEqual(calls[0]["resize_mode"], "智能缩放")
|
|
self.assertEqual(calls[0]["background"], "transparent")
|
|
self.assertNotIn("moderation", calls[0])
|
|
|
|
def test_single_node_forwards_25_quality_api_values(self):
|
|
calls = []
|
|
generated = torch.ones((1, 4, 6, 3), dtype=torch.float32)
|
|
|
|
class FakeClient:
|
|
def __init__(self):
|
|
self.base_url = ""
|
|
self.response_log_enabled = True
|
|
self.poll_log_enabled = True
|
|
|
|
def generate_image_async_sync(self, **kwargs):
|
|
calls.append(kwargs)
|
|
return [Image.new("RGB", (6, 4), "green")]
|
|
|
|
@staticmethod
|
|
def _pil_list_to_tensor(_images):
|
|
return generated
|
|
|
|
with (
|
|
patch("comfyui_o1key.nodes.gpt_image.GptImageClient", FakeClient),
|
|
patch(
|
|
"comfyui_o1key.nodes.gpt_image.get_base_url_by_route",
|
|
return_value="https://example.invalid",
|
|
),
|
|
):
|
|
for display_quality, api_quality in (("超高", "xhigh"), ("最高", "max")):
|
|
O1keyGPTImage.generate(
|
|
prompt="测试",
|
|
模型="gpt-image-2.5-sunburst",
|
|
质量=display_quality,
|
|
)
|
|
|
|
self.assertEqual([call["quality"] for call in calls], ["xhigh", "max"])
|
|
|
|
|
|
class GPTImageBatchExecutionTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_batch_task_forwards_request_controls_without_removed_options(self):
|
|
calls = []
|
|
reference = Image.new("RGB", (6, 4), "red")
|
|
generated = Image.new("RGB", (6, 4), "green")
|
|
pair = (ImageInfo(reference, "source", ".png", ""),)
|
|
|
|
class FakeClient:
|
|
async def generate_image_async(self, **kwargs):
|
|
calls.append(kwargs)
|
|
return [generated]
|
|
|
|
with (
|
|
patch.object(
|
|
O1keyGPTImageBatch,
|
|
"_save_images",
|
|
return_value=["output.png"],
|
|
) as save_images,
|
|
patch("builtins.print"),
|
|
):
|
|
result = await O1keyGPTImageBatch._run_task(
|
|
FakeClient(), pair, "测试", 0, 1,
|
|
"gpt-image-2-c-sp", "auto", "1024x1024", 1, 0, None, "webp",
|
|
"transparent", "智能缩放",
|
|
"output", "和原始图片名保持一致", asyncio.Lock(), None,
|
|
)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(calls[0]["resize_mode"], "智能缩放")
|
|
self.assertEqual(calls[0]["background"], "transparent")
|
|
self.assertNotIn("moderation", calls[0])
|
|
self.assertIs(save_images.call_args.args[0][0], generated)
|
|
|
|
async def test_batch_starts_all_tasks_together(self):
|
|
task_count = 4
|
|
started = []
|
|
all_started = asyncio.Event()
|
|
|
|
async def fake_run_task(cls, *_args):
|
|
task_index = _args[3]
|
|
started.append(task_index)
|
|
if len(started) == task_count:
|
|
all_started.set()
|
|
await asyncio.wait_for(all_started.wait(), timeout=1)
|
|
return {
|
|
"task_index": task_index,
|
|
"success": True,
|
|
"generated_count": 1,
|
|
"saved_files": [],
|
|
"error": None,
|
|
}
|
|
|
|
task_defs = [(index, tuple(), f"prompt {index}") for index in range(task_count)]
|
|
with patch.object(O1keyGPTImageBatch, "_run_task", classmethod(fake_run_task)), patch("builtins.print"):
|
|
results = await O1keyGPTImageBatch._process_async(
|
|
object(), task_defs, "model", "auto", "auto", 1, 0,
|
|
None, "png", "auto", "智能缩放", "output",
|
|
"自然数字", None,
|
|
)
|
|
|
|
self.assertEqual(sorted(started), list(range(task_count)))
|
|
self.assertEqual(len(results), task_count)
|
|
self.assertTrue(all(item["success"] for item in results))
|
|
|
|
|
|
class GPTImageDownloadRetryTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_result_download_retries_transient_http_status_without_logging_url(self):
|
|
buffer = BytesIO()
|
|
Image.new("RGB", (3, 2), "purple").save(buffer, format="PNG")
|
|
image_bytes = buffer.getvalue()
|
|
|
|
class Content:
|
|
def __init__(self, body):
|
|
self.body = body
|
|
|
|
async def iter_chunked(self, _chunk_size):
|
|
yield self.body
|
|
|
|
class Response:
|
|
http_version = "HTTP/1.1"
|
|
|
|
def __init__(self, status, body):
|
|
self.status = status
|
|
self.headers = {"Content-Length": str(len(body))}
|
|
self.content = Content(body)
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, _exc_type, _exc, _tb):
|
|
return None
|
|
|
|
class Session:
|
|
def __init__(self):
|
|
self.responses = [Response(503, b"busy"), Response(200, image_bytes)]
|
|
self.calls = 0
|
|
|
|
def get(self, _url, **_kwargs):
|
|
response = self.responses[self.calls]
|
|
self.calls += 1
|
|
return response
|
|
|
|
client = object.__new__(GptImageClient)
|
|
session = Session()
|
|
signed_url = "https://example.invalid/private?signature=secret"
|
|
with (
|
|
patch(
|
|
"comfyui_o1key.clients.gpt_image_client.asyncio.sleep",
|
|
new=AsyncMock(),
|
|
),
|
|
patch("builtins.print") as print_mock,
|
|
):
|
|
image, byte_count, _elapsed = await client._download_image_with_response_retry(
|
|
session,
|
|
signed_url,
|
|
"第 1 张图片",
|
|
)
|
|
|
|
self.assertEqual(session.calls, 2)
|
|
self.assertEqual(image.size, (3, 2))
|
|
self.assertEqual(byte_count, len(image_bytes))
|
|
rendered = " ".join(str(value) for call in print_mock.call_args_list for value in call.args)
|
|
self.assertNotIn(signed_url, rendered)
|
|
self.assertNotIn("signature=secret", rendered)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|