safetensors: Torch SD-based models tensor invalid for input size

There might be a slight discrepancy between the loading and saving process in safetensors. When loading a SD-based model like sd-1.4 packaged into a PyTorch checkpoint, we’ll call it sd-v1-4.ckpt. We can package its state_dict and discard the torch format.

Packaging as safe_tensors

    sf_filename = "sd-v1-4.safetensors"
    filename = "sd-v1-4.ckpt"

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

    # appears to pop nothing in this case
    shared = shared_pointers(loaded)
    for shared_weights in shared:
        for name in shared_weights[1:]:
            loaded.pop(name)

    loaded = {k: v.contiguous() for k, v in loaded.items()}

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

    check_file_size(local, filename)

Loading the tensors

load_file('sd-v1-4.safetensors', device='cpu')

Results in error:

File “venv\lib\site-packages\safetensors\torch.py”, line 99, in load_file result[k] = f.get_tensor(k) RuntimeError: shape ‘[1280, 1280, 3, 3]’ is invalid for input of size 7290352

Expected behaviour: safetensors fails while trying to save unexpected tensor data or creates tensors which can be loaded Affected version: safetensors=2.4.0, torch=1.12.1+cu113

ckpt size: 3.97 GB 4,265,381,888 bytes (4,265,380,512 bytes) safetensor size: 3.97 GB 4,265,148,416 bytes (4,265,146,304 bytes) SHA fe4efff1e174c627256e44ec2991ba279b3816e364b49f9be2abc0b3ff3f8556

Using pruned version of CompVis/stable-diffusion-v-1-4-original

Apologies if this is already fixed with addition of more dtypes. Will try to get more info by running through check output and debug info of this specific tensor