using_safetensors.md 5.27 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
<!--Copyright 2023 The HuggingFace Team. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
specific language governing permissions and limitations under the License.
-->

Steven Liu's avatar
Steven Liu committed
13
# Load safetensors
14

15
16
[[open-in-colab]]

Steven Liu's avatar
Steven Liu committed
17
[safetensors](https://github.com/huggingface/safetensors) is a safe and fast file format for storing and loading tensors. Typically, PyTorch model weights are saved or *pickled* into a `.bin` file with Python's [`pickle`](https://docs.python.org/3/library/pickle.html) utility. However, `pickle` is not secure and pickled files may contain malicious code that can be executed. safetensors is a secure alternative to `pickle`, making it ideal for sharing model weights.
18

Steven Liu's avatar
Steven Liu committed
19
This guide will show you how you load `.safetensor` files, and how to convert Stable Diffusion model weights stored in other formats to `.safetensor`. Before you start, make sure you have safetensors installed:
20

21
22
23
```py
# uncomment to install the necessary libraries in Colab
#!pip install safetensors
24
25
```

Steven Liu's avatar
Steven Liu committed
26
If you look at the [`runwayml/stable-diffusion-v1-5`](https://huggingface.co/runwayml/stable-diffusion-v1-5/tree/main) repository, you'll see weights inside the `text_encoder`, `unet` and `vae` subfolders are stored in the `.safetensors` format. By default, 🤗 Diffusers automatically loads these `.safetensors` files from their subfolders if they're available in the model repository.
27

Steven Liu's avatar
Steven Liu committed
28
For more explicit control, you can optionally set `use_safetensors=True` (if `safetensors` is not installed, you'll get an error message asking you to install it):
29

Steven Liu's avatar
Steven Liu committed
30
31
```py
from diffusers import DiffusionPipeline
32

Steven Liu's avatar
Steven Liu committed
33
pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_safetensors=True)
34
35
```

Patrick von Platen's avatar
Patrick von Platen committed
36
However, model weights are not necessarily stored in separate subfolders like in the example above. Sometimes, all the weights are stored in a single `.safetensors` file. In this case, if the weights are Stable Diffusion weights, you can load the file directly with the [`~diffusers.loaders.FromSingleFileMixin.from_single_file`] method:
37

Steven Liu's avatar
Steven Liu committed
38
39
```py
from diffusers import StableDiffusionPipeline
40

Patrick von Platen's avatar
Patrick von Platen committed
41
pipeline = StableDiffusionPipeline.from_single_file(
Steven Liu's avatar
Steven Liu committed
42
43
44
    "https://huggingface.co/WarriorMama777/OrangeMixs/blob/main/Models/AbyssOrangeMix/AbyssOrangeMix.safetensors"
)
```
45

Steven Liu's avatar
Steven Liu committed
46
## Convert to safetensors
47

Steven Liu's avatar
Steven Liu committed
48
Not all weights on the Hub are available in the `.safetensors` format, and you may encounter weights stored as `.bin`. In this case, use the [Convert Space](https://huggingface.co/spaces/diffusers/convert) to convert the weights to `.safetensors`. The Convert Space downloads the pickled weights, converts them, and opens a Pull Request to upload the newly converted `.safetensors` file on the Hub. This way, if there is any malicious code contained in the pickled files, they're uploaded to the Hub - which has a [security scanner](https://huggingface.co/docs/hub/security-pickle#hubs-security-scanner) to detect unsafe files and suspicious pickle imports - instead of your computer.
49

Steven Liu's avatar
Steven Liu committed
50
You can use the model with the new `.safetensors` weights by specifying the reference to the Pull Request in the `revision` parameter (you can also test it in this [Check PR](https://huggingface.co/spaces/diffusers/check_pr) Space on the Hub), for example `refs/pr/22`:
51

Steven Liu's avatar
Steven Liu committed
52
53
```py
from diffusers import DiffusionPipeline
54

55
56
57
pipeline = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1", revision="refs/pr/22", use_safetensors=True
)
58
59
```

Steven Liu's avatar
Steven Liu committed
60
## Why use safetensors?
61

Steven Liu's avatar
Steven Liu committed
62
There are several reasons for using safetensors:
63

Steven Liu's avatar
Steven Liu committed
64
65
- Safety is the number one reason for using safetensors. As open-source and model distribution grows, it is important to be able to trust the model weights you downloaded don't contain any malicious code. The current size of the header in safetensors prevents parsing extremely large JSON files.
- Loading speed between switching models is another reason to use safetensors, which performs zero-copy of the tensors. It is especially fast compared to `pickle` if you're loading the weights to CPU (the default case), and just as fast if not faster when directly loading the weights to GPU. You'll only notice the performance difference if the model is already loaded, and not if you're downloading the weights or loading the model for the first time.
66

Steven Liu's avatar
Steven Liu committed
67
	The time it takes to load the entire pipeline:
68

Steven Liu's avatar
Steven Liu committed
69
	```py
70
 	from diffusers import StableDiffusionPipeline
71

72
73
74
 	pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", use_safetensors=True)
 	"Loaded in safetensors 0:00:02.033658"
 	"Loaded in PyTorch 0:00:02.663379"
Steven Liu's avatar
Steven Liu committed
75
	```
76

Steven Liu's avatar
Steven Liu committed
77
	But the actual time it takes to load 500MB of the model weights is only:
78

Steven Liu's avatar
Steven Liu committed
79
80
81
82
	```bash
	safetensors: 3.4873ms
	PyTorch: 172.7537ms
	```
83

Steven Liu's avatar
Steven Liu committed
84
- Lazy loading is also supported in safetensors, which is useful in distributed settings to only load some of the tensors. This format allowed the [BLOOM](https://huggingface.co/bigscience/bloom) model to be loaded in 45 seconds on 8 GPUs instead of 10 minutes with regular PyTorch weights.