Skip to content

Commit

Permalink
tests/fix: include all models required for tests as fixtures
Browse files Browse the repository at this point in the history
  • Loading branch information
tazlin committed Oct 3, 2024
1 parent 1404626 commit cd9497b
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 3 deletions.
16 changes: 16 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,13 +122,17 @@ def isolated_comfy_horde_instance(init_horde) -> Comfy_Horde:
_sdxl_refined_model_name = "AlbedoBase XL (SDXL)"
_stable_cascade_base_model_name = "Stable Cascade 1.0"
_flux1_schnell_fp8_base_model_name = "Flux.1-Schnell fp8 (Compact)"
_am_pony_xl_model_name = "AMPonyXL"
_rev_animated_model_name = "Rev Animated"

_all_model_names = [
_testing_model_name,
_sdxl_1_0_model_name,
_sdxl_refined_model_name,
_stable_cascade_base_model_name,
_flux1_schnell_fp8_base_model_name,
_am_pony_xl_model_name,
_rev_animated_model_name,
]

# !!!!
Expand Down Expand Up @@ -166,6 +170,18 @@ def flux1_schnell_fp8_base_model_name(shared_model_manager: type[SharedModelMana
return _flux1_schnell_fp8_base_model_name


@pytest.fixture(scope="session")
def am_pony_xl_model_name(shared_model_manager: type[SharedModelManager]) -> str:
"""The default AMPonyXL model name used for testing."""
return _am_pony_xl_model_name


@pytest.fixture(scope="session")
def rev_animated_model_name(shared_model_manager: type[SharedModelManager]) -> str:
"""The default Rev Animated model name used for testing."""
return _rev_animated_model_name


# !!!!
# If you're adding a model name, follow the pattern and **add it to `_all_model_names`**
# !!!!
Expand Down
9 changes: 6 additions & 3 deletions tests/test_horde_lora.py
Original file line number Diff line number Diff line change
Expand Up @@ -390,6 +390,7 @@ def test_text_to_image_lora_character_hires_fix(
self,
shared_model_manager: type[SharedModelManager],
hordelib_instance: HordeLib,
rev_animated_model_name: str,
):
assert shared_model_manager.manager.lora

Expand Down Expand Up @@ -435,7 +436,7 @@ def test_text_to_image_lora_character_hires_fix(
],
"ddim_steps": 30,
"n_iter": 1,
"model": "Rev Animated",
"model": rev_animated_model_name,
}

pil_image = hordelib_instance.basic_inference_single_image(data).image
Expand All @@ -454,6 +455,7 @@ def test_text_to_image_lora_character_hires_fix_sdxl(
self,
shared_model_manager: type[SharedModelManager],
hordelib_instance: HordeLib,
am_pony_xl_model_name: str,
):
assert shared_model_manager.manager.lora

Expand Down Expand Up @@ -488,7 +490,7 @@ def test_text_to_image_lora_character_hires_fix_sdxl(
],
"ddim_steps": 12,
"n_iter": 1,
"model": "AMPonyXL",
"model": am_pony_xl_model_name,
}

pil_image = hordelib_instance.basic_inference_single_image(data).image
Expand Down Expand Up @@ -1010,6 +1012,7 @@ def test_login_gated_lora(
self,
shared_model_manager: type[SharedModelManager],
hordelib_instance: HordeLib,
sdxl_1_0_base_model_name: str,
):
assert shared_model_manager.manager.lora

Expand Down Expand Up @@ -1044,7 +1047,7 @@ def test_login_gated_lora(
"loras": [{"name": download_gated_lora_version_id, "model": 1.0, "clip": 1.0, "is_version": True}],
"ddim_steps": 25,
"n_iter": 1,
"model": "SDXL 1.0",
"model": sdxl_1_0_base_model_name,
}
ret = hordelib_instance.basic_inference_single_image(data)
assert isinstance(ret, ResultingImageReturn)
Expand Down

0 comments on commit cd9497b

Please sign in to comment.