From 67c8cf4bc29342c57ae979ebd69ccc9a6fcb262b Mon Sep 17 00:00:00 2001 From: Sergey Borisov Date: Sat, 8 Jul 2023 14:26:25 +0300 Subject: [PATCH] Controlnet model detection --- .../backend/model_management/models/controlnet.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/invokeai/backend/model_management/models/controlnet.py b/invokeai/backend/model_management/models/controlnet.py index 9563f87afd..3b73097799 100644 --- a/invokeai/backend/model_management/models/controlnet.py +++ b/invokeai/backend/model_management/models/controlnet.py @@ -13,6 +13,7 @@ from .base import ( calc_model_size_by_fs, calc_model_size_by_data, classproperty, + InvalidModelException, ) class ControlNetModelFormat(str, Enum): @@ -73,10 +74,18 @@ class ControlNetModel(ModelBase): @classmethod def detect_format(cls, path: str): + if not os.path.exists(path): + raise ModelNotFoundException() + if os.path.isdir(path): - return ControlNetModelFormat.Diffusers - else: - return ControlNetModelFormat.Checkpoint + if os.path.exists(os.path.join(path, "config.json")): + return ControlNetModelFormat.Diffusers + + if os.path.isfile(path): + if any([path.endswith(f".{ext}") for ext in ["safetensors", "ckpt", "pt", "pth"]]): + return ControlNetModelFormat.Checkpoint + + raise InvalidModelException(f"Not a valid model: {path}") @classmethod def convert_if_required(