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): test VaeFolderProbe

Browse Source
This commit is contained in:
Kevin Turner
2023-09-20 12:48:59 -07:00
parent e487bcd0f7
commit 2c39aec22d
5 changed files with 153 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
tests/test_model_probe.py Normal file
View File

@ -0,0 +1,19 @@
from pathlib import Path
import pytest
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)
])
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)
base_type = probe.get_base_type()
assert base_type == expected_type