diff --git a/environments-and-requirements/requirements-base.txt b/environments-and-requirements/requirements-base.txt index b8937605e3..d76531a133 100644 --- a/environments-and-requirements/requirements-base.txt +++ b/environments-and-requirements/requirements-base.txt @@ -30,7 +30,7 @@ test-tube>=0.7.5 torch-fidelity torchmetrics transformers==4.21.* -picklescan==0.0.5 +picklescan git+https://github.com/openai/CLIP.git@main#egg=clip git+https://github.com/Birch-san/k-diffusion.git@mps#egg=k-diffusion git+https://github.com/invoke-ai/clipseg.git@relaxed-python-requirement#egg=clipseg diff --git a/installer/requirements.in b/installer/requirements.in index cb1b7090ff..ff84af7a58 100644 --- a/installer/requirements.in +++ b/installer/requirements.in @@ -19,7 +19,7 @@ torch-fidelity torchvision==0.13.1 ; platform_system == 'Darwin' torchvision==0.13.1+cu116 ; platform_system == 'Linux' or platform_system == 'Windows' transformers -picklescan==0.0.5 +picklescan https://github.com/openai/CLIP/archive/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1.zip https://github.com/TencentARC/GFPGAN/archive/2eac2033893ca7f427f4035d80fe95b92649ac56.zip https://github.com/invoke-ai/k-diffusion/archive/7f16b2c33411f26b3eae78d10648d625cb0c1095.zip diff --git a/ldm/invoke/model_cache.py b/ldm/invoke/model_cache.py index a0130e740b..18837531cc 100644 --- a/ldm/invoke/model_cache.py +++ b/ldm/invoke/model_cache.py @@ -12,13 +12,15 @@ import time import gc import hashlib import psutil +import sys import transformers import traceback import os -from sys import getrefcount from omegaconf import OmegaConf from omegaconf.errors import ConfigAttributeError from ldm.util import instantiate_from_config +from picklescan.scanner import scan_file_path + DEFAULT_MAX_MODELS=2 @@ -201,7 +203,7 @@ class ModelCache(object): height = mconfig.height # scan model - self.scan_model(model_name, weights) + self._scan_model(model_name, weights) print(f'>> Loading {model_name} from {weights}') @@ -279,10 +281,8 @@ class ModelCache(object): if self._has_cuda(): torch.cuda.empty_cache() - def scan_model(self, model_name, checkpoint): + def _scan_model(self, model_name, checkpoint): # scan model - from picklescan.scanner import scan_file_path - import sys print(f'>> Scanning Model: {model_name}') scan_result = scan_file_path(checkpoint) if scan_result.infected_files != 0: