Alexander Garzon
e8fdd92f21
Manual testing in ComfyUI surfaced two defects in the gpt-image-2.5 work.
background=transparent was silently useless: the API returned a correct RGBA
PNG, but pil_to_tensor ran convert("RGB") on every response and composited the
alpha onto black. Both nodes now output (IMAGE, MASK), with the alpha carried on
the mask as 1 - alpha to match the LoadImage convention, so Join Image with
Alpha consumes it directly. The three response decode helpers collapsed into
response_item_to_pil, since the old path already decoded each item to a PIL
image and then discarded the alpha -- routing through a single decode was
smaller than adding a second one.
CUSTOM_DIM_MIN was 1024 while validate_custom_dimensions() accepts a short edge
down to 480 (1440x480 at the 3:1 aspect limit). The widget was a second, stricter
gate that disagreed with the validator, so legal sizes like 1536x864 snapped up
to 1024 and were unreachable from the UI. The floor is now derived from what the
validator accepts.
The existing tests missed both: they assert on request kwargs and never checked
that alpha survives the response, and they called run_generate() directly,
bypassing the widget bounds entirely.
Verified against the live API (transparent output recomposes to a fully
transparent corner pixel) and by manual testing in ComfyUI.
287 lines
11 KiB
Python
287 lines
11 KiB
Python
import base64
|
|
import io
|
|
import types
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
import imagent.client as client
|
|
from imagent.build import run_generate, run_edit
|
|
|
|
|
|
def _b64(color=(1, 2, 3), size=(16, 16)):
|
|
buf = io.BytesIO()
|
|
Image.new("RGB", size, color).save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode()
|
|
|
|
|
|
class FakeImages:
|
|
def __init__(self, recorder):
|
|
self._rec = recorder
|
|
|
|
def generate(self, **kwargs):
|
|
self._rec["generate"] = kwargs
|
|
item = types.SimpleNamespace(b64_json=_b64(), revised_prompt="a revised prompt")
|
|
return types.SimpleNamespace(data=[item for _ in range(kwargs.get("n", 1))])
|
|
|
|
def edit(self, **kwargs):
|
|
self._rec["edit"] = kwargs
|
|
item = types.SimpleNamespace(b64_json=_b64(), revised_prompt=None)
|
|
return types.SimpleNamespace(data=[item for _ in range(kwargs.get("n", 1))])
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self, recorder):
|
|
self.images = FakeImages(recorder)
|
|
|
|
|
|
def test_generate_missing_key_returns_empty(monkeypatch):
|
|
monkeypatch.setattr(client, "get_client", lambda: None)
|
|
img, _mask, info = run_generate(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1)
|
|
assert img.shape == (1, 512, 512, 3)
|
|
assert info.startswith("Error")
|
|
|
|
|
|
def test_generate_sends_expected_params_and_decodes(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
img, _mask, info = run_generate(
|
|
prompt="a cat", model="gpt-image-2", size="1024x1024", quality="high",
|
|
background="opaque", output_format="png", n=2)
|
|
sent = rec["generate"]
|
|
assert sent["model"] == "gpt-image-2"
|
|
assert sent["prompt"] == "a cat"
|
|
assert sent["n"] == 2
|
|
assert sent["size"] == "1024x1024"
|
|
assert sent["quality"] == "high"
|
|
assert sent["background"] == "opaque"
|
|
assert img.shape[0] == 2
|
|
assert info == "a revised prompt"
|
|
|
|
|
|
def test_generate_drops_transparent_background_for_gpt_image_2(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
run_generate(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="transparent", output_format="png", n=1)
|
|
assert rec["generate"]["background"] == "auto" # gpt-image-2 can't do transparent
|
|
|
|
|
|
def test_generate_extended_quality_gated_to_gpt_image_25(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
run_generate(
|
|
prompt="x", model="gpt-image-2.5-sunburst", size="auto", quality="max",
|
|
background="auto", output_format="png", n=1)
|
|
assert rec["generate"]["quality"] == "max"
|
|
rec.clear()
|
|
run_generate(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="max",
|
|
background="auto", output_format="png", n=1)
|
|
assert rec["generate"]["quality"] == "high"
|
|
|
|
|
|
def test_generate_keeps_transparent_background_for_gpt_image_25(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
run_generate(
|
|
prompt="x", model="gpt-image-2.5-flare", size="auto", quality="auto",
|
|
background="transparent", output_format="png", n=1)
|
|
assert rec["generate"]["background"] == "transparent"
|
|
|
|
|
|
def test_generate_adds_compression_only_for_jpeg(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
run_generate(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="jpeg", n=1, output_compression=80)
|
|
assert rec["generate"]["output_compression"] == 80
|
|
rec.clear()
|
|
run_generate(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1, output_compression=80)
|
|
assert "output_compression" not in rec["generate"]
|
|
|
|
|
|
def test_generate_moderation_only_when_not_auto(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
run_generate(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1, moderation="low")
|
|
assert rec["generate"]["moderation"] == "low"
|
|
rec.clear()
|
|
run_generate(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1, moderation="auto")
|
|
assert "moderation" not in rec["generate"]
|
|
|
|
|
|
def test_n_max_is_8():
|
|
from imagent import params
|
|
assert params.N_MAX == 8
|
|
|
|
|
|
def test_generate_api_error_returns_empty(monkeypatch):
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient({}))
|
|
monkeypatch.setattr(
|
|
FakeImages, "generate",
|
|
lambda self, **kw: (_ for _ in ()).throw(RuntimeError("network error")))
|
|
img, _mask, info = run_generate(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1)
|
|
assert img.shape == (1, 512, 512, 3)
|
|
assert info.startswith("Error")
|
|
|
|
|
|
def test_generate_invalid_custom_size_returns_error(monkeypatch):
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient({}))
|
|
img, _mask, info = run_generate(
|
|
prompt="x", model="gpt-image-2", size="custom", quality="auto",
|
|
background="auto", output_format="png", n=1,
|
|
custom_width=1000, custom_height=1000) # 1000 is not a multiple of 16
|
|
assert img.shape == (1, 512, 512, 3)
|
|
assert info.startswith("Error")
|
|
|
|
|
|
def test_generate_custom_dimensions_for_gpt_image_2(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
run_generate(
|
|
prompt="x", model="gpt-image-2", size="custom", quality="auto",
|
|
background="auto", output_format="png", n=1,
|
|
custom_width=2048, custom_height=1152)
|
|
assert rec["generate"]["size"] == "2048x1152"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# OpenAIImageEdit tests
|
|
# ---------------------------------------------------------------------------
|
|
import torch
|
|
|
|
|
|
def _img_tensor(h=16, w=16):
|
|
return torch.full((1, h, w, 3), 0.5, dtype=torch.float32)
|
|
|
|
|
|
def test_edit_sends_image_list_and_decodes(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
img, _mask, info = run_edit(
|
|
prompt="make it snowy", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1,
|
|
images=[_img_tensor(), _img_tensor()])
|
|
sent = rec["edit"]
|
|
assert sent["model"] == "gpt-image-2"
|
|
assert isinstance(sent["image"], list) and len(sent["image"]) == 2
|
|
assert "mask" not in sent
|
|
assert img.shape[0] == 1
|
|
|
|
|
|
def test_edit_includes_mask_when_provided(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
mask = torch.zeros((16, 16), dtype=torch.float32)
|
|
mask[0:8, 0:8] = 1.0
|
|
run_edit(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1,
|
|
images=[_img_tensor()], mask=mask)
|
|
assert "mask" in rec["edit"]
|
|
|
|
|
|
def test_edit_moderation_only_when_not_auto(monkeypatch):
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
run_edit(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1,
|
|
images=[_img_tensor()], moderation="low")
|
|
assert rec["edit"]["moderation"] == "low"
|
|
rec.clear()
|
|
run_edit(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1,
|
|
images=[_img_tensor()], moderation="auto")
|
|
assert "moderation" not in rec["edit"]
|
|
|
|
|
|
def test_edit_downscales_large_reference(monkeypatch):
|
|
from imagent import image_io
|
|
rec = {}
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec))
|
|
big = torch.full((1, 2100, 2100, 3), 0.5, dtype=torch.float32)
|
|
run_edit(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1, images=[big])
|
|
sent_img = rec["edit"]["image"][0]
|
|
w, h = Image.open(sent_img).size
|
|
assert w * h <= image_io.MAX_INPUT_PIXELS
|
|
|
|
|
|
def test_edit_missing_image_returns_error(monkeypatch):
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient({}))
|
|
img, _mask, info = run_edit(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1, images=[])
|
|
assert info.startswith("Error")
|
|
assert img.shape == (1, 512, 512, 3)
|
|
|
|
|
|
def test_edit_api_error_returns_empty(monkeypatch):
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient({}))
|
|
monkeypatch.setattr(
|
|
FakeImages, "edit",
|
|
lambda self, **kw: (_ for _ in ()).throw(RuntimeError("network error")))
|
|
img, _mask, info = run_edit(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1, images=[_img_tensor()])
|
|
assert img.shape == (1, 512, 512, 3)
|
|
assert info.startswith("Error")
|
|
|
|
|
|
def test_edit_mask_with_multiple_refs_returns_error(monkeypatch):
|
|
monkeypatch.setattr(client, "get_client", lambda: FakeClient({}))
|
|
mask = torch.zeros((16, 16), dtype=torch.float32)
|
|
img, _mask, info = run_edit(
|
|
prompt="x", model="gpt-image-2", size="auto", quality="auto",
|
|
background="auto", output_format="png", n=1,
|
|
images=[_img_tensor(), _img_tensor()], mask=mask)
|
|
assert img.shape == (1, 512, 512, 3)
|
|
assert info.startswith("Error")
|
|
|
|
|
|
def test_model_widget_matches_catalog_and_gpt_image_25_options():
|
|
import pytest
|
|
pytest.importorskip("comfy_api")
|
|
from imagent.node_edit import OpenAIImageEdit
|
|
from imagent.node_generate import OpenAIImageGenerate
|
|
for node in (OpenAIImageGenerate, OpenAIImageEdit):
|
|
model_input = next(i for i in node.define_schema().inputs if i.id == "model")
|
|
assert [o.key for o in model_input.options] == client.MODELS
|
|
for key in ("gpt-image-2.5-flare", "gpt-image-2.5-sunburst"):
|
|
subs = {i.id: i for i in next(o for o in model_input.options if o.key == key).inputs}
|
|
assert "transparent" in subs["background"].options
|
|
assert "custom" in [o.key for o in subs["size"].options]
|
|
|
|
|
|
def test_extension_registers_both_nodes():
|
|
# Registration uses the comfy_api V3 entrypoint, so this needs ComfyUI present.
|
|
import pytest
|
|
pytest.importorskip("comfy_api")
|
|
import asyncio
|
|
import importlib.util
|
|
root = Path(__file__).resolve().parent.parent
|
|
spec = importlib.util.spec_from_file_location("imagent_root_init", root / "__init__.py")
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
nodes = asyncio.run(mod.ImagentExtension().get_node_list())
|
|
schemas = {n.define_schema().node_id: n.define_schema().display_name for n in nodes}
|
|
assert set(schemas) == {"ImagentOpenAIImage", "ImagentOpenAIImageEdit"}
|
|
assert schemas["ImagentOpenAIImage"] == "🤖 Imagent: OpenAI Image"
|
|
assert schemas["ImagentOpenAIImageEdit"] == "🤖 Imagent: OpenAI Image Edit"
|