118 lines
3.8 KiB
Python
118 lines
3.8 KiB
Python
import base64
|
|
import io
|
|
import math
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
|
|
import gates
|
|
from gates import image_chooser
|
|
|
|
|
|
def test_normalize_selection_deduplicates_and_preserves_batch_order():
|
|
assert image_chooser.normalize_selection("[3, 1, 3]", 4) == (1, 3)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"selection",
|
|
[
|
|
[],
|
|
"[]",
|
|
"not json",
|
|
"1",
|
|
[True],
|
|
[1.0],
|
|
["1"],
|
|
[-1],
|
|
[3],
|
|
],
|
|
)
|
|
def test_normalize_selection_rejects_invalid_choices(selection):
|
|
with pytest.raises(ValueError):
|
|
image_chooser.normalize_selection(selection, 3)
|
|
|
|
|
|
def test_select_batch_keeps_batch_dimension_and_original_pixels():
|
|
images = torch.arange(4 * 2 * 3 * 3, dtype=torch.float32).reshape(4, 2, 3, 3)
|
|
|
|
one = image_chooser.select_batch(images, [2])
|
|
many = image_chooser.select_batch(images, [3, 0, 3])
|
|
|
|
assert one.shape == (1, 2, 3, 3)
|
|
assert torch.equal(one[0], images[2])
|
|
assert many.shape == (2, 2, 3, 3)
|
|
assert torch.equal(many, images[[0, 3]])
|
|
assert many.dtype == images.dtype
|
|
|
|
|
|
def test_encode_previews_returns_small_jpegs_and_source_dimensions():
|
|
images = torch.zeros((2, 80, 40, 3), dtype=torch.float32)
|
|
images[1] = 1.0
|
|
|
|
previews = image_chooser.encode_previews(images, max_side=32)
|
|
|
|
assert [preview["index"] for preview in previews] == [0, 1]
|
|
assert [(preview["width"], preview["height"]) for preview in previews] == [
|
|
(40, 80),
|
|
(40, 80),
|
|
]
|
|
decoded = [
|
|
Image.open(io.BytesIO(base64.b64decode(preview["image"])))
|
|
for preview in previews
|
|
]
|
|
assert all(image.format == "JPEG" and image.mode == "RGB" for image in decoded)
|
|
assert [image.size for image in decoded] == [(16, 32), (16, 32)]
|
|
|
|
|
|
def test_encode_previews_accepts_bfloat16_images():
|
|
images = torch.zeros((1, 8, 8, 3), dtype=torch.bfloat16)
|
|
assert len(image_chooser.encode_previews(images)) == 1
|
|
|
|
|
|
def test_run_waits_for_token_scoped_selection(monkeypatch):
|
|
fake_server = types.ModuleType("gates.gate_server")
|
|
|
|
def send_image_choices(node_id, token, images):
|
|
assert images.shape[0] == 3
|
|
assert image_chooser.gate_bus.GateBus.put_token_payload(
|
|
node_id, token, [2, 0]
|
|
)
|
|
|
|
fake_server.send_image_choices = send_image_choices
|
|
monkeypatch.setitem(sys.modules, "gates.gate_server", fake_server)
|
|
monkeypatch.setattr(gates, "gate_server", fake_server, raising=False)
|
|
|
|
class InterruptProcessingException(Exception):
|
|
pass
|
|
|
|
fake_mm = types.ModuleType("comfy.model_management")
|
|
fake_mm.processing_interrupted = lambda: False
|
|
fake_mm.InterruptProcessingException = InterruptProcessingException
|
|
fake_comfy = types.ModuleType("comfy")
|
|
fake_comfy.model_management = fake_mm
|
|
monkeypatch.setitem(sys.modules, "comfy", fake_comfy)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", fake_mm)
|
|
|
|
images = torch.arange(3, dtype=torch.float32).reshape(3, 1, 1, 1)
|
|
selected, = image_chooser.ImageChooserGate().run(images, unique_id="12")
|
|
|
|
assert torch.equal(selected.flatten(), torch.tensor([0.0, 2.0]))
|
|
assert "12" not in image_chooser.gate_bus.GateBus.active_tokens
|
|
|
|
|
|
def test_image_chooser_node_contract():
|
|
inputs = image_chooser.ImageChooserGate.INPUT_TYPES()
|
|
|
|
assert inputs["required"] == {"images": ("IMAGE",)}
|
|
assert inputs["hidden"] == {"unique_id": "UNIQUE_ID"}
|
|
assert image_chooser.ImageChooserGate.RETURN_TYPES == ("IMAGE",)
|
|
assert image_chooser.ImageChooserGate.RETURN_NAMES == ("images",)
|
|
assert image_chooser.ImageChooserGate.FUNCTION == "run"
|
|
assert image_chooser.ImageChooserGate.CATEGORY == "Dataset Gates"
|
|
assert math.isnan(image_chooser.ImageChooserGate.IS_CHANGED(images=None))
|
|
assert image_chooser.NODE_CLASS_MAPPINGS["ImageChooserGate"] \
|
|
is image_chooser.ImageChooserGate
|