Repository: AUTOMATIC1111/stable-diffusion-webui

Issue:
[Bug]:  Slow .safetensors loading
### Is there an existing issue for this?

- [X] I have searched the existing issues and checked the recent builds/commits

### What happened?

System specs: 5800X3D, RTX 3080 12GB, 64GB DDR4-3600, Windows 11 22H2 22621.963, latest Nvidia driver (527.56), latest BIOS update.

This is most likely a specific issue related only to my PC, but I've seen a couple of comments about it on other sites. safetensors load significantly slower than the same model in .ckpt, around 2-3 minutes to load a safetensors compared to .ckpt loading in less than 10 seconds.

However, safetensors do load fast when doing merges, both with the inbuilt merger and with merge extensions. I think this is because they mostly only load into RAM and not fully into VRAM, but switching to either of the models used for a merge after the merge is done also makes them load instantly.

Troubleshooting I've tried so far:
-launching with optimizations and without optimizations
-Fresh UI reinstall
-deleted and rebuilt venv
-set SAFETENSORS_FAST_GPU=1 and without the set parameter
-Python 3.10.6, 3.10.8 and 3.10.9
-No extensions
-No antivirus
-No GPU undervolt
-setting python.exe to Max Performance in Nvidia panel
-different browsers
-browser HW acceleration on and off

Videocard works as it should for everything else.

The most notable thing I've noticed is that when loading a .ckpt, both python.exe and System show reads in Resource Manager. When loading a .safetensors, only python.exe shows reads.

ckpt loading:
![ckpt](https://user-images.githubusercontent.com/66507343/208746830-6cc6a8e8-ca53-42c8-83e4-bce66b0bcb9a.png)
safetensors loading:
![safetensors](https://user-images.githubusercontent.com/66507343/208746828-893a666a-7a3c-4e62-8367-4d5325188db0.png)

### Steps to reproduce the problem

1. Go to Stable Diffusion checkpoint dropdown
2. Load safetensors model
3. Takes 2-3 minutes to load compared to the same model but in .ckpt


### What should have happened?

1. Go to Stable Diffusion checkpoint dropdown
2. Load safetensors model
3. Load in less than 10 seconds like the same model but in .ckpt does.

### Commit where the problem happens

 685f9631b56ff8bd43bce24ff5ce0f9a0e9af490

### What platforms do you use to access UI ?

Windows

### What browsers do you use to access the UI ?

Brave

### Command Line Arguments

```Shell
Optimized .bat:

--xformers --deepdanbooru --gradio-img2img-tool color-sketch --gradio-inpaint-tool=color-sketch --opt-split-attention --precision autocast --opt-channelslast --api

Non optimized:
--deepdanbooru 

Both show the same issue.
```


### Additional information, context and logs

_No response_


Discussion:
have the same observation, thought it should load faster as advertised.
cc @Narsil
Hi thanks for the ping, this definitely shouldn't happen.

Do you mind sharing:

- The actual files you're trying to load ? (Part of it could come from incorrect precision in the files, since you use `--precision autocast` it could cause some issues).
   - I am guessing this : https://huggingface.co/riffusion/riffusion-model-v1/blob/main/riffusion-model-v1.ckpt
   - And this : https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/refs%2Fpr%2F46/v1-5-pruned-emaonly.safetensors
- The exact PyTorch version
- Can you confirm you're loading on the GPU correct ?
- The type of Disk your are using (SSD, HDD).
- Both files are stored on the same disk I guess ?
- You don't have an antivirus that might be inspecting the file slowing down its load ?


I need to try and reproduce, but so far I have failed (I must confess I'm only using cloud Windows since I don't own a Windows machine anymore).

> they mostly only load into RAM and not fully into VRAM

Loading on CPU is always going to be faster than loading onto GPU yes. 
However it's very surprising that the read speeds are vastly different for both files.

Could you isolate the issue maybe ?


```python
import torch
import datetime

start = datetime.datetime.now()
weights = torch.load("riffusion-model-v1.ckpt", device="cuda:0")
print(f"Loaded PT in {datetime.datetime.now() - start}")
```

```python
from safetensors.torch import load_file
import datetime

start = datetime.datetime.now()
weights = load_file("v1-5-pruned-emaonly.safetensors", device="cuda:0")
print(f"Loaded SF in {datetime.datetime.now() - start}")
```

If SF is indeed slower, we can also try this: https://gist.github.com/Narsil/3edeec2669a5e94e4707aa0f901d2282 to check if the slowness is indeed in the lib or not (It needs a bit adaptation to load the checkpoint you want onto GPU).

And report if you see the same things (just trying to ignore webui if possible.

In the meantime, I will keep trying to reproduce.

NB: The two files are not the same size it seems, but the SF one is smaller, so it should be faster to load.

```
Loading weights [3aafa6fe] from C:\Users\Administrator\src\stable-diffusion-webui\models\Stable-diffusion\riffusion-model-v1.ckpt
Loaded C:\Users\Administrator\src\stable-diffusion-webui\models\Stable-diffusion\riffusion-model-v1.ckpt in 0:01:15.861647
Applying cross attention optimization (Doggettx).
Weights loaded. in 0:01:19.734654
Loading weights [d7049739] from C:\Users\Administrator\src\stable-diffusion-webui\models\Stable-diffusion\v1-5-pruned-emaonly.safetensors
Loaded C:\Users\Administrator\src\stable-diffusion-webui\models\Stable-diffusion\v1-5-pruned-emaonly.safetensors in 0:00:00.308988
Applying cross attention optimization (Doggettx).
Weights loaded. in 0:00:44.164628
```

Here is what I get on a brand new Windows (Server 2022) with cuda support.
I merely added 2 timings within `sd-webui` one just after loading the file, and the other one where the classic prompt is

So it's indeed spending a bunch of time doing something AFTER loading the weights, instead of loading correctly the correct weights for SD-1.5 directly in a proper form.

I'll investigate a bit more to see what's going on for SD-1.5
Ok, I was able to investigate. It seems `map_location or shared.weight_load_location` resolves to `None` meaning the tensors are loaded on CPU (which is extremely fast), but than are sent to VRAM with `load_state_dict` (which is painfully slow and I don't know why.

https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/sd_models.py#L170

I changed that line and forced `device="cuda:0"` and then I got those results: 

```
Loaded C:\Users\Administrator\src\stable-diffusion-webui\models\Stable-diffusion\v1-5-pruned-emaonly.safetensors in 0:00:05.154992
Using map_location None
Loaded state dict 0:00:02.660000
Model half 0:00:02.677993
Model first stage to 0:00:02.696012
Applying cross attention optimization (Doggettx).
HIGHJACK stuff in 0:00:08.760007
callback stuff in 0:00:08.760007
load on device cuda in 0:00:09.488031
Weights loaded. in 0:00:09.488031
```

So the load time is slower (it loads directly on GPU) but it seems faster overall (because there's no need to reallocate on GPU afterwards).

Could that explain the issue you're having ?

As for a fix, I'm not sure what's the proper fix here. 
IMO the `device=` should always be properly set (both for PT and SF) instead of using `None` (the device is stored in the pickle themselves, but that's not really portable, since if some weights where saved on GPU they can't be loaded on a machine without GPU anymore).

Resolving patch:
diff --git a/modules/sd_models.py b/modules/sd_models.py
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -168,7 +168,10 @@ def get_state_dict_from_checkpoint(pl_sd):
 def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):
     _, extension = os.path.splitext(checkpoint_file)
     if extension.lower() == ".safetensors":
-        pl_sd = safetensors.torch.load_file(checkpoint_file, device=map_location or shared.weight_load_location)
+        device = map_location or shared.weight_load_location
+        if device is None:
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)
     else:
         pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)
 


Test patch: