imaginAIry/tests/img_processors/test_control_modes.py
Bryce 9b95e8b0b6 perf: improve cli startup time
- do not provide automatically imported api functions and objects in `imaginairy` root module
- horrible hack to overcome horrible design choices by easy_install/setuptools

The hack modifies the installed script to remove the __import__ pkg_resources line

If we don't do this then the scripts will be slow to start up because of
pkg_resources.require() which is called by setuptools to ensure the
"correct" version of the package is installed.

before modification example:
```
__requires__ = 'imaginAIry==14.0.0b5'
__import__('pkg_resources').require('imaginAIry==14.0.0b5')
__file__ = '/home/user/projects/imaginairy/imaginairy/bin/aimg'
with open(__file__) as f:
    exec(compile(f.read(), __file__, 'exec'))
```
2023-12-12 20:54:39 -08:00

31 lines
1.1 KiB
Python

import pytest
from lightning_fabric import seed_everything
from imaginairy.img_processors.control_modes import CONTROL_MODES
from imaginairy.img_utils import pillow_img_to_torch_image, torch_img_to_pillow_img
from imaginairy.schema import LazyLoadingImage
from tests import TESTS_FOLDER
from tests.utils import assert_image_similar_to_expectation
def control_img_to_pillow_img(img_t):
return torch_img_to_pillow_img((img_t - 0.5) * 2)
control_mode_params = list(CONTROL_MODES.items())
@pytest.mark.parametrize(("control_name", "control_func"), control_mode_params)
def test_control_images(filename_base_for_outputs, control_func, control_name):
seed_everything(42)
img = LazyLoadingImage(filepath=f"{TESTS_FOLDER}/data/bench2.png")
img_t = pillow_img_to_torch_image(img)
if control_name == "inpaint":
control_t = control_func(img_t.clone(), img_t.clone())
else:
control_t = control_func(img_t.clone())
control_img = control_img_to_pillow_img(control_t)
img_path = f"{filename_base_for_outputs}.png"
assert_image_similar_to_expectation(control_img, img_path, threshold=8000)