stable-diffusion-webui icon indicating copy to clipboard operation
stable-diffusion-webui copied to clipboard

[Feature Request]: Safely loading models using safetensors 🥒

Open pattontim opened this issue 1 year ago • 0 comments

Is there an existing issue for this?

  • [X] I have searched the existing issues and checked the recent builds/commits

What would your feature do ?

November 14th is National Pickle Day so I thought I'd catalog the current state of pickle obsolescence technology. Naturally, the current way models are loaded via python pickle is potentially unsafe, relying on a user's trust of the model. Rudimentary pickle scanning is in the UI to protect from basic attacks, yet techniques exist which may be able to circumvent this sort of implementation. Recognizing the risk, Huggingface appears to have assigned a dev to work on a new format supporting zero copy loading of models without arbitrary code execution, the safetensors library. Currently, the 0.2.4 python bindings for safetensors supports basic creation of and loading of models via numpy or Torch. The benefits and implementation including attack surface and conversion are discussed in the repo:

https://github.com/huggingface/safetensors

An ideal implementation of safetensors would silently load files in any of the model subfolders with a .safetensors suffix, calling safetensors.torch.load_file and using the model as normal. Similarly, model merges would have an option to Save as .safetensors format. Depending on the implementation, further changes may be needed to support models embedding extra information in more advanced model files.

Proposed workflow

  1. Start the webui with models/StableDiffusion/modelA.safetensors present
  2. Weights are loaded from the safetensors model as if normal.
  3. User generates an image, switches to modelB (.ckpt)
  4. Switches to checkpoint merger
  5. Selects modelA and modelB to merge
  6. Hits Save and modelC.ckpt is created
  7. Selects Use safetensors format checkbox
  8. Hits Save and modelC.safetensors is created
  9. The same process can be done for training.

Additional information

There is a Space for converting models, including code for converting models to the new safetensors format. There are a few already converted models, which can be found and downloaded in the PRs attached to each discussion here: https://huggingface.co/datasets/safetensors/conversions. There was a file size limit but it appears to have been sorted out recently?

Org page: https://huggingface.co/safetensors

🥒🥒🥒🥒🥒🥒🥒🥒🥒🥒🥒

pattontim avatar Nov 14 '22 17:11 pattontim