Skip to content

BUG: pickle_scan stops if it encounters a non-picke file #14

@phizch

Description

@phizch

When I tried scanning my downloads directory the scanning stopped before all files had been scanned.

The reason is that result = torch.load(path.as_posix(), pickle_module=pickle_inspector.pickle) threw an exception because one of the files in my download directory wasn't a valid checkpoint. (I think it was a misnamed safetensors file).

Python does produce an error message, but it's not saved in the scan_outputs.txt file, and it's not picked up on by the cmd script when it's reporting on potentially malicious files.

How to reproduce:

  1. Create an empty file named "0.bin" in one of the directories to be scanned.
  2. Scan

Expected behavior:

All files should be scanned.

Actual behavior:

The scan stops after trying to scan the empty file with an error like this:

...L:/StableDiffusion/Downloads/0.bin
Traceback (most recent call last):
  File "L:\StableDiffusion\stable-diffusion-pickle-scanner-main\pickle_scan.py", line 24, in <module>
    result = torch.load(path.as_posix(), pickle_module=pickle_inspector.pickle)
  File "L:\StableDiffusion\stable-diffusion-webui\venv\lib\site-packages\torch\serialization.py", line 713, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "L:\StableDiffusion\stable-diffusion-webui\venv\lib\site-packages\torch\serialization.py", line 920, in _legacy_load
    magic_number = pickle_module.load(f, **pickle_load_args)

Fix:

Do the scanning in try/except blocks.

for path in BASE_DIR.glob(r'**/*'):
  if path.suffix in EXTENSIONS:
    print("")
    print("..." + path.as_posix())
    try:
      result = torch.load(path.as_posix(), pickle_module=pickle_inspector.pickle)
      ..
    except:
      print("NOT A VAILD CHECKPOINT FILE")

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions