skainet-io-safetensors/sk.ainet.io.safetensors/StorageAwareSafeTensorsLoader

StorageAwareSafeTensorsLoader

class StorageAwareSafeTensorsLoader(sourceProvider: () -> RandomAccessSource, filePath: String? = null, onProgress: (current: Long, total: Long, tensorName: String?) -> Unit = { _, _, _ -> })(source)

SafeTensors loader that produces TensorStorage descriptors with zero-copy file-backed handles where possible.

Unlike SafeTensorsParametersLoader which always decodes into typed arrays, this loader returns raw TensorStorage descriptors that can be:

File-backed (zero-copy): When a file path is provided, tensors reference the original file via BufferHandle.FileBacked. No heap allocation occurs for the tensor data itself.
Borrowed (single allocation): When no file path is available, tensor bytes are loaded into a single ByteArray and wrapped as BufferHandle.Borrowed.

Usage:

// Zero-copy: tensors reference the file directly
val loader = StorageAwareSafeTensorsLoader(sourceProvider, filePath = "/models/model.safetensors")
val tensors = loader.loadAll()
// tensors[0].isFileBacked == true

// Heap-loaded: tensors are borrowed byte arrays
val loader = StorageAwareSafeTensorsLoader(sourceProvider)
val tensors = loader.loadAll()
// tensors[0].ownership == Ownership.BORROWED

Constructors

StorageAwareSafeTensorsLoader

constructor(sourceProvider: () -> RandomAccessSource, filePath: String? = null, onProgress: (current: Long, total: Long, tensorName: String?) -> Unit = { _, _, _ -> })

Functions

listTensors

fun listTensors(): List<StreamingSafeTensorInfo>

List all tensor names and their metadata without loading data.

load

fun load(name: String): TensorStorage

Load a single tensor by name as TensorStorage.

loadAll

fun loadAll(): Map<String, TensorStorage>

Load all tensors as TensorStorage descriptors.