mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
test(model management): test VaeFolderProbe
This commit is contained in:
parent
e487bcd0f7
commit
2c39aec22d
19
tests/test_model_probe.py
Normal file
19
tests/test_model_probe.py
Normal 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
|
29
tests/test_model_probe/vae/sd-vae-ft-mse/config.json
Normal file
29
tests/test_model_probe/vae/sd-vae-ft-mse/config.json
Normal file
@ -0,0 +1,29 @@
|
|||||||
|
{
|
||||||
|
"_class_name": "AutoencoderKL",
|
||||||
|
"_diffusers_version": "0.4.2",
|
||||||
|
"act_fn": "silu",
|
||||||
|
"block_out_channels": [
|
||||||
|
128,
|
||||||
|
256,
|
||||||
|
512,
|
||||||
|
512
|
||||||
|
],
|
||||||
|
"down_block_types": [
|
||||||
|
"DownEncoderBlock2D",
|
||||||
|
"DownEncoderBlock2D",
|
||||||
|
"DownEncoderBlock2D",
|
||||||
|
"DownEncoderBlock2D"
|
||||||
|
],
|
||||||
|
"in_channels": 3,
|
||||||
|
"latent_channels": 4,
|
||||||
|
"layers_per_block": 2,
|
||||||
|
"norm_num_groups": 32,
|
||||||
|
"out_channels": 3,
|
||||||
|
"sample_size": 256,
|
||||||
|
"up_block_types": [
|
||||||
|
"UpDecoderBlock2D",
|
||||||
|
"UpDecoderBlock2D",
|
||||||
|
"UpDecoderBlock2D",
|
||||||
|
"UpDecoderBlock2D"
|
||||||
|
]
|
||||||
|
}
|
31
tests/test_model_probe/vae/sdxl-vae/config.json
Normal file
31
tests/test_model_probe/vae/sdxl-vae/config.json
Normal file
@ -0,0 +1,31 @@
|
|||||||
|
{
|
||||||
|
"_class_name": "AutoencoderKL",
|
||||||
|
"_diffusers_version": "0.18.0.dev0",
|
||||||
|
"_name_or_path": ".",
|
||||||
|
"act_fn": "silu",
|
||||||
|
"block_out_channels": [
|
||||||
|
128,
|
||||||
|
256,
|
||||||
|
512,
|
||||||
|
512
|
||||||
|
],
|
||||||
|
"down_block_types": [
|
||||||
|
"DownEncoderBlock2D",
|
||||||
|
"DownEncoderBlock2D",
|
||||||
|
"DownEncoderBlock2D",
|
||||||
|
"DownEncoderBlock2D"
|
||||||
|
],
|
||||||
|
"in_channels": 3,
|
||||||
|
"latent_channels": 4,
|
||||||
|
"layers_per_block": 2,
|
||||||
|
"norm_num_groups": 32,
|
||||||
|
"out_channels": 3,
|
||||||
|
"sample_size": 1024,
|
||||||
|
"scaling_factor": 0.13025,
|
||||||
|
"up_block_types": [
|
||||||
|
"UpDecoderBlock2D",
|
||||||
|
"UpDecoderBlock2D",
|
||||||
|
"UpDecoderBlock2D",
|
||||||
|
"UpDecoderBlock2D"
|
||||||
|
]
|
||||||
|
}
|
37
tests/test_model_probe/vae/taesd/config.json
Normal file
37
tests/test_model_probe/vae/taesd/config.json
Normal file
@ -0,0 +1,37 @@
|
|||||||
|
{
|
||||||
|
"_class_name": "AutoencoderTiny",
|
||||||
|
"_diffusers_version": "0.20.0.dev0",
|
||||||
|
"act_fn": "relu",
|
||||||
|
"decoder_block_out_channels": [
|
||||||
|
64,
|
||||||
|
64,
|
||||||
|
64,
|
||||||
|
64
|
||||||
|
],
|
||||||
|
"encoder_block_out_channels": [
|
||||||
|
64,
|
||||||
|
64,
|
||||||
|
64,
|
||||||
|
64
|
||||||
|
],
|
||||||
|
"force_upcast": false,
|
||||||
|
"in_channels": 3,
|
||||||
|
"latent_channels": 4,
|
||||||
|
"latent_magnitude": 3,
|
||||||
|
"latent_shift": 0.5,
|
||||||
|
"num_decoder_blocks": [
|
||||||
|
3,
|
||||||
|
3,
|
||||||
|
3,
|
||||||
|
1
|
||||||
|
],
|
||||||
|
"num_encoder_blocks": [
|
||||||
|
1,
|
||||||
|
3,
|
||||||
|
3,
|
||||||
|
3
|
||||||
|
],
|
||||||
|
"out_channels": 3,
|
||||||
|
"scaling_factor": 1.0,
|
||||||
|
"upsampling_scaling_factor": 2
|
||||||
|
}
|
37
tests/test_model_probe/vae/taesdxl/config.json
Normal file
37
tests/test_model_probe/vae/taesdxl/config.json
Normal file
@ -0,0 +1,37 @@
|
|||||||
|
{
|
||||||
|
"_class_name": "AutoencoderTiny",
|
||||||
|
"_diffusers_version": "0.20.0.dev0",
|
||||||
|
"act_fn": "relu",
|
||||||
|
"decoder_block_out_channels": [
|
||||||
|
64,
|
||||||
|
64,
|
||||||
|
64,
|
||||||
|
64
|
||||||
|
],
|
||||||
|
"encoder_block_out_channels": [
|
||||||
|
64,
|
||||||
|
64,
|
||||||
|
64,
|
||||||
|
64
|
||||||
|
],
|
||||||
|
"force_upcast": false,
|
||||||
|
"in_channels": 3,
|
||||||
|
"latent_channels": 4,
|
||||||
|
"latent_magnitude": 3,
|
||||||
|
"latent_shift": 0.5,
|
||||||
|
"num_decoder_blocks": [
|
||||||
|
3,
|
||||||
|
3,
|
||||||
|
3,
|
||||||
|
1
|
||||||
|
],
|
||||||
|
"num_encoder_blocks": [
|
||||||
|
1,
|
||||||
|
3,
|
||||||
|
3,
|
||||||
|
3
|
||||||
|
],
|
||||||
|
"out_channels": 3,
|
||||||
|
"scaling_factor": 1.0,
|
||||||
|
"upsampling_scaling_factor": 2
|
||||||
|
}
|
Loading…
Reference in New Issue
Block a user