diff --git a/tests/test_model_probe.py b/tests/test_model_probe.py index e0f2cff4e5..248b7d602f 100644 --- a/tests/test_model_probe.py +++ b/tests/test_model_probe.py @@ -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)