security: prefer safetensors over torch dumps#204
security: prefer safetensors over torch dumps#204coldwaterq wants to merge 2 commits intonikopueringer:mainfrom
Conversation
and provide a conversion script
|
It's not related to this change so I didn't include it in the PR, but I also believe that HF_CHECKPOINT_FILENAME = "CorridorKey.pth"is supposed to be HF_CHECKPOINT_FILENAME = "CorridorKey_v1.0.pth" |
|
Seems like a prudent change. @HYP3R00T What are your thoughts? There are a fair amount of new lines of code. Do you think we should merge it? |
|
Discussed this, and it seems like it might be better for me to convert the official file to safetensors and avoid .pth altogether! Going to look into that. |
|
Yeah, if that's easier for you it would also be better. The conversion complicates this merge but it was just an attempt to make it easier for you. |
|
@coldwaterq Hey there! Wanted to let you know that I am now hosting a .safetensors version of CorridorKey on HF, alongside the old .pth for now. I also merged a PR to import the .safetensor file by default, along with some old legacy fallback to .pth for the people still running the old model, and a conversion script as well. Much of that was directly informed by what you did with this PR. You laid out the groundwork for this, so thank you for doing so! Hopefully this will be a big step in getting CK to be more secure. Since I believe we have effectively duplicated the actions of this PR, I'm going to go ahead and close it, but please let me know if I am doing this in error! Thank you again for your help! |
What does this change?
Add a conversion script and update the loading to make it easy to swap to safetensors without breaking backwards compatability. Safetesnors ensures that pickles aren't used to save the models at all, since pytorch keeps having security advisories this ensures that you don't have to worry if you are running the latest version of pytorch and know that even if the model is corrupted, an attacker can't get a shell on your system.
How was it tested?
The conversion script was tested on the model in huggingface and then that converted file was loaded to ensure the safetensor load calls are setup correctly and nothing the model requires was lost.
ran
uv run pytestChecklist
uv run pytestpassesuv run ruff checkpassesuv run ruff format --checkpasses