Implementation:Ollama Ollama Imagegen Safetensors

Knowledge Sources	Ollama
Domains	Image Generation, Model Loading
Last Updated	2025-02-15 00:00 GMT

Overview

Manages safetensors file indexing and tensor loading with native MLX mmap support for zero-copy weight access.

Description

The safetensors.go file provides the ModelWeights struct for managing weights from multiple safetensors files. It parses JSON headers (SafetensorHeader/TensorInfo with dtype, shape, data_offsets) from safetensors files to build a tensor index without loading data, then provides lazy loading through native MLX mmap (mlx_load_safetensors) for zero-copy GPU access. The ModelWeights struct supports loading from both directories (LoadModelWeights) and explicit file paths (LoadModelWeightsFromPaths), dtype conversion, and tensor caching. dtypeFromString converts safetensors dtype strings (F32, BF16, etc.) to MLX Dtype constants.

Usage

Used as the primary weight loading mechanism for standalone model directories containing .safetensors files.

Code Reference

Source Location

• Repository: Ollama
• File: x/imagegen/safetensors/safetensors.go
• Lines: 1-318

Signature

type SafetensorHeader map[string]TensorInfo

type TensorInfo struct {
	Dtype       string  `json:"dtype"`
	Shape       []int32 `json:"shape"`
	DataOffsets [2]int  `json:"data_offsets"`
}

type ModelWeights struct {
	dir         string
	tensorFiles map[string]string
	tensorInfo  map[string]TensorInfo
	nativeCache map[string]*mlx.SafetensorsFile
	cache       map[string]*mlx.Array
}

func LoadModelWeights(dir string) (*ModelWeights, error)
func LoadModelWeightsFromPaths(paths []string) (*ModelWeights, error)
func (mw *ModelWeights) Load(dtype mlx.Dtype) error
func (mw *ModelWeights) Get(name string) *mlx.Array
func (mw *ModelWeights) ReleaseAll()

Import

import "github.com/ollama/ollama/x/imagegen/safetensors"

I/O Contract

Inputs

Name	Type	Required	Description
dir	string	Yes	Directory containing .safetensors files
dtype	mlx.Dtype	No	Target dtype for conversion (0 = original dtype)

Outputs

Name	Type	Description
*ModelWeights	*ModelWeights	Indexed weight manager with lazy loading
error	error	Error if no safetensors files found or parsing fails

Usage Examples

weights, err := safetensors.LoadModelWeights("/path/to/model")
if err != nil {
    return err
}

// Load all tensors as bfloat16
if err := weights.Load(mlx.DtypeBFloat16); err != nil {
    return err
}

// Get specific tensor
embedding := weights.Get("model.embed_tokens.weight")

Related Pages

• Principle:Ollama_Ollama_ImageGeneration