Safetensors

Calling patch_safetensors() replaces safetensors.torch.safe_open with a SafeOpen wrapper that decompresses marked tensors on read. After a single call, the rest of the program uses the safetensors API unchanged.

Basic usage

from ptwm import patch_safetensors
patch_safetensors()

import safetensors.torch
tensors = safetensors.torch.load_file("model.safetensors")
# Will decompress on the fly if a sibling .ptwm exists.

Compressing safetensors models

Tensor-by-tensor with the CLI:

ptwm compress model.safetensors
# produces model.safetensors.ptwm

Or programmatically:

from ptwm.integrations import compress_safetensors_file

compress_safetensors_file(
    "model.safetensors",
    "model.safetensors.ptwm",
    mode="b",                   # "a" = native, "b" = .safetensors shell
)

Mode "a" writes a native .ptwm directory with one blob per shard; mode "b" writes a .safetensors shell that wraps a .ptwm payload so existing loaders still recognise the filename. See .ptwm container.

Process-safe patching

patch_safetensors() uses MultiProcessPatcher from ptwm.utils._patch, so it survives Python's multi-process import dance (for example, in DataLoader worker processes). Calling it more than once is idempotent.

Unpatching

PTWM ships no unpatch entry point. The patcher's contract is simple: the override sticks for the lifetime of the process. To get a fresh interpreter without it, exit and re-launch.