Skip to content

Commit

Permalink
tests: warm loaded models now are unloaded at test run start
Browse files Browse the repository at this point in the history
  • Loading branch information
@tazlin
tazlin committed Jul 21, 2023
1 parent 4074167 commit 749e577
Show file tree
Hide file tree
Showing 8 changed files with 82 additions and 70 deletions.
16 changes: 12 additions & 4 deletions tests/conftest.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from typing import Generator

import PIL.Image
from loguru import logger
import pytest

from hordelib.comfy_horde import Comfy_Horde
Expand All @@ -20,6 +21,7 @@ def init_horde():
hordelib.initialise()
from hordelib.settings import UserSettings

UserSettings.set_ram_to_leave_free_mb("100%")
UserSettings.set_vram_to_leave_free_mb("90%")


Expand Down Expand Up @@ -50,20 +52,26 @@ def shared_model_manager(hordelib_instance: HordeLib) -> Generator[type[SharedMo
assert SharedModelManager.manager.blip is not None
assert SharedModelManager.manager.clip is not None

for model_availible in SharedModelManager.manager.compvis.get_loaded_models():
assert SharedModelManager.manager.unload_model(model_availible)

yield SharedModelManager

SharedModelManager._instance = None # type: ignore
SharedModelManager.manager = None # type: ignore


_testing_model_name = "Deliberate"


@pytest.fixture(scope="class")
def stable_diffusion_modelname_for_testing(shared_model_manager: type[SharedModelManager]) -> str:
def stable_diffusion_model_name_for_testing(shared_model_manager: type[SharedModelManager]) -> str:
"""Loads the stable diffusion model for testing. This model is used by many tests.
This fixture returns the model name as string."""
shared_model_manager.load_model_managers([CompVisModelManager])
model_name = "Deliberate"
assert shared_model_manager.manager.load(model_name)
return model_name

assert shared_model_manager.manager.load(_testing_model_name)
return _testing_model_name


@pytest.fixture(scope="session")
Expand Down
20 changes: 12 additions & 8 deletions tests/test_horde_inference.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ class TestHordeInference:
def test_text_to_image(
self,
hordelib_instance: HordeLib,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
):
data = {
"sampler_name": "k_dpmpp_2m",
Expand All @@ -33,10 +33,11 @@ def test_text_to_image(
"prompt": "an ancient llamia monster",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
}
pil_image = hordelib_instance.basic_inference(data)
assert pil_image is not None
assert isinstance(pil_image, Image.Image)

img_filename = "text_to_image.png"
pil_image.save(f"images/{img_filename}", quality=100)
Expand All @@ -49,7 +50,7 @@ def test_text_to_image(
def test_text_to_image_small(
self,
hordelib_instance: HordeLib,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
):
data = {
"sampler_name": "k_dpmpp_2m",
Expand All @@ -68,10 +69,11 @@ def test_text_to_image_small(
"prompt": "a photo of cute dinosaur ### painting, drawing, artwork, red border",
"ddim_steps": 20,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
}
pil_image = hordelib_instance.basic_inference(data)
assert pil_image is not None
assert isinstance(pil_image, Image.Image)

img_filename = "text_to_image_small.png"
pil_image.save(f"images/{img_filename}", quality=100)
Expand All @@ -84,7 +86,7 @@ def test_text_to_image_small(
def test_text_to_image_clip_skip_2(
self,
hordelib_instance: HordeLib,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
):
data = {
"sampler_name": "k_dpmpp_2m",
Expand All @@ -103,10 +105,11 @@ def test_text_to_image_clip_skip_2(
"prompt": "an ancient llamia monster",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
}
pil_image = hordelib_instance.basic_inference(data)
assert pil_image is not None
assert isinstance(pil_image, Image.Image)

img_filename = "text_to_image_clip_skip_2.png"
pil_image.save(f"images/{img_filename}", quality=100)
Expand All @@ -119,7 +122,7 @@ def test_text_to_image_clip_skip_2(
def test_text_to_image_hires_fix(
self,
hordelib_instance: HordeLib,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
):
data = {
"sampler_name": "k_dpmpp_2m",
Expand All @@ -138,10 +141,11 @@ def test_text_to_image_hires_fix(
"prompt": "an ancient llamia monster",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
}
pil_image = hordelib_instance.basic_inference(data)
assert pil_image is not None
assert isinstance(pil_image, Image.Image)

img_filename = "text_to_image_hires_fix.png"
pil_image.save(f"images/{img_filename}", quality=100)
Expand Down
20 changes: 10 additions & 10 deletions tests/test_horde_inference_controlnet.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ def test_controlnet_sd1(
self,
shared_model_manager: type[SharedModelManager],
hordelib_instance: HordeLib,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
):
data = {
"sampler_name": "k_dpmpp_2m",
Expand All @@ -40,7 +40,7 @@ def test_controlnet_sd1(
"prompt": "a man walking in the snow",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_db0.jpg"),
"source_processing": "img2img",
}
Expand Down Expand Up @@ -80,7 +80,7 @@ def test_controlnet_sd1(
def test_controlnet_fake_cn(
self,
hordelib_instance: HordeLib,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
db0_test_image: Image.Image,
):
data = {
Expand All @@ -100,7 +100,7 @@ def test_controlnet_fake_cn(
"prompt": "a man walking in the snow",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": db0_test_image,
"source_processing": "img2img",
}
Expand All @@ -111,7 +111,7 @@ def test_controlnet_fake_cn(
def test_controlnet_strength(
self,
hordelib_instance: HordeLib,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
):
data = {
"sampler_name": "k_dpmpp_2m",
Expand All @@ -130,7 +130,7 @@ def test_controlnet_strength(
"prompt": "a man walking on the moon",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_db0.jpg"),
"source_processing": "img2img",
}
Expand All @@ -157,7 +157,7 @@ def test_controlnet_strength(
def test_controlnet_hires_fix(
self,
hordelib_instance: HordeLib,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
):
data = {
"sampler_name": "k_dpmpp_2m",
Expand All @@ -177,7 +177,7 @@ def test_controlnet_hires_fix(
"prompt": "a man walking in the jungle",
"ddim_steps": 15,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_db0.jpg"),
"source_processing": "img2img",
}
Expand All @@ -198,7 +198,7 @@ def test_controlnet_hires_fix(
def test_controlnet_image_is_control(
self,
hordelib_instance: HordeLib,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
):
data = {
"sampler_name": "k_dpmpp_2m",
Expand All @@ -217,7 +217,7 @@ def test_controlnet_image_is_control(
"prompt": "a woman standing in the snow",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_image_is_control.png"),
"source_processing": "img2img",
}
Expand Down
32 changes: 16 additions & 16 deletions tests/test_horde_inference_img2img.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
class TestHordeInference:
def test_image_to_image(
self,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
hordelib_instance: HordeLib,
):
data = {
Expand All @@ -33,7 +33,7 @@ def test_image_to_image(
"prompt": "a dinosaur",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_db0.jpg"),
"source_processing": "img2img",
}
Expand All @@ -51,7 +51,7 @@ def test_image_to_image(

def test_image_to_image_hires_fix_small(
self,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
hordelib_instance: HordeLib,
):
data = {
Expand All @@ -71,7 +71,7 @@ def test_image_to_image_hires_fix_small(
"prompt": "a dinosaur",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_db0.jpg"),
"source_processing": "img2img",
}
Expand All @@ -89,7 +89,7 @@ def test_image_to_image_hires_fix_small(

def test_image_to_image_hires_fix_large(
self,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
hordelib_instance: HordeLib,
):
data = {
Expand All @@ -109,7 +109,7 @@ def test_image_to_image_hires_fix_large(
"prompt": "a dinosaur",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_db0.jpg"),
"source_processing": "img2img",
}
Expand All @@ -126,7 +126,7 @@ def test_image_to_image_hires_fix_large(

def test_img2img_masked_denoise_1(
self,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
hordelib_instance: HordeLib,
):
data = {
Expand All @@ -141,7 +141,7 @@ def test_img2img_masked_denoise_1(
"prompt": "a mecha robot sitting on a bench",
"ddim_steps": 20,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_img2img_alpha.png"),
"source_processing": "img2img",
}
Expand All @@ -159,7 +159,7 @@ def test_img2img_masked_denoise_1(

def test_img2img_masked_denoise_high(
self,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
hordelib_instance: HordeLib,
):
data = {
Expand All @@ -174,7 +174,7 @@ def test_img2img_masked_denoise_high(
"prompt": "a mecha robot sitting on a bench",
"ddim_steps": 20,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_img2img_alpha.png"),
"source_processing": "img2img",
}
Expand All @@ -192,7 +192,7 @@ def test_img2img_masked_denoise_high(

def test_img2img_masked_denoise_mid(
self,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
hordelib_instance: HordeLib,
):
data = {
Expand All @@ -207,7 +207,7 @@ def test_img2img_masked_denoise_mid(
"prompt": "a mecha robot sitting on a bench",
"ddim_steps": 20,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_img2img_alpha.png"),
"source_processing": "img2img",
}
Expand All @@ -225,7 +225,7 @@ def test_img2img_masked_denoise_mid(

def test_img2img_masked_denoise_low(
self,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
hordelib_instance: HordeLib,
):
data = {
Expand All @@ -240,7 +240,7 @@ def test_img2img_masked_denoise_low(
"prompt": "a mecha robot sitting on a bench",
"ddim_steps": 20,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": Image.open("images/test_img2img_alpha.png"),
"source_processing": "img2img",
}
Expand All @@ -258,7 +258,7 @@ def test_img2img_masked_denoise_low(

def test_image_to_faulty_source_image(
self,
stable_diffusion_modelname_for_testing: str,
stable_diffusion_model_name_for_testing: str,
hordelib_instance: HordeLib,
):
data = {
Expand All @@ -278,7 +278,7 @@ def test_image_to_faulty_source_image(
"prompt": "an ancient llamia monster",
"ddim_steps": 25,
"n_iter": 1,
"model": stable_diffusion_modelname_for_testing,
"model": stable_diffusion_model_name_for_testing,
"source_image": "THIS SHOULD FAILOVER TO TEXT2IMG",
"source_processing": "img2img",
}
Expand Down
Loading

0 comments on commit 749e577

Please sign in to comment.