maple-diffusion
maple-diffusion copied to clipboard
./maple-convert.py should ignore non-tensor checkpoint state
Some weight checkpoints have extra non-tensor data in their state_dict. maple-convert.py should ignore this extra data rather than crashing.
# model weights
-for k in ckpt["state_dict"]:
- if "first_stage_model.encoder" in k: continue
- ckpt["state_dict"][k].numpy().astype('float16').tofile(outpath / (k + ".bin"))
+for k, v in ckpt["state_dict"].items():
+ if "first_stage_model.encoder" in k:
+ continue
+ if not hasattr(v, "numpy"):
+ continue
+ v.numpy().astype('float16').tofile(outpath / (k + ".bin"))