mlxrunner: read model manifests through the manifest package

The runner and its weight loader read a model's manifest through
x/imagegen/manifest, the last piece of the removed image generation
engine. It was a hand-rolled copy of the manifest package: its own model
name parser with the default registry and namespace spelled out, its own
blob path builder, and re-spelled media types, plus a model_index.json
reader and other helpers that nothing has called since the engine went.

The manifest package gains the three lookups the runner needs, a config
layer by path, its contents, and the tensor layers, and ReadConfigJSON is
built on the second of them. The weight loader resolves the model name
with the shared parser, which fills in the same defaults the copy did, and
locates blobs with BlobsPath. The architectures' calls to read their
config.json compile unchanged. x/imagegen is gone.

manifest/manifest.go

@@ -40,22 +40,47 @@ func (m *Manifest) FileInfo() os.FileInfo {
 	return m.fi
 }
 
-// ReadConfigJSON reads and unmarshals a config layer as JSON.
-func (m *Manifest) ReadConfigJSON(configPath string, v any) error {
+// ConfigLayer returns the JSON layer stored under configPath.
+func (m *Manifest) ConfigLayer(configPath string) (Layer, bool) {
 	for _, layer := range m.Layers {
 		if layer.MediaType == "application/vnd.ollama.image.json" && layer.Name == configPath {
-			blobPath, err := BlobsPath(layer.Digest)
-			if err != nil {
-				return err
-			}
-			data, err := os.ReadFile(blobPath)
-			if err != nil {
-				return err
-			}
-			return json.Unmarshal(data, v)
+			return layer, true
 		}
 	}
-	return fmt.Errorf("config %q not found in manifest", configPath)
+	return Layer{}, false
+}
+
+// ReadConfig returns the contents of the JSON layer stored under configPath.
+func (m *Manifest) ReadConfig(configPath string) ([]byte, error) {
+	layer, ok := m.ConfigLayer(configPath)
+	if !ok {
+		return nil, fmt.Errorf("config %q not found in manifest", configPath)
+	}
+	blobPath, err := BlobsPath(layer.Digest)
+	if err != nil {
+		return nil, err
+	}
+	return os.ReadFile(blobPath)
+}
+
+// ReadConfigJSON reads and unmarshals a config layer as JSON.
+func (m *Manifest) ReadConfigJSON(configPath string, v any) error {
+	data, err := m.ReadConfig(configPath)
+	if err != nil {
+		return err
+	}
+	return json.Unmarshal(data, v)
+}
+
+// TensorLayers returns the layers holding tensor data.
+func (m *Manifest) TensorLayers() []Layer {
+	var layers []Layer
+	for _, layer := range m.Layers {
+		if layer.MediaType == MediaTypeImageTensor {
+			layers = append(layers, layer)
+		}
+	}
+	return layers
 }
 
 func (m *Manifest) Remove() error {

mlxrunner/client.go

@@ -25,9 +25,10 @@ import (
 	"github.com/ollama/ollama/envconfig"
 	"github.com/ollama/ollama/format"
 	"github.com/ollama/ollama/llm"
+	"github.com/ollama/ollama/manifest"
 	"github.com/ollama/ollama/ml"
 	"github.com/ollama/ollama/mlx"
-	"github.com/ollama/ollama/x/imagegen/manifest"
+	"github.com/ollama/ollama/types/model"
 )
 
 // Client wraps an MLX runner subprocess to implement llm.LlamaServer for LLM models.
@@ -62,11 +63,15 @@ func NewClient(modelName string, softContextLength int) (*Client, error) {
 		client:            http.DefaultClient,
 	}
 
-	modelManifest, err := manifest.LoadManifest(modelName)
+	modelManifest, err := manifest.ParseNamedManifest(model.ParseName(modelName))
 	if err != nil {
 		return nil, err
 	}
-	c.memory.Store(uint64(modelManifest.TotalTensorSize()))
+	var tensorBytes int64
+	for _, layer := range modelManifest.TensorLayers() {
+		tensorBytes += layer.Size
+	}
+	c.memory.Store(uint64(tensorBytes))
 
 	return c, nil
 }

mlxrunner/model/root.go

@@ -10,8 +10,8 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/ollama/ollama/manifest"
 	modeltypes "github.com/ollama/ollama/types/model"
-	"github.com/ollama/ollama/x/imagegen/manifest"
 )
 
 // TensorQuantInfo describes per-tensor quantization metadata.
@@ -20,9 +20,9 @@ type TensorQuantInfo struct {
 	GroupSize int
 }
 
-// Root wraps a ModelManifest with pre-scanned quantization metadata.
+// Root wraps a model's manifest with pre-scanned quantization metadata.
 type Root struct {
-	Manifest *manifest.ModelManifest
+	Manifest *manifest.Manifest
 	Draft    *modeltypes.Draft
 
 	// Backwards-compatible model-level quant metadata (first tensor blob).
@@ -36,7 +36,7 @@ type Root struct {
 // Open loads a manifest for the given model name and scans tensor blobs for
 // quantization metadata.
 func Open(modelName string) (*Root, error) {
-	m, err := manifest.LoadManifest(modelName)
+	m, err := manifest.ParseNamedManifest(modeltypes.ParseName(modelName))
 	if err != nil {
 		return nil, err
 	}
@@ -47,8 +47,11 @@ func Open(modelName string) (*Root, error) {
 	}
 	root.Draft = readDraftConfig(m)
 
-	for _, layer := range m.GetTensorLayers("") {
-		blobPath := m.BlobPath(layer.Digest)
+	for _, layer := range m.TensorLayers() {
+		blobPath, err := manifest.BlobsPath(layer.Digest)
+		if err != nil {
+			return nil, err
+		}
 
 		infos, blobQuantType, blobGroupSize, err := readBlobTensorQuantInfo(blobPath)
 		if err != nil {
@@ -71,12 +74,16 @@ func Open(modelName string) (*Root, error) {
 	return root, nil
 }
 
-func readDraftConfig(m *manifest.ModelManifest) *modeltypes.Draft {
-	if m == nil || m.Manifest == nil || m.Manifest.Config.Digest == "" {
+func readDraftConfig(m *manifest.Manifest) *modeltypes.Draft {
+	if m == nil || m.Config.Digest == "" {
 		return nil
 	}
 
-	data, err := os.ReadFile(m.BlobPath(m.Manifest.Config.Digest))
+	configPath, err := manifest.BlobsPath(m.Config.Digest)
+	if err != nil {
+		return nil
+	}
+	data, err := os.ReadFile(configPath)
 	if err != nil {
 		return nil
 	}
@@ -89,7 +96,7 @@ func readDraftConfig(m *manifest.ModelManifest) *modeltypes.Draft {
 		return cfg.Draft
 	}
 
-	if m.GetConfigLayer("draft/config.json") != nil {
+	if _, ok := m.ConfigLayer("draft/config.json"); ok {
 		return &modeltypes.Draft{
 			ModelFormat:  "safetensors",
 			TensorPrefix: "draft.",

mlxrunner/runner.go

@@ -13,6 +13,7 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/manifest"
 	"github.com/ollama/ollama/mlx"
 	"github.com/ollama/ollama/mlx/mlxthread"
 	"github.com/ollama/ollama/mlxrunner/batch"
@@ -235,12 +236,15 @@ func loadTensorsFromManifest(root *model.Root) (map[string]*mlx.Array, error) {
 	// Phase 1: Load all tensors raw from all blobs
 	rawTensors := make(map[string]*mlx.Array)
 	seen := make(map[string]bool)
-	for _, layer := range root.Manifest.GetTensorLayers("") {
+	for _, layer := range root.Manifest.TensorLayers() {
 		if seen[layer.Digest] {
 			continue
 		}
 		seen[layer.Digest] = true
-		blobPath := root.Manifest.BlobPath(layer.Digest)
+		blobPath, err := manifest.BlobsPath(layer.Digest)
+		if err != nil {
+			return nil, err
+		}
 		for name, arr := range mlx.Load(blobPath) {
 			rawTensors[name] = arr
 		}

x/imagegen/manifest/manifest.go

@@ -1,307 +0,0 @@
-package manifest
-
-import (
-	"encoding/binary"
-	"encoding/json"
-	"fmt"
-	"io"
-	"os"
-	"path/filepath"
-	"sort"
-	"strings"
-
-	"github.com/ollama/ollama/envconfig"
-)
-
-// ManifestLayer represents a layer in the manifest.
-type ManifestLayer struct {
-	MediaType string `json:"mediaType"`
-	Digest    string `json:"digest"`
-	Size      int64  `json:"size"`
-	Name      string `json:"name,omitempty"` // Path-style name: "component/tensor" or "path/to/config.json"
-}
-
-// Manifest represents the manifest JSON structure.
-type Manifest struct {
-	SchemaVersion int             `json:"schemaVersion"`
-	MediaType     string          `json:"mediaType"`
-	Config        ManifestLayer   `json:"config"`
-	Layers        []ManifestLayer `json:"layers"`
-}
-
-// ModelManifest holds a parsed manifest with helper methods.
-type ModelManifest struct {
-	Manifest *Manifest
-	BlobDir  string
-}
-
-func DefaultBlobDir() string {
-	return filepath.Join(envconfig.Models(), "blobs")
-}
-
-// DefaultManifestDir returns the manifest storage directory.
-// Respects OLLAMA_MODELS.
-
-func DefaultManifestDir() string {
-	return filepath.Join(envconfig.Models(), "manifests")
-}
-
-// LoadManifest loads a manifest for the given model name.
-// Model name format: "modelname" or "modelname:tag" or "host/namespace/name:tag"
-func LoadManifest(modelName string) (*ModelManifest, error) {
-	manifestPath := resolveManifestPath(modelName)
-
-	data, err := os.ReadFile(manifestPath)
-	if err != nil {
-		return nil, fmt.Errorf("read manifest: %w", err)
-	}
-
-	var manifest Manifest
-	if err := json.Unmarshal(data, &manifest); err != nil {
-		return nil, fmt.Errorf("parse manifest: %w", err)
-	}
-
-	return &ModelManifest{
-		Manifest: &manifest,
-		BlobDir:  DefaultBlobDir(),
-	}, nil
-}
-
-// resolveManifestPath converts a model name to a manifest file path.
-func resolveManifestPath(modelName string) string {
-	// Parse model name into components
-	// Default: registry.ollama.ai/library/<name>/<tag>
-	host := "registry.ollama.ai"
-	namespace := "library"
-	name := modelName
-	tag := "latest"
-
-	// Handle explicit tag
-	if idx := strings.LastIndex(name, ":"); idx != -1 {
-		tag = name[idx+1:]
-		name = name[:idx]
-	}
-
-	// Handle full path like "host/namespace/name"
-	parts := strings.Split(name, "/")
-	switch len(parts) {
-	case 3:
-		host = parts[0]
-		namespace = parts[1]
-		name = parts[2]
-	case 2:
-		namespace = parts[0]
-		name = parts[1]
-	}
-
-	return filepath.Join(DefaultManifestDir(), host, namespace, name, tag)
-}
-
-// BlobPath returns the full path to a blob given its digest.
-func (m *ModelManifest) BlobPath(digest string) string {
-	// Convert "sha256:abc123" to "sha256-abc123"
-	blobName := strings.Replace(digest, ":", "-", 1)
-	return filepath.Join(m.BlobDir, blobName)
-}
-
-// GetTensorLayers returns tensor layers, optionally filtered by component.
-// If component is empty, returns all tensor layers (for LLM models).
-// If component is specified (e.g., "text_encoder", "transformer", "vae"),
-// returns only layers with that prefix.
-func (m *ModelManifest) GetTensorLayers(component string) []ManifestLayer {
-	var layers []ManifestLayer
-	for _, layer := range m.Manifest.Layers {
-		if layer.MediaType != "application/vnd.ollama.image.tensor" {
-			continue
-		}
-		if component == "" || strings.HasPrefix(layer.Name, component+"/") {
-			layers = append(layers, layer)
-		}
-	}
-	return layers
-}
-
-// GetConfigLayer returns the config layer for a given path.
-func (m *ModelManifest) GetConfigLayer(configPath string) *ManifestLayer {
-	for _, layer := range m.Manifest.Layers {
-		if layer.MediaType == "application/vnd.ollama.image.json" && layer.Name == configPath {
-			return &layer
-		}
-	}
-	return nil
-}
-
-// ReadConfig reads and returns the content of a config file.
-func (m *ModelManifest) ReadConfig(configPath string) ([]byte, error) {
-	layer := m.GetConfigLayer(configPath)
-	if layer == nil {
-		return nil, fmt.Errorf("config %q not found in manifest", configPath)
-	}
-
-	blobPath := m.BlobPath(layer.Digest)
-	return os.ReadFile(blobPath)
-}
-
-// ReadConfigJSON reads and unmarshals a config file.
-func (m *ModelManifest) ReadConfigJSON(configPath string, v any) error {
-	data, err := m.ReadConfig(configPath)
-	if err != nil {
-		return err
-	}
-	return json.Unmarshal(data, v)
-}
-
-// OpenBlob opens a blob for reading.
-func (m *ModelManifest) OpenBlob(digest string) (io.ReadCloser, error) {
-	return os.Open(m.BlobPath(digest))
-}
-
-// HasTensorLayers returns true if the manifest has any tensor layers.
-func (m *ModelManifest) HasTensorLayers() bool {
-	for _, layer := range m.Manifest.Layers {
-		if layer.MediaType == "application/vnd.ollama.image.tensor" {
-			return true
-		}
-	}
-	return false
-}
-
-// TotalTensorSize returns the total size in bytes of all tensor layers.
-func (m *ModelManifest) TotalTensorSize() int64 {
-	var total int64
-	for _, layer := range m.Manifest.Layers {
-		if layer.MediaType == "application/vnd.ollama.image.tensor" {
-			total += layer.Size
-		}
-	}
-	return total
-}
-
-// ModelInfo contains metadata about an image generation model.
-type ModelInfo struct {
-	Architecture   string
-	ParameterCount int64
-	Quantization   string
-}
-
-// GetModelInfo returns metadata about an image generation model.
-func GetModelInfo(modelName string) (*ModelInfo, error) {
-	manifest, err := LoadManifest(modelName)
-	if err != nil {
-		return nil, fmt.Errorf("failed to load manifest: %w", err)
-	}
-
-	info := &ModelInfo{}
-
-	// Read model_index.json for architecture, parameter count, and quantization
-	if data, err := manifest.ReadConfig("model_index.json"); err == nil {
-		var index struct {
-			Architecture   string `json:"architecture"`
-			ParameterCount int64  `json:"parameter_count"`
-			Quantization   string `json:"quantization"`
-		}
-		if json.Unmarshal(data, &index) == nil {
-			info.Architecture = index.Architecture
-			info.ParameterCount = index.ParameterCount
-			info.Quantization = index.Quantization
-		}
-	}
-
-	// Fallback: detect quantization from first tensor blob's __metadata__
-	if info.Quantization == "" {
-		info.Quantization = detectQuantizationFromBlobs(manifest)
-	}
-	if info.Quantization == "" {
-		info.Quantization = "BF16"
-	}
-
-	// Fallback: estimate parameter count if not in config
-	if info.ParameterCount == 0 {
-		var totalSize int64
-		for _, layer := range manifest.Manifest.Layers {
-			if layer.MediaType == "application/vnd.ollama.image.tensor" {
-				totalSize += layer.Size
-			}
-		}
-		// Assume BF16 (2 bytes/param) as rough estimate
-		info.ParameterCount = totalSize / 2
-	}
-
-	return info, nil
-}
-
-// detectQuantizationFromBlobs reads __metadata__ from the first tensor blob
-// to detect quantization type.
-func detectQuantizationFromBlobs(manifest *ModelManifest) string {
-	for _, layer := range manifest.Manifest.Layers {
-		if layer.MediaType != "application/vnd.ollama.image.tensor" {
-			continue
-		}
-		data, err := readBlobHeader(manifest.BlobPath(layer.Digest))
-		if err != nil {
-			continue
-		}
-		var header map[string]json.RawMessage
-		if json.Unmarshal(data, &header) != nil {
-			continue
-		}
-		if metaRaw, ok := header["__metadata__"]; ok {
-			var meta map[string]string
-			if json.Unmarshal(metaRaw, &meta) == nil {
-				if qt, ok := meta["quant_type"]; ok && qt != "" {
-					return strings.ToUpper(qt)
-				}
-			}
-		}
-		// Only check the first tensor blob
-		break
-	}
-	return ""
-}
-
-// ParseBlobTensorNames reads a safetensors blob and returns all "main" tensor names.
-// Filters out __metadata__, .scale, and .bias entries to return only primary weight tensors.
-func ParseBlobTensorNames(path string) ([]string, error) {
-	data, err := readBlobHeader(path)
-	if err != nil {
-		return nil, err
-	}
-
-	var header map[string]json.RawMessage
-	if err := json.Unmarshal(data, &header); err != nil {
-		return nil, err
-	}
-
-	var names []string
-	for k := range header {
-		if k == "__metadata__" || strings.HasSuffix(k, ".scale") || strings.HasSuffix(k, ".bias") {
-			continue
-		}
-		names = append(names, k)
-	}
-
-	sort.Strings(names)
-	return names, nil
-}
-
-// readBlobHeader reads the JSON header bytes from a safetensors blob file.
-func readBlobHeader(path string) ([]byte, error) {
-	f, err := os.Open(path)
-	if err != nil {
-		return nil, err
-	}
-	defer f.Close()
-
-	var headerSize uint64
-	if err := binary.Read(f, binary.LittleEndian, &headerSize); err != nil {
-		return nil, err
-	}
-	if headerSize > 1024*1024 {
-		return nil, fmt.Errorf("header too large: %d", headerSize)
-	}
-	data := make([]byte, headerSize)
-	if _, err := io.ReadFull(f, data); err != nil {
-		return nil, err
-	}
-	return data, nil
-}

x/imagegen/manifest/manifest_test.go

@@ -1,57 +0,0 @@
-package manifest
-
-import (
-	"path/filepath"
-	"testing"
-)
-
-func TestTotalTensorSize(t *testing.T) {
-	m := &ModelManifest{
-		Manifest: &Manifest{
-			Layers: []ManifestLayer{
-				{MediaType: "application/vnd.ollama.image.tensor", Size: 1000},
-				{MediaType: "application/vnd.ollama.image.tensor", Size: 2000},
-				{MediaType: "application/vnd.ollama.image.json", Size: 500}, // not a tensor
-				{MediaType: "application/vnd.ollama.image.tensor", Size: 3000},
-			},
-		},
-	}
-
-	got := m.TotalTensorSize()
-	want := int64(6000)
-	if got != want {
-		t.Errorf("TotalTensorSize() = %d, want %d", got, want)
-	}
-}
-
-func TestTotalTensorSizeEmpty(t *testing.T) {
-	m := &ModelManifest{
-		Manifest: &Manifest{
-			Layers: []ManifestLayer{},
-		},
-	}
-
-	if got := m.TotalTensorSize(); got != 0 {
-		t.Errorf("TotalTensorSize() = %d, want 0", got)
-	}
-}
-
-func TestManifestAndBlobDirsRespectOLLAMAModels(t *testing.T) {
-	modelsDir := filepath.Join(t.TempDir(), "models")
-
-	// Simulate packaged/systemd environment
-	t.Setenv("OLLAMA_MODELS", modelsDir)
-	t.Setenv("HOME", "/usr/share/ollama")
-
-	// Manifest dir must respect OLLAMA_MODELS
-	wantManifest := filepath.Join(modelsDir, "manifests")
-	if got := DefaultManifestDir(); got != wantManifest {
-		t.Fatalf("DefaultManifestDir() = %q, want %q", got, wantManifest)
-	}
-
-	// Blob dir must respect OLLAMA_MODELS
-	wantBlobs := filepath.Join(modelsDir, "blobs")
-	if got := DefaultBlobDir(); got != wantBlobs {
-		t.Fatalf("DefaultBlobDir() = %q, want %q", got, wantBlobs)
-	}
-}