[diffusers]仅读取unet | 模型格式转化diffusion_pytorch_model.safetensors转为diffusion_pytorch_model.bin

问题描述

SError: Error no file named config.json found in directory /share/huanggao/zjc/code_mid/model_yw/bk-base-2m/unet.

无法读取unet

下载模型

我们从网上下载的模型
通过
pipe = StableDiffusionPipeline.from_pretrained("nota-ai/bk-sdm-base-2m", torch_dtype=torch.float16)
指令下载的unet是
diffusion_pytorch_model.safetensors格式的

我从网上直接下载的模型路径
/root/.cache/huggingface/hub/models--nota-ai--bk-sdm-base-2m/snapshots/e8b5597155c5b2c77585570b99113f1c77b97338/unet
转移到
/share/huanggao/zjc/code_mid/model_yw/bk-base-2m/unet
文件格式
config.json 
diffusion_pytorch_model.safetensors
（这种方法无法直接用
config.json文件会显示找不到）

smallsd路径
~/.cache/huggingface/hub/models--OFA-Sys--small-stable-diffusion-v0/snapshots/38e10e5e71e8fbf717a47a81e7543cd01c1a8140/unet
格式是（diffusers格式）
文件格式
config.json 
diffusion_pytorch_model.bin

转化策略（不应该转化）

转化文件路径
python ./scripts/convert_original_stable_diffusion_to_diffusers.py

转化文件
convert_original_stable_diffusion_to_diffusers.py

转化方法
python ./scripts/convert_original_stable_diffusion_to_diffusers.py
--checkpoint_path xxx.safetensors
--dump_path save_dir
--from_safetensors

python ./convert_original_stable_diffusion_to_diffusers.py  --checkpoint_path  /root/.cache/huggingface/hub/models--nota-ai--bk-sdm-base-2m/snapshots/e8b5597155c5b2c77585570b99113f1c77b97338/unet/diffusion_pytorch_model.safetensors  --dump_path   /share/huanggao/zjc/code_mid    --from_safetensors

转化过程会报错

KeyError: 'time_embed.0.weight'
报错截图

解决办法

1 如果使用整个模型，如下指令

pipe = StableDiffusionPipeline.from_pretrained("nota-ai/bk-sdm-base-2m",  torch_dtype=torch.float16)

pipe = StableDiffusionPipeline.from_pretrained("nota-ai/bk-sdm-base-2m",  torch_dtype=torch.float16)

2 如果只用其中的Unet

unet_smallsd = NewUNet.from_pretrained(
        args.smallsd_model_name_or_path, subfolder="unet", use_safetensors=True, revision=args.non_ema_revision
    )  

unet_smallsd = NewUNet.from_pretrained(
        args.smallsd_model_name_or_path, subfolder="unet", use_safetensors=True, revision=args.non_ema_revision
    )  # -> small
    if accelerator.is_main_process:  print("Unet_student already Loaded! >>> ")

其中args.smallsd_model_name_or_path为
export smallsd="nota-ai/bk-sdm-base-2m"

我的理解

当采用了，这个指令来读取模型
unet_smallsd = NewUNet.from_pretrained(
        args.smallsd_model_name_or_path, subfolder="unet", use_safetensors=True, revision=args.non_ema_revision
文件会保存到下面路径
~/.cache/huggingface/hub/models--nota-ai--bk-sdm-base-2m/snapshots/e8b5597155c5b2c77585570b99113f1c77b97338

models--nota-ai--bk-sdm-base-2m/
就等同于路径“nota-ai/bk-sdm-base-2m”

解决办法（待尝试）
 

通过下面方法读取和保存
pt_state_dict = safetensors.torch.load_file("/xx/xxx.safetensors", device="cpu")
torch.save(pt_state_dict, "pt_state_dict.bin")