154 lines
4.7 KiB
Python
154 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
PLUGIN_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_models_module():
|
|
module_path = PLUGIN_ROOT / "lib" / "models.py"
|
|
assert module_path.exists(), f"Expected module at {module_path}"
|
|
spec = importlib.util.spec_from_file_location("comfyui_kieapi_models", module_path)
|
|
assert spec is not None and spec.loader is not None
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def load_image_utils_module():
|
|
module_path = PLUGIN_ROOT / "lib" / "image_utils.py"
|
|
assert module_path.exists(), f"Expected module at {module_path}"
|
|
spec = importlib.util.spec_from_file_location("comfyui_kieapi_image_utils", module_path)
|
|
assert spec is not None and spec.loader is not None
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def test_upload_request_serializes_expected_fields():
|
|
module = load_models_module()
|
|
|
|
request = module.KieBase64UploadRequest(
|
|
base64Data="data:image/png;base64,abc123",
|
|
uploadPath="images/base64",
|
|
fileName="reference.png",
|
|
)
|
|
|
|
payload = request.model_dump()
|
|
assert payload["base64Data"] == "data:image/png;base64,abc123"
|
|
assert payload["uploadPath"] == "images/base64"
|
|
assert payload["fileName"] == "reference.png"
|
|
|
|
|
|
def test_create_task_request_serializes_expected_model_and_input():
|
|
module = load_models_module()
|
|
|
|
request = module.KieCreateTaskRequest(
|
|
input=module.KieNanoBanana2Input(
|
|
prompt="banana ship",
|
|
aspect_ratio="auto",
|
|
resolution="2K",
|
|
output_format="jpg",
|
|
google_search=False,
|
|
image_input=["https://example.com/reference.png"],
|
|
)
|
|
)
|
|
|
|
payload = request.model_dump()
|
|
assert payload["model"] == "nano-banana-2"
|
|
assert payload["input"]["prompt"] == "banana ship"
|
|
assert payload["input"]["image_input"] == ["https://example.com/reference.png"]
|
|
|
|
|
|
def test_task_detail_parses_result_urls_from_result_json():
|
|
module = load_models_module()
|
|
|
|
response = module.KieTaskRecordResponse.model_validate(
|
|
{
|
|
"code": 200,
|
|
"msg": "success",
|
|
"data": {
|
|
"taskId": "task_123",
|
|
"state": "success",
|
|
"resultJson": "{\"resultUrls\":[\"https://example.com/1.jpg\",\"https://example.com/2.jpg\"]}",
|
|
"failCode": "",
|
|
"failMsg": "",
|
|
},
|
|
}
|
|
)
|
|
|
|
assert response.data.result_urls() == ["https://example.com/1.jpg", "https://example.com/2.jpg"]
|
|
|
|
|
|
def test_task_detail_exposes_failure_information():
|
|
module = load_models_module()
|
|
|
|
response = module.KieTaskRecordResponse.model_validate(
|
|
{
|
|
"code": 200,
|
|
"msg": "success",
|
|
"data": {
|
|
"taskId": "task_123",
|
|
"state": "fail",
|
|
"resultJson": "",
|
|
"failCode": "content_blocked",
|
|
"failMsg": "prompt rejected",
|
|
},
|
|
}
|
|
)
|
|
|
|
assert response.data.failure_reason() == "content_blocked: prompt rejected"
|
|
|
|
|
|
def test_tensor_to_base64_png_returns_data_uri():
|
|
module = load_image_utils_module()
|
|
image = torch.tensor([[[[1.0, 0.0, 0.0]]]], dtype=torch.float32)
|
|
|
|
encoded = module.tensor_to_base64_png(image)
|
|
|
|
assert encoded.startswith("data:image/png;base64,")
|
|
|
|
|
|
def test_tensor_batch_to_base64_pngs_returns_one_data_uri_per_image():
|
|
module = load_image_utils_module()
|
|
image_batch = torch.tensor(
|
|
[
|
|
[[[1.0, 0.0, 0.0]]],
|
|
[[[0.0, 1.0, 0.0]]],
|
|
[[[0.0, 0.0, 1.0]]],
|
|
],
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
encoded = module.tensor_batch_to_base64_pngs(image_batch)
|
|
|
|
assert len(encoded) == 3
|
|
assert all(item.startswith("data:image/png;base64,") for item in encoded)
|
|
|
|
|
|
def test_decode_image_bytes_returns_single_image_tensor():
|
|
module = load_image_utils_module()
|
|
image = torch.tensor([[[[0.0, 1.0, 0.0]]]], dtype=torch.float32)
|
|
encoded = module.tensor_to_base64_png(image)
|
|
raw_bytes = module.base64_data_uri_to_bytes(encoded)
|
|
|
|
decoded = module.decode_image_bytes(raw_bytes)
|
|
|
|
assert tuple(decoded.shape) == (1, 1, 1, 3)
|
|
assert torch.isclose(decoded[0, 0, 0, 1], torch.tensor(1.0))
|
|
|
|
|
|
def test_stack_image_tensors_batches_multiple_results():
|
|
module = load_image_utils_module()
|
|
first = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
|
second = torch.ones((1, 2, 2, 3), dtype=torch.float32)
|
|
|
|
batch = module.stack_image_tensors([first, second])
|
|
|
|
assert tuple(batch.shape) == (2, 2, 2, 3)
|
|
assert torch.equal(batch[0], first[0])
|
|
assert torch.equal(batch[1], second[0])
|