Remove SafeTensors source weights after a converted GGUF loads

internal/convert/source_weights.go

@@ -0,0 +1,77 @@
+package convert
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+// sourceWeightIndexName is the shard index that names the SafeTensors files
+// llama.cpp reads; it is only useful alongside those shards.
+const sourceWeightIndexName = "model.safetensors.index.json"
+
+// ConvertedSourceWeightFiles returns the top-level SafeTensors files the GGUF
+// converter consumed in modelDir, plus their shard index. llama.cpp's
+// convert_hf_to_gguf.py reads every "model*.safetensors" file at the top
+// level, so companions such as head.safetensors or adapter directories are
+// not listed and stay in place.
+func ConvertedSourceWeightFiles(modelDir string) ([]string, error) {
+	entries, err := os.ReadDir(modelDir)
+	if err != nil {
+		return nil, err
+	}
+	var names []string
+	for _, entry := range entries {
+		if !entry.Type().IsRegular() {
+			continue
+		}
+		name := entry.Name()
+		lower := strings.ToLower(name)
+		if lower == sourceWeightIndexName || (strings.HasPrefix(lower, "model") && strings.HasSuffix(lower, ".safetensors")) {
+			names = append(names, name)
+		}
+	}
+	return names, nil
+}
+
+// RemoveConvertedSourceWeights deletes the SafeTensors weights in modelDir
+// once a converted GGUF serves the model, so the checkpoint is not stored
+// twice. It keeps them when no weight GGUF exists, or when a vision model is
+// still missing its mmproj, since a later load converts the projector from
+// the same weights. It returns the removed file names and their total size.
+func RemoveConvertedSourceWeights(modelDir string) ([]string, int64, error) {
+	if _, ok, err := FindGGUFForDType(modelDir, ""); err != nil || !ok {
+		return nil, 0, err
+	}
+	if hasVisionConfig(modelDir) {
+		if _, ok, err := FindMMProjForDType(modelDir, ""); err != nil || !ok {
+			return nil, 0, err
+		}
+	}
+	names, err := ConvertedSourceWeightFiles(modelDir)
+	if err != nil {
+		return nil, 0, err
+	}
+	var (
+		removed []string
+		freed   int64
+		errs    []error
+	)
+	for _, name := range names {
+		path := filepath.Join(modelDir, name)
+		info, statErr := os.Stat(path)
+		if err := os.Remove(path); err != nil {
+			if !errors.Is(err, os.ErrNotExist) {
+				errs = append(errs, fmt.Errorf("removing %s: %w", name, err))
+			}
+			continue
+		}
+		removed = append(removed, name)
+		if statErr == nil {
+			freed += info.Size()
+		}
+	}
+	return removed, freed, errors.Join(errs...)
+}

internal/convert/source_weights_test.go

@@ -0,0 +1,111 @@
+package convert
+
+import (
+	"os"
+	"path/filepath"
+	"reflect"
+	"sort"
+	"testing"
+)
+
+func writeSourceWeightFixture(t *testing.T, dir string, files map[string]string) {
+	t.Helper()
+	for name, content := range files {
+		path := filepath.Join(dir, name)
+		if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
+			t.Fatalf("mkdir %s: %v", name, err)
+		}
+		if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
+			t.Fatalf("write %s: %v", name, err)
+		}
+	}
+}
+
+func remainingFiles(t *testing.T, dir string) []string {
+	t.Helper()
+	var names []string
+	err := filepath.WalkDir(dir, func(path string, d os.DirEntry, err error) error {
+		if err != nil || d.IsDir() {
+			return err
+		}
+		rel, _ := filepath.Rel(dir, path)
+		names = append(names, filepath.ToSlash(rel))
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("walk: %v", err)
+	}
+	sort.Strings(names)
+	return names
+}
+
+func TestRemoveConvertedSourceWeightsKeepsCompanions(t *testing.T) {
+	dir := filepath.Join(t.TempDir(), "JEV-9B")
+	writeSourceWeightFixture(t, dir, map[string]string{
+		"config.json":                       `{"architectures":["Qwen3ForCausalLM"]}`,
+		"model-00001-of-00002.safetensors":  "aaaa",
+		"model-00002-of-00002.safetensors":  "bb",
+		"model.safetensors.index.json":      "{}",
+		"head.safetensors":                  "head",
+		"adapter/adapter_model.safetensors": "lora",
+		"JEV-9B-q8_0.gguf":                  "gguf",
+	})
+
+	removed, freed, err := RemoveConvertedSourceWeights(dir)
+	if err != nil {
+		t.Fatalf("RemoveConvertedSourceWeights: %v", err)
+	}
+	wantRemoved := []string{"model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors", "model.safetensors.index.json"}
+	if !reflect.DeepEqual(removed, wantRemoved) {
+		t.Fatalf("removed = %v, want %v", removed, wantRemoved)
+	}
+	if freed != 8 {
+		t.Fatalf("freed = %d, want 8", freed)
+	}
+	want := []string{"JEV-9B-q8_0.gguf", "adapter/adapter_model.safetensors", "config.json", "head.safetensors"}
+	if got := remainingFiles(t, dir); !reflect.DeepEqual(got, want) {
+		t.Fatalf("remaining = %v, want %v", got, want)
+	}
+}
+
+func TestRemoveConvertedSourceWeightsKeepsWeightsWithoutGGUF(t *testing.T) {
+	dir := filepath.Join(t.TempDir(), "demo")
+	writeSourceWeightFixture(t, dir, map[string]string{
+		"config.json":       `{"architectures":["LlamaForCausalLM"]}`,
+		"model.safetensors": "weights",
+	})
+
+	removed, _, err := RemoveConvertedSourceWeights(dir)
+	if err != nil {
+		t.Fatalf("RemoveConvertedSourceWeights: %v", err)
+	}
+	if len(removed) != 0 {
+		t.Fatalf("removed = %v, want none", removed)
+	}
+}
+
+func TestRemoveConvertedSourceWeightsKeepsVisionWeightsUntilMMProjExists(t *testing.T) {
+	dir := filepath.Join(t.TempDir(), "vlm")
+	writeSourceWeightFixture(t, dir, map[string]string{
+		"config.json":       `{"architectures":["Qwen2VLForConditionalGeneration"],"vision_config":{}}`,
+		"model.safetensors": "weights",
+		"vlm-f16.gguf":      "gguf",
+	})
+
+	removed, _, err := RemoveConvertedSourceWeights(dir)
+	if err != nil {
+		t.Fatalf("RemoveConvertedSourceWeights: %v", err)
+	}
+	if len(removed) != 0 {
+		t.Fatalf("removed = %v before mmproj exists, want none", removed)
+	}
+
+	writeSourceWeightFixture(t, dir, map[string]string{"mmproj-vlm-f16.gguf": "mmproj"})
+	removed, _, err = RemoveConvertedSourceWeights(dir)
+	if err != nil {
+		t.Fatalf("RemoveConvertedSourceWeights: %v", err)
+	}
+	if !reflect.DeepEqual(removed, []string{"model.safetensors"}) {
+		t.Fatalf("removed = %v after mmproj exists, want [model.safetensors]", removed)
+	}
+}

internal/inference/engine.go

@@ -134,6 +134,28 @@ func loadEngineWithProgressMode(modelDir string, lm *model.LocalModel, progress
 		return "", nil
 	}
 
+	// startConverted serves a GGUF this load just converted. A GGUF that
+	// llama-server rejects is removed so the next load converts again; once
+	// one loads, the SafeTensors source is no longer needed.
+	startConverted := func(ggufPath string) (Engine, error) {
+		mmproj, err := resolveMMProj()
+		if err != nil {
+			return nil, err
+		}
+		var eng Engine
+		if embedding {
+			eng, err = newLlamaEmbeddingEngine(ggufPath, lm.FullName(), verbose, progress, numCtx, numParallel, nGPULayers, cacheTypeK, cacheTypeV, mmproj)
+		} else {
+			eng, err = newLlamaSpeculativeEngine(ggufPath, lm.FullName(), verbose, progress, numCtx, numParallel, nGPULayers, cacheTypeK, cacheTypeV, speculative, mmproj)
+		}
+		if err != nil {
+			removeConvertedGGUFIfInvalid(ggufPath, err)
+			return nil, err
+		}
+		go removeConvertedSourceWeights(modelDir)
+		return eng, nil
+	}
+
 	if normalizedDType != "" {
 		if ggufPath, ok, err := convert.FindGGUFForDType(modelDir, normalizedDType); err != nil {
 			return nil, err
@@ -157,24 +179,7 @@ func loadEngineWithProgressMode(modelDir string, lm *model.LocalModel, progress
 			if err != nil {
 				return nil, fmt.Errorf("auto-converting SafeTensors to GGUF: %w", err)
 			}
-			mmproj, err := resolveMMProj()
-			if err != nil {
-				return nil, err
-			}
-			if embedding {
-				eng, err := newLlamaEmbeddingEngine(ggufPath, lm.FullName(), verbose, progress, numCtx, numParallel, nGPULayers, cacheTypeK, cacheTypeV, mmproj)
-				if err != nil {
-					removeConvertedGGUFIfInvalid(ggufPath, err)
-					return nil, err
-				}
-				return eng, nil
-			}
-			eng, err := newLlamaSpeculativeEngine(ggufPath, lm.FullName(), verbose, progress, numCtx, numParallel, nGPULayers, cacheTypeK, cacheTypeV, speculative, mmproj)
-			if err != nil {
-				removeConvertedGGUFIfInvalid(ggufPath, err)
-				return nil, err
-			}
-			return eng, nil
+			return startConverted(ggufPath)
 		}
 	}
 
@@ -210,24 +215,7 @@ func loadEngineWithProgressMode(modelDir string, lm *model.LocalModel, progress
 		if err != nil {
 			return nil, fmt.Errorf("auto-converting SafeTensors to GGUF: %w", err)
 		}
-		mmproj, err := resolveMMProj()
-		if err != nil {
-			return nil, err
-		}
-		if embedding {
-			eng, err := newLlamaEmbeddingEngine(ggufPath, lm.FullName(), verbose, progress, numCtx, numParallel, nGPULayers, cacheTypeK, cacheTypeV, mmproj)
-			if err != nil {
-				removeConvertedGGUFIfInvalid(ggufPath, err)
-				return nil, err
-			}
-			return eng, nil
-		}
-		eng, err := newLlamaSpeculativeEngine(ggufPath, lm.FullName(), verbose, progress, numCtx, numParallel, nGPULayers, cacheTypeK, cacheTypeV, speculative, mmproj)
-		if err != nil {
-			removeConvertedGGUFIfInvalid(ggufPath, err)
-			return nil, err
-		}
-		return eng, nil
+		return startConverted(ggufPath)
 
 	default:
 		return nil, fmt.Errorf("%w: %s", ErrUnsupportedFormat, format)
@@ -300,6 +288,23 @@ func convertSafeTensors(modelDir string, progress ConvertProgressFunc, dtype str
 	return path, nil
 }
 
+// removeConvertedSourceWeights deletes the SafeTensors a loaded GGUF was
+// converted from and rewrites the manifest to list the GGUF instead. It runs
+// in the background because refreshing the manifest hashes the new GGUF.
+func removeConvertedSourceWeights(modelDir string) {
+	removed, freed, err := convert.RemoveConvertedSourceWeights(modelDir)
+	if err != nil {
+		log.Printf("warning: could not remove converted SafeTensors in %s: %v", modelDir, err)
+	}
+	if len(removed) == 0 {
+		return
+	}
+	log.Printf("CONVERT: removed %d converted SafeTensors file(s) from %s, freed %d bytes", len(removed), modelDir, freed)
+	if err := model.RefreshManifestFiles(modelDir); err != nil {
+		log.Printf("warning: could not refresh manifest in %s: %v", modelDir, err)
+	}
+}
+
 func removeConvertedGGUFIfInvalid(ggufPath string, err error) {
 	if !shouldRemoveConvertedGGUF(err) {
 		log.Printf("keeping converted GGUF after llama-server load failure: %s", ggufPath)

internal/model/files.go

@@ -192,3 +192,21 @@ func fileSHA256(filePath string) (string, error) {
 	}
 	return hex.EncodeToString(sum.Sum(nil)), nil
 }
+
+// RefreshManifestFiles rescans modelDir after files were added or removed in
+// place, such as a GGUF conversion replacing its SafeTensors source, and
+// rewrites the manifest's file list, size and format to match the disk.
+func RefreshManifestFiles(modelDir string) error {
+	lm, err := LoadManifestInDir(modelDir)
+	if err != nil {
+		return err
+	}
+	if _, err := EnsureLocalModelFiles(modelDir, lm); err != nil {
+		return err
+	}
+	lm.Size = localModelEntriesSize(lm.FileEntries)
+	if format := DetectFormat(lm.Files); format != FormatUnknown {
+		lm.Format = format
+	}
+	return SaveManifestInDir(modelDir, lm)
+}

internal/model/files_test.go

@@ -3,6 +3,7 @@ package model
 import (
 	"os"
 	"path/filepath"
+	"reflect"
 	"testing"
 )
 
@@ -67,3 +68,49 @@ func TestEnsureLocalModelFiles_BackfillsMissingEntries(t *testing.T) {
 		t.Fatal("changed = true on second pass, want false")
 	}
 }
+
+func TestRefreshManifestFiles_ReplacesConvertedSourceWithGGUF(t *testing.T) {
+	baseDir := t.TempDir()
+	modelDir := ModelDir(baseDir, "Acme", "demo")
+	if err := os.MkdirAll(modelDir, 0o755); err != nil {
+		t.Fatalf("mkdir model dir: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(modelDir, "config.json"), []byte(`{}`), 0o644); err != nil {
+		t.Fatalf("write config.json: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(modelDir, "demo-f16.gguf"), []byte("gguf-data"), 0o644); err != nil {
+		t.Fatalf("write gguf: %v", err)
+	}
+	lm := &LocalModel{
+		Namespace: "Acme",
+		Name:      "demo",
+		Format:    FormatSafeTensors,
+		Size:      100,
+		Files:     []string{"config.json", "model.safetensors"},
+		FileEntries: []LocalModelFile{
+			{Path: "config.json", Size: 2},
+			{Path: "model.safetensors", Size: 98},
+		},
+	}
+	if err := SaveManifest(baseDir, lm); err != nil {
+		t.Fatalf("save manifest: %v", err)
+	}
+
+	if err := RefreshManifestFiles(modelDir); err != nil {
+		t.Fatalf("RefreshManifestFiles: %v", err)
+	}
+
+	got, err := LoadManifestInDir(modelDir)
+	if err != nil {
+		t.Fatalf("load manifest: %v", err)
+	}
+	if got.Format != FormatGGUF {
+		t.Fatalf("Format = %q, want %q", got.Format, FormatGGUF)
+	}
+	if want := []string{"config.json", "demo-f16.gguf"}; !reflect.DeepEqual(got.Files, want) {
+		t.Fatalf("Files = %v, want %v", got.Files, want)
+	}
+	if got.Size != int64(len("{}")+len("gguf-data")) {
+		t.Fatalf("Size = %d, want %d", got.Size, len("{}")+len("gguf-data"))
+	}
+}