Unpin picklescan req and cleanup

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

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

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: