github/InvokeAI
1
1
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-04-23 03:00:31 -04:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into stripped-models

Browse Source
This commit is contained in:
Billy
2025-03-19 07:52:56 +11:00
parent b9972be7f1 7fe4d4c21a
commit 8e14f9d971
3 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
invokeai/app/services/model_load/model_load_default.py
View File

@@ -85,8 +85,11 @@ class ModelLoadService(ModelLoadServiceBase):
def torch_load_file(checkpoint: Path) -> AnyModel:
scan_result = scan_file_path(checkpoint)
if scan_result.infected_files != 0 or scan_result.scan_err:
raise Exception("The model at {checkpoint} is potentially infected by malware. Aborting load.")
if scan_result.infected_files != 0:
raise Exception(f"The model at {checkpoint} is potentially infected by malware. Aborting load.")
if scan_result.scan_err:
raise Exception(f"Error scanning model at {checkpoint} for malware. Aborting load.")
result = torch_load(checkpoint, map_location="cpu")
return result

6
invokeai/backend/model_manager/legacy_probe.py
View File

@@ -484,8 +484,10 @@ class ModelProbe(object):
"""
# scan model
scan_result = pscan.scan_file_path(checkpoint)
if scan_result.infected_files != 0 or scan_result.scan_err:
raise Exception("The model {model_name} is potentially infected by malware. Aborting import.")
if scan_result.infected_files != 0:
raise Exception(f"The model {model_name} is potentially infected by malware. Aborting import.")
if scan_result.scan_err:
raise Exception(f"Error scanning model {model_name} for malware. Aborting import.")
# Probing utilities

7
invokeai/backend/model_manager/util/model_util.py
View File

@@ -58,8 +58,11 @@ def read_checkpoint_meta(path: Union[str, Path], scan: bool = True) -> Dict[str,
else:
if scan:
scan_result = pscan.scan_file_path(path)
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 scan_result.infected_files != 0:
raise Exception(f"The model at {path} is potentially infected by malware. Aborting import.")
if scan_result.scan_err:
raise Exception(f"Error scanning model at {path} for malware. Aborting import.")
checkpoint = torch.load(path, map_location=torch.device("meta"))
return checkpoint