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.
1559 lines
62 KiB
Python
1559 lines
62 KiB
Python
import asyncio
|
||
import os
|
||
from pathlib import Path
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
from io import BytesIO
|
||
import unittest
|
||
import uuid
|
||
from unittest.mock import AsyncMock, patch
|
||
|
||
from PIL import Image
|
||
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
CUSTOM_NODES_ROOT = ROOT.parent
|
||
COMFY_ROOT = CUSTOM_NODES_ROOT.parent
|
||
sys.path.insert(0, str(COMFY_ROOT))
|
||
sys.path.insert(0, str(CUSTOM_NODES_ROOT))
|
||
|
||
from comfyui_o1key.utils import o1key_image_jobs as JOBS # noqa: E402
|
||
|
||
|
||
def _payload(batch_id=None, generator_id=1, save_id=2):
|
||
return {
|
||
"batch_id": batch_id or str(uuid.uuid4()),
|
||
"generator_node_id": generator_id,
|
||
"save_node_id": save_id,
|
||
"prompt": "一只宇航猫",
|
||
"model": "Nano Banana 2",
|
||
"model_route": "畅速",
|
||
"thinking_level": "高",
|
||
"resolution": "2K",
|
||
"aspect_ratio": "1:1",
|
||
"image_count": 1,
|
||
"seed": 42,
|
||
"references": [],
|
||
}
|
||
|
||
|
||
class PayloadTests(unittest.TestCase):
|
||
def test_missing_or_smart_resolution_uses_provider_default(self):
|
||
missing = _payload()
|
||
missing.pop("resolution")
|
||
self.assertEqual(
|
||
JOBS.normalize_job_payload(missing)["resolution"],
|
||
"智能",
|
||
)
|
||
|
||
for model in ("Nano Banana 2", "gpt-image-2", "Seedream 5.0 Pro"):
|
||
with self.subTest(model=model):
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": model,
|
||
"resolution": "智能",
|
||
})
|
||
self.assertEqual(job["resolution"], "智能")
|
||
|
||
def test_normalizes_an_immutable_batch_identity_and_reference_order(self):
|
||
payload = _payload()
|
||
payload["references"] = [
|
||
{"name": "one.jpg", "subfolder": "refs", "type": "input"},
|
||
{"name": "two.png", "subfolder": "", "type": "input"},
|
||
]
|
||
job = JOBS.normalize_job_payload(payload)
|
||
payload["references"].reverse()
|
||
|
||
self.assertEqual(job["batch_id"], payload["batch_id"])
|
||
self.assertEqual([item["name"] for item in job["references"]], ["one.jpg", "two.png"])
|
||
self.assertEqual(job["generator_node_id"], 1)
|
||
self.assertEqual(job["save_node_id"], 2)
|
||
|
||
def test_missing_thinking_level_defaults_to_low(self):
|
||
payload = _payload()
|
||
payload.pop("thinking_level")
|
||
job = JOBS.normalize_job_payload(payload)
|
||
self.assertEqual(job["thinking_level"], "低")
|
||
self.assertEqual(job["resize_mode"], "不缩放")
|
||
self.assertNotIn("color_correction", job)
|
||
self.assertEqual(job["filename_prefix"], "o1key")
|
||
self.assertEqual(job["save_format"], "原始")
|
||
self.assertEqual(job["save_location"], "")
|
||
self.assertEqual(job["naming_rule"], "自定义前缀")
|
||
self.assertEqual(job["task_prompts"], ["一只宇航猫"])
|
||
self.assertEqual(job["total_task_count"], 1)
|
||
self.assertNotIn("output_format", job)
|
||
self.assertNotIn("background", job)
|
||
self.assertNotIn("moderation", job)
|
||
self.assertNotIn("google_search", job)
|
||
|
||
smart_job = JOBS.normalize_job_payload({**_payload(), "resize_mode": "智能缩放"})
|
||
self.assertEqual(smart_job["resize_mode"], "智能缩放")
|
||
|
||
with self.assertRaisesRegex(ValueError, "缩放图片"):
|
||
JOBS.normalize_job_payload({**_payload(), "resize_mode": "拉伸"})
|
||
|
||
legacy_color = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"color_correction": "legacy-value-is-ignored",
|
||
})
|
||
self.assertNotIn("color_correction", legacy_color)
|
||
|
||
def test_rejects_invalid_uuid_and_model_combinations(self):
|
||
payload = _payload(batch_id="not-a-uuid")
|
||
with self.assertRaisesRegex(ValueError, "批次 ID"):
|
||
JOBS.normalize_job_payload(payload)
|
||
|
||
payload = _payload()
|
||
payload.update({"model": "Nano Banana Pro", "resolution": "512"})
|
||
with self.assertRaisesRegex(ValueError, "Nano Banana 2"):
|
||
JOBS.normalize_job_payload(payload)
|
||
|
||
def test_unified_limits_and_gpt_parameters(self):
|
||
payload = _payload()
|
||
payload.update({
|
||
"model": "gpt-image-2",
|
||
"resolution": "2K",
|
||
"aspect_ratio": "3:2",
|
||
"quality": "高",
|
||
"output_format": "webp",
|
||
"background": "transparent",
|
||
"moderation": "low",
|
||
"image_count": 9,
|
||
"references": [
|
||
{"name": f"{index}.png", "subfolder": "", "type": "input"}
|
||
for index in range(10)
|
||
],
|
||
})
|
||
job = JOBS.normalize_job_payload(payload)
|
||
self.assertEqual(job["image_count"], 9)
|
||
self.assertEqual(len(job["references"]), 10)
|
||
self.assertEqual(job["quality"], "高")
|
||
self.assertEqual(job["output_format"], "webp")
|
||
self.assertEqual(job["background"], "transparent")
|
||
self.assertNotIn("moderation", job)
|
||
|
||
default_gpt = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_id": str(uuid.uuid4()),
|
||
"model": "gpt-image-2",
|
||
"resolution": "1K",
|
||
"aspect_ratio": "1:1",
|
||
"save_format": "png",
|
||
})
|
||
self.assertEqual(default_gpt["output_format"], "png")
|
||
self.assertEqual(default_gpt["resize_mode"], "智能缩放")
|
||
self.assertEqual(default_gpt["save_format"], "原始")
|
||
|
||
for model in ("gpt-image-2.5-sunburst", "gpt-image-2.5-flare"):
|
||
with self.subTest(model=model):
|
||
gpt_25 = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_id": str(uuid.uuid4()),
|
||
"model": model,
|
||
"resolution": "4K",
|
||
"aspect_ratio": "9:16",
|
||
})
|
||
self.assertEqual(gpt_25["model"], model)
|
||
self.assertEqual(gpt_25["resolution"], "4K")
|
||
self.assertEqual(gpt_25["aspect_ratio"], "9:16")
|
||
self.assertEqual(gpt_25["output_format"], "png")
|
||
self.assertEqual(gpt_25["save_format"], "原始")
|
||
|
||
expanded = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_id": str(uuid.uuid4()),
|
||
"model": model,
|
||
"image_count": 3,
|
||
"quality": "超高",
|
||
})
|
||
self.assertEqual(expanded["image_count"], 3)
|
||
self.assertEqual(expanded["quality"], "超高")
|
||
|
||
with self.assertRaisesRegex(ValueError, "仅支持 GPT Image 2.5"):
|
||
JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_id": str(uuid.uuid4()),
|
||
"model": "gpt-image-2",
|
||
"quality": "最高",
|
||
})
|
||
|
||
with self.assertRaisesRegex(ValueError, "透明背景仅支持 PNG 或 WebP"):
|
||
JOBS.normalize_job_payload({
|
||
**payload,
|
||
"batch_id": str(uuid.uuid4()),
|
||
"output_format": "jpeg",
|
||
"background": "transparent",
|
||
})
|
||
|
||
legacy_payload = {
|
||
**payload,
|
||
"batch_id": str(uuid.uuid4()),
|
||
"resolution": "3648x2048（2K 横版 16:9）",
|
||
"aspect_ratio": "1:1",
|
||
}
|
||
legacy_payload.pop("moderation", None)
|
||
legacy_job = JOBS.normalize_job_payload(legacy_payload)
|
||
self.assertEqual(legacy_job["resolution"], "3648x2048（2K 横版 16:9）")
|
||
self.assertNotIn("moderation", legacy_job)
|
||
|
||
self.assertNotIn("moderation", JOBS.normalize_job_payload({
|
||
**payload,
|
||
"batch_id": str(uuid.uuid4()),
|
||
"moderation": "high",
|
||
}))
|
||
|
||
payload["references"].append(
|
||
{"name": "overflow.png", "subfolder": "", "type": "input"}
|
||
)
|
||
with self.assertRaisesRegex(ValueError, "最多支持 10 张"):
|
||
JOBS.normalize_job_payload(payload)
|
||
|
||
payload["references"] = []
|
||
payload["mask"] = {"name": "mask.png", "subfolder": "", "type": "input"}
|
||
with self.assertRaisesRegex(ValueError, "未提供参考图"):
|
||
JOBS.normalize_job_payload(payload)
|
||
|
||
def test_nano_ignores_gpt_only_output_parameters(self):
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"output_format": "webp",
|
||
"background": "transparent",
|
||
"moderation": "low",
|
||
})
|
||
self.assertNotIn("output_format", job)
|
||
self.assertNotIn("background", job)
|
||
self.assertNotIn("moderation", job)
|
||
|
||
def test_seedream_keeps_output_format_and_forces_original_save(self):
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "Seedream 5.0 Pro",
|
||
"resolution": "2K",
|
||
"aspect_ratio": "21:9",
|
||
"output_format": "png",
|
||
"background": "transparent",
|
||
"moderation": "low",
|
||
"save_format": "webp",
|
||
})
|
||
self.assertEqual(job["output_format"], "png")
|
||
self.assertEqual(job["save_format"], "原始")
|
||
self.assertNotIn("background", job)
|
||
self.assertNotIn("moderation", job)
|
||
|
||
with self.assertRaisesRegex(ValueError, "png 或 jpeg"):
|
||
JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "Seedream 5.0 Pro",
|
||
"resolution": "2K",
|
||
"aspect_ratio": "1:1",
|
||
"output_format": "webp",
|
||
})
|
||
|
||
with self.assertRaisesRegex(ValueError, "分辨率"):
|
||
JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "Seedream 5.0 Pro",
|
||
"resolution": "4K",
|
||
"aspect_ratio": "1:1",
|
||
"output_format": "png",
|
||
})
|
||
|
||
def test_seedream_layer_decomposition_has_strict_preflight_validation(self):
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"prompt": "",
|
||
"model": "Seedream 5.0 Pro",
|
||
"resolution": "auto",
|
||
"output_format": "png",
|
||
"layer_decomposition": True,
|
||
"references": [{"name": "poster.png", "subfolder": "", "type": "input"}],
|
||
})
|
||
self.assertIs(job["layer_decomposition"], True)
|
||
self.assertEqual(job["task_prompts"], [""])
|
||
|
||
with self.assertRaisesRegex(ValueError, "必须且只能上传1张"):
|
||
JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"prompt": "",
|
||
"model": "Seedream 5.0 Pro",
|
||
"resolution": "2K",
|
||
"output_format": "png",
|
||
"layer_decomposition": True,
|
||
})
|
||
with self.assertRaisesRegex(ValueError, "不支持批量出图"):
|
||
JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "Seedream 5.0 Pro",
|
||
"resolution": "2K",
|
||
"output_format": "png",
|
||
"layer_decomposition": True,
|
||
"batch_enabled": True,
|
||
"references": [{"name": "poster.png", "subfolder": "", "type": "input"}],
|
||
})
|
||
|
||
def test_google_search_is_kept_only_for_nano_banana_2_when_true(self):
|
||
enabled = JOBS.normalize_job_payload({**_payload(), "google_search": True})
|
||
self.assertIs(enabled["google_search"], True)
|
||
|
||
disabled = JOBS.normalize_job_payload({**_payload(), "google_search": False})
|
||
self.assertNotIn("google_search", disabled)
|
||
|
||
unsupported = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "Nano Banana Pro",
|
||
"google_search": True,
|
||
})
|
||
self.assertNotIn("google_search", unsupported)
|
||
|
||
def test_batch_prompts_expand_prompt_major_and_enforce_total_limit(self):
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"prompt": "第一条\n---\n第二条\n---\n第三条",
|
||
"image_count": 2,
|
||
})
|
||
self.assertEqual(
|
||
job["task_prompts"],
|
||
["第一条", "第一条", "第二条", "第二条", "第三条", "第三条"],
|
||
)
|
||
self.assertEqual(job["total_task_count"], 6)
|
||
|
||
too_many_prompts = "\n---\n".join(
|
||
f"提示词 {index}" for index in range(JOBS.MAX_UNIFIED_IMAGE_TASKS + 1)
|
||
)
|
||
with self.assertRaisesRegex(ValueError, "最多支持 1000 个"):
|
||
JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"prompt": too_many_prompts,
|
||
})
|
||
|
||
def test_batch_pairing_expands_ten_and_one_hundred_tasks(self):
|
||
outfits = [
|
||
{"name": f"outfit-{index}.png", "subfolder": "", "type": "input"}
|
||
for index in range(10)
|
||
]
|
||
models = [
|
||
{"name": f"model-{index}.png", "subfolder": "", "type": "input"}
|
||
for index in range(10)
|
||
]
|
||
grouped = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "一组搭配+多模特",
|
||
"references": outfits[:9],
|
||
"model_references": models,
|
||
})
|
||
self.assertEqual(grouped["total_task_count"], 10)
|
||
self.assertEqual(grouped["tasks"][0]["reference_indices"], tuple(range(9)))
|
||
self.assertEqual(grouped["tasks"][9]["model_reference_indices"], (9,))
|
||
|
||
cartesian = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "全部搭配×全部模特",
|
||
"references": outfits,
|
||
"model_references": models,
|
||
})
|
||
self.assertEqual(cartesian["total_task_count"], 100)
|
||
self.assertEqual(cartesian["tasks"][0]["reference_indices"], (0,))
|
||
self.assertEqual(cartesian["tasks"][9]["model_reference_indices"], (9,))
|
||
self.assertEqual(cartesian["tasks"][10]["reference_indices"], (1,))
|
||
|
||
single_references = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "单图素材批量",
|
||
"references": outfits[:3],
|
||
"model_references": models[:1],
|
||
})
|
||
self.assertEqual(single_references["total_task_count"], 3)
|
||
self.assertEqual(
|
||
[task["reference_indices"] for task in single_references["tasks"]],
|
||
[(0,), (1,), (2,)],
|
||
)
|
||
self.assertTrue(
|
||
all(
|
||
task["model_reference_indices"] == ()
|
||
for task in single_references["tasks"]
|
||
)
|
||
)
|
||
self.assertEqual(single_references["model_references"], [])
|
||
|
||
fifty_references = [
|
||
{"name": f"single-{index}.png", "subfolder": "", "type": "input"}
|
||
for index in range(50)
|
||
]
|
||
fifty_models = [
|
||
{"name": f"target-{index}.png", "subfolder": "", "type": "input"}
|
||
for index in range(50)
|
||
]
|
||
single_fifty = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "单图素材批量",
|
||
"references": fifty_references,
|
||
})
|
||
self.assertEqual(single_fifty["total_task_count"], 50)
|
||
|
||
grouped_fifty_targets = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "一组搭配+多模特",
|
||
"references": outfits[:1],
|
||
"model_references": fifty_models,
|
||
})
|
||
self.assertEqual(grouped_fifty_targets["total_task_count"], 50)
|
||
|
||
with self.assertRaisesRegex(ValueError, "最多支持 50 张"):
|
||
JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "单图素材批量",
|
||
"references": [
|
||
*fifty_references,
|
||
{"name": "overflow.png", "subfolder": "", "type": "input"},
|
||
],
|
||
})
|
||
|
||
with self.assertRaisesRegex(ValueError, "最多支持 9 张素材图"):
|
||
JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "一组搭配+多模特",
|
||
"references": outfits,
|
||
"model_references": models[:1],
|
||
})
|
||
|
||
with self.assertRaisesRegex(ValueError, "暂不支持蒙版"):
|
||
JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "gpt-image-2",
|
||
"resolution": "1K",
|
||
"batch_enabled": True,
|
||
"references": outfits[:1],
|
||
"model_references": models[:1],
|
||
"mask": {"name": "mask.png", "subfolder": "", "type": "input"},
|
||
})
|
||
|
||
|
||
class NodeErrorFormattingTests(unittest.IsolatedAsyncioTestCase):
|
||
async def test_error_wrappers_apply_after_all_model_dispatches(self):
|
||
raw = "status_code=403, insufficient balance"
|
||
expected = "上游额度不足！"
|
||
cases = (
|
||
("Nano Banana 2", "_execute_nano_banana_job"),
|
||
("gpt-image-2", "_execute_gpt_image_job"),
|
||
("Seedream 5.0 Pro", "_execute_seedream_job"),
|
||
)
|
||
|
||
for model, executor_name in cases:
|
||
with self.subTest(model=model), patch.object(
|
||
JOBS,
|
||
executor_name,
|
||
new=AsyncMock(side_effect=RuntimeError(raw)),
|
||
):
|
||
with self.assertRaisesRegex(RuntimeError, expected):
|
||
await JOBS.execute_image_job(
|
||
{"model": model},
|
||
"unused",
|
||
)
|
||
|
||
|
||
class SeedreamExecutionTests(unittest.IsolatedAsyncioTestCase):
|
||
async def test_background_job_uses_exact_size_and_selected_output_format(self):
|
||
calls = []
|
||
session = object()
|
||
|
||
class FakeSessionContext:
|
||
async def __aenter__(self):
|
||
return session
|
||
|
||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||
return False
|
||
|
||
class FakeClient:
|
||
def __init__(self, **_kwargs):
|
||
pass
|
||
|
||
async def generate_async(self, **kwargs):
|
||
calls.append(kwargs)
|
||
return [Image.new("RGB", (6, 4), "green")], {}
|
||
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "Seedream 5.0 Pro",
|
||
"resolution": "2K",
|
||
"aspect_ratio": "16:9",
|
||
"output_format": "jpeg",
|
||
})
|
||
with (
|
||
tempfile.TemporaryDirectory() as temp_dir,
|
||
patch(
|
||
"comfyui_o1key.clients.seedream_image_client.SeedreamImageClient",
|
||
FakeClient,
|
||
),
|
||
patch(
|
||
"comfyui_o1key.utils.config.get_api_key_or_raise",
|
||
return_value="secret",
|
||
),
|
||
patch(
|
||
"comfyui_o1key.utils.config.get_base_url_by_route",
|
||
return_value="https://cf-api.o1key.com",
|
||
),
|
||
patch(
|
||
"comfyui_o1key.utils.http2_client.create_http_client",
|
||
return_value=FakeSessionContext(),
|
||
),
|
||
):
|
||
result = await JOBS._execute_seedream_job(job, temp_dir)
|
||
smart_job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "Seedream 5.0 Pro",
|
||
"resolution": "智能",
|
||
"output_format": "jpeg",
|
||
})
|
||
await JOBS._execute_seedream_job(smart_job, temp_dir)
|
||
|
||
self.assertEqual(len(result["images"]), 1)
|
||
self.assertEqual(result["warnings"], [])
|
||
self.assertEqual(len(calls), 2)
|
||
self.assertIs(calls[0]["session"], session)
|
||
self.assertEqual(calls[0]["model"], "dola-seedream-5-0-pro-260628-ep")
|
||
self.assertEqual(calls[0]["size"], "2816x1584")
|
||
self.assertEqual(calls[0]["output_format"], "jpeg")
|
||
self.assertIsNone(calls[1]["size"])
|
||
|
||
async def test_layer_job_saves_every_result_with_layer_metadata(self):
|
||
calls = []
|
||
|
||
class FakeSessionContext:
|
||
async def __aenter__(self):
|
||
return object()
|
||
|
||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||
return False
|
||
|
||
class FakeClient:
|
||
def __init__(self, **_kwargs):
|
||
pass
|
||
|
||
async def generate_async(self, **kwargs):
|
||
calls.append(kwargs)
|
||
base = Image.new("RGB", (6, 4), "white")
|
||
layer = Image.new("RGBA", (3, 2), (10, 20, 30, 64))
|
||
setattr(base, "_o1key_seedream_layer", {"z_index": 0, "name": "底图"})
|
||
setattr(layer, "_o1key_seedream_layer", {"z_index": 1, "name": "主体"})
|
||
return [base, layer], {}
|
||
|
||
with tempfile.TemporaryDirectory() as temp_dir:
|
||
reference_path = os.path.join(temp_dir, "reference.png")
|
||
Image.new("RGB", (6, 4), "red").save(reference_path)
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"prompt": "",
|
||
"model": "Seedream 5.0 Pro",
|
||
"resolution": "1.5K",
|
||
"output_format": "png",
|
||
"layer_decomposition": True,
|
||
"references": [{"name": "poster.png", "subfolder": "", "type": "input"}],
|
||
})
|
||
job["reference_paths"] = (reference_path,)
|
||
job["model_reference_paths"] = ()
|
||
with (
|
||
patch(
|
||
"comfyui_o1key.clients.seedream_image_client.SeedreamImageClient",
|
||
FakeClient,
|
||
),
|
||
patch("comfyui_o1key.utils.config.get_api_key_or_raise", return_value="secret"),
|
||
patch("comfyui_o1key.utils.config.get_base_url_by_route", return_value="https://cf-api.o1key.com"),
|
||
patch(
|
||
"comfyui_o1key.utils.http2_client.create_http_client",
|
||
return_value=FakeSessionContext(),
|
||
),
|
||
):
|
||
result = await JOBS._execute_seedream_job(job, temp_dir)
|
||
|
||
self.assertEqual(len(result["images"]), 2)
|
||
self.assertEqual([item["result_index"] for item in result["images"]], [1, 2])
|
||
self.assertEqual(result["images"][1]["layer"]["name"], "主体")
|
||
with Image.open(os.path.join(temp_dir, result["images"][1]["filename"])) as saved:
|
||
self.assertEqual(saved.mode, "RGBA")
|
||
self.assertEqual(saved.getpixel((0, 0))[3], 64)
|
||
self.assertEqual(calls[0]["size"], "1.5K")
|
||
self.assertIs(calls[0]["layer_decomposition"], True)
|
||
|
||
|
||
class GptExecutionTests(unittest.IsolatedAsyncioTestCase):
|
||
async def test_single_reference_batch_bounds_live_reference_tensors(self):
|
||
opened = 0
|
||
maximum_opened = 0
|
||
started = 0
|
||
release = asyncio.Event()
|
||
|
||
class TrackingImage:
|
||
def __init__(self, path):
|
||
nonlocal opened, maximum_opened
|
||
self.path = path
|
||
self.closed = False
|
||
opened += 1
|
||
maximum_opened = max(maximum_opened, opened)
|
||
|
||
def close(self):
|
||
nonlocal opened
|
||
if not self.closed:
|
||
self.closed = True
|
||
opened -= 1
|
||
|
||
class FakeClient:
|
||
def __init__(self):
|
||
self.base_url = ""
|
||
self.response_log_enabled = True
|
||
self.poll_log_enabled = True
|
||
|
||
async def generate_image_async(self, **_kwargs):
|
||
nonlocal started
|
||
started += 1
|
||
if started == JOBS.MAX_CONCURRENT_REQUESTS_PER_JOB:
|
||
release.set()
|
||
await release.wait()
|
||
return [Image.new("RGB", (5, 4), "purple")]
|
||
|
||
references = [
|
||
{"name": f"reference-{index}.png", "subfolder": "", "type": "input"}
|
||
for index in range(12)
|
||
]
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "gpt-image-2",
|
||
"resolution": "1K",
|
||
"batch_enabled": True,
|
||
"batch_mode": "单图素材批量",
|
||
"references": references,
|
||
})
|
||
job["reference_paths"] = tuple(item["name"] for item in references)
|
||
job["model_reference_paths"] = ()
|
||
job["mask_path"] = None
|
||
with tempfile.TemporaryDirectory() as output_dir, patch(
|
||
"comfyui_o1key.clients.gpt_image_client.GptImageClient",
|
||
FakeClient,
|
||
), patch(
|
||
"comfyui_o1key.utils.config.get_base_url_by_route",
|
||
return_value="https://example.invalid",
|
||
), patch.object(
|
||
JOBS,
|
||
"_load_reference_image",
|
||
side_effect=TrackingImage,
|
||
), patch(
|
||
"comfyui_o1key.utils.image_utils.pil_to_tensor",
|
||
side_effect=lambda images: ("tensor", images[0].path),
|
||
):
|
||
result = await JOBS.execute_image_job(job, output_dir)
|
||
|
||
self.assertEqual(len(result["images"]), 12)
|
||
self.assertEqual(opened, 0)
|
||
self.assertLessEqual(maximum_opened, JOBS.MAX_CONCURRENT_REQUESTS_PER_JOB)
|
||
|
||
async def test_cartesian_batch_uses_task_specific_outfit_model_pairs(self):
|
||
calls = []
|
||
|
||
class FakeClient:
|
||
def __init__(self):
|
||
self.base_url = ""
|
||
self.response_log_enabled = True
|
||
self.poll_log_enabled = True
|
||
|
||
async def generate_image_async(self, **kwargs):
|
||
signature = tuple(
|
||
tuple(round(float(channel), 2) for channel in tensor.mean(dim=(0, 1, 2)))
|
||
for tensor in kwargs["image_tensor"]
|
||
)
|
||
calls.append(signature)
|
||
return [Image.new("RGB", (5, 4), "purple")]
|
||
|
||
with tempfile.TemporaryDirectory() as source_dir, tempfile.TemporaryDirectory() as output_dir:
|
||
outfit_paths = []
|
||
model_paths = []
|
||
for name, color in (("outfit-red.png", "red"), ("outfit-green.png", "green")):
|
||
path = os.path.join(source_dir, name)
|
||
Image.new("RGB", (5, 4), color).save(path)
|
||
outfit_paths.append(path)
|
||
for name, color in (("model-blue.png", "blue"), ("model-white.png", "white")):
|
||
path = os.path.join(source_dir, name)
|
||
Image.new("RGB", (5, 4), color).save(path)
|
||
model_paths.append(path)
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "gpt-image-2",
|
||
"resolution": "1K",
|
||
"batch_enabled": True,
|
||
"batch_mode": "全部搭配×全部模特",
|
||
"references": [
|
||
{"name": os.path.basename(path), "subfolder": "", "type": "input"}
|
||
for path in outfit_paths
|
||
],
|
||
"model_references": [
|
||
{"name": os.path.basename(path), "subfolder": "", "type": "input"}
|
||
for path in model_paths
|
||
],
|
||
})
|
||
job["reference_paths"] = tuple(outfit_paths)
|
||
job["model_reference_paths"] = tuple(model_paths)
|
||
job["mask_path"] = None
|
||
with patch(
|
||
"comfyui_o1key.clients.gpt_image_client.GptImageClient",
|
||
FakeClient,
|
||
), patch(
|
||
"comfyui_o1key.utils.config.get_base_url_by_route",
|
||
return_value="https://example.invalid",
|
||
):
|
||
result = await JOBS.execute_image_job(
|
||
job,
|
||
output_dir,
|
||
)
|
||
|
||
self.assertEqual(len(result["images"]), 4)
|
||
self.assertEqual(
|
||
calls,
|
||
[
|
||
((1.0, 0.0, 0.0), (0.0, 0.0, 1.0)),
|
||
((1.0, 0.0, 0.0), (1.0, 1.0, 1.0)),
|
||
((0.0, 0.5, 0.0), (0.0, 0.0, 1.0)),
|
||
((0.0, 0.5, 0.0), (1.0, 1.0, 1.0)),
|
||
],
|
||
)
|
||
|
||
async def test_each_requested_image_is_a_concurrent_n_one_call(self):
|
||
calls = []
|
||
active = 0
|
||
maximum_active = 0
|
||
|
||
class FakeClient:
|
||
def __init__(self):
|
||
self.base_url = ""
|
||
self.response_log_enabled = True
|
||
self.poll_log_enabled = True
|
||
|
||
async def generate_image_async(self, **kwargs):
|
||
nonlocal active, maximum_active
|
||
calls.append(kwargs)
|
||
active += 1
|
||
maximum_active = max(maximum_active, active)
|
||
await asyncio.sleep(0.01)
|
||
active -= 1
|
||
buffer = BytesIO()
|
||
Image.new("RGB", (5, 4), "purple").save(
|
||
buffer,
|
||
format="WEBP",
|
||
lossless=True,
|
||
)
|
||
raw = buffer.getvalue()
|
||
image = Image.open(BytesIO(raw))
|
||
image.load()
|
||
image._o1key_original_format = "WEBP"
|
||
image._o1key_original_bytes = raw
|
||
return [image]
|
||
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "gpt-image-2.5-sunburst",
|
||
"resolution": "4K",
|
||
"aspect_ratio": "智能",
|
||
"prompt": "紫色方块\n---\n绿色圆形",
|
||
"image_count": 2,
|
||
"quality": "超高",
|
||
"output_format": "webp",
|
||
"background": "transparent",
|
||
"moderation": "low",
|
||
})
|
||
job["reference_paths"] = ()
|
||
job["mask_path"] = None
|
||
with tempfile.TemporaryDirectory() as output_dir, patch(
|
||
"comfyui_o1key.clients.gpt_image_client.GptImageClient",
|
||
FakeClient,
|
||
), patch(
|
||
"comfyui_o1key.utils.config.get_base_url_by_route",
|
||
return_value="https://example.invalid",
|
||
):
|
||
result = await JOBS.execute_image_job(
|
||
job,
|
||
output_dir,
|
||
)
|
||
|
||
self.assertEqual(len(result["images"]), 4)
|
||
self.assertEqual(len(calls), 4)
|
||
self.assertEqual(
|
||
[call["prompt"] for call in calls],
|
||
["紫色方块", "紫色方块", "绿色圆形", "绿色圆形"],
|
||
)
|
||
self.assertGreater(maximum_active, 1)
|
||
self.assertTrue(all(call["n"] == 1 for call in calls))
|
||
self.assertTrue(all(call["quality"] == "xhigh" for call in calls))
|
||
self.assertTrue(all(call["model"] == "gpt-image-2.5-sunburst-sp" for call in calls))
|
||
self.assertTrue(all(call["special_price_parallel"] is False for call in calls))
|
||
self.assertTrue(all(call["size"] == "2880x2880" for call in calls))
|
||
self.assertTrue(all(call["output_format"] == "webp" for call in calls))
|
||
self.assertTrue(all(call["background"] == "transparent" for call in calls))
|
||
self.assertTrue(all("moderation" not in call for call in calls))
|
||
self.assertTrue(all(call["resize_mode"] == "智能缩放" for call in calls))
|
||
self.assertTrue(all(item["filename"].endswith(".webp") for item in result["images"]))
|
||
self.assertTrue(all(item["type"] == "temp" for item in result["images"]))
|
||
|
||
async def test_partial_failures_keep_the_exact_request_slot(self):
|
||
class FakeClient:
|
||
def __init__(self):
|
||
self.base_url = ""
|
||
self.response_log_enabled = True
|
||
self.poll_log_enabled = True
|
||
|
||
async def generate_image_async(self, **kwargs):
|
||
if kwargs["prompt"] == "失败项":
|
||
raise RuntimeError("第二张生成失败")
|
||
return [Image.new("RGB", (5, 4), "purple")]
|
||
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "gpt-image-2",
|
||
"resolution": "1K",
|
||
"prompt": "成功项\n---\n失败项",
|
||
"quality": "自动",
|
||
"output_format": "png",
|
||
"background": "auto",
|
||
})
|
||
job["reference_paths"] = ()
|
||
job["mask_path"] = None
|
||
with tempfile.TemporaryDirectory() as output_dir, patch(
|
||
"comfyui_o1key.clients.gpt_image_client.GptImageClient",
|
||
FakeClient,
|
||
), patch(
|
||
"comfyui_o1key.utils.config.get_base_url_by_route",
|
||
return_value="https://example.invalid",
|
||
):
|
||
result = await JOBS.execute_image_job(
|
||
job,
|
||
output_dir,
|
||
)
|
||
|
||
self.assertEqual(len(result["images"]), 1)
|
||
self.assertEqual(result["images"][0]["request_index"], 1)
|
||
self.assertEqual(
|
||
result["failed_items"],
|
||
[{"request_index": 2, "error": "第二张生成失败"}],
|
||
)
|
||
|
||
|
||
class NanoExecutionTests(unittest.IsolatedAsyncioTestCase):
|
||
async def test_single_reference_batch_prepares_one_source_at_a_time(self):
|
||
opened = 0
|
||
maximum_opened = 0
|
||
published = []
|
||
|
||
def progress_callback(_value):
|
||
pass
|
||
|
||
progress_callback.publish_images = lambda images: published.extend(images)
|
||
|
||
class TrackingImage:
|
||
def __init__(self, path):
|
||
nonlocal opened, maximum_opened
|
||
self.path = path
|
||
self.closed = False
|
||
opened += 1
|
||
maximum_opened = max(maximum_opened, opened)
|
||
|
||
def close(self):
|
||
nonlocal opened
|
||
if not self.closed:
|
||
self.closed = True
|
||
opened -= 1
|
||
|
||
class FakeHttpClient:
|
||
async def __aenter__(self):
|
||
return object()
|
||
|
||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||
return False
|
||
|
||
async def fake_prepare(images, **_kwargs):
|
||
await asyncio.sleep(0)
|
||
return [{"signature": images[0].path}]
|
||
|
||
async def fake_generate(**_kwargs):
|
||
await asyncio.sleep(0)
|
||
return [Image.new("RGB", (5, 4), "purple")], {}
|
||
|
||
references = [
|
||
{"name": f"reference-{index}.png", "subfolder": "", "type": "input"}
|
||
for index in range(12)
|
||
]
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "单图素材批量",
|
||
"references": references,
|
||
})
|
||
job["reference_paths"] = tuple(item["name"] for item in references)
|
||
job["model_reference_paths"] = ()
|
||
with tempfile.TemporaryDirectory() as output_dir, patch(
|
||
"comfyui_o1key.utils.config.get_api_key_or_raise",
|
||
return_value="secret",
|
||
), patch(
|
||
"comfyui_o1key.utils.config.get_base_url_by_route",
|
||
return_value="https://example.invalid",
|
||
), patch(
|
||
"comfyui_o1key.utils.http2_client.create_http_client",
|
||
return_value=FakeHttpClient(),
|
||
), patch(
|
||
"comfyui_o1key.utils.nano_banana_async.prepare_nano_banana_inline_images",
|
||
new=fake_prepare,
|
||
), patch(
|
||
"comfyui_o1key.utils.nano_banana_async.generate_nano_banana_async",
|
||
new=fake_generate,
|
||
), patch.object(
|
||
JOBS,
|
||
"_load_reference_image",
|
||
side_effect=TrackingImage,
|
||
):
|
||
result = await JOBS.execute_image_job(job, output_dir, progress_callback)
|
||
|
||
self.assertEqual(len(result["images"]), 12)
|
||
self.assertEqual(len(published), 12)
|
||
self.assertEqual(
|
||
sorted(item["request_index"] for item in published),
|
||
list(range(1, 13)),
|
||
)
|
||
self.assertEqual(opened, 0)
|
||
self.assertEqual(maximum_opened, 1)
|
||
|
||
async def test_cartesian_batch_encodes_only_each_task_outfit_model_pair(self):
|
||
calls = []
|
||
|
||
class FakeHttpClient:
|
||
async def __aenter__(self):
|
||
return object()
|
||
|
||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||
return False
|
||
|
||
async def fake_prepare(images, **_kwargs):
|
||
return [{"signature": image.getpixel((0, 0))} for image in images]
|
||
|
||
async def fake_generate(**kwargs):
|
||
calls.append(tuple(item["signature"] for item in kwargs["inline_images"]))
|
||
return [Image.new("RGB", (5, 4), "purple")], {}
|
||
|
||
with tempfile.TemporaryDirectory() as source_dir, tempfile.TemporaryDirectory() as output_dir:
|
||
outfit_paths = []
|
||
model_paths = []
|
||
for name, color in (("outfit-red.png", "red"), ("outfit-green.png", "green")):
|
||
path = os.path.join(source_dir, name)
|
||
Image.new("RGB", (5, 4), color).save(path)
|
||
outfit_paths.append(path)
|
||
for name, color in (("model-blue.png", "blue"), ("model-white.png", "white")):
|
||
path = os.path.join(source_dir, name)
|
||
Image.new("RGB", (5, 4), color).save(path)
|
||
model_paths.append(path)
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "全部搭配×全部模特",
|
||
"references": [
|
||
{"name": os.path.basename(path), "subfolder": "", "type": "input"}
|
||
for path in outfit_paths
|
||
],
|
||
"model_references": [
|
||
{"name": os.path.basename(path), "subfolder": "", "type": "input"}
|
||
for path in model_paths
|
||
],
|
||
})
|
||
job["reference_paths"] = tuple(outfit_paths)
|
||
job["model_reference_paths"] = tuple(model_paths)
|
||
with patch(
|
||
"comfyui_o1key.utils.config.get_api_key_or_raise",
|
||
return_value="secret",
|
||
), patch(
|
||
"comfyui_o1key.utils.config.get_base_url_by_route",
|
||
return_value="https://example.invalid",
|
||
), patch(
|
||
"comfyui_o1key.utils.http2_client.create_http_client",
|
||
return_value=FakeHttpClient(),
|
||
), patch(
|
||
"comfyui_o1key.utils.nano_banana_async.prepare_nano_banana_inline_images",
|
||
new=fake_prepare,
|
||
), patch(
|
||
"comfyui_o1key.utils.nano_banana_async.generate_nano_banana_async",
|
||
new=fake_generate,
|
||
):
|
||
result = await JOBS.execute_image_job(
|
||
job,
|
||
output_dir,
|
||
)
|
||
|
||
self.assertEqual(len(result["images"]), 4)
|
||
self.assertEqual(calls, [
|
||
((255, 0, 0), (0, 0, 255)),
|
||
((255, 0, 0), (255, 255, 255)),
|
||
((0, 128, 0), (0, 0, 255)),
|
||
((0, 128, 0), (255, 255, 255)),
|
||
])
|
||
|
||
async def test_references_are_reused_as_exact_inline_data_without_urls(self):
|
||
calls = []
|
||
|
||
class FakeHttpClient:
|
||
async def __aenter__(self):
|
||
return object()
|
||
|
||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||
return False
|
||
|
||
async def fake_generate(**kwargs):
|
||
calls.append(kwargs)
|
||
await asyncio.sleep(0)
|
||
return [Image.new("RGB", (5, 4), "purple")], {}
|
||
|
||
with tempfile.TemporaryDirectory() as source_dir, tempfile.TemporaryDirectory() as output_dir:
|
||
reference_path = os.path.join(source_dir, "reference.jpg")
|
||
Image.new("RGB", (5, 4), "orange").save(reference_path, format="JPEG")
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"prompt": "第一条\n---\n第二条",
|
||
"image_count": 2,
|
||
"google_search": True,
|
||
"color_correction": "智能纠正",
|
||
"references": [
|
||
{"name": "reference.jpg", "subfolder": "", "type": "input"},
|
||
],
|
||
})
|
||
job["reference_paths"] = (reference_path,)
|
||
|
||
with patch(
|
||
"comfyui_o1key.utils.config.get_api_key_or_raise",
|
||
return_value="secret",
|
||
), patch(
|
||
"comfyui_o1key.utils.config.get_base_url_by_route",
|
||
return_value="https://example.invalid",
|
||
), patch(
|
||
"comfyui_o1key.utils.http2_client.create_http_client",
|
||
return_value=FakeHttpClient(),
|
||
), patch(
|
||
"comfyui_o1key.utils.nano_banana_async.generate_nano_banana_async",
|
||
new=fake_generate,
|
||
):
|
||
result = await JOBS.execute_image_job(
|
||
job,
|
||
output_dir,
|
||
)
|
||
|
||
self.assertEqual(len(result["images"]), 4)
|
||
self.assertEqual(len(calls), 4)
|
||
self.assertTrue(all(call["download_semaphore"] is None for call in calls))
|
||
self.assertTrue(all(call["google_search"] is True for call in calls))
|
||
self.assertEqual(
|
||
[call["prompt"] for call in calls],
|
||
["第一条", "第一条", "第二条", "第二条"],
|
||
)
|
||
self.assertNotIn("image_urls", calls[0])
|
||
self.assertIs(calls[0]["inline_images"], calls[1]["inline_images"])
|
||
inline_data = calls[0]["inline_images"][0]["inlineData"]
|
||
self.assertEqual(inline_data["mimeType"], "image/jpeg")
|
||
self.assertNotIn("data:", inline_data["data"])
|
||
|
||
|
||
class IsolationTests(unittest.TestCase):
|
||
def test_terminal_history_survives_a_new_store_instance_without_secrets(self):
|
||
with tempfile.TemporaryDirectory() as temp_dir:
|
||
path = os.path.join(temp_dir, "o1key", JOBS.JOB_HISTORY_FILENAME)
|
||
batch_id = str(uuid.uuid4())
|
||
store = JOBS.PersistentJobHistory(path)
|
||
store.upsert({
|
||
"batch_id": batch_id,
|
||
"state": "completed",
|
||
"generator_node_id": 10,
|
||
"save_node_id": 11,
|
||
"images": [{
|
||
"filename": "saved.png",
|
||
"subfolder": "project",
|
||
"type": "output",
|
||
"batch_id": batch_id,
|
||
"request_index": 1,
|
||
"signed_url": "https://secret.invalid/token",
|
||
"base64": "private-pixels",
|
||
"layer": {
|
||
"z_index": 1,
|
||
"name": "主体",
|
||
"url": "https://secret.invalid/layer-token",
|
||
},
|
||
}],
|
||
"prompt": "must-not-be-stored",
|
||
"total_count": 1,
|
||
"create_time": 1000,
|
||
"execution_start_time": 1100,
|
||
"execution_end_time": 1200,
|
||
})
|
||
|
||
restored = JOBS.PersistentJobHistory(path).list()
|
||
self.assertEqual(len(restored), 1)
|
||
self.assertEqual(restored[0]["batch_id"], batch_id)
|
||
self.assertEqual(restored[0]["images"][0]["filename"], "saved.png")
|
||
self.assertEqual(restored[0]["images"][0]["layer"], {
|
||
"z_index": 1,
|
||
"name": "主体",
|
||
})
|
||
raw = Path(path).read_text(encoding="utf-8")
|
||
self.assertNotIn("must-not-be-stored", raw)
|
||
self.assertNotIn("secret.invalid", raw)
|
||
self.assertNotIn("private-pixels", raw)
|
||
|
||
JOBS.PersistentJobHistory(path).delete(batch_id)
|
||
self.assertEqual(JOBS.PersistentJobHistory(path).list(), [])
|
||
|
||
def test_layer_job_status_distinguishes_one_request_from_multiple_results(self):
|
||
batch_id = str(uuid.uuid4())
|
||
|
||
async def executor(_job, _progress):
|
||
return {}
|
||
|
||
async def sender(_event, _payload):
|
||
return None
|
||
|
||
manager = JOBS.ParallelImageJobManager(executor, sender)
|
||
record = JOBS.JobRecord(job={
|
||
"batch_id": batch_id,
|
||
"generator_node_id": 1,
|
||
"save_node_id": 2,
|
||
"total_task_count": 1,
|
||
"layer_decomposition": True,
|
||
"submitted_at": 0,
|
||
})
|
||
record.state = "completed"
|
||
record.images = [
|
||
{"filename": "base.png", "type": "temp"},
|
||
{"filename": "layer.png", "type": "temp"},
|
||
]
|
||
manager.jobs[batch_id] = record
|
||
|
||
status = manager.status(batch_id)
|
||
self.assertEqual(status["request_count"], 1)
|
||
self.assertEqual(status["result_count"], 2)
|
||
self.assertEqual(status["total_count"], 2)
|
||
|
||
def test_batch_snapshot_keeps_outfit_and_model_files_separate(self):
|
||
with tempfile.TemporaryDirectory() as input_dir, tempfile.TemporaryDirectory() as temp_dir:
|
||
Image.new("RGB", (5, 4), "red").save(os.path.join(input_dir, "outfit.png"))
|
||
Image.new("RGB", (5, 4), "blue").save(os.path.join(input_dir, "model.png"))
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"batch_enabled": True,
|
||
"batch_mode": "全部搭配×全部模特",
|
||
"references": [{"name": "outfit.png", "subfolder": "", "type": "input"}],
|
||
"model_references": [{"name": "model.png", "subfolder": "", "type": "input"}],
|
||
})
|
||
JOBS.snapshot_reference_files(job, input_dir, temp_dir)
|
||
self.assertEqual(len(job["reference_paths"]), 1)
|
||
self.assertEqual(len(job["model_reference_paths"]), 1)
|
||
self.assertIn(f"{os.sep}references{os.sep}", job["reference_paths"][0])
|
||
self.assertIn(f"{os.sep}models{os.sep}", job["model_reference_paths"][0])
|
||
self.assertEqual(JOBS._task_reference_paths(job, 0), (
|
||
job["reference_paths"][0],
|
||
job["model_reference_paths"][0],
|
||
))
|
||
JOBS.cleanup_job_snapshot(job, temp_dir)
|
||
|
||
def test_seedream_snapshot_rejects_invalid_dimensions_before_copying(self):
|
||
with tempfile.TemporaryDirectory() as input_dir, tempfile.TemporaryDirectory() as temp_dir:
|
||
Image.new("RGB", (14, 15), "red").save(os.path.join(input_dir, "tiny.png"))
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "Seedream 5.0 Pro",
|
||
"references": [{"name": "tiny.png", "subfolder": "", "type": "input"}],
|
||
})
|
||
|
||
with self.assertRaisesRegex(ValueError, "宽和高都必须大于 14px"):
|
||
JOBS.snapshot_reference_files(job, input_dir, temp_dir)
|
||
|
||
self.assertFalse(os.path.exists(os.path.join(
|
||
temp_dir,
|
||
"o1key_image_jobs",
|
||
job["batch_id"],
|
||
)))
|
||
|
||
def test_save_accepts_batch_prompt_results_above_legacy_eighty_one_limit(self):
|
||
job = JOBS.normalize_job_payload(_payload())
|
||
descriptors = []
|
||
with tempfile.TemporaryDirectory() as temp_dir:
|
||
for index in range(82):
|
||
filename = JOBS._result_filename(job, index, 0, "png")
|
||
Path(temp_dir, filename).write_bytes(b"result")
|
||
descriptors.append({
|
||
"filename": filename,
|
||
"subfolder": "",
|
||
"type": "temp",
|
||
"batch_id": job["batch_id"],
|
||
})
|
||
|
||
paths = JOBS._resolve_temp_result_paths(
|
||
job["batch_id"], descriptors, temp_dir,
|
||
)
|
||
|
||
self.assertEqual(len(paths), 82)
|
||
|
||
def test_save_route_accepts_only_matching_batch_temp_descriptors(self):
|
||
job = JOBS.normalize_job_payload(_payload())
|
||
filename = JOBS._result_filename(job, 0, 0, "jpg")
|
||
with tempfile.TemporaryDirectory() as temp_dir:
|
||
Path(temp_dir, filename).write_bytes(b"result")
|
||
descriptor = {
|
||
"filename": filename,
|
||
"subfolder": "",
|
||
"type": "temp",
|
||
"batch_id": job["batch_id"],
|
||
}
|
||
self.assertEqual(
|
||
JOBS._resolve_temp_result_paths(job["batch_id"], [descriptor], temp_dir),
|
||
[os.path.join(temp_dir, filename)],
|
||
)
|
||
with self.assertRaisesRegex(ValueError, "临时图片文件名"):
|
||
JOBS._resolve_temp_result_paths(
|
||
job["batch_id"],
|
||
[{**descriptor, "filename": "../outside.jpg"}],
|
||
temp_dir,
|
||
)
|
||
with self.assertRaisesRegex(ValueError, "批次 ID"):
|
||
JOBS._resolve_temp_result_paths(
|
||
job["batch_id"],
|
||
[{**descriptor, "batch_id": str(uuid.uuid4())}],
|
||
temp_dir,
|
||
)
|
||
|
||
def test_same_reference_is_copied_to_distinct_batch_snapshots(self):
|
||
with tempfile.TemporaryDirectory() as input_dir, tempfile.TemporaryDirectory() as temp_dir:
|
||
source = os.path.join(input_dir, "same-name.jpg")
|
||
Image.new("RGB", (5, 4), "red").save(source, format="JPEG")
|
||
|
||
first = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"references": [{"name": "same-name.jpg", "subfolder": "", "type": "input"}],
|
||
})
|
||
second = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"references": [{"name": "same-name.jpg", "subfolder": "", "type": "input"}],
|
||
})
|
||
JOBS.snapshot_reference_files(first, input_dir, temp_dir)
|
||
JOBS.snapshot_reference_files(second, input_dir, temp_dir)
|
||
|
||
first_path = first["reference_paths"][0]
|
||
second_path = second["reference_paths"][0]
|
||
self.assertNotEqual(first_path, second_path)
|
||
self.assertEqual(Path(first_path).read_bytes(), Path(second_path).read_bytes())
|
||
|
||
Image.new("RGB", (5, 4), "blue").save(source, format="JPEG")
|
||
self.assertEqual(Path(first_path).read_bytes(), Path(second_path).read_bytes())
|
||
|
||
JOBS.cleanup_job_snapshot(first, temp_dir)
|
||
JOBS.cleanup_job_snapshot(second, temp_dir)
|
||
self.assertFalse(os.path.exists(first["snapshot_root"]))
|
||
self.assertFalse(os.path.exists(second["snapshot_root"]))
|
||
|
||
def test_gpt_mask_is_copied_into_the_batch_snapshot(self):
|
||
with tempfile.TemporaryDirectory() as input_dir, tempfile.TemporaryDirectory() as temp_dir:
|
||
Image.new("RGB", (5, 4), "red").save(os.path.join(input_dir, "reference.png"))
|
||
Image.new("L", (5, 4), 255).save(os.path.join(input_dir, "mask.png"))
|
||
job = JOBS.normalize_job_payload({
|
||
**_payload(),
|
||
"model": "gpt-image-2",
|
||
"resolution": "1K",
|
||
"references": [
|
||
{"name": "reference.png", "subfolder": "", "type": "input"},
|
||
],
|
||
"mask": {"name": "mask.png", "subfolder": "", "type": "input"},
|
||
})
|
||
|
||
JOBS.snapshot_reference_files(job, input_dir, temp_dir)
|
||
self.assertTrue(os.path.isfile(job["mask_path"]))
|
||
self.assertEqual(
|
||
os.path.commonpath([job["snapshot_root"], job["mask_path"]]),
|
||
job["snapshot_root"],
|
||
)
|
||
self.assertEqual(Path(job["mask_path"]).read_bytes(), Path(input_dir, "mask.png").read_bytes())
|
||
|
||
JOBS.cleanup_job_snapshot(job, temp_dir)
|
||
self.assertFalse(os.path.exists(job["snapshot_root"]))
|
||
|
||
def test_provider_results_stay_in_temp_until_the_save_node_runs(self):
|
||
with tempfile.TemporaryDirectory() as temp_dir, tempfile.TemporaryDirectory() as output_dir:
|
||
first = JOBS.normalize_job_payload(_payload())
|
||
second = JOBS.normalize_job_payload(_payload())
|
||
first_path, first_subfolder = JOBS._prepare_output_directory(first, temp_dir)
|
||
second_path, second_subfolder = JOBS._prepare_output_directory(second, temp_dir)
|
||
self.assertEqual(first_path, os.path.abspath(temp_dir))
|
||
self.assertEqual(second_path, os.path.abspath(temp_dir))
|
||
self.assertEqual(first_subfolder, "")
|
||
self.assertEqual(second_subfolder, "")
|
||
self.assertEqual(list(Path(temp_dir).iterdir()), [])
|
||
|
||
first_buffer = BytesIO()
|
||
Image.new("RGB", (3, 3), "red").save(first_buffer, format="JPEG", quality=90)
|
||
first_raw = first_buffer.getvalue()
|
||
first_image = Image.open(BytesIO(first_raw))
|
||
first_image.load()
|
||
first_image._o1key_original_bytes = first_raw
|
||
first_image._o1key_original_format = "JPEG"
|
||
second_image = Image.new("RGB", (3, 3), "blue")
|
||
|
||
first_name = JOBS._result_filename(first, 0, 0, JOBS._provider_result_format(first_image))
|
||
second_name = JOBS._result_filename(second, 0, 0, JOBS._provider_result_format(second_image))
|
||
first_file = os.path.join(first_path, first_name)
|
||
second_file = os.path.join(second_path, second_name)
|
||
self.assertNotEqual(first_name, second_name)
|
||
JOBS._write_provider_result(first_image, first_file)
|
||
JOBS._write_provider_result(second_image, second_file)
|
||
first_image.close()
|
||
second_image.close()
|
||
self.assertEqual(Path(first_file).read_bytes(), first_raw)
|
||
self.assertNotEqual(Path(first_file).read_bytes(), Path(second_file).read_bytes())
|
||
self.assertFalse(os.path.exists(first_file + ".tmp"))
|
||
self.assertFalse(os.path.exists(second_file + ".tmp"))
|
||
self.assertEqual(list(Path(output_dir).iterdir()), [])
|
||
|
||
recovered = JOBS.find_completed_job_outputs(
|
||
first["batch_id"],
|
||
output_dir,
|
||
first["generator_node_id"],
|
||
first["save_node_id"],
|
||
temp_dir,
|
||
)
|
||
self.assertEqual(recovered["state"], "completed")
|
||
self.assertTrue(recovered["recovered_from_disk"])
|
||
self.assertEqual(recovered["images"][0]["batch_id"], first["batch_id"])
|
||
self.assertEqual(recovered["images"][0]["subfolder"], "")
|
||
self.assertEqual(recovered["images"][0]["type"], "temp")
|
||
self.assertFalse(Path(output_dir, "o1key_parallel").exists())
|
||
|
||
Path(output_dir, first_name).write_bytes(first_raw)
|
||
recovered_output = JOBS.find_completed_job_outputs(
|
||
first["batch_id"],
|
||
output_dir,
|
||
first["generator_node_id"],
|
||
first["save_node_id"],
|
||
temp_dir,
|
||
)
|
||
self.assertEqual(recovered_output["images"][0]["type"], "output")
|
||
|
||
def test_saved_job_record_returns_the_same_output_after_browser_refresh(self):
|
||
job = JOBS.normalize_job_payload(_payload())
|
||
record = JOBS.JobRecord(job=job, state="completed", images=[{
|
||
"filename": f"o1key_{job['batch_id']}_0001_01.png",
|
||
"type": "temp",
|
||
"batch_id": job["batch_id"],
|
||
}])
|
||
self.assertEqual(JOBS._saved_record_images(record, job["batch_id"]), [])
|
||
|
||
record.images = [{
|
||
"filename": "external-preview.png",
|
||
"subfolder": "o1key_external_preview/token",
|
||
"type": "temp",
|
||
"external_saved": True,
|
||
"batch_id": job["batch_id"],
|
||
}]
|
||
self.assertEqual(
|
||
JOBS._saved_record_images(record, job["batch_id"]),
|
||
record.images,
|
||
)
|
||
|
||
record.images = [{
|
||
"filename": "saved.png",
|
||
"subfolder": "",
|
||
"type": "output",
|
||
"batch_id": job["batch_id"],
|
||
"request_index": 1,
|
||
"result_index": 1,
|
||
}]
|
||
recovered = JOBS._saved_record_images(record, job["batch_id"])
|
||
self.assertEqual(recovered, record.images)
|
||
recovered[0]["filename"] = "mutated.png"
|
||
self.assertEqual(record.images[0]["filename"], "saved.png")
|
||
|
||
|
||
class SchedulerTests(unittest.IsolatedAsyncioTestCase):
|
||
async def test_manager_emits_safe_partial_images_while_batch_is_running(self):
|
||
partial_seen = asyncio.Event()
|
||
events = []
|
||
descriptor = {
|
||
"filename": "partial.png",
|
||
"subfolder": "",
|
||
"type": "temp",
|
||
"request_index": 1,
|
||
"result_index": 1,
|
||
"signed_url": "https://secret.invalid/result",
|
||
}
|
||
|
||
async def executor(job, progress_callback):
|
||
partial = {**descriptor, "batch_id": job["batch_id"]}
|
||
progress_callback.publish_images([partial])
|
||
await asyncio.wait_for(partial_seen.wait(), timeout=1)
|
||
return {"images": [partial], "warnings": []}
|
||
|
||
async def sender(event, payload):
|
||
events.append((event, dict(payload)))
|
||
if payload["state"] == "running" and payload["images"]:
|
||
partial_seen.set()
|
||
|
||
manager = JOBS.ParallelImageJobManager(executor, sender)
|
||
job = JOBS.normalize_job_payload(_payload())
|
||
await manager.submit(job)
|
||
await manager.jobs[job["batch_id"]].task
|
||
|
||
partial_events = [
|
||
payload for event, payload in events
|
||
if event == "o1key.image_job"
|
||
and payload["state"] == "running"
|
||
and payload["images"]
|
||
]
|
||
self.assertTrue(partial_events)
|
||
self.assertEqual(partial_events[0]["images"][0]["filename"], "partial.png")
|
||
self.assertNotIn("signed_url", partial_events[0]["images"][0])
|
||
self.assertEqual(partial_events[0]["succeeded_count"], 1)
|
||
|
||
async def test_manager_persists_terminal_summary_for_process_restart(self):
|
||
async def executor(job, _progress_callback):
|
||
return {
|
||
"images": [{
|
||
"filename": "result.png",
|
||
"subfolder": "",
|
||
"type": "output",
|
||
"batch_id": job["batch_id"],
|
||
"request_index": 1,
|
||
"result_index": 1,
|
||
}],
|
||
"warnings": [],
|
||
}
|
||
|
||
async def sender(_event, _payload):
|
||
return None
|
||
|
||
with tempfile.TemporaryDirectory() as temp_dir:
|
||
path = os.path.join(temp_dir, JOBS.JOB_HISTORY_FILENAME)
|
||
manager = JOBS.ParallelImageJobManager(
|
||
executor,
|
||
sender,
|
||
history_store=JOBS.PersistentJobHistory(path),
|
||
)
|
||
job = JOBS.normalize_job_payload(_payload())
|
||
await manager.submit(job)
|
||
await manager.jobs[job["batch_id"]].task
|
||
|
||
restored = JOBS.PersistentJobHistory(path).list()
|
||
self.assertEqual(restored[0]["batch_id"], job["batch_id"])
|
||
self.assertEqual(restored[0]["state"], "completed")
|
||
self.assertGreater(restored[0]["execution_end_time"], 0)
|
||
|
||
async def test_queue_positions_and_cancellation_cover_waiting_and_running_batches(self):
|
||
started = asyncio.Event()
|
||
release = asyncio.Event()
|
||
events = []
|
||
|
||
async def executor(_job, _progress_callback):
|
||
started.set()
|
||
await release.wait()
|
||
return {"images": [], "warnings": []}
|
||
|
||
async def sender(event, payload):
|
||
events.append((event, dict(payload)))
|
||
|
||
manager = JOBS.ParallelImageJobManager(executor, sender, max_concurrent=1)
|
||
first = JOBS.normalize_job_payload(_payload(generator_id=300, save_id=301))
|
||
second = JOBS.normalize_job_payload(_payload(generator_id=300, save_id=302))
|
||
await manager.submit(first)
|
||
await started.wait()
|
||
submitted_second = await manager.submit(second)
|
||
await asyncio.sleep(0)
|
||
|
||
self.assertEqual(submitted_second["state"], "queued")
|
||
self.assertEqual(submitted_second["queue_position"], 1)
|
||
self.assertEqual(submitted_second["total_count"], 1)
|
||
self.assertEqual(manager.status(second["batch_id"])["queue_position"], 1)
|
||
|
||
cancelled_second = await manager.cancel(second["batch_id"])
|
||
self.assertEqual(cancelled_second["state"], "cancelled")
|
||
cancelled_first = await manager.cancel(first["batch_id"])
|
||
self.assertEqual(cancelled_first["state"], "cancelled")
|
||
self.assertTrue(any(
|
||
event == "o1key.image_job"
|
||
and payload["batch_id"] == second["batch_id"]
|
||
and payload["state"] == "cancelled"
|
||
for event, payload in events
|
||
))
|
||
|
||
async def test_caps_at_ten_batches_and_keeps_every_result_on_its_batch(self):
|
||
release = asyncio.Event()
|
||
active = 0
|
||
maximum_active = 0
|
||
events = []
|
||
|
||
async def executor(job, progress_callback):
|
||
nonlocal active, maximum_active
|
||
active += 1
|
||
maximum_active = max(maximum_active, active)
|
||
progress_callback(0.42)
|
||
await release.wait()
|
||
progress_callback(1.0)
|
||
active -= 1
|
||
return {
|
||
"images": [{
|
||
"filename": f"{job['batch_id']}.png",
|
||
"subfolder": job["batch_id"],
|
||
"type": "output",
|
||
"batch_id": job["batch_id"],
|
||
}],
|
||
"warnings": [],
|
||
}
|
||
|
||
async def sender(event, payload):
|
||
events.append((event, dict(payload)))
|
||
|
||
manager = JOBS.ParallelImageJobManager(executor, sender, max_concurrent=10)
|
||
jobs = []
|
||
for index in range(12):
|
||
job = JOBS.normalize_job_payload(
|
||
_payload(generator_id=100, save_id=200 + index)
|
||
)
|
||
jobs.append(job)
|
||
await manager.submit(job)
|
||
|
||
for _ in range(100):
|
||
if maximum_active == 10:
|
||
break
|
||
await asyncio.sleep(0.01)
|
||
|
||
self.assertEqual(maximum_active, 10)
|
||
self.assertEqual(sum(record.state == "queued" for record in manager.jobs.values()), 2)
|
||
release.set()
|
||
await asyncio.gather(*(record.task for record in manager.jobs.values()))
|
||
|
||
self.assertEqual(maximum_active, 10)
|
||
completed = [payload for event, payload in events if event == "o1key.image_job" and payload["state"] == "completed"]
|
||
running_progress = [
|
||
payload for event, payload in events
|
||
if event == "o1key.image_job"
|
||
and payload["state"] == "running"
|
||
and payload["progress_percent"] > 0
|
||
]
|
||
self.assertTrue(running_progress)
|
||
self.assertEqual(len(completed), 12)
|
||
for payload in completed:
|
||
self.assertEqual(payload["progress_percent"], 100)
|
||
self.assertEqual(payload["total_count"], 1)
|
||
self.assertEqual(payload["succeeded_count"], 1)
|
||
self.assertEqual(payload["failed_count"], 0)
|
||
self.assertEqual(payload["images"][0]["batch_id"], payload["batch_id"])
|
||
self.assertEqual(payload["images"][0]["subfolder"], payload["batch_id"])
|
||
self.assertEqual(
|
||
payload["save_node_id"],
|
||
next(job["save_node_id"] for job in jobs if job["batch_id"] == payload["batch_id"]),
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main(verbosity=2)
|