234 lines
8.3 KiB
Python
234 lines
8.3 KiB
Python
import unittest
|
|
|
|
from core.contracts import ResolvedModel
|
|
from core.path_templates import infer_download_path_templates
|
|
from core.settings import calculate_template_subfolder
|
|
|
|
BASE_MODELS = {
|
|
"base_models": [
|
|
{
|
|
"name": "Pony",
|
|
"aliases": ["pony", "ponyxl", "pony diffusion"],
|
|
},
|
|
{
|
|
"name": "SDXL 1.0",
|
|
"aliases": ["sdxl", "sdxl10"],
|
|
},
|
|
{
|
|
"name": "Illustrious",
|
|
"aliases": ["illustrious", "illustriousxl"],
|
|
},
|
|
{
|
|
"name": "Flux.1 Krea",
|
|
"aliases": ["Flux.1 Krea", "flux 1 krea"],
|
|
},
|
|
{
|
|
"name": "Krea 2",
|
|
"aliases": ["krea", "krea 2"],
|
|
},
|
|
{
|
|
"name": "Anima",
|
|
"aliases": ["anima"],
|
|
},
|
|
{
|
|
"name": "Ideogram 4.0",
|
|
"aliases": ["Ideogram 4.0", "ideogram", "ideogram 4 0"],
|
|
},
|
|
{
|
|
"name": "Upscaler",
|
|
"aliases": ["upscaler"],
|
|
},
|
|
]
|
|
}
|
|
|
|
|
|
def _typed_models(items):
|
|
return [ResolvedModel.from_mapping(item) for item in items]
|
|
|
|
|
|
class PathTemplateInferenceTests(unittest.TestCase):
|
|
def test_rejects_non_model_collections(self):
|
|
with self.assertRaises(TypeError):
|
|
infer_download_path_templates(0, BASE_MODELS)
|
|
|
|
def test_detects_base_model_and_first_tag_layout(self):
|
|
result = infer_download_path_templates(
|
|
_typed_models(
|
|
[
|
|
{
|
|
"category": "loras",
|
|
"relative_path": r"ponyxl\style\model-a.safetensors",
|
|
"path": r"C:\models\loras\ponyxl\style\model-a.safetensors",
|
|
},
|
|
{
|
|
"category": "lora",
|
|
"relative_path": r"ponyxl\character\model-b.safetensors",
|
|
"path": r"C:\models\loras\ponyxl\character\model-b.safetensors",
|
|
},
|
|
{
|
|
"category": "loras",
|
|
"relative_path": r"sdxl\style\model-c.safetensors",
|
|
"path": r"C:\models\loras\sdxl\style\model-c.safetensors",
|
|
},
|
|
]
|
|
),
|
|
BASE_MODELS,
|
|
)
|
|
|
|
self.assertEqual("{base_model}/{first_tag}", result["templates"]["loras"])
|
|
self.assertEqual("ponyxl", result["base_model_path_mappings"]["Pony"])
|
|
self.assertEqual("sdxl", result["base_model_path_mappings"]["SDXL 1.0"])
|
|
self.assertTrue(result["categories"]["loras"]["apply"])
|
|
|
|
def test_detects_flat_layout(self):
|
|
result = infer_download_path_templates(
|
|
_typed_models(
|
|
[
|
|
{
|
|
"category": "vae",
|
|
"relative_path": "vae-a.safetensors",
|
|
"path": r"C:\models\vae\vae-a.safetensors",
|
|
},
|
|
{
|
|
"category": "vae",
|
|
"relative_path": "vae-b.safetensors",
|
|
"path": r"C:\models\vae\vae-b.safetensors",
|
|
},
|
|
]
|
|
),
|
|
BASE_MODELS,
|
|
)
|
|
|
|
self.assertEqual("", result["templates"]["vae"])
|
|
self.assertTrue(result["categories"]["vae"]["apply"])
|
|
|
|
def test_detects_nested_checkpoint_base_model_mappings(self):
|
|
result = infer_download_path_templates(
|
|
_typed_models(
|
|
[
|
|
{
|
|
"category": "checkpoints",
|
|
"relative_path": r"SDXL\Pony\pony-realism.safetensors",
|
|
"path": r"C:\models\checkpoints\SDXL\Pony\pony-realism.safetensors",
|
|
},
|
|
{
|
|
"category": "checkpoints",
|
|
"relative_path": r"SDXL\Illustrious\wai.safetensors",
|
|
"path": r"C:\models\checkpoints\SDXL\Illustrious\wai.safetensors",
|
|
},
|
|
{
|
|
"category": "checkpoints",
|
|
"relative_path": r"SDXL\Illustrious\realistic\real.safetensors",
|
|
"path": r"C:\models\checkpoints\SDXL\Illustrious\realistic\real.safetensors",
|
|
},
|
|
{
|
|
"category": "checkpoints",
|
|
"relative_path": r"SDXL\SDXL\juggernaut.safetensors",
|
|
"path": r"C:\models\checkpoints\SDXL\SDXL\juggernaut.safetensors",
|
|
},
|
|
]
|
|
),
|
|
BASE_MODELS,
|
|
)
|
|
|
|
self.assertEqual("{base_model}", result["templates"]["checkpoints"])
|
|
self.assertEqual("SDXL/Pony", result["base_model_path_mappings"]["Pony"])
|
|
self.assertEqual(
|
|
"SDXL/Illustrious",
|
|
result["base_model_path_mappings"]["Illustrious"],
|
|
)
|
|
self.assertEqual(
|
|
"SDXL/SDXL",
|
|
result["base_model_path_mappings"]["SDXL 1.0"],
|
|
)
|
|
|
|
def test_base_model_mapping_matches_longer_metadata_name(self):
|
|
subfolder = calculate_template_subfolder(
|
|
"checkpoints",
|
|
{"base_model": "Pony Diffusion V6 XL"},
|
|
{
|
|
"download_path_templates": {"checkpoints": "{base_model}"},
|
|
"base_model_path_mappings": {"Pony": "SDXL/Pony"},
|
|
},
|
|
)
|
|
|
|
self.assertEqual("SDXL/Pony", subfolder)
|
|
|
|
def test_detect_does_not_map_krea2_to_flux_krea_folder(self):
|
|
result = infer_download_path_templates(
|
|
_typed_models(
|
|
[
|
|
{
|
|
"category": "loras",
|
|
"relative_path": r"FLUX\KREA\concept\snofs_krea_v1.safetensors",
|
|
"path": r"C:\models\loras\FLUX\KREA\concept\snofs_krea_v1.safetensors",
|
|
},
|
|
{
|
|
"category": "loras",
|
|
"relative_path": r"FLUX\KREA\style\another_flux_krea.safetensors",
|
|
"path": r"C:\models\loras\FLUX\KREA\style\another_flux_krea.safetensors",
|
|
},
|
|
]
|
|
),
|
|
BASE_MODELS,
|
|
)
|
|
|
|
self.assertNotIn("Krea 2", result["base_model_path_mappings"])
|
|
self.assertEqual(
|
|
"FLUX/KREA",
|
|
result["base_model_path_mappings"]["Flux.1 Krea"],
|
|
)
|
|
|
|
def test_detect_does_not_map_anima_to_wan_animate_folder(self):
|
|
result = infer_download_path_templates(
|
|
_typed_models(
|
|
[
|
|
{
|
|
"category": "checkpoints",
|
|
"relative_path": r"WAN\WAN2.2ANIMATE\model-a.safetensors",
|
|
"path": r"C:\models\checkpoints\WAN\WAN2.2ANIMATE\model-a.safetensors",
|
|
},
|
|
{
|
|
"category": "checkpoints",
|
|
"relative_path": r"WAN\WAN2.2ANIMATE\model-b.safetensors",
|
|
"path": r"C:\models\checkpoints\WAN\WAN2.2ANIMATE\model-b.safetensors",
|
|
},
|
|
]
|
|
),
|
|
BASE_MODELS,
|
|
)
|
|
|
|
self.assertNotIn("Anima", result["base_model_path_mappings"])
|
|
|
|
def test_detect_accepts_version_and_plural_alias_suffixes(self):
|
|
result = infer_download_path_templates(
|
|
_typed_models(
|
|
[
|
|
{
|
|
"category": "diffusion_models",
|
|
"relative_path": r"IDEOGRAM4\model-a.safetensors",
|
|
"path": r"C:\models\diffusion_models\IDEOGRAM4\model-a.safetensors",
|
|
},
|
|
{
|
|
"category": "upscale_models",
|
|
"relative_path": r"ESRGAN\Upscalers\photo\model-a.pth",
|
|
"path": r"C:\models\upscale_models\ESRGAN\Upscalers\photo\model-a.pth",
|
|
},
|
|
]
|
|
),
|
|
BASE_MODELS,
|
|
)
|
|
|
|
self.assertEqual(
|
|
"IDEOGRAM4",
|
|
result["base_model_path_mappings"]["Ideogram 4.0"],
|
|
)
|
|
self.assertEqual(
|
|
"ESRGAN/Upscalers",
|
|
result["base_model_path_mappings"]["Upscaler"],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|