From 756008dc5899081c5aa51e5bd8f24c1b3975a59e Mon Sep 17 00:00:00 2001 From: Brandon Rising Date: Mon, 25 Nov 2024 14:03:36 -0500 Subject: [PATCH] fix: Fail scan on InvalidMagicError in picklescan, update default for read_checkpoint_meta to scan unless explicitly told not to --- invokeai/app/services/model_load/model_load_default.py | 2 +- invokeai/backend/model_manager/probe.py | 2 +- invokeai/backend/model_manager/util/model_util.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/invokeai/app/services/model_load/model_load_default.py b/invokeai/app/services/model_load/model_load_default.py index be2cc2478af..21b25aaab78 100644 --- a/invokeai/app/services/model_load/model_load_default.py +++ b/invokeai/app/services/model_load/model_load_default.py @@ -86,7 +86,7 @@ def load_model_from_path( def torch_load_file(checkpoint: Path) -> AnyModel: scan_result = scan_file_path(checkpoint) - if scan_result.infected_files != 0: + if scan_result.infected_files != 0 or scan_result.scan_err: raise Exception("The model at {checkpoint} is potentially infected by malware. Aborting load.") result = torch_load(checkpoint, map_location="cpu") return result diff --git a/invokeai/backend/model_manager/probe.py b/invokeai/backend/model_manager/probe.py index b1ff2fd080f..333d81b7eff 100644 --- a/invokeai/backend/model_manager/probe.py +++ b/invokeai/backend/model_manager/probe.py @@ -469,7 +469,7 @@ def _scan_model(cls, model_name: str, checkpoint: Path) -> None: """ # scan model scan_result = scan_file_path(checkpoint) - if scan_result.infected_files != 0: + if scan_result.infected_files != 0 or scan_result.scan_err: raise Exception("The model {model_name} is potentially infected by malware. Aborting import.") diff --git a/invokeai/backend/model_manager/util/model_util.py b/invokeai/backend/model_manager/util/model_util.py index e218124fb8a..dd9dea67ef2 100644 --- a/invokeai/backend/model_manager/util/model_util.py +++ b/invokeai/backend/model_manager/util/model_util.py @@ -44,7 +44,7 @@ def _fast_safetensors_reader(path: str) -> Dict[str, torch.Tensor]: return checkpoint -def read_checkpoint_meta(path: Union[str, Path], scan: bool = False) -> Dict[str, torch.Tensor]: +def read_checkpoint_meta(path: Union[str, Path], scan: bool = True) -> Dict[str, torch.Tensor]: if str(path).endswith(".safetensors"): try: path_str = path.as_posix() if isinstance(path, Path) else path @@ -55,7 +55,7 @@ def read_checkpoint_meta(path: Union[str, Path], scan: bool = False) -> Dict[str else: if scan: scan_result = scan_file_path(path) - if scan_result.infected_files != 0: + if scan_result.infected_files != 0 or scan_result.scan_err: raise Exception(f'The model file "{path}" is potentially infected by malware. Aborting import.') if str(path).endswith(".gguf"): # The GGUF reader used here uses numpy memmap, so these tensors are not loaded into memory during this function