247 lines
9.5 KiB
Python
247 lines
9.5 KiB
Python
import pytest
|
|
|
|
from comfy_api_nodes.nodes_openai import (
|
|
OpenAIGPTImage1,
|
|
OpenAIGPTImage2,
|
|
_GPT_IMAGE_2_SIZES,
|
|
_resolve_gpt_image_2_size,
|
|
calculate_tokens_price_image_1,
|
|
calculate_tokens_price_image_1_5,
|
|
calculate_tokens_price_image_2,
|
|
)
|
|
from comfy_api_nodes.apis.openai import OpenAIImageGenerationResponse, Usage
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_response(input_tokens: int, output_tokens: int) -> OpenAIImageGenerationResponse:
|
|
return OpenAIImageGenerationResponse(
|
|
data=[],
|
|
usage=Usage(input_tokens=input_tokens, output_tokens=output_tokens),
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Price extractor tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_price_image_1_formula():
|
|
response = _make_response(input_tokens=1_000_000, output_tokens=1_000_000)
|
|
assert calculate_tokens_price_image_1(response) == pytest.approx(50.0)
|
|
|
|
|
|
def test_price_image_1_5_formula():
|
|
response = _make_response(input_tokens=1_000_000, output_tokens=1_000_000)
|
|
assert calculate_tokens_price_image_1_5(response) == pytest.approx(40.0)
|
|
|
|
|
|
def test_price_image_2_formula():
|
|
response = _make_response(input_tokens=1_000_000, output_tokens=1_000_000)
|
|
assert calculate_tokens_price_image_2(response) == pytest.approx(38.0)
|
|
|
|
|
|
def test_price_image_2_cheaper_than_1():
|
|
response = _make_response(input_tokens=500, output_tokens=196)
|
|
assert calculate_tokens_price_image_2(response) < calculate_tokens_price_image_1(response)
|
|
|
|
|
|
def test_price_image_2_cheaper_output_than_1_5():
|
|
# gpt-image-2 output rate ($30/1M) is lower than gpt-image-1.5 ($32/1M)
|
|
response = _make_response(input_tokens=0, output_tokens=1_000_000)
|
|
assert calculate_tokens_price_image_2(response) < calculate_tokens_price_image_1_5(response)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _resolve_gpt_image_2_size tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_resolve_preset_passthrough_when_custom_zero():
|
|
# 0/0 means "use size preset"
|
|
assert _resolve_gpt_image_2_size("1024x1024", 0, 0) == "1024x1024"
|
|
assert _resolve_gpt_image_2_size("auto", 0, 0) == "auto"
|
|
assert _resolve_gpt_image_2_size("3840x2160", 0, 0) == "3840x2160"
|
|
|
|
|
|
def test_resolve_preset_passthrough_when_only_one_dim_set():
|
|
# only one dimension set → still use preset
|
|
assert _resolve_gpt_image_2_size("auto", 1024, 0) == "auto"
|
|
assert _resolve_gpt_image_2_size("auto", 0, 1024) == "auto"
|
|
|
|
|
|
def test_resolve_custom_overrides_preset():
|
|
assert _resolve_gpt_image_2_size("auto", 1024, 1024) == "1024x1024"
|
|
assert _resolve_gpt_image_2_size("1024x1024", 2048, 1152) == "2048x1152"
|
|
assert _resolve_gpt_image_2_size("auto", 3840, 2160) == "3840x2160"
|
|
|
|
|
|
def test_resolve_custom_rejects_edge_too_large():
|
|
with pytest.raises(ValueError, match="3840"):
|
|
_resolve_gpt_image_2_size("auto", 4096, 1024)
|
|
|
|
|
|
def test_resolve_custom_rejects_non_multiple_of_16():
|
|
with pytest.raises(ValueError, match="multiple of 16"):
|
|
_resolve_gpt_image_2_size("auto", 1025, 1024)
|
|
|
|
|
|
def test_resolve_custom_rejects_bad_ratio():
|
|
with pytest.raises(ValueError, match="ratio"):
|
|
_resolve_gpt_image_2_size("auto", 3840, 1024) # 3.75:1 > 3:1
|
|
|
|
|
|
def test_resolve_custom_rejects_too_few_pixels():
|
|
with pytest.raises(ValueError, match="Total pixels"):
|
|
_resolve_gpt_image_2_size("auto", 16, 16)
|
|
|
|
|
|
def test_resolve_custom_rejects_too_many_pixels():
|
|
# 3840x2176 exceeds 8,294,400
|
|
with pytest.raises(ValueError, match="Total pixels"):
|
|
_resolve_gpt_image_2_size("auto", 3840, 2176)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# OpenAIGPTImage1 schema tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestOpenAIGPTImage1Schema:
|
|
def setup_method(self):
|
|
self.schema = OpenAIGPTImage1.define_schema()
|
|
|
|
def test_node_id(self):
|
|
assert self.schema.node_id == "OpenAIGPTImage1"
|
|
|
|
def test_display_name(self):
|
|
assert self.schema.display_name == "OpenAI GPT Image 1 & 1.5"
|
|
|
|
def test_model_options_exclude_gpt_image_2(self):
|
|
model_input = next(i for i in self.schema.inputs if i.name == "model")
|
|
assert "gpt-image-2" not in model_input.options
|
|
|
|
def test_model_options_include_legacy_models(self):
|
|
model_input = next(i for i in self.schema.inputs if i.name == "model")
|
|
assert "gpt-image-1" in model_input.options
|
|
assert "gpt-image-1.5" in model_input.options
|
|
|
|
def test_has_background_with_transparent(self):
|
|
bg_input = next(i for i in self.schema.inputs if i.name == "background")
|
|
assert "transparent" in bg_input.options
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# OpenAIGPTImage2 schema tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestOpenAIGPTImage2Schema:
|
|
def setup_method(self):
|
|
self.schema = OpenAIGPTImage2.define_schema()
|
|
|
|
def test_node_id(self):
|
|
assert self.schema.node_id == "OpenAIGPTImage2"
|
|
|
|
def test_display_name(self):
|
|
assert self.schema.display_name == "OpenAI GPT Image 2"
|
|
|
|
def test_category(self):
|
|
assert "OpenAI" in self.schema.category
|
|
|
|
def test_no_transparent_background(self):
|
|
bg_input = next(i for i in self.schema.inputs if i.name == "background")
|
|
assert "transparent" not in bg_input.options
|
|
|
|
def test_background_options(self):
|
|
bg_input = next(i for i in self.schema.inputs if i.name == "background")
|
|
assert set(bg_input.options) == {"auto", "opaque"}
|
|
|
|
def test_quality_options(self):
|
|
quality_input = next(i for i in self.schema.inputs if i.name == "quality")
|
|
assert set(quality_input.options) == {"auto", "low", "medium", "high"}
|
|
|
|
def test_quality_default_is_auto(self):
|
|
quality_input = next(i for i in self.schema.inputs if i.name == "quality")
|
|
assert quality_input.default == "auto"
|
|
|
|
def test_all_popular_sizes_present(self):
|
|
size_input = next(i for i in self.schema.inputs if i.name == "size")
|
|
for size in ["1024x1024", "1536x1024", "1024x1536", "2048x2048", "2048x1152", "3840x2160", "2160x3840"]:
|
|
assert size in size_input.options, f"Missing size: {size}"
|
|
|
|
def test_no_custom_size_option(self):
|
|
size_input = next(i for i in self.schema.inputs if i.name == "size")
|
|
assert "custom" not in size_input.options
|
|
|
|
def test_size_default_is_auto(self):
|
|
size_input = next(i for i in self.schema.inputs if i.name == "size")
|
|
assert size_input.default == "auto"
|
|
|
|
def test_custom_width_and_height_inputs_exist(self):
|
|
input_names = [i.name for i in self.schema.inputs]
|
|
assert "custom_width" in input_names
|
|
assert "custom_height" in input_names
|
|
|
|
def test_custom_width_height_default_zero(self):
|
|
width_input = next(i for i in self.schema.inputs if i.name == "custom_width")
|
|
height_input = next(i for i in self.schema.inputs if i.name == "custom_height")
|
|
assert width_input.default == 0
|
|
assert height_input.default == 0
|
|
|
|
def test_custom_width_height_step_is_16(self):
|
|
width_input = next(i for i in self.schema.inputs if i.name == "custom_width")
|
|
height_input = next(i for i in self.schema.inputs if i.name == "custom_height")
|
|
assert width_input.step == 16
|
|
assert height_input.step == 16
|
|
|
|
def test_custom_width_height_max_is_3840(self):
|
|
width_input = next(i for i in self.schema.inputs if i.name == "custom_width")
|
|
height_input = next(i for i in self.schema.inputs if i.name == "custom_height")
|
|
assert width_input.max == 3840
|
|
assert height_input.max == 3840
|
|
|
|
def test_uses_num_images_not_n(self):
|
|
input_names = [i.name for i in self.schema.inputs]
|
|
assert "num_images" in input_names
|
|
assert "n" not in input_names
|
|
|
|
def test_model_input_shows_gpt_image_2(self):
|
|
model_input = next(i for i in self.schema.inputs if i.name == "model")
|
|
assert model_input.options == ["gpt-image-2"]
|
|
assert model_input.default == "gpt-image-2"
|
|
|
|
def test_has_image_and_mask_inputs(self):
|
|
input_names = [i.name for i in self.schema.inputs]
|
|
assert "image" in input_names
|
|
assert "mask" in input_names
|
|
|
|
def test_is_api_node(self):
|
|
assert self.schema.is_api_node is True
|
|
|
|
def test_sizes_match_constant(self):
|
|
size_input = next(i for i in self.schema.inputs if i.name == "size")
|
|
assert size_input.options == _GPT_IMAGE_2_SIZES
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# OpenAIGPTImage2 execute validation tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_raises_on_empty_prompt():
|
|
with pytest.raises(Exception):
|
|
await OpenAIGPTImage2.execute(prompt=" ")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_raises_mask_without_image():
|
|
import torch
|
|
mask = torch.ones(1, 64, 64)
|
|
with pytest.raises(ValueError, match="mask without an input image"):
|
|
await OpenAIGPTImage2.execute(prompt="test", mask=mask)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_raises_invalid_custom_size():
|
|
with pytest.raises(ValueError):
|
|
await OpenAIGPTImage2.execute(prompt="test", custom_width=4096, custom_height=1024)
|