Publish current ComfyUI O1Key code baseline
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
This commit is contained in:
Jony
@@ -0,0 +1,307 @@
|
||||
"""Offline tests for the Seedream image transport."""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
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.clients import seedream_image_client as CLIENT # noqa: E402
|
||||
|
||||
|
||||
class SeedreamBodyTests(unittest.TestCase):
|
||||
def test_model_routes_resolve_to_the_documented_identifier(self):
|
||||
for route in ("畅速", "直连", "专线"):
|
||||
self.assertEqual(
|
||||
CLIENT.resolve_seedream_model("Seedream 5.0 Pro", route),
|
||||
CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
)
|
||||
|
||||
def test_body_keeps_reference_order_and_fixed_watermark(self):
|
||||
body = CLIENT.build_seedream_submit_body(
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
prompt=" 一只宇航猫 ",
|
||||
size="2368x1776",
|
||||
output_format="PNG",
|
||||
image_urls=[
|
||||
"https://example.invalid/one.png",
|
||||
"https://example.invalid/two.jpg",
|
||||
],
|
||||
)
|
||||
|
||||
self.assertEqual(body, {
|
||||
"model": CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
"prompt": "一只宇航猫",
|
||||
"n": 1,
|
||||
"size": "2368x1776",
|
||||
"output_format": "png",
|
||||
"watermark": False,
|
||||
"images": [
|
||||
"https://example.invalid/one.png",
|
||||
"https://example.invalid/two.jpg",
|
||||
],
|
||||
})
|
||||
|
||||
def test_body_rejects_webp_and_non_https_references(self):
|
||||
with self.assertRaisesRegex(ValueError, "png 或 jpeg"):
|
||||
CLIENT.build_seedream_submit_body(
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
prompt="test",
|
||||
size="1024x1024",
|
||||
output_format="webp",
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "HTTPS URL"):
|
||||
CLIENT.build_seedream_submit_body(
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
prompt="test",
|
||||
size="1024x1024",
|
||||
output_format="jpeg",
|
||||
image_urls=["http://example.invalid/reference.png"],
|
||||
)
|
||||
|
||||
def test_body_rejects_sizes_outside_the_documented_matrix(self):
|
||||
with self.assertRaisesRegex(ValueError, "图片尺寸无效"):
|
||||
CLIENT.build_seedream_submit_body(
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
prompt="test",
|
||||
size="1920x1080",
|
||||
output_format="jpeg",
|
||||
)
|
||||
|
||||
def test_body_omits_size_for_smart_resolution(self):
|
||||
body = CLIENT.build_seedream_submit_body(
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
prompt="test",
|
||||
size=None,
|
||||
output_format="jpeg",
|
||||
)
|
||||
|
||||
self.assertNotIn("size", body)
|
||||
|
||||
def test_reference_dimensions_match_current_volcengine_limits(self):
|
||||
CLIENT.validate_seedream_reference_dimensions(15, 15)
|
||||
CLIENT.validate_seedream_reference_dimensions(240, 15)
|
||||
CLIENT.validate_seedream_reference_dimensions(15, 240)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "宽和高都必须大于 14px"):
|
||||
CLIENT.validate_seedream_reference_dimensions(14, 240)
|
||||
with self.assertRaisesRegex(ValueError, "宽高比必须在 1:16～16:1"):
|
||||
CLIENT.validate_seedream_reference_dimensions(241, 15)
|
||||
with self.assertRaisesRegex(ValueError, "总像素不能超过"):
|
||||
CLIENT.validate_seedream_reference_dimensions(6001, 6000)
|
||||
|
||||
def test_layer_reference_uses_its_documented_total_pixel_floor(self):
|
||||
CLIENT.validate_seedream_reference_dimensions(
|
||||
512,
|
||||
512,
|
||||
layer_decomposition=True,
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "总像素必须在 512×512"):
|
||||
CLIENT.validate_seedream_reference_dimensions(
|
||||
511,
|
||||
512,
|
||||
layer_decomposition=True,
|
||||
)
|
||||
|
||||
def test_reference_file_size_is_checked_against_exact_upload_payload(self):
|
||||
class OversizedPayload:
|
||||
def __len__(self):
|
||||
return CLIENT.SEEDREAM_REFERENCE_MAX_BYTES + 1
|
||||
|
||||
source = Image.new("RGB", (15, 15), "red")
|
||||
with (
|
||||
patch.object(
|
||||
CLIENT,
|
||||
"image_to_upload_payload",
|
||||
return_value=(OversizedPayload(), ".png", "image/png"),
|
||||
),
|
||||
self.assertRaisesRegex(ValueError, "文件不能超过 30MB"),
|
||||
):
|
||||
CLIENT.validate_seedream_reference_image(source)
|
||||
|
||||
def test_layer_decomposition_allows_blank_prompt_and_requires_one_png_reference(self):
|
||||
body = CLIENT.build_seedream_submit_body(
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
prompt="",
|
||||
size="1.5K",
|
||||
output_format="png",
|
||||
image_urls=["https://example.invalid/poster.png"],
|
||||
layer_decomposition=True,
|
||||
)
|
||||
self.assertNotIn("prompt", body)
|
||||
self.assertEqual(body["size"], "1.5K")
|
||||
self.assertIs(body["layer_decomposition"], True)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "必须且只能提供1张"):
|
||||
CLIENT.build_seedream_submit_body(
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
prompt="",
|
||||
size="auto",
|
||||
output_format="png",
|
||||
layer_decomposition=True,
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "仅支持 png"):
|
||||
CLIENT.build_seedream_submit_body(
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
prompt="",
|
||||
size="2K",
|
||||
output_format="jpeg",
|
||||
image_urls=["https://example.invalid/poster.png"],
|
||||
layer_decomposition=True,
|
||||
)
|
||||
|
||||
def test_layer_metadata_is_sanitized_without_result_urls(self):
|
||||
metadata = CLIENT.extract_seedream_layer_metadata({
|
||||
"data": {
|
||||
"images": [{
|
||||
"url": "https://signed.invalid/secret.png",
|
||||
"z_index": "1",
|
||||
"size": "1408x1780",
|
||||
"output_format": "png",
|
||||
"bounding_box": {
|
||||
"absolute": [320, 180, 1728, 1960],
|
||||
"normalized": [156, 88, 844, 957],
|
||||
},
|
||||
"name": "主体",
|
||||
"description": "画面中的主要人物",
|
||||
}],
|
||||
},
|
||||
})
|
||||
self.assertEqual(metadata[0]["z_index"], 1)
|
||||
self.assertEqual(metadata[0]["bounding_box"]["absolute"], [320, 180, 1728, 1960])
|
||||
self.assertNotIn("url", metadata[0])
|
||||
|
||||
|
||||
class SeedreamLifecycleTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_upload_submit_poll_and_parse_use_the_shared_async_lifecycle(self):
|
||||
source = Image.new("RGB", (15, 15), "red")
|
||||
result = Image.new("RGB", (2, 2), "blue")
|
||||
completed = {
|
||||
"task_id": "task-seedream",
|
||||
"status": "SUCCESS",
|
||||
"data": {"images": [{"url": "https://example.invalid/result.png"}]},
|
||||
}
|
||||
metrics = {
|
||||
"download_bytes": 123,
|
||||
"download_seconds": 0.2,
|
||||
"download_wall_seconds": 0.1,
|
||||
"inline_images": 0,
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
CLIENT,
|
||||
"upload_images_to_temp_urls",
|
||||
new=AsyncMock(return_value=["https://example.invalid/reference.png"]),
|
||||
) as upload,
|
||||
patch.object(
|
||||
CLIENT,
|
||||
"submit_async_image_task",
|
||||
new=AsyncMock(return_value="task-seedream"),
|
||||
) as submit,
|
||||
patch.object(
|
||||
CLIENT,
|
||||
"poll_async_image_task",
|
||||
new=AsyncMock(return_value=completed),
|
||||
) as poll,
|
||||
patch.object(
|
||||
CLIENT,
|
||||
"parse_completed_async_image_task",
|
||||
new=AsyncMock(return_value=(completed, ([result], metrics))),
|
||||
) as parse,
|
||||
):
|
||||
client = CLIENT.SeedreamImageClient(
|
||||
base_url="https://cf-api.o1key.com/",
|
||||
api_key="secret",
|
||||
)
|
||||
images, timing = await client.generate_async(
|
||||
session=object(),
|
||||
prompt="test",
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
size="1024x1024",
|
||||
output_format="jpeg",
|
||||
images=[source],
|
||||
)
|
||||
|
||||
self.assertEqual(images, [result])
|
||||
upload.assert_awaited_once()
|
||||
self.assertEqual(upload.await_args.kwargs["base_url"], "https://cf-api.o1key.com")
|
||||
body = submit.await_args.args[3]
|
||||
self.assertEqual(body["images"], ["https://example.invalid/reference.png"])
|
||||
self.assertEqual(body["output_format"], "jpeg")
|
||||
self.assertIs(body["watermark"], False)
|
||||
poll.assert_awaited_once()
|
||||
parse.assert_awaited_once()
|
||||
self.assertEqual(timing["task_id"], "task-seedream")
|
||||
self.assertEqual(timing["download_bytes"], 123)
|
||||
|
||||
async def test_layer_results_are_sorted_by_z_index_and_keep_safe_metadata(self):
|
||||
top = Image.new("RGBA", (2, 2), (0, 0, 255, 64))
|
||||
base = Image.new("RGB", (2, 2), "white")
|
||||
completed = {
|
||||
"status": "SUCCESS",
|
||||
"data": {"images": [
|
||||
{"url": "https://signed.invalid/top.png", "z_index": 2, "name": "文字"},
|
||||
{"url": "https://signed.invalid/base.png", "z_index": 0, "name": "底图"},
|
||||
]},
|
||||
}
|
||||
with (
|
||||
patch.object(CLIENT, "upload_images_to_temp_urls", new=AsyncMock(
|
||||
return_value=["https://example.invalid/reference.png"]
|
||||
)),
|
||||
patch.object(CLIENT, "submit_async_image_task", new=AsyncMock(return_value="task-layer")),
|
||||
patch.object(CLIENT, "poll_async_image_task", new=AsyncMock(return_value=completed)),
|
||||
patch.object(CLIENT, "parse_completed_async_image_task", new=AsyncMock(
|
||||
return_value=(completed, ([top, base], {
|
||||
"download_bytes": 1,
|
||||
"download_seconds": 0.0,
|
||||
"download_wall_seconds": 0.0,
|
||||
"inline_images": 0,
|
||||
}))
|
||||
)),
|
||||
):
|
||||
images, timing = await CLIENT.SeedreamImageClient(
|
||||
base_url="https://cf-api.o1key.com",
|
||||
api_key="secret",
|
||||
).generate_async(
|
||||
session=object(),
|
||||
prompt="",
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
size="auto",
|
||||
output_format="png",
|
||||
images=[Image.new("RGB", (512, 512), "red")],
|
||||
layer_decomposition=True,
|
||||
)
|
||||
|
||||
self.assertEqual([item["z_index"] for item in timing["result_metadata"]], [0, 2])
|
||||
self.assertEqual(getattr(images[0], "_o1key_seedream_layer")["name"], "底图")
|
||||
|
||||
async def test_invalid_reference_is_rejected_before_upload(self):
|
||||
upload = AsyncMock(return_value=["https://example.invalid/reference.png"])
|
||||
with patch.object(CLIENT, "upload_images_to_temp_urls", new=upload):
|
||||
with self.assertRaisesRegex(ValueError, "宽和高都必须大于 14px"):
|
||||
await CLIENT.SeedreamImageClient(
|
||||
base_url="https://cf.invalid",
|
||||
api_key="secret",
|
||||
).generate_async(
|
||||
session=object(),
|
||||
prompt="test",
|
||||
model=CLIENT.SEEDREAM_API_MODEL_ID,
|
||||
size="1024x1024",
|
||||
output_format="jpeg",
|
||||
images=[Image.new("RGB", (14, 15), "red")],
|
||||
)
|
||||
|
||||
upload.assert_not_awaited()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
Reference in New Issue
Block a user