diff --git a/manifest/manifest.go b/manifest/manifest.go index d010ead73..f84270705 100644 --- a/manifest/manifest.go +++ b/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 { diff --git a/mlxrunner/client.go b/mlxrunner/client.go index 9787a01cc..8c0288ec0 100644 --- a/mlxrunner/client.go +++ b/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 } diff --git a/mlxrunner/model/root.go b/mlxrunner/model/root.go index 4a89ca302..166c95a0f 100644 --- a/mlxrunner/model/root.go +++ b/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.", diff --git a/mlxrunner/runner.go b/mlxrunner/runner.go index 1cbd57679..a129c15af 100644 --- a/mlxrunner/runner.go +++ b/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 } diff --git a/x/imagegen/manifest/manifest.go b/x/imagegen/manifest/manifest.go deleted file mode 100644 index 4de66644c..000000000 --- a/x/imagegen/manifest/manifest.go +++ /dev/null @@ -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// - 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 -} diff --git a/x/imagegen/manifest/manifest_test.go b/x/imagegen/manifest/manifest_test.go deleted file mode 100644 index 03361c6df..000000000 --- a/x/imagegen/manifest/manifest_test.go +++ /dev/null @@ -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) - } -}