add safetensors

vae.py

@@ -1,6 +1,7 @@
 # Adopted from LDM's KL-VAE: https://github.com/CompVis/latent-diffusion
 import torch
 import torch.nn as nn
+from safetensors.torch import load_file
 
 import numpy as np
 
@@ -461,7 +462,7 @@ class AutoencoderKL(nn.Module):
             self.init_from_ckpt(ckpt_path)
 
     def init_from_ckpt(self, path):
-        sd = torch.load(path, map_location="cpu")["model"]
+        sd = load_file(path)
         msg = self.load_state_dict(sd, strict=False)
         print("Loading pre-trained KL-VAE")
         print("Missing keys:")