mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-10-31 03:20:40 +00:00
365 lines
11 KiB
Python
365 lines
11 KiB
Python
import os.path
|
|
|
|
import pytest
|
|
|
|
from imaginairy import LazyLoadingImage
|
|
from imaginairy.api import imagine, imagine_image_files
|
|
from imaginairy.img_processors.control_modes import CONTROL_MODES
|
|
from imaginairy.img_utils import pillow_fit_image_within
|
|
from imaginairy.schema import ControlNetInput, ImaginePrompt
|
|
from imaginairy.utils import get_device
|
|
|
|
from . import TESTS_FOLDER
|
|
from .utils import assert_image_similar_to_expectation
|
|
|
|
|
|
def test_imagine(sampler_type, filename_base_for_outputs):
|
|
prompt_text = "a scenic old-growth forest with diffuse light poking through the canopy. high resolution nature photography"
|
|
prompt = ImaginePrompt(
|
|
prompt_text, width=512, height=512, steps=20, seed=1, sampler_type=sampler_type
|
|
)
|
|
result = next(imagine(prompt))
|
|
|
|
threshold_lookup = {"k_dpm_2_a": 26000}
|
|
threshold = threshold_lookup.get(sampler_type, 10000)
|
|
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(
|
|
result.img, img_path=img_path, threshold=threshold
|
|
)
|
|
|
|
|
|
compare_prompts = [
|
|
"a photo of a bowl of fruit",
|
|
"a headshot photo of a happy couple smiling at the camera",
|
|
"a painting of a beautiful cloudy sunset at the beach",
|
|
"a photo of a dog",
|
|
"a photo of a handshake",
|
|
"a photo of an astronaut riding a horse on the moon. the earth visible in the background",
|
|
]
|
|
|
|
|
|
@pytest.mark.skipif(get_device() != "cuda", reason="Too slow to run on CPU or MPS")
|
|
@pytest.mark.parametrize(
|
|
"model_version", ["SD-1.4", "SD-1.5", "SD-2.0", "SD-2.0-v", "SD-2.1", "SD-2.1-v"]
|
|
)
|
|
def test_model_versions(filename_base_for_orig_outputs, model_version):
|
|
"""Test that we can switch between model versions."""
|
|
prompts = []
|
|
for prompt_text in compare_prompts:
|
|
prompts.append(
|
|
ImaginePrompt(
|
|
prompt_text,
|
|
seed=1,
|
|
model=model_version,
|
|
)
|
|
)
|
|
|
|
threshold = 35000
|
|
|
|
for i, result in enumerate(imagine(prompts)):
|
|
img_path = f"{filename_base_for_orig_outputs}_{result.prompt.prompt_text}_{result.prompt.model}.png"
|
|
result.img.save(img_path)
|
|
|
|
for i, result in enumerate(imagine(prompts)):
|
|
img_path = f"{filename_base_for_orig_outputs}_{result.prompt.prompt_text}_{result.prompt.model}.png"
|
|
assert_image_similar_to_expectation(
|
|
result.img, img_path=img_path, threshold=threshold
|
|
)
|
|
|
|
|
|
def test_img2img_beach_to_sunset(
|
|
sampler_type, filename_base_for_outputs, filename_base_for_orig_outputs
|
|
):
|
|
img = LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/beach_at_sainte_adresse.jpg")
|
|
prompt = ImaginePrompt(
|
|
"a painting of beautiful cloudy sunset at the beach",
|
|
init_image=img,
|
|
init_image_strength=0.5,
|
|
prompt_strength=15,
|
|
mask_prompt="(sky|clouds) AND !(buildings|trees)",
|
|
mask_mode="replace",
|
|
width=512,
|
|
height=512,
|
|
steps=40 * 2,
|
|
seed=1,
|
|
sampler_type=sampler_type,
|
|
)
|
|
result = next(imagine(prompt))
|
|
|
|
pillow_fit_image_within(img).save(f"{filename_base_for_orig_outputs}__orig.jpg")
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=2900)
|
|
|
|
|
|
def test_img_to_img_from_url_cats(
|
|
sampler_type,
|
|
filename_base_for_outputs,
|
|
mocked_responses,
|
|
filename_base_for_orig_outputs,
|
|
):
|
|
with open(
|
|
os.path.join(TESTS_FOLDER, "data", "val2017-000000039769-cococats.jpg"), "rb"
|
|
) as f:
|
|
img_data = f.read()
|
|
mocked_responses.get(
|
|
"http://images.cocodataset.org/val2017/000000039769.jpg",
|
|
body=img_data,
|
|
status=200,
|
|
content_type="image/jpeg",
|
|
)
|
|
img = LazyLoadingImage(url="http://images.cocodataset.org/val2017/000000039769.jpg")
|
|
|
|
prompt = ImaginePrompt(
|
|
"dogs lying on a hot pink couch",
|
|
init_image=img,
|
|
init_image_strength=0.5,
|
|
width=512,
|
|
height=512,
|
|
steps=50,
|
|
seed=1,
|
|
sampler_type=sampler_type,
|
|
)
|
|
|
|
result = next(imagine(prompt))
|
|
|
|
img = pillow_fit_image_within(img)
|
|
img.save(f"{filename_base_for_orig_outputs}__orig.jpg")
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=17000)
|
|
|
|
|
|
def test_img2img_low_noise(
|
|
filename_base_for_outputs,
|
|
sampler_type,
|
|
):
|
|
fruit_path = os.path.join(TESTS_FOLDER, "data", "bowl_of_fruit.jpg")
|
|
img = LazyLoadingImage(filepath=fruit_path)
|
|
|
|
prompt = ImaginePrompt(
|
|
"a white bowl filled with gold coins",
|
|
prompt_strength=12,
|
|
init_image=img,
|
|
init_image_strength=0.5,
|
|
mask_prompt="(fruit{*2} OR stem{*10} OR fruit stem{*3})",
|
|
mask_mode="replace",
|
|
# steps=40,
|
|
seed=1,
|
|
sampler_type=sampler_type,
|
|
)
|
|
|
|
result = next(imagine(prompt))
|
|
|
|
threshold_lookup = {
|
|
"k_dpm_2_a": 26000,
|
|
"k_euler_a": 18000,
|
|
"k_dpm_adaptive": 13000,
|
|
}
|
|
threshold = threshold_lookup.get(sampler_type, 14000)
|
|
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(
|
|
result.img, img_path=img_path, threshold=threshold
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("init_strength", [0, 0.05, 0.2, 1])
|
|
def test_img_to_img_fruit_2_gold(
|
|
filename_base_for_outputs,
|
|
sampler_type,
|
|
init_strength,
|
|
filename_base_for_orig_outputs,
|
|
):
|
|
img = LazyLoadingImage(
|
|
filepath=os.path.join(TESTS_FOLDER, "data", "bowl_of_fruit.jpg")
|
|
)
|
|
target_steps = 25
|
|
needed_steps = 25 if init_strength >= 1 else int(target_steps / (1 - init_strength))
|
|
prompt = ImaginePrompt(
|
|
"a white bowl filled with gold coins",
|
|
prompt_strength=12,
|
|
init_image=img,
|
|
init_image_strength=init_strength,
|
|
mask_prompt="(fruit{*2} OR stem{*10} OR fruit stem{*3})",
|
|
mask_mode="replace",
|
|
steps=needed_steps,
|
|
seed=1,
|
|
sampler_type=sampler_type,
|
|
)
|
|
|
|
result = next(imagine(prompt))
|
|
|
|
threshold_lookup = {
|
|
"k_dpm_2_a": 32000,
|
|
"k_euler_a": 18000,
|
|
"k_dpm_adaptive": 13000,
|
|
"k_dpmpp_2s": 16000,
|
|
}
|
|
threshold = threshold_lookup.get(sampler_type, 16000)
|
|
|
|
pillow_fit_image_within(img).save(f"{filename_base_for_orig_outputs}__orig.jpg")
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(
|
|
result.img, img_path=img_path, threshold=threshold
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
|
|
def test_img_to_img_fruit_2_gold_repeat():
|
|
img = LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/bowl_of_fruit.jpg")
|
|
run_count = 1
|
|
|
|
kwargs = {
|
|
"prompt": "a white bowl filled with gold coins. sharp focus",
|
|
"prompt_strength": 12,
|
|
"init_image": img,
|
|
"init_image_strength": 0.2,
|
|
"mask_prompt": "(fruit OR stem{*5} OR fruit stem)",
|
|
"mask_mode": "replace",
|
|
"steps": 20,
|
|
"seed": 946188797,
|
|
"sampler_type": "plms",
|
|
"fix_faces": True,
|
|
"upscale": True,
|
|
}
|
|
prompts = [
|
|
ImaginePrompt(**kwargs),
|
|
ImaginePrompt(**kwargs),
|
|
ImaginePrompt(**kwargs),
|
|
]
|
|
for result in imagine(prompts, debug_img_callback=None):
|
|
result.img.save(
|
|
f"{TESTS_FOLDER}/test_output/img2img_fruit_2_gold_plms_{get_device()}_run-{run_count:02}.jpg"
|
|
)
|
|
run_count += 1
|
|
|
|
|
|
@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
|
|
def test_img_to_file():
|
|
prompt = ImaginePrompt(
|
|
"an old growth forest, diffuse light poking through the canopy. high-resolution, nature photography, nat geo photo",
|
|
width=512 + 64,
|
|
height=512 - 64,
|
|
steps=20,
|
|
seed=2,
|
|
sampler_type="PLMS",
|
|
upscale=True,
|
|
)
|
|
out_folder = f"{TESTS_FOLDER}/test_output"
|
|
imagine_image_files(prompt, outdir=out_folder)
|
|
|
|
|
|
@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
|
|
def test_inpainting_bench(filename_base_for_outputs, filename_base_for_orig_outputs):
|
|
img = LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/bench2.png")
|
|
prompt = ImaginePrompt(
|
|
"a wise old man",
|
|
init_image=img,
|
|
init_image_strength=0.4,
|
|
mask_image=LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/bench2_mask.png"),
|
|
width=512,
|
|
height=512,
|
|
steps=40,
|
|
seed=1,
|
|
sampler_type="plms",
|
|
)
|
|
result = next(imagine(prompt))
|
|
|
|
pillow_fit_image_within(img).save(f"{filename_base_for_orig_outputs}_orig.jpg")
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=2800)
|
|
|
|
|
|
@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
|
|
def test_cliptext_inpainting_pearl_doctor(
|
|
filename_base_for_outputs, filename_base_for_orig_outputs
|
|
):
|
|
img = LazyLoadingImage(
|
|
filepath=f"{TESTS_FOLDER}/data/girl_with_a_pearl_earring.jpg"
|
|
)
|
|
prompt = ImaginePrompt(
|
|
"a female doctor in the hospital",
|
|
prompt_strength=12,
|
|
init_image=img,
|
|
init_image_strength=0.2,
|
|
mask_prompt="face AND NOT (bandana OR hair OR blue fabric){*5}",
|
|
mask_mode=ImaginePrompt.MaskMode.KEEP,
|
|
width=512,
|
|
height=512,
|
|
steps=40,
|
|
sampler_type="plms",
|
|
seed=181509347,
|
|
)
|
|
result = next(imagine(prompt))
|
|
|
|
pillow_fit_image_within(img).save(f"{filename_base_for_orig_outputs}_orig.jpg")
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=32000)
|
|
|
|
|
|
@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
|
|
def test_tile_mode(filename_base_for_outputs):
|
|
prompt_text = "gold coins"
|
|
prompt = ImaginePrompt(
|
|
prompt_text,
|
|
width=400,
|
|
height=400,
|
|
steps=15,
|
|
seed=1,
|
|
tile_mode="xy",
|
|
)
|
|
result = next(imagine(prompt))
|
|
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=25000)
|
|
|
|
|
|
control_modes = list(CONTROL_MODES.keys())
|
|
|
|
|
|
@pytest.mark.parametrize("control_mode", control_modes)
|
|
@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
|
|
def test_controlnet(filename_base_for_outputs, control_mode):
|
|
prompt_text = "a photo of a woman sitting on a bench"
|
|
control_input = ControlNetInput(
|
|
mode=control_mode,
|
|
image=LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/bench2.png"),
|
|
)
|
|
|
|
prompt = ImaginePrompt(
|
|
prompt_text,
|
|
width=512,
|
|
height=512,
|
|
steps=15,
|
|
seed=0,
|
|
control_inputs=[control_input],
|
|
fix_faces=True,
|
|
)
|
|
prompt.steps = 1
|
|
prompt.width = 256
|
|
prompt.height = 256
|
|
result = next(imagine(prompt))
|
|
prompt.steps = 15
|
|
prompt.width = 512
|
|
prompt.height = 512
|
|
result = next(imagine(prompt))
|
|
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=24000)
|
|
|
|
|
|
@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
|
|
def test_large_image(filename_base_for_outputs):
|
|
prompt_text = "a stormy ocean. oil painting"
|
|
prompt = ImaginePrompt(
|
|
prompt_text,
|
|
width=1920,
|
|
height=1080,
|
|
steps=15,
|
|
seed=0,
|
|
)
|
|
result = next(imagine(prompt))
|
|
|
|
img_path = f"{filename_base_for_outputs}.png"
|
|
assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=35000)
|