97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
import importlib.util
|
|
import pathlib
|
|
import unittest
|
|
|
|
|
|
ROOT = pathlib.Path(__file__).resolve().parents[1]
|
|
MODEL_INFO_PATH = ROOT / "utils" / "model_info.py"
|
|
|
|
|
|
def load_model_info_module():
|
|
spec = importlib.util.spec_from_file_location("lls_model_info_test", MODEL_INFO_PATH)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class FluxClipStub:
|
|
def tokenize(self, _text):
|
|
return {"text": "", "t5xxl": ["token"], "l": ["token"]}
|
|
|
|
|
|
class SdxlClipStub:
|
|
def tokenize(self, _text):
|
|
return {"text": "", "g": ["token"], "l": ["token"]}
|
|
|
|
|
|
class Sd15ClipStub:
|
|
def tokenize(self, _text):
|
|
return {"text": ""}
|
|
|
|
|
|
class BrokenClipStub:
|
|
def tokenize(self, _text):
|
|
raise RuntimeError("tokenizer unavailable")
|
|
|
|
|
|
class ModelWithFamily:
|
|
_lls_family = "FLUX"
|
|
|
|
|
|
class ModelWithType:
|
|
model_type = "flux-dev"
|
|
|
|
|
|
class LatentFormatStub:
|
|
def __init__(self, latent_channels):
|
|
self.latent_channels = latent_channels
|
|
|
|
|
|
class ModelWithLatentFormat:
|
|
def __init__(self, latent_channels):
|
|
self.latent_channels = latent_channels
|
|
|
|
def get_model_object(self, name):
|
|
if name != "latent_format":
|
|
raise AssertionError(f"unexpected model object: {name}")
|
|
return LatentFormatStub(self.latent_channels)
|
|
|
|
|
|
class PlainModel:
|
|
pass
|
|
|
|
|
|
class TestModelInfoInference(unittest.TestCase):
|
|
def test_infer_family_from_clip_uses_tokenizer_shape(self):
|
|
model_info = load_model_info_module()
|
|
|
|
self.assertEqual(model_info.infer_family_from_clip(FluxClipStub()), "FLUX_DEV")
|
|
self.assertEqual(model_info.infer_family_from_clip(SdxlClipStub()), "SDXL")
|
|
self.assertEqual(model_info.infer_family_from_clip(Sd15ClipStub()), "SD1.5")
|
|
self.assertEqual(model_info.infer_family_from_clip(BrokenClipStub()), "SD1.5")
|
|
|
|
def test_infer_family_from_model_prioritizes_lls_family_marker(self):
|
|
model_info = load_model_info_module()
|
|
|
|
self.assertEqual(model_info.infer_family_from_model(ModelWithFamily()), "FLUX_DEV")
|
|
|
|
def test_infer_family_from_model_falls_back_to_model_type_and_latent_format(self):
|
|
model_info = load_model_info_module()
|
|
|
|
self.assertEqual(model_info.infer_family_from_model(ModelWithType()), "FLUX_DEV")
|
|
self.assertEqual(model_info.infer_family_from_model(ModelWithLatentFormat(128)), "FLUX_DEV")
|
|
self.assertEqual(model_info.infer_family_from_model(ModelWithLatentFormat(4)), "SD1.5")
|
|
self.assertEqual(model_info.infer_family_from_model(PlainModel()), "SD1.5")
|
|
|
|
def test_infer_task_mode_from_latent_uses_source_field(self):
|
|
model_info = load_model_info_module()
|
|
|
|
self.assertEqual(model_info.infer_task_mode_from_latent({"source": "image_encode"}), "img2img")
|
|
self.assertEqual(model_info.infer_task_mode_from_latent({"source": "empty_latent"}), "txt2img")
|
|
self.assertEqual(model_info.infer_task_mode_from_latent({}), "txt2img")
|
|
self.assertEqual(model_info.infer_task_mode_from_latent(None), "txt2img")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|