Add basic safetensors SD model loading support

[pattontim]

Addresses requirements for #4714

For now, will automatically use and load .safetensor files in the SD model subdirectory. I tested it with trinart but should work for all pruned sd-based models. This basic support will allow the implementation of other features mentioned in the issue such as optional fast GPU load and zero copy (see #684 ) and converting existing models to safetensors format.

Any extraneous information like global step in the file may be lost when creating safetensors files using convert.py, simply edit https://github.com/huggingface/safetensors/blob/main/bindings/python/convert.py to get:

import os
from collections import defaultdict
from inspect import signature
from typing import Dict, List, Optional

import torch

from safetensors.torch import save_file

def shared_pointers(tensors):
    ptrs = defaultdict(list)
    #print(tensors.keys())
    #print(dir(tensors))
    for k, v in tensors.items():
        ptrs[v.data_ptr()].append(k)
    failing = []
    for ptr, names in ptrs.items():
        if len(names) > 1:
            failing.append(names)
    return failing


def check_file_size(sf_filename: str, pt_filename: str):
    sf_size = os.stat(sf_filename).st_size
    pt_size = os.stat(pt_filename).st_size

    if (sf_size - pt_size) / pt_size > 0.01:
        raise RuntimeError(
            f"""The file size different is more than 1%:
         - {sf_filename}: {sf_size}
         - {pt_filename}: {pt_size}
         """
        )

def convert_single():
    sf_filename = "model.safetensors"
    filename = "C:\\convert_loc\\model_to_convert.ckpt"

    loaded = torch.load(filename)
    loaded = loaded['state_dict']

    local = os.path.join("C:\\convert_loc\\convert\\", sf_filename)
    shared = shared_pointers(loaded)
    for shared_weights in shared:
        for name in shared_weights[1:]:
            loaded.pop(name)

    # For tensors to be contiguous
    loaded = {k: v.contiguous() for k, v in loaded.items()}

    save_file(loaded, local, metadata={"format": "pt"})

    check_file_size(local, filename)

    operations = None
    return operations

if __name__ == "__main__":
    convert_single()

[pattontim]

To maintain merge functionality between file types needs calls to safetensors.torch.loadfile in extras.py.

Edit: commit soon

[pattontim]

The convert script will likely need to be modified also because other models don't seem to work when loading the tensors out of them. Only trinart_60k appears to work out of the box. sd1.4 didn't work

[pattontim]

Me and the safetensors dev have gone back and forth about a bug with loading models via CPU. Unfortunately it means for now safetensor models can only be loaded to GPU. Otherwise everything is now working including models. But again, cause of that bug it means essentially only GPUs can use safetensors and model swaps between safetensors will be on gpu only.

[0xdevalias]

https://huggingface.co/docs/safetensors/speed

Safetensors is really fast.

on CPU, safetensors is faster than pytorch by: 76.6 X This speedup is due to the fact that this library avoids unnecessary copies by mapping the file directly.

on GPU, safetensors is faster than pytorch by: 2.1 X The speedup works because this library is able to skip unecessary CPU allocations.

Originally posted by @0xdevalias in https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/4714#issuecomment-1325986115

[AUTOMATIC1111]

wouldn't that mean that for model merging you'd have to load models into GPU memory, which will ultimately require at least 3x the amount of memory a single model takes?

[AUTOMATIC1111]

there are some changes i don't agree in this, and the bug you're referring to has been fixed it seems so I ended up adding my own implementation using your code as reference

[pattontim]

I ended up adding my own implementation using your code as reference

This is my first time contributing to a major project, so it would have been nice if the work was based on my branch in order to build contribution history. I understand you are working with a lot of PRs and it was more practical to simply build off master. Yet there is no I or Y in OUR and this is a hobby unlicensed project so I will take more time before judging your character

Merging now works so I'm closing this PR.