From 5feb62d440101ef50983f796442e98e94cdd6a74 Mon Sep 17 00:00:00 2001 From: Brandon Rising Date: Fri, 15 Mar 2024 12:10:39 -0400 Subject: [PATCH] Allow users to specify model type and skip detection step of probe --- invokeai/backend/model_manager/probe.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/invokeai/backend/model_manager/probe.py b/invokeai/backend/model_manager/probe.py index 33e886b964..af26d0a985 100644 --- a/invokeai/backend/model_manager/probe.py +++ b/invokeai/backend/model_manager/probe.py @@ -132,12 +132,13 @@ class ModelProbe(object): format_type = ModelFormat.Diffusers if model_path.is_dir() else ModelFormat.Checkpoint model_info = None - model_type = None - if format_type is ModelFormat.Diffusers: - model_type = cls.get_model_type_from_folder(model_path) - else: - model_type = cls.get_model_type_from_checkpoint(model_path) - format_type = ModelFormat.ONNX if model_type == ModelType.ONNX else format_type + model_type = fields['type'] if 'type' in fields else None + if not model_type: + if format_type is ModelFormat.Diffusers: + model_type = cls.get_model_type_from_folder(model_path) + else: + model_type = cls.get_model_type_from_checkpoint(model_path) + format_type = ModelFormat.ONNX if model_type == ModelType.ONNX else format_type probe_class = cls.PROBES[format_type].get(model_type) if not probe_class: