55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
# Integration test for the Batch Image Loader nodes — part of ComfyUI-AI2Go-Utils. GPL-3.0.
|
|
import json
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("comfy_api") # only runs inside a ComfyUI environment
|
|
import torch # noqa: E402
|
|
from PIL import Image # noqa: E402
|
|
|
|
|
|
@pytest.fixture
|
|
def input_dir(tmp_path, monkeypatch):
|
|
import folder_paths
|
|
monkeypatch.setattr(folder_paths, "get_input_directory", lambda: str(tmp_path))
|
|
return tmp_path
|
|
|
|
|
|
def test_advanced_images_masks_filenames(input_dir):
|
|
from nodes.batch_image_loader import _run_image
|
|
Image.new("RGBA", (32, 16), (255, 0, 0, 128)).save(input_dir / "a.png")
|
|
Image.new("RGB", (8, 8), (0, 255, 0)).save(input_dir / "b.png")
|
|
out = _run_image(json.dumps([{"name": "a.png"}, {"name": "b.png"}]), "off", 1200, advanced=True)
|
|
|
|
assert len(out) == 25 and out[0] == 2
|
|
img1, mask1, name1 = out[1], out[2], out[3]
|
|
assert img1.shape == (1, 16, 32, 3) # [1,H,W,C]
|
|
assert mask1.shape == (1, 16, 32) # inverted alpha, image-sized
|
|
assert torch.allclose(mask1, torch.full_like(mask1, 1.0 - 128 / 255), atol=1e-3)
|
|
assert name1 == "a.png"
|
|
img2, mask2 = out[4], out[5]
|
|
assert img2.shape == (1, 8, 8, 3)
|
|
assert mask2.shape == (1, 64, 64) and mask2.max() == 0 # no alpha -> stock 64x64 zeros
|
|
assert out[7] is None # padding starts after file 2's group
|
|
|
|
|
|
def test_simple_layout(input_dir):
|
|
from nodes.batch_image_loader import _run_image
|
|
Image.new("RGB", (8, 8), (0, 0, 255)).save(input_dir / "c.png")
|
|
out = _run_image(json.dumps([{"name": "c.png"}]), "off", 1200, advanced=False)
|
|
assert len(out) == 9 and out[0] == 1
|
|
assert out[1].shape == (1, 8, 8, 3) and out[2] is None
|
|
|
|
|
|
def test_downscale_fit_applies(input_dir):
|
|
from nodes.batch_image_loader import _run_image
|
|
Image.new("RGB", (2400, 1200), (9, 9, 9)).save(input_dir / "big.png")
|
|
out = _run_image(json.dumps([{"name": "big.png"}]), "fit", 1200, advanced=False)
|
|
assert out[1].shape == (1, 600, 1200, 3)
|
|
|
|
|
|
def test_missing_file_raises(input_dir):
|
|
from nodes.batch_image_loader import _run_image
|
|
with pytest.raises(ValueError, match="gone.png"):
|
|
_run_image(json.dumps([{"name": "gone.png"}]), "off", 1200, advanced=False)
|