Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(model_management): factor out name strings

Browse Source
This commit is contained in:
Kevin Turner 2023-08-05 15:46:46 -07:00
parent 80876bbbd1
commit 7f4c387080
2 changed files with 9 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tests/test_model_manager.py
View File

@ -5,6 +5,9 @@ import pytest
from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.config import InvokeAIAppConfig
from invokeai.backend import ModelManager, BaseModelType, ModelType, SubModelType from invokeai.backend import ModelManager, BaseModelType, ModelType, SubModelType
BASIC_MODEL_NAME = ("SDXL base", BaseModelType.StableDiffusionXL, ModelType.Main)
VAE_OVERRIDE_MODEL_NAME = ("SDXL with VAE", BaseModelType.StableDiffusionXL, ModelType.Main)
@pytest.fixture @pytest.fixture
def model_manager(datadir) -> ModelManager: def model_manager(datadir) -> ModelManager:
@ -14,14 +17,11 @@ def model_manager(datadir) -> ModelManager:
def test_get_model_names(model_manager: ModelManager): def test_get_model_names(model_manager: ModelManager):
names = model_manager.model_names() names = model_manager.model_names()
assert names[:2] == [ assert names[:2] == [BASIC_MODEL_NAME, VAE_OVERRIDE_MODEL_NAME]
("SDXL base", BaseModelType.StableDiffusionXL, ModelType.Main),
("SDXL with VAE", BaseModelType.StableDiffusionXL, ModelType.Main),
]
def test_get_model_path_for_diffusers(model_manager: ModelManager, datadir: Path): def test_get_model_path_for_diffusers(model_manager: ModelManager, datadir: Path):
model_config = model_manager._get_model_config(BaseModelType.StableDiffusionXL, "SDXL base", ModelType.Main) model_config = model_manager._get_model_config(BASIC_MODEL_NAME[1], BASIC_MODEL_NAME[0], BASIC_MODEL_NAME[2])
top_model_path, is_override = model_manager._get_model_path(model_config) top_model_path, is_override = model_manager._get_model_path(model_config)
expected_model_path = datadir / "models" / "sdxl" / "main" / "SDXL base 1_0" expected_model_path = datadir / "models" / "sdxl" / "main" / "SDXL base 1_0"
assert top_model_path == expected_model_path assert top_model_path == expected_model_path
@ -29,7 +29,9 @@ def test_get_model_path_for_diffusers(model_manager: ModelManager, datadir: Path
def test_get_model_path_for_overridden_vae(model_manager: ModelManager, datadir: Path): def test_get_model_path_for_overridden_vae(model_manager: ModelManager, datadir: Path):
model_config = model_manager._get_model_config(BaseModelType.StableDiffusionXL, "SDXL with VAE", ModelType.Main) model_config = model_manager._get_model_config(
VAE_OVERRIDE_MODEL_NAME[1], VAE_OVERRIDE_MODEL_NAME[0], VAE_OVERRIDE_MODEL_NAME[2]
)
vae_model_path, is_override = model_manager._get_model_path(model_config, SubModelType.Vae) vae_model_path, is_override = model_manager._get_model_path(model_config, SubModelType.Vae)
expected_vae_path = datadir / "models" / "sdxl" / "vae" / "sdxl-vae-fp16-fix" expected_vae_path = datadir / "models" / "sdxl" / "vae" / "sdxl-vae-fp16-fix"
assert vae_model_path == expected_vae_path assert vae_model_path == expected_vae_path

2
tests/test_model_manager/configs/relative_sub.models.yaml
View File

@ -9,7 +9,7 @@ sdxl/main/SDXL base:
sdxl/main/SDXL with VAE: sdxl/main/SDXL with VAE:
path: sdxl/main/SDXL base 1_0 path: sdxl/main/SDXL base 1_0
description: SDXL base v1.0 description: SDXL with customized VAE
vae: sdxl/vae/sdxl-vae-fp16-fix/ vae: sdxl/vae/sdxl-vae-fp16-fix/
variant: normal variant: normal
format: diffusers format: diffusers