diff --git a/invokeai/backend/model_management/model_probe.py b/invokeai/backend/model_management/model_probe.py index 145c56c273..3045849065 100644 --- a/invokeai/backend/model_management/model_probe.py +++ b/invokeai/backend/model_management/model_probe.py @@ -481,9 +481,19 @@ class ControlNetFolderProbe(FolderProbeBase): with open(config_file, "r") as file: config = json.load(file) # no obvious way to distinguish between sd2-base and sd2-768 - return ( - BaseModelType.StableDiffusion1 if config["cross_attention_dim"] == 768 else BaseModelType.StableDiffusion2 + dimension = config["cross_attention_dim"] + base_model = ( + BaseModelType.StableDiffusion1 + if dimension == 768 + else BaseModelType.StableDiffusion2 + if dimension == 1024 + else BaseModelType.StableDiffusionXL + if dimension == 2048 + else None ) + if not base_model: + raise InvalidModelException(f"Unable to determine model base for {self.folder_path}") + return base_model class LoRAFolderProbe(FolderProbeBase):