From ee042ab76d75e9292332dd0feea99056ee468767 Mon Sep 17 00:00:00 2001 From: Sergey Borisov Date: Wed, 5 Jul 2023 14:18:30 +0300 Subject: [PATCH] Fix ckpt scanning on conversion --- .../backend/model_management/convert_ckpt_to_diffusers.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py index 1eeee92fb7..e3e64940de 100644 --- a/invokeai/backend/model_management/convert_ckpt_to_diffusers.py +++ b/invokeai/backend/model_management/convert_ckpt_to_diffusers.py @@ -29,7 +29,7 @@ import invokeai.backend.util.logging as logger from invokeai.app.services.config import InvokeAIAppConfig from .model_manager import ModelManager -from .model_cache import ModelCache +from picklescan.scanner import scan_file_path from .models import BaseModelType, ModelVariantType try: @@ -1014,7 +1014,10 @@ def load_pipeline_from_original_stable_diffusion_ckpt( checkpoint = load_file(checkpoint_path) else: if scan_needed: - ModelCache.scan_model(checkpoint_path, checkpoint_path) + # scan model + scan_result = scan_file_path(checkpoint_path) + if scan_result.infected_files != 0: + raise "The model {checkpoint_path} is potentially infected by malware. Aborting import." checkpoint = torch.load(checkpoint_path) # sometimes there is a state_dict key and sometimes not