Add safetensors checkpoint loading support - #599
johnwbrisbin wants to merge 1 commit into
Conversation
All four checkpoint-loading paths (build_sam3_video_model, build_sam3_multiplex_video_model, build_sam3_multiplex_video_predictor, and the internal _load_checkpoint helper) now detect a .safetensors extension and use safetensors.torch.load_file instead of torch.load. The official Comfy-Org/sam3.1 HuggingFace distribution ships sam3.1_multiplex_fp16.safetensors; torch.load fails on this format with "unsupported pickle protocol: 104" (the safetensors magic bytes are misread as a pickle header). .pt / .pth checkpoints continue to use torch.load, now with weights_only=False so that legacy checkpoints using older pickle protocols still load without raising UnpicklingError.
|
Hi @johnwbrisbin! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
Summary
All four checkpoint-loading paths now detect a
.safetensorsextension and usesafetensors.torch.load_fileinstead oftorch.load:_load_checkpointhelperbuild_sam3_video_modelbuild_sam3_multiplex_video_modelbuild_sam3_multiplex_video_predictorMotivation
The official Comfy-Org/sam3.1 HuggingFace distribution ships
sam3.1_multiplex_fp16.safetensors. Passing this file to any of the above builders fails with:The safetensors magic bytes are misread as a legacy pickle header by
torch.load..pt/.pthcheckpoints continue to usetorch.load, now withweights_only=Falseso that legacy checkpoints using older pickle protocols still load without raisingUnpicklingError.Test plan
sam3.1_multiplex_fp16.safetensorsviabuild_sam3_multiplex_video_model— previously crashed, now loads correctly.ptcheckpoint viabuild_sam3_video_model— still loads correctly withweights_only=False