81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
import unittest
|
|
|
|
try:
|
|
from .test_pro_edit_helpers import FakeModel, import_plugin_submodule, load_plugin_package
|
|
except ImportError:
|
|
from test_pro_edit_helpers import FakeModel, import_plugin_submodule, load_plugin_package
|
|
|
|
|
|
class TestModelProfiles(unittest.TestCase):
|
|
def setUp(self):
|
|
self.plugin = load_plugin_package()
|
|
self.registry = import_plugin_submodule(self.plugin, "model_profiles.registry")
|
|
|
|
def test_resolves_sdxl_base_profile_from_family(self):
|
|
profile = self.registry.resolve_model_profile(
|
|
model=None,
|
|
model_info={"checkpoint_name": "plain-sdxl-base.safetensors", "family": "SDXL"},
|
|
)
|
|
|
|
self.assertEqual(profile["profile_id"], "sdxl_base")
|
|
self.assertEqual(profile["backend_type"], "none")
|
|
self.assertEqual(profile["sampler_strategy"], "standard_k")
|
|
self.assertEqual(profile["loader_strategy"], "sdxl_checkpoint")
|
|
|
|
def test_resolves_flux_kontext_profile_from_name(self):
|
|
profile = self.registry.resolve_model_profile(
|
|
model=None,
|
|
model_info={"checkpoint_name": "demo-flux-kontext-dev.safetensors", "family": "FLUX_DEV"},
|
|
)
|
|
|
|
self.assertEqual(profile["profile_id"], "flux_edit")
|
|
self.assertEqual(profile["role"], "edit")
|
|
self.assertEqual(profile["backend_type"], "flux_edit")
|
|
self.assertEqual(profile["sampler_strategy"], "flux_guided")
|
|
self.assertEqual(profile["loader_strategy"], "flux_split_or_bundle")
|
|
|
|
def test_model_info_override_has_highest_priority(self):
|
|
profile = self.registry.resolve_model_profile(
|
|
model=None,
|
|
model_info={
|
|
"checkpoint_name": "plain-sdxl-base.safetensors",
|
|
"family": "SDXL",
|
|
"profile_id": "sdxl_edit",
|
|
"role": "edit",
|
|
"backend_type": "sdxl_native",
|
|
"sampler_strategy": "standard_k",
|
|
"supports_inpaint_native": True,
|
|
"supports_image_edit_native": True,
|
|
"preferred_edit_backend": "sdxl",
|
|
},
|
|
)
|
|
|
|
self.assertEqual(profile["profile_id"], "sdxl_edit")
|
|
self.assertEqual(profile["role"], "edit")
|
|
self.assertEqual(profile["backend_type"], "sdxl_native")
|
|
self.assertTrue(profile["supports_image_edit_native"])
|
|
|
|
def test_legacy_tags_upgrade_to_matching_profile(self):
|
|
model = FakeModel(
|
|
family="FLUX_DEV",
|
|
model_role="edit",
|
|
supports_inpaint_native=False,
|
|
supports_image_edit_native=True,
|
|
preferred_edit_backend="flux",
|
|
profile_id="",
|
|
backend_type="",
|
|
sampler_strategy="",
|
|
loader_strategy="",
|
|
model_name="legacy-flux-edit.safetensors",
|
|
)
|
|
|
|
profile = self.registry.resolve_model_profile(model=model, model_info=None)
|
|
|
|
self.assertEqual(profile["profile_id"], "flux_edit")
|
|
self.assertEqual(profile["backend_type"], "flux_edit")
|
|
self.assertEqual(profile["sampler_strategy"], "flux_guided")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|