Skip to content

security: prefer safetensors over torch dumps#204

Open
coldwaterq wants to merge 2 commits intonikopueringer:mainfrom
coldwaterq:safetensors-as-alternative
Open

security: prefer safetensors over torch dumps#204
coldwaterq wants to merge 2 commits intonikopueringer:mainfrom
coldwaterq:safetensors-as-alternative

Conversation

@coldwaterq
Copy link

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 pytest

Checklist

  • [ X ] uv run pytest passes
  • [ X ] uv run ruff check passes
  • [ X ] uv run ruff format --check passes

@coldwaterq
Copy link
Author

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"

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant