ollama source for Momentry Core verification

x/imagegen/safetensors/loader.go

@@ -0,0 +1,429 @@
+package safetensors
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+	"github.com/ollama/ollama/x/imagegen/nn"
+)
+
+// WeightSource is an interface for loading weights.
+// Both ModelWeights (directory-based) and ManifestWeights (blob-based) implement this.
+type WeightSource interface {
+	GetTensor(name string) (*mlx.Array, error)
+	ListTensors() []string
+	HasTensor(name string) bool
+	Quantization() string // Returns "NVFP4", "INT4", "INT8", or ""
+	GroupSize() int       // Returns quantization group size, or 0 if not specified
+}
+
+// QuantizationParams returns groupSize, bits, mode for a quantization type.
+// MLX quantization modes:
+//   - "affine": scale + zero-point bias, group_size=32/64/128
+//   - "nvfp4": NVIDIA FP4 with E4M3 scales, group_size=16 (no bias)
+//   - "mxfp8": Microsoft MX FP8 with E4M3 scales, group_size=32 (no bias)
+func QuantizationParams(quantization string) (groupSize, bits int, mode string) {
+	switch strings.ToUpper(quantization) {
+	case "NVFP4":
+		// NVIDIA FP4: group_size=16, bits=4, E4M3 scales (no qbias)
+		return 16, 4, "nvfp4"
+	case "FP4", "Q4", "INT4":
+		// 4-bit quantization with affine mode (scale + qbias)
+		return 32, 4, "affine"
+	case "MXFP8":
+		// Microsoft MX FP8: group_size=32, bits=8, E4M3 scales (no qbias)
+		return 32, 8, "mxfp8"
+	case "FP8", "Q8", "INT8":
+		// 8-bit quantization with affine mode (default for quantized models)
+		return 64, 8, "affine"
+	case "":
+		return 0, 0, ""
+	default:
+		return 32, 8, "affine" // Default to affine
+	}
+}
+
+// Transformer allows structs to transform weight arrays before assignment.
+// Implement this to apply operations like transpose during loading.
+type Transformer interface {
+	Transform(field string, arr *mlx.Array) *mlx.Array
+}
+
+// LoadModule loads weights into a struct using reflection and struct tags.
+//
+// Struct tags use the format: `weight:"path[,optional]"`
+//   - path: the weight name suffix (appended to prefix)
+//   - optional: if present, missing weights don't cause errors
+//   - "-": skip this field entirely
+//   - no tag on struct pointer: recurse with current prefix
+//   - no tag on *mlx.Array: skip (computed fields don't need loading)
+//
+// For slices of struct pointers, the loader iterates with .0, .1, .2... suffixes.
+// The slice must be pre-allocated to the correct length.
+//
+// Example:
+//
+//	type Attention struct {
+//	    QProj *nn.Linear  `weight:"self_attn.q_proj"`
+//	    KProj *nn.Linear  `weight:"self_attn.k_proj"`
+//	    Cache *mlx.Array  // no tag = skipped (computed field)
+//	}
+//
+//	err := LoadModule(&attn, weights, "model.layers.0")
+func LoadModule(dst any, weights WeightSource, prefix string) error {
+	v := reflect.ValueOf(dst)
+	if v.Kind() != reflect.Ptr || v.IsNil() {
+		return fmt.Errorf("LoadModule: dst must be a non-nil pointer")
+	}
+	v = v.Elem()
+	if v.Kind() != reflect.Struct {
+		return fmt.Errorf("LoadModule: dst must be a pointer to struct, got %v", v.Kind())
+	}
+
+	var errs []string
+	loadStruct(v, weights, prefix, &errs, false)
+
+	if len(errs) > 0 {
+		return fmt.Errorf("LoadModule: missing weights:\n  %s", strings.Join(errs, "\n  "))
+	}
+	return nil
+}
+
+// loadStruct recursively loads weights into a struct value.
+func loadStruct(v reflect.Value, weights WeightSource, prefix string, errs *[]string, parentOptional bool) {
+	t := v.Type()
+
+	for i := 0; i < t.NumField(); i++ {
+		field := t.Field(i)
+		fieldVal := v.Field(i)
+
+		// Skip unexported fields
+		if !fieldVal.CanSet() {
+			continue
+		}
+
+		// Parse tag
+		tag, hasTag := field.Tag.Lookup("weight")
+		if tag == "-" {
+			continue
+		}
+
+		// Parse tag options
+		optional := parentOptional
+		weightPath := tag
+		if idx := strings.Index(tag, ","); idx != -1 {
+			weightPath = tag[:idx]
+			if strings.Contains(tag[idx+1:], "optional") {
+				optional = true
+			}
+		}
+
+		// Build full path
+		fullPath := joinPath(prefix, weightPath)
+
+		// For struct pointers without a tag, recurse with current prefix
+		if !hasTag && fieldVal.Kind() == reflect.Ptr {
+			elemType := fieldVal.Type().Elem()
+			if elemType.Kind() == reflect.Struct && elemType != reflect.TypeOf(mlx.Array{}) {
+				if fieldVal.IsNil() {
+					fieldVal.Set(reflect.New(elemType))
+				}
+				loadStruct(fieldVal.Elem(), weights, prefix, errs, optional)
+				continue
+			}
+		}
+
+		// Handle nn.LinearLayer interface fields specially
+		linearLayerType := reflect.TypeOf((*nn.LinearLayer)(nil)).Elem()
+		if field.Type == linearLayerType {
+			if !hasTag {
+				continue // no tag = skip
+			}
+			layer, err := LoadLinearLayer(weights, fullPath)
+			if err != nil {
+				if !optional {
+					*errs = append(*errs, fullPath+": "+err.Error())
+				}
+				continue
+			}
+			fieldVal.Set(reflect.ValueOf(layer))
+			continue
+		}
+
+		// Handle nn.MultiLinearLayer interface fields specially
+		multiLinearLayerType := reflect.TypeOf((*nn.MultiLinearLayer)(nil)).Elem()
+		if field.Type == multiLinearLayerType {
+			if !hasTag {
+				continue // no tag = skip
+			}
+			layer, err := LoadMultiLinearLayer(weights, fullPath)
+			if err != nil {
+				if !optional {
+					*errs = append(*errs, fullPath+": "+err.Error())
+				}
+				continue
+			}
+			fieldVal.Set(reflect.ValueOf(layer))
+			continue
+		}
+
+		// Handle by kind
+		switch fieldVal.Kind() {
+		case reflect.Ptr:
+			elemType := fieldVal.Type().Elem()
+
+			// *mlx.Array - load directly (but skip if no tag - computed fields)
+			if fieldVal.Type() == reflect.TypeOf((*mlx.Array)(nil)) {
+				if !hasTag {
+					continue // no tag on *mlx.Array = computed field, skip
+				}
+				arr, err := weights.GetTensor(fullPath)
+				if err != nil {
+					if !optional {
+						*errs = append(*errs, fullPath)
+					}
+					continue
+				}
+				// Transform before assigning if parent implements Transformer
+				if t, ok := v.Addr().Interface().(Transformer); ok {
+					arr = t.Transform(field.Name, arr)
+				}
+				fieldVal.Set(reflect.ValueOf(arr))
+				continue
+			}
+
+			// Pointer to struct - allocate and recurse
+			if elemType.Kind() == reflect.Struct {
+				if optional && !hasWeightsWithPrefix(weights, fullPath) {
+					continue
+				}
+				if fieldVal.IsNil() {
+					fieldVal.Set(reflect.New(elemType))
+				}
+				loadStruct(fieldVal.Elem(), weights, fullPath, errs, optional)
+			}
+
+		case reflect.Slice:
+			elemType := fieldVal.Type().Elem()
+			if elemType.Kind() == reflect.Ptr && elemType.Elem().Kind() == reflect.Struct {
+				loadSlice(fieldVal, weights, fullPath, errs)
+			}
+		}
+	}
+}
+
+// hasWeightsWithPrefix checks if any weights exist with the given prefix.
+func hasWeightsWithPrefix(weights WeightSource, prefix string) bool {
+	for _, name := range weights.ListTensors() {
+		if strings.HasPrefix(name, prefix+".") || name == prefix {
+			return true
+		}
+	}
+	return false
+}
+
+// loadSlice loads weights into each element of a slice of struct pointers.
+func loadSlice(v reflect.Value, weights WeightSource, prefix string, errs *[]string) {
+	elemStructType := v.Type().Elem().Elem()
+
+	for i := 0; i < v.Len(); i++ {
+		elem := v.Index(i)
+		if elem.IsNil() {
+			elem.Set(reflect.New(elemStructType))
+		}
+		loadStruct(elem.Elem(), weights, fmt.Sprintf("%s.%d", prefix, i), errs, false)
+	}
+}
+
+// joinPath joins path segments with dots, handling empty segments.
+func joinPath(prefix, suffix string) string {
+	if prefix == "" {
+		return suffix
+	}
+	if suffix == "" {
+		return prefix
+	}
+	return prefix + "." + suffix
+}
+
+// LoadMultiLinearLayer loads a per-head linear layer from weights.
+// Weight shape should be [num_heads, output_dims, input_dims].
+// If quantized, always dequantizes since batched quantized matmul isn't supported.
+func LoadMultiLinearLayer(weights WeightSource, path string) (nn.MultiLinearLayer, error) {
+	// Check if this is a quantized layer by looking for scale tensor
+	scalePath := path + ".weight_scale"
+	hasScale := weights.HasTensor(scalePath)
+
+	weight, err := weights.GetTensor(path + ".weight")
+	if err != nil {
+		return nil, fmt.Errorf("failed to load weight %s: %w", path, err)
+	}
+
+	if hasScale {
+		scales, err := weights.GetTensor(scalePath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load scales %s: %w", scalePath, err)
+		}
+
+		var qbiases *mlx.Array
+		qbiasPath := path + ".weight_qbias"
+		if weights.HasTensor(qbiasPath) {
+			qbiases, _ = weights.GetTensor(qbiasPath)
+		}
+
+		// Always dequantize for MultiLinear - no batched quantized matmul support
+		// Detect bits from tensor shapes (supports mixed-precision Q4/Q8)
+		weightShape := weight.Shape()
+		scalesShape := scales.Shape()
+		weightCols := int(weightShape[len(weightShape)-1])
+		scalesCols := int(scalesShape[len(scalesShape)-1])
+
+		// Detect quantization from tensor shapes
+		// groupSize = weightCols * packFactor / scalesCols
+		// Note: groupSize4 = 2 * groupSize8 always, so ambiguous cases need metadata
+		groupSize4 := weightCols * 8 / scalesCols
+		groupSize8 := weightCols * 4 / scalesCols
+
+		var bits, groupSize int
+		// Use metadata to help disambiguate when shapes are ambiguous
+		// (e.g., Q4 with group_size=64 has same shapes as Q8 with group_size=32)
+		quantType := strings.ToUpper(weights.Quantization())
+		isQ8Type := quantType == "Q8" || quantType == "FP8" || quantType == "INT8"
+
+		if groupSize4 == 32 {
+			// Unambiguous: Q4 with group_size=32
+			bits = 4
+			groupSize = 32
+		} else if groupSize8 == 64 {
+			// Unambiguous: Q8 with group_size=64
+			bits = 8
+			groupSize = 64
+		} else if groupSize4 == 64 && groupSize8 == 32 {
+			// Ambiguous: could be Q4/gs=64 or Q8/gs=32, use metadata
+			if isQ8Type {
+				bits = 8
+				groupSize = 32
+			} else {
+				bits = 4
+				groupSize = 64
+			}
+		} else {
+			// Fallback: use global quantization params
+			_, bits, _ = QuantizationParams(weights.Quantization())
+			packFactor := 32 / bits
+			groupSize = weightCols * packFactor / scalesCols
+		}
+		weight = mlx.Dequantize(weight, scales, qbiases, groupSize, bits, "affine")
+	}
+
+	return nn.NewMultiLinear(weight), nil
+}
+
+// LoadLinearLayer loads a linear layer from weights, automatically detecting if it's quantized.
+// If {path}.weight_scale exists, creates a QuantizedLinear layer (or dequantizes if no kernel support).
+func LoadLinearLayer(weights WeightSource, path string) (nn.LinearLayer, error) {
+	// Check if this is a quantized layer by looking for scale tensor
+	scalePath := path + ".weight_scale"
+	hasScale := weights.HasTensor(scalePath)
+	if hasScale {
+		weight, err := weights.GetTensor(path + ".weight")
+		if err != nil {
+			return nil, fmt.Errorf("failed to load quantized weight %s: %w", path, err)
+		}
+
+		scales, err := weights.GetTensor(scalePath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load scales %s: %w", scalePath, err)
+		}
+
+		// Bias is optional
+		var bias *mlx.Array
+		biasPath := path + ".bias"
+		if weights.HasTensor(biasPath) {
+			bias, _ = weights.GetTensor(biasPath)
+		}
+
+		var qbiases *mlx.Array
+		qbiasPath := path + ".weight_qbias"
+		if weights.HasTensor(qbiasPath) {
+			qbiases, _ = weights.GetTensor(qbiasPath)
+		}
+
+		// Detect bits from tensor shapes (supports mixed-precision Q4/Q8)
+		weightShape := weight.Shape()
+		scalesShape := scales.Shape()
+		weightCols := int(weightShape[len(weightShape)-1])
+		scalesCols := int(scalesShape[len(scalesShape)-1])
+
+		// Detect quantization from tensor shapes
+		// groupSize = weightCols * packFactor / scalesCols
+		// Note: groupSize4 = 2 * groupSize8 always, so ambiguous cases need metadata
+		groupSize4 := weightCols * 8 / scalesCols
+		groupSize8 := weightCols * 4 / scalesCols
+
+		var bits, groupSize int
+		mode := "affine"
+		// Use metadata to help disambiguate when shapes are ambiguous
+		// (e.g., Q4 with group_size=64 has same shapes as Q8 with group_size=32)
+		quantType := strings.ToUpper(weights.Quantization())
+		isQ8Type := quantType == "Q8" || quantType == "FP8" || quantType == "INT8"
+
+		if groupSize4 == 32 {
+			// Unambiguous: Q4 with group_size=32
+			bits = 4
+			groupSize = 32
+		} else if groupSize8 == 64 {
+			// Unambiguous: Q8 with group_size=64
+			bits = 8
+			groupSize = 64
+		} else if groupSize4 == 64 && groupSize8 == 32 {
+			// Ambiguous: could be Q4/gs=64 or Q8/gs=32, use metadata
+			if isQ8Type {
+				bits = 8
+				groupSize = 32
+			} else {
+				bits = 4
+				groupSize = 64
+			}
+		} else {
+			// Fallback: use global quantization params
+			_, bits, mode = QuantizationParams(weights.Quantization())
+			packFactor := 32 / bits
+			groupSize = weightCols * packFactor / scalesCols
+		}
+
+		// NVFP4 and MXFP8 don't have native quantized matmul kernels in MLX,
+		// so we always dequantize at load time. Affine modes (FP4, FP8) have kernel support.
+		if mlx.MetalIsAvailable() && mode != "nvfp4" && mode != "mxfp8" {
+			return &nn.QuantizedLinear{
+				Weight:    weight,
+				Scales:    scales,
+				QBiases:   qbiases,
+				Bias:      bias,
+				GroupSize: groupSize,
+				Bits:      bits,
+				Mode:      mode,
+			}, nil
+		}
+
+		dequantized := mlx.Dequantize(weight, scales, qbiases, groupSize, bits, mode)
+		return nn.NewLinear(dequantized, bias), nil
+	}
+
+	// Load as regular Linear
+	weight, err := weights.GetTensor(path + ".weight")
+	if err != nil {
+		return nil, fmt.Errorf("failed to load weight %s: %w", path, err)
+	}
+
+	// Bias is optional
+	var bias *mlx.Array
+	biasPath := path + ".bias"
+	if weights.HasTensor(biasPath) {
+		bias, _ = weights.GetTensor(biasPath)
+	}
+
+	return nn.NewLinear(weight, bias), nil
+}

x/imagegen/safetensors/safetensors.go

@@ -0,0 +1,318 @@
+package safetensors
+
+import (
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+// SafetensorHeader represents the JSON header of a safetensors file
+type SafetensorHeader map[string]TensorInfo
+
+// TensorInfo contains metadata about a tensor
+type TensorInfo struct {
+	Dtype       string  `json:"dtype"`
+	Shape       []int32 `json:"shape"`
+	DataOffsets [2]int  `json:"data_offsets"`
+}
+
+// parseSafetensorHeader reads only the JSON header from a safetensors file.
+func parseSafetensorHeader(path string) (SafetensorHeader, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open file: %w", err)
+	}
+	defer f.Close()
+
+	var headerSize uint64
+	if err := binary.Read(f, binary.LittleEndian, &headerSize); err != nil {
+		return nil, fmt.Errorf("failed to read header size: %w", err)
+	}
+
+	headerBytes := make([]byte, headerSize)
+	if _, err := f.Read(headerBytes); err != nil {
+		return nil, fmt.Errorf("failed to read header: %w", err)
+	}
+
+	var header SafetensorHeader
+	if err := json.Unmarshal(headerBytes, &header); err != nil {
+		return nil, fmt.Errorf("failed to parse header: %w", err)
+	}
+
+	delete(header, "__metadata__")
+	return header, nil
+}
+
+// dtypeFromString converts safetensors dtype string to mlx.Dtype
+func dtypeFromString(s string) mlx.Dtype {
+	switch strings.ToUpper(s) {
+	case "F32", "FLOAT32":
+		return mlx.DtypeFloat32
+	case "F16", "FLOAT16":
+		return mlx.DtypeFloat16
+	case "BF16", "BFLOAT16":
+		return mlx.DtypeBFloat16
+	case "I32", "INT32":
+		return mlx.DtypeInt32
+	case "I64", "INT64":
+		return mlx.DtypeInt64
+	case "U8", "UINT8":
+		return mlx.DtypeUint8
+	case "F8_E4M3", "F8_E5M2", "F8_E4M3FN", "F8_E5M2FNUZ":
+		return mlx.DtypeUint8 // FP8 types stored as raw uint8 bytes
+	default:
+		return mlx.DtypeFloat32
+	}
+}
+
+// ModelWeights manages weights from multiple safetensor files.
+type ModelWeights struct {
+	dir         string                          // Model directory
+	tensorFiles map[string]string               // tensor name -> file path
+	tensorInfo  map[string]TensorInfo           // tensor name -> metadata
+	nativeCache map[string]*mlx.SafetensorsFile // file path -> loaded native handle
+	cache       map[string]*mlx.Array           // tensor name -> array (after Load)
+}
+
+// LoadModelWeights scans safetensor files and builds a tensor index.
+// This only reads JSON headers, not tensor data.
+func LoadModelWeights(dir string) (*ModelWeights, error) {
+	mw := &ModelWeights{
+		dir:         dir,
+		tensorFiles: make(map[string]string),
+		tensorInfo:  make(map[string]TensorInfo),
+		nativeCache: make(map[string]*mlx.SafetensorsFile),
+	}
+
+	entries, err := os.ReadDir(dir)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read directory: %w", err)
+	}
+
+	for _, entry := range entries {
+		if strings.HasSuffix(entry.Name(), ".safetensors") {
+			path := filepath.Join(dir, entry.Name())
+
+			header, err := parseSafetensorHeader(path)
+			if err != nil {
+				return nil, fmt.Errorf("failed to parse %s: %w", entry.Name(), err)
+			}
+
+			for name, info := range header {
+				mw.tensorFiles[name] = path
+				mw.tensorInfo[name] = info
+			}
+		}
+	}
+
+	if len(mw.tensorFiles) == 0 {
+		return nil, fmt.Errorf("no safetensor files found in %s", dir)
+	}
+
+	return mw, nil
+}
+
+// LoadModelWeightsFromPaths loads weights from specific safetensor file paths.
+// Used for loading from blob storage where files are not in a directory.
+func LoadModelWeightsFromPaths(paths []string) (*ModelWeights, error) {
+	mw := &ModelWeights{
+		tensorFiles: make(map[string]string),
+		tensorInfo:  make(map[string]TensorInfo),
+		nativeCache: make(map[string]*mlx.SafetensorsFile),
+	}
+
+	for _, path := range paths {
+		header, err := parseSafetensorHeader(path)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse %s: %w", path, err)
+		}
+
+		for name, info := range header {
+			mw.tensorFiles[name] = path
+			mw.tensorInfo[name] = info
+		}
+	}
+
+	if len(mw.tensorFiles) == 0 {
+		return nil, fmt.Errorf("no tensors found in provided paths")
+	}
+
+	return mw, nil
+}
+
+// Load loads all tensors into cache with the specified dtype.
+// If dtype is 0, tensors are loaded in their original dtype.
+// Automatically uses streaming (memory-efficient) when dtype conversion is needed,
+// or native loading when tensors are already in the target dtype.
+func (mw *ModelWeights) Load(dtype mlx.Dtype) error {
+	if dtype == 0 {
+		return mw.loadNative()
+	}
+
+	// Check if any tensor needs conversion
+	needsConversion := false
+	for name := range mw.tensorFiles {
+		info := mw.tensorInfo[name]
+		if dtypeFromString(info.Dtype) != dtype {
+			needsConversion = true
+			break
+		}
+	}
+
+	if needsConversion {
+		return mw.loadStreaming(dtype)
+	}
+	return mw.loadNative()
+}
+
+// loadNative loads all tensors using the native memory-mapped loader.
+func (mw *ModelWeights) loadNative() error {
+	mw.cache = make(map[string]*mlx.Array)
+
+	fileToTensors := make(map[string][]string)
+	for name, path := range mw.tensorFiles {
+		fileToTensors[path] = append(fileToTensors[path], name)
+	}
+
+	for path, names := range fileToTensors {
+		native, err := mlx.LoadSafetensorsNative(path)
+		if err != nil {
+			return fmt.Errorf("failed to load %s: %w", path, err)
+		}
+
+		for _, name := range names {
+			arr := native.Get(name)
+			if arr == nil {
+				native.Free()
+				return fmt.Errorf("tensor %q not found in %s", name, path)
+			}
+			mw.cache[name] = arr
+		}
+
+		mw.nativeCache[path] = native
+	}
+
+	return nil
+}
+
+// loadStreaming loads tensors with dtype conversion.
+// Uses the same pattern as Python: replace each entry in the map after conversion,
+// so the original tensor loses its reference and can be freed.
+func (mw *ModelWeights) loadStreaming(dtype mlx.Dtype) error {
+	mw.cache = make(map[string]*mlx.Array)
+
+	fileToTensors := make(map[string][]string)
+	for name, path := range mw.tensorFiles {
+		fileToTensors[path] = append(fileToTensors[path], name)
+	}
+
+	for path, names := range fileToTensors {
+		native, err := mlx.LoadSafetensorsNative(path)
+		if err != nil {
+			return fmt.Errorf("failed to load %s: %w", path, err)
+		}
+
+		for _, name := range names {
+			src := native.Get(name)
+			if src == nil {
+				native.Free()
+				return fmt.Errorf("tensor %q not found in %s", name, path)
+			}
+
+			dst := mlx.AsType(src, dtype)
+			mlx.Eval(dst)
+			native.Set(name, dst)
+			mw.cache[name] = dst
+		}
+
+		native.Free()
+	}
+
+	return nil
+}
+
+// Get returns a tensor from cache. Call Load() first.
+func (mw *ModelWeights) Get(name string) (*mlx.Array, error) {
+	if mw.cache == nil {
+		return nil, fmt.Errorf("cache not initialized: call Load() first")
+	}
+	arr, ok := mw.cache[name]
+	if !ok {
+		return nil, fmt.Errorf("tensor %q not found in cache", name)
+	}
+	return arr, nil
+}
+
+// GetTensor loads a tensor using the native loader without caching.
+// For bulk loading, use Load() + Get() instead.
+func (mw *ModelWeights) GetTensor(name string) (*mlx.Array, error) {
+	if mw.cache != nil {
+		if arr, ok := mw.cache[name]; ok {
+			return arr, nil
+		}
+	}
+
+	path, ok := mw.tensorFiles[name]
+	if !ok {
+		return nil, fmt.Errorf("tensor %q not found", name)
+	}
+
+	native, ok := mw.nativeCache[path]
+	if !ok {
+		var err error
+		native, err = mlx.LoadSafetensorsNative(path)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load %s: %w", path, err)
+		}
+		mw.nativeCache[path] = native
+	}
+
+	return native.Get(name), nil
+}
+
+// GetTensorInfo returns metadata about a tensor without loading it.
+func (mw *ModelWeights) GetTensorInfo(name string) (TensorInfo, bool) {
+	info, ok := mw.tensorInfo[name]
+	return info, ok
+}
+
+// ListTensors returns all tensor names.
+func (mw *ModelWeights) ListTensors() []string {
+	names := make([]string, 0, len(mw.tensorFiles))
+	for name := range mw.tensorFiles {
+		names = append(names, name)
+	}
+	sort.Strings(names)
+	return names
+}
+
+// HasTensor checks if a tensor exists.
+func (mw *ModelWeights) HasTensor(name string) bool {
+	_, ok := mw.tensorFiles[name]
+	return ok
+}
+
+// Quantization returns empty string for directory-based weights (not quantized).
+func (mw *ModelWeights) Quantization() string {
+	return ""
+}
+
+// GroupSize returns 0 for directory-based weights (use default).
+func (mw *ModelWeights) GroupSize() int {
+	return 0
+}
+
+// ReleaseAll releases all cached native file handles.
+func (mw *ModelWeights) ReleaseAll() {
+	for path, native := range mw.nativeCache {
+		native.Free()
+		delete(mw.nativeCache, path)
+	}
+}
+

x/imagegen/safetensors/safetensors_test.go

@@ -0,0 +1,165 @@
+package safetensors
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/ollama/ollama/x/imagegen/mlx"
+)
+
+func TestLoadModelWeights(t *testing.T) {
+	// Skip if no model available
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	mw, err := LoadModelWeights(modelDir)
+	if err != nil {
+		t.Fatalf("LoadModelWeights: %v", err)
+	}
+	defer mw.ReleaseAll()
+
+	// Check we found tensors
+	tensors := mw.ListTensors()
+	if len(tensors) == 0 {
+		t.Fatal("no tensors found")
+	}
+	t.Logf("found %d tensors", len(tensors))
+
+	// Check HasTensor
+	if !mw.HasTensor(tensors[0]) {
+		t.Errorf("HasTensor(%q) = false", tensors[0])
+	}
+	if mw.HasTensor("nonexistent.weight") {
+		t.Error("HasTensor returned true for nonexistent tensor")
+	}
+}
+
+func TestGetTensor(t *testing.T) {
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	mw, err := LoadModelWeights(modelDir)
+	if err != nil {
+		t.Fatalf("LoadModelWeights: %v", err)
+	}
+	defer mw.ReleaseAll()
+
+	tensors := mw.ListTensors()
+	if len(tensors) == 0 {
+		t.Skip("no tensors")
+	}
+
+	// Load first tensor
+	arr, err := mw.GetTensor(tensors[0])
+	if err != nil {
+		t.Fatalf("GetTensor(%q): %v", tensors[0], err)
+	}
+
+	// Verify it has a shape
+	shape := arr.Shape()
+	if len(shape) == 0 {
+		t.Error("tensor has no shape")
+	}
+	t.Logf("%s: shape=%v dtype=%v", tensors[0], shape, arr.Dtype())
+}
+
+func TestLoadWithDtype(t *testing.T) {
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	mw, err := LoadModelWeights(modelDir)
+	if err != nil {
+		t.Fatalf("LoadModelWeights: %v", err)
+	}
+	defer mw.ReleaseAll()
+
+	// Load all tensors as bfloat16
+	if err := mw.Load(mlx.DtypeBFloat16); err != nil {
+		t.Fatalf("Load: %v", err)
+	}
+
+	// Get a tensor from cache
+	tensors := mw.ListTensors()
+	arr, err := mw.Get(tensors[0])
+	if err != nil {
+		t.Fatalf("Get: %v", err)
+	}
+
+	// Verify dtype (unless it was already bf16)
+	t.Logf("%s: dtype=%v", tensors[0], arr.Dtype())
+}
+
+func TestLookupTensor(t *testing.T) {
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	mw, err := LoadModelWeights(modelDir)
+	if err != nil {
+		t.Fatalf("LoadModelWeights: %v", err)
+	}
+	defer mw.ReleaseAll()
+
+	// HasTensor returns false for nonexistent
+	if mw.HasTensor("nonexistent") {
+		t.Error("HasTensor should return false for nonexistent")
+	}
+
+	// HasTensor returns true for existing tensor
+	tensors := mw.ListTensors()
+	if !mw.HasTensor(tensors[0]) {
+		t.Error("HasTensor should return true for existing tensor")
+	}
+}
+
+func TestParseSafetensorHeader(t *testing.T) {
+	modelDir := "../weights/gpt-oss-20b"
+	if _, err := os.Stat(modelDir); os.IsNotExist(err) {
+		t.Skip("model weights not available")
+	}
+
+	// Find a safetensors file
+	entries, err := os.ReadDir(modelDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var stFile string
+	for _, e := range entries {
+		if filepath.Ext(e.Name()) == ".safetensors" {
+			stFile = filepath.Join(modelDir, e.Name())
+			break
+		}
+	}
+	if stFile == "" {
+		t.Skip("no safetensors file found")
+	}
+
+	header, err := parseSafetensorHeader(stFile)
+	if err != nil {
+		t.Fatalf("parseSafetensorHeader: %v", err)
+	}
+
+	if len(header) == 0 {
+		t.Error("header is empty")
+	}
+
+	// Check a tensor has valid info
+	for name, info := range header {
+		if info.Dtype == "" {
+			t.Errorf("%s: empty dtype", name)
+		}
+		if len(info.Shape) == 0 {
+			t.Errorf("%s: empty shape", name)
+		}
+		break // just check one
+	}
+}