Transformers库读取.safetensors文件报错

报错问题

在load一些pretrained模型的时候会报错：

if metadata.get("format") == "pt": 
AttributeError: 'NoneType' object has no attribute 'get'

原因是在保存为.safetensors文件时没有保存metadata。

解决方法

引用自：GitHub链接
可以使用以下代码将.safetensors文件加上metadata重新保存：

from safetensors import safe_open
from safetensors.torch import save_file

safetensor_path = "ckpt/sd/text_encoder/model.safetensors" # 此处替换为需要重新保存的文件位置

fname, ext = safetensor_path.split("/")[-1].split(".")
# ext = 'safetensors' # 扩展名
# fname = 'model' # 文件名

tensors = dict()
with safe_open(safetensor_path, framework="pt", device="cpu") as f:
    for key in f.keys():
        tensors[key] = f.get_tensor(key)

save_file(tensors, f"ckpt/sd/text_encoder/{fname}-with-format.{ext}", metadata={"format": "pt"})