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

lint

Browse Source
This commit is contained in:
Kevin Turner 2023-09-20 12:53:25 -07:00
parent 2c39aec22d
commit 6392098961
1 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
tests/test_model_probe.py
View File

@ -6,12 +6,15 @@ from invokeai.backend import BaseModelType
from invokeai.backend.model_management.model_probe import VaeFolderProbe
@pytest.mark.parametrize("vae_path,expected_type", [
("sd-vae-ft-mse", BaseModelType.StableDiffusion1),
("sdxl-vae", BaseModelType.StableDiffusionXL),
("taesd", BaseModelType.StableDiffusion1),
("taesdxl", BaseModelType.StableDiffusionXL)
])
@pytest.mark.parametrize(
"vae_path,expected_type",
[
("sd-vae-ft-mse", BaseModelType.StableDiffusion1),
("sdxl-vae", BaseModelType.StableDiffusionXL),
("taesd", BaseModelType.StableDiffusion1),
("taesdxl", BaseModelType.StableDiffusionXL),
],
)
def test_get_base_type(vae_path: str, expected_type: BaseModelType, datadir: Path):
sd1_vae_path = datadir / "vae" / vae_path
probe = VaeFolderProbe(sd1_vae_path)