Skip to content

Commit 5e74675

Browse files
[docs] Load safetensors (#3333)
* safetensors * apply feedback * apply feedback * Apply suggestions from code review --------- Co-authored-by: Patrick von Platen <[email protected]>
1 parent c49e9ed commit 5e74675

File tree

2 files changed

+48
-61
lines changed

2 files changed

+48
-61
lines changed

docs/source/en/_toctree.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
title: Load and compare different schedulers
2727
- local: using-diffusers/custom_pipeline_overview
2828
title: Load community pipelines
29+
- local: using-diffusers/using_safetensors
30+
title: Load safetensors
2931
- local: using-diffusers/kerascv
3032
title: Load KerasCV Stable Diffusion checkpoints
3133
title: Loading & Hub
@@ -50,8 +52,6 @@
5052
title: Community pipelines
5153
- local: using-diffusers/contribute_pipeline
5254
title: How to contribute a community pipeline
53-
- local: using-diffusers/using_safetensors
54-
title: Using safetensors
5555
- local: using-diffusers/stable_diffusion_jax_how_to
5656
title: Stable Diffusion in JAX/Flax
5757
- local: using-diffusers/weighted_prompts
Lines changed: 46 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,87 +1,74 @@
1-
# What is safetensors ?
1+
# Load safetensors
22

3-
[safetensors](https://github.com/huggingface/safetensors) is a different format
4-
from the classic `.bin` which uses Pytorch which uses pickle. It contains the
5-
exact same data, which is just the model weights (or tensors).
3+
[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.
64

7-
Pickle is notoriously unsafe which allow any malicious file to execute arbitrary code.
8-
The hub itself tries to prevent issues from it, but it's not a silver bullet.
5+
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:
96

10-
`safetensors` first and foremost goal is to make loading machine learning models *safe*
11-
in the sense that no takeover of your computer can be done.
12-
13-
Hence the name.
14-
15-
# Why use safetensors ?
16-
17-
**Safety** can be one reason, if you're attempting to use a not well known model and
18-
you're not sure about the source of the file.
19-
20-
And a secondary reason, is **the speed of loading**. Safetensors can load models much faster
21-
than regular pickle files. If you spend a lot of times switching models, this can be
22-
a huge timesave.
23-
24-
Numbers taken AMD EPYC 7742 64-Core Processor
7+
```bash
8+
!pip install safetensors
259
```
26-
from diffusers import StableDiffusionPipeline
2710

28-
pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
11+
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.
2912

30-
# Loaded in safetensors 0:00:02.033658
31-
# Loaded in Pytorch 0:00:02.663379
32-
```
13+
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):
3314

34-
This is for the entire loading time, the actual weights loading time to load 500MB:
15+
```py
16+
from diffusers import DiffusionPipeline
3517

36-
```
37-
Safetensors: 3.4873ms
38-
PyTorch: 172.7537ms
18+
pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_safetensors=True)
3919
```
4020

41-
Performance in general is a tricky business, and there are a few things to understand:
21+
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.FromCkptMixin.from_ckpt`] method:
4222

43-
- If you're using the model for the first time from the hub, you will have to download the weights.
44-
That's extremely likely to be much slower than any loading method, therefore you will not see any difference
45-
- If you're loading the model for the first time (let's say after a reboot) then your machine will have to
46-
actually read the disk. It's likely to be as slow in both cases. Again the speed difference may not be as visible (this depends on hardware and the actual model).
47-
- The best performance benefit is when the model was already loaded previously on your computer and you're switching from one model to another. Your OS, is trying really hard not to read from disk, since this is slow, so it will keep the files around in RAM, making it loading again much faster. Since safetensors is doing zero-copy of the tensors, reloading will be faster than pytorch since it has at least once extra copy to do.
23+
```py
24+
from diffusers import StableDiffusionPipeline
4825

49-
# How to use safetensors ?
26+
pipeline = StableDiffusionPipeline.from_ckpt(
27+
"https://huggingface.co/WarriorMama777/OrangeMixs/blob/main/Models/AbyssOrangeMix/AbyssOrangeMix.safetensors"
28+
)
29+
```
5030

51-
If you have `safetensors` installed, and all the weights are available in `safetensors` format, \
52-
then by default it will use that instead of the pytorch weights.
31+
## Convert to safetensors
5332

54-
If you are really paranoid about this, the ultimate weapon would be disabling `torch.load`:
55-
```python
56-
import torch
33+
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 Space below 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.
5734

35+
<iframe
36+
src="https://safetensors-convert.hf.space"
37+
frameborder="0"
38+
width="850"
39+
height="450"
40+
></iframe>
5841

59-
def _raise():
60-
raise RuntimeError("I don't want to use pickle")
42+
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`:
6143

44+
```py
45+
from diffusers import DiffusionPipeline
6246

63-
torch.load = lambda *args, **kwargs: _raise()
47+
pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", revision="refs/pr/22")
6448
```
6549

66-
# I want to use model X but it doesn't have safetensors weights.
50+
## Why use safetensors?
6751

68-
Just go to this [space](https://huggingface.co/spaces/diffusers/convert).
69-
This will create a new PR with the weights, let's say `refs/pr/22`.
52+
There are several reasons for using safetensors:
7053

71-
This space will download the pickled version, convert it, and upload it on the hub as a PR.
72-
If anything bad is contained in the file, it's Huggingface hub that will get issues, not your own computer.
73-
And we're equipped with dealing with it.
54+
- 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.
55+
- 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.
7456

75-
Then in order to use the model, even before the branch gets accepted by the original author you can do:
57+
The time it takes to load the entire pipeline:
7658

77-
```python
78-
from diffusers import DiffusionPipeline
59+
```py
60+
from diffusers import StableDiffusionPipeline
7961

80-
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", revision="refs/pr/22")
81-
```
62+
pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
63+
"Loaded in safetensors 0:00:02.033658"
64+
"Loaded in PyTorch 0:00:02.663379"
65+
```
8266

83-
or you can test it directly online with this [space](https://huggingface.co/spaces/diffusers/check_pr).
67+
But the actual time it takes to load 500MB of the model weights is only:
8468

85-
And that's it !
69+
```bash
70+
safetensors: 3.4873ms
71+
PyTorch: 172.7537ms
72+
```
8673

87-
Anything unclear, concerns, or found a bugs ? [Open an issue](https://github.com/huggingface/diffusers/issues/new/choose)
74+
- 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.

0 commit comments

Comments
 (0)