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

tests: add test for probing diffusers model variant type

Browse Source
This commit is contained in:
psychedelicious
2024-04-23 19:47:58 +10:00
committed by Kent Keirsey
parent bbf919a933
commit 2537d260e3
6 changed files with 133 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/test_model_probe.py
View File

@ -4,7 +4,7 @@ import pytest
from torch import tensor
from invokeai.backend.model_manager import BaseModelType, ModelRepoVariant
from invokeai.backend.model_manager.config import InvalidModelConfigException
from invokeai.backend.model_manager.config import InvalidModelConfigException, MainDiffusersConfig, ModelVariantType
from invokeai.backend.model_manager.probe import (
CkptType,
ModelProbe,
@ -78,3 +78,11 @@ def test_probe_handles_state_dict_with_integer_keys():
}
with pytest.raises(InvalidModelConfigException):
ModelProbe.get_model_type_from_checkpoint(Path("embedding.pt"), state_dict_with_integer_keys)
def test_probe_sd1_diffusers_inpainting(datadir: Path):
config = ModelProbe.probe(datadir / "sd-1/main/dreamshaper-8-inpainting")
assert isinstance(config, MainDiffusersConfig)
assert config.base is BaseModelType.StableDiffusion1
assert config.variant is ModelVariantType.Inpaint
assert config.repo_variant is ModelRepoVariant.FP16