K-【学习Diffusers 四】 读取模型参数 bin格式、safetensors格式_safetensors文件

该操作多用于推理

safetensors格式的参数读取方法

1 拿到pipeline中的unet的办法
unet = pipeline.pipe.unet

2 safetensors格式文件的参数读取方法
state_dict = safetensors.torch.load_file(args.model_id, device="cpu")
unet.load_state_dict(state_dict)  # 读入模型

3 args.model_id = "xxxx/unet/diffusion_pytorch_model.safetensors"
unet模型路径

4 将修改的unet放回pipeline

pipeline.pipe.unet = unet.half().to(args.device)

对于bin格式

1 读取参数
state_dict=torch.load('../train-output/'+ args.model_name_or_path.split('/')[-1] +'/unet/diffusion_pytorch_model.bin')
2 送入uent
unet.load_state_dict(state_dict)
3 送回pipe
pipe.unet = unet
4 送入GPU
pipe.to("cuda")