328 lines
8.4 KiB
Python
328 lines
8.4 KiB
Python
import pytest
|
|
|
|
from .template import (
|
|
APITestTemplate,
|
|
portrait_imgs,
|
|
girl_img,
|
|
mask_img,
|
|
disable_in_cq,
|
|
get_model,
|
|
console_log_context,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("gen_type", ["img2img", "txt2img"])
|
|
def test_no_unit(gen_type):
|
|
assert APITestTemplate(
|
|
f"test_no_unit{gen_type}",
|
|
gen_type,
|
|
payload_overrides={},
|
|
unit_overrides=[],
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
@pytest.mark.parametrize("gen_type", ["img2img", "txt2img"])
|
|
def test_unrecognized_param(gen_type):
|
|
assert APITestTemplate(
|
|
f"test_unrecognized_param_{gen_type}",
|
|
gen_type,
|
|
payload_overrides={},
|
|
unit_overrides={
|
|
"foo": True,
|
|
"is_ui": False,
|
|
},
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
@pytest.mark.parametrize("gen_type", ["img2img", "txt2img"])
|
|
def test_multiple_iter(gen_type):
|
|
assert APITestTemplate(
|
|
f"test_multiple_iter{gen_type}",
|
|
gen_type,
|
|
payload_overrides={"n_iter": 2},
|
|
unit_overrides={},
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
@pytest.mark.parametrize("gen_type", ["img2img", "txt2img"])
|
|
def test_batch_size(gen_type):
|
|
assert APITestTemplate(
|
|
f"test_batch_size{gen_type}",
|
|
gen_type,
|
|
payload_overrides={"batch_size": 2},
|
|
unit_overrides={},
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
@pytest.mark.parametrize("gen_type", ["img2img", "txt2img"])
|
|
def test_2_units(gen_type):
|
|
assert APITestTemplate(
|
|
f"test_2_units{gen_type}",
|
|
gen_type,
|
|
payload_overrides={},
|
|
unit_overrides=[{}, {}],
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
@pytest.mark.parametrize("gen_type", ["img2img", "txt2img"])
|
|
def test_preprocessor(gen_type):
|
|
assert APITestTemplate(
|
|
f"test_preprocessor{gen_type}",
|
|
gen_type,
|
|
payload_overrides={},
|
|
unit_overrides={"module": "canny"},
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
@pytest.mark.parametrize("param_name", ("processor_res", "threshold_a", "threshold_b"))
|
|
@pytest.mark.parametrize("gen_type", ["img2img", "txt2img"])
|
|
def test_invalid_param(gen_type, param_name):
|
|
with console_log_context() as log_context:
|
|
assert APITestTemplate(
|
|
f"test_invalid_param{(gen_type, param_name)}",
|
|
gen_type,
|
|
payload_overrides={},
|
|
unit_overrides={param_name: -100},
|
|
input_image=girl_img,
|
|
).exec()
|
|
number = "-100" if param_name == "processor_res" else "-100.0"
|
|
assert log_context.is_in_console_logs(
|
|
[
|
|
f"[canny.{param_name}] Invalid value({number}), using default value",
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("save_map", [True, False])
|
|
@pytest.mark.parametrize("gen_type", ["img2img", "txt2img"])
|
|
def test_save_map(gen_type, save_map):
|
|
assert APITestTemplate(
|
|
f"test_save_map{(gen_type, save_map)}",
|
|
gen_type,
|
|
payload_overrides={},
|
|
unit_overrides={"save_detected_map": save_map},
|
|
input_image=girl_img,
|
|
).exec(expected_output_num=2 if save_map else 1)
|
|
|
|
|
|
def test_ip_adapter_face():
|
|
assert APITestTemplate(
|
|
"test_ip_adapter_face",
|
|
"txt2img",
|
|
payload_overrides={},
|
|
unit_overrides={
|
|
"module": "ip-adapter_clip_sd15",
|
|
"model": get_model("ip-adapter-plus-face_sd15"),
|
|
},
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
def test_ip_adapter_fullface():
|
|
assert APITestTemplate(
|
|
"test_ip_adapter_fullface",
|
|
"txt2img",
|
|
payload_overrides={},
|
|
unit_overrides={
|
|
"module": "ip-adapter_clip_sd15",
|
|
"model": get_model("ip-adapter-full-face_sd15"),
|
|
},
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
def test_control_lora():
|
|
assert APITestTemplate(
|
|
"test_control_lora",
|
|
"txt2img",
|
|
payload_overrides={},
|
|
unit_overrides={
|
|
"module": "canny",
|
|
"model": get_model("control_lora_rank128_v11p_sd15_canny"),
|
|
},
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
def test_t2i_adapter():
|
|
assert APITestTemplate(
|
|
"test_t2i_adapter",
|
|
"txt2img",
|
|
payload_overrides={},
|
|
unit_overrides={
|
|
"module": "canny",
|
|
"model": get_model("t2iadapter_canny_sd15v2"),
|
|
},
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
def test_reference():
|
|
assert APITestTemplate(
|
|
"test_reference",
|
|
"txt2img",
|
|
payload_overrides={},
|
|
unit_overrides={
|
|
"module": "reference_only",
|
|
"model": "None",
|
|
},
|
|
input_image=girl_img,
|
|
).exec(result_only=False)
|
|
|
|
|
|
def test_advanced_weighting():
|
|
assert APITestTemplate(
|
|
"test_advanced_weighting",
|
|
"txt2img",
|
|
payload_overrides={},
|
|
unit_overrides={"advanced_weighting": [0.75] * 13}, # SD1.5
|
|
input_image=girl_img,
|
|
).exec()
|
|
|
|
|
|
def test_hr_option():
|
|
assert APITestTemplate(
|
|
"test_hr_option",
|
|
"txt2img",
|
|
payload_overrides={
|
|
"enable_hr": True,
|
|
"denoising_strength": 0.75,
|
|
},
|
|
unit_overrides={"hr_option": "Both"},
|
|
input_image=girl_img,
|
|
).exec(expected_output_num=3)
|
|
|
|
|
|
def test_hr_option_default():
|
|
"""In non-hr run, hr_option should be ignored."""
|
|
assert APITestTemplate(
|
|
"test_hr_option_default",
|
|
"txt2img",
|
|
payload_overrides={"enable_hr": False},
|
|
unit_overrides={"hr_option": "Both"},
|
|
input_image=girl_img,
|
|
).exec(expected_output_num=2)
|
|
|
|
|
|
@disable_in_cq
|
|
def test_masked_controlnet_txt2img():
|
|
assert APITestTemplate(
|
|
f"test_masked_controlnet_txt2img",
|
|
"txt2img",
|
|
payload_overrides={},
|
|
unit_overrides={
|
|
"image": girl_img,
|
|
"mask_image": mask_img,
|
|
},
|
|
).exec()
|
|
|
|
|
|
@disable_in_cq
|
|
def test_masked_controlnet_img2img():
|
|
assert APITestTemplate(
|
|
f"test_masked_controlnet_img2img",
|
|
"img2img",
|
|
payload_overrides={
|
|
"init_images": [girl_img],
|
|
},
|
|
# Note: Currently you must give ControlNet unit input image to specify
|
|
# mask.
|
|
# TODO: Fix this for img2img.
|
|
unit_overrides={
|
|
"image": girl_img,
|
|
"mask_image": mask_img,
|
|
},
|
|
).exec()
|
|
|
|
|
|
@disable_in_cq
|
|
def test_txt2img_inpaint():
|
|
assert APITestTemplate(
|
|
"txt2img_inpaint",
|
|
"txt2img",
|
|
payload_overrides={},
|
|
unit_overrides={
|
|
"image": girl_img,
|
|
"mask_image": mask_img,
|
|
"model": get_model("v11p_sd15_inpaint"),
|
|
"module": "inpaint_only",
|
|
},
|
|
).exec()
|
|
|
|
|
|
@disable_in_cq
|
|
def test_img2img_inpaint():
|
|
assert APITestTemplate(
|
|
"img2img_inpaint",
|
|
"img2img",
|
|
payload_overrides={
|
|
"init_images": [girl_img],
|
|
"mask": mask_img,
|
|
},
|
|
unit_overrides={
|
|
"model": get_model("v11p_sd15_inpaint"),
|
|
"module": "inpaint_only",
|
|
},
|
|
).exec()
|
|
|
|
|
|
@disable_in_cq
|
|
def test_lama_outpaint():
|
|
assert APITestTemplate(
|
|
"txt2img_lama_outpaint",
|
|
"txt2img",
|
|
payload_overrides={
|
|
"width": 768,
|
|
"height": 768,
|
|
},
|
|
# Outpaint should not need a mask.
|
|
unit_overrides={
|
|
"image": girl_img,
|
|
"model": get_model("v11p_sd15_inpaint"),
|
|
"module": "inpaint_only+lama",
|
|
"resize_mode": "Resize and Fill", # OUTER_FIT
|
|
},
|
|
).exec(result_only=False)
|
|
|
|
|
|
@disable_in_cq
|
|
def test_ip_adapter_auto():
|
|
with console_log_context() as log_context:
|
|
assert APITestTemplate(
|
|
"txt2img_ip_adapter_auto",
|
|
"txt2img",
|
|
payload_overrides={},
|
|
unit_overrides={
|
|
"image": girl_img,
|
|
"model": get_model("ip-adapter_sd15"),
|
|
"module": "ip-adapter-auto",
|
|
},
|
|
).exec()
|
|
|
|
assert log_context.is_in_console_logs(["ip-adapter-auto => ip-adapter_clip_h"])
|
|
|
|
|
|
@disable_in_cq
|
|
@pytest.mark.parametrize("img_index", [i for i, _ in enumerate(portrait_imgs)])
|
|
def test_pulid(img_index: int):
|
|
"""PuLID should not memory leak."""
|
|
assert APITestTemplate(
|
|
f"txt2img_pulid_{img_index}",
|
|
"txt2img",
|
|
payload_overrides={
|
|
"width": 768,
|
|
"height": 768,
|
|
},
|
|
unit_overrides={
|
|
"image": portrait_imgs[img_index],
|
|
"model": get_model("ip-adapter_pulid_sdxl_fp16"),
|
|
"module": "ip-adapter_pulid",
|
|
},
|
|
).exec()
|