madebyollin / maple-diffusion

Stable Diffusion inference on iOS / macOS using MPSGraph
https://madebyoll.in/posts/maple_diffusion/
MIT License
793 stars 51 forks source link

./maple-convert.py should ignore non-tensor checkpoint state #17

Closed jackpal closed 1 year ago

jackpal commented 1 year ago

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"))
madebyollin commented 1 year ago

Ah, sure! Added the hasattr-based skip, hopefully it's more robust now.

jackpal commented 1 year ago

Great! Thanks!