Unify Model Builder with Strategy Pattern

cmd/cli/commands/integration_test.go

@@ -264,7 +264,7 @@ func createAndPushTestModel(t *testing.T, registryURL, modelRef string, contextS
 
 	// Create a builder from the GGUF file
 	t.Logf("Creating test model %s from %s", modelRef, absPath)
-	pkg, err := builder.FromGGUF(absPath)
+	pkg, err := builder.FromPath(absPath)
 	require.NoError(t, err)
 
 	// Set context size if specified

cmd/cli/commands/package.go

@@ -241,7 +241,7 @@ func initializeBuilder(cmd *cobra.Command, opts packageOptions) (*builderInitRes
 		}
 	} else if opts.ggufPath != "" {
 		cmd.PrintErrf("Adding GGUF file from %q\n", opts.ggufPath)
-		pkg, err := builder.FromGGUF(opts.ggufPath)
+		pkg, err := builder.FromPath(opts.ggufPath)
 		if err != nil {
 			return nil, fmt.Errorf("add gguf file: %w", err)
 		}
@@ -262,7 +262,7 @@ func initializeBuilder(cmd *cobra.Command, opts packageOptions) (*builderInitRes
 		}
 
 		cmd.PrintErrf("Found %d safetensors file(s)\n", len(safetensorsPaths))
-		pkg, err := builder.FromSafetensors(safetensorsPaths)
+		pkg, err := builder.FromPaths(safetensorsPaths)
 		if err != nil {
 			return nil, fmt.Errorf("create safetensors model: %w", err)
 		}

cmd/mdltool/main.go

@@ -286,7 +286,7 @@ func cmdPackage(args []string) int {
 	var b *builder.Builder
 	if isSafetensors {
 		fmt.Println("Creating safetensors model")
-		b, err = builder.FromSafetensors(safetensorsPaths)
+		b, err = builder.FromPaths(safetensorsPaths)
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "Error creating model from safetensors: %v\n", err)
 			return 1
@@ -302,7 +302,7 @@ func cmdPackage(args []string) int {
 			}
 		}
 	} else {
-		b, err = builder.FromGGUF(source)
+		b, err = builder.FromPath(source)
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "Error creating model from gguf: %v\n", err)
 			return 1

pkg/distribution/builder/builder.go

@@ -4,11 +4,11 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"time"
 
-	"github.com/docker/model-runner/pkg/distribution/internal/gguf"
+	"github.com/docker/model-runner/pkg/distribution/format"
 	"github.com/docker/model-runner/pkg/distribution/internal/mutate"
 	"github.com/docker/model-runner/pkg/distribution/internal/partial"
-	"github.com/docker/model-runner/pkg/distribution/internal/safetensors"
 	"github.com/docker/model-runner/pkg/distribution/oci"
 	"github.com/docker/model-runner/pkg/distribution/types"
 )
@@ -19,23 +19,86 @@ type Builder struct {
 	originalLayers []oci.Layer // Snapshot of layers when created from existing model
 }
 
-// FromGGUF returns a *Builder that builds a model artifacts from a GGUF file
-func FromGGUF(path string) (*Builder, error) {
-	mdl, err := gguf.NewModel(path)
+// FromPath returns a *Builder that builds model artifacts from a file path.
+// It auto-detects the model format (GGUF or Safetensors) and discovers any shards.
+// This is the preferred entry point for creating models from local files.
+func FromPath(path string) (*Builder, error) {
+	// Auto-detect format from file extension
+	f, err := format.DetectFromPath(path)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("detect format: %w", err)
 	}
-	return &Builder{
-		model: mdl,
-	}, nil
+
+	// Discover all shards if this is a sharded model
+	paths, err := f.DiscoverShards(path)
+	if err != nil {
+		return nil, fmt.Errorf("discover shards: %w", err)
+	}
+
+	// Create model using the format abstraction
+	return fromFormat(f, paths)
+}
+
+// FromPaths returns a *Builder that builds model artifacts from multiple file paths.
+// All paths must be of the same format. Use this when you already have the list of files.
+func FromPaths(paths []string) (*Builder, error) {
+	if len(paths) == 0 {
+		return nil, fmt.Errorf("at least one path is required")
+	}
+
+	// Detect and verify format from all paths
+	f, err := format.DetectFromPaths(paths)
+	if err != nil {
+		return nil, fmt.Errorf("detect format: %w", err)
+	}
+
+	// Create model using the format abstraction
+	return fromFormat(f, paths)
 }
 
-// FromSafetensors returns a *Builder that builds model artifacts from safetensors files
-func FromSafetensors(safetensorsPaths []string) (*Builder, error) {
-	mdl, err := safetensors.NewModel(safetensorsPaths)
+// fromFormat creates a Builder using the unified format abstraction.
+// This is the internal implementation that creates layers and config.
+func fromFormat(f format.Format, paths []string) (*Builder, error) {
+	// Create layers from paths
+	layers := make([]oci.Layer, len(paths))
+	diffIDs := make([]oci.Hash, len(paths))
+
+	mediaType := f.MediaType()
+	for i, path := range paths {
+		layer, err := partial.NewLayer(path, mediaType)
+		if err != nil {
+			return nil, fmt.Errorf("create layer from %q: %w", path, err)
+		}
+		diffID, err := layer.DiffID()
+		if err != nil {
+			return nil, fmt.Errorf("get diffID for %q: %w", path, err)
+		}
+		layers[i] = layer
+		diffIDs[i] = diffID
+	}
+
+	// Extract config metadata using format-specific logic
+	config, err := f.ExtractConfig(paths)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("extract config: %w", err)
 	}
+
+	// Build the model
+	created := time.Now()
+	mdl := &partial.BaseModel{
+		ModelConfigFile: types.ConfigFile{
+			Config: config,
+			Descriptor: types.Descriptor{
+				Created: &created,
+			},
+			RootFS: oci.RootFS{
+				Type:    "rootfs",
+				DiffIDs: diffIDs,
+			},
+		},
+		LayerList: layers,
+	}
+
 	return &Builder{
 		model: mdl,
 	}, nil

pkg/distribution/builder/builder_test.go

@@ -15,7 +15,7 @@ import (
 
 func TestBuilder(t *testing.T) {
 	// Create a builder from a GGUF file
-	b, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf"))
+	b, err := builder.FromPath(filepath.Join("..", "assets", "dummy.gguf"))
 	if err != nil {
 		t.Fatalf("Failed to create builder from GGUF: %v", err)
 	}
@@ -63,7 +63,7 @@ func TestBuilder(t *testing.T) {
 
 func TestWithMultimodalProjectorInvalidPath(t *testing.T) {
 	// Create a builder from a GGUF file
-	b, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf"))
+	b, err := builder.FromPath(filepath.Join("..", "assets", "dummy.gguf"))
 	if err != nil {
 		t.Fatalf("Failed to create builder from GGUF: %v", err)
 	}
@@ -77,7 +77,7 @@ func TestWithMultimodalProjectorInvalidPath(t *testing.T) {
 
 func TestWithMultimodalProjectorChaining(t *testing.T) {
 	// Create a builder from a GGUF file
-	b, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf"))
+	b, err := builder.FromPath(filepath.Join("..", "assets", "dummy.gguf"))
 	if err != nil {
 		t.Fatalf("Failed to create builder from GGUF: %v", err)
 	}
@@ -147,7 +147,7 @@ func TestWithMultimodalProjectorChaining(t *testing.T) {
 
 func TestFromModel(t *testing.T) {
 	// Step 1: Create an initial model from GGUF with context size 2048
-	initialBuilder, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf"))
+	initialBuilder, err := builder.FromPath(filepath.Join("..", "assets", "dummy.gguf"))
 	if err != nil {
 		t.Fatalf("Failed to create initial builder from GGUF: %v", err)
 	}
@@ -230,7 +230,7 @@ func TestFromModel(t *testing.T) {
 
 func TestFromModelWithAdditionalLayers(t *testing.T) {
 	// Create an initial model from GGUF
-	initialBuilder, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf"))
+	initialBuilder, err := builder.FromPath(filepath.Join("..", "assets", "dummy.gguf"))
 	if err != nil {
 		t.Fatalf("Failed to create initial builder from GGUF: %v", err)
 	}

pkg/distribution/distribution/bundle_test.go

@@ -6,7 +6,7 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/docker/model-runner/pkg/distribution/internal/gguf"
+	"github.com/docker/model-runner/pkg/distribution/builder"
 	"github.com/docker/model-runner/pkg/distribution/internal/mutate"
 	"github.com/docker/model-runner/pkg/distribution/internal/partial"
 	"github.com/docker/model-runner/pkg/distribution/types"
@@ -23,10 +23,11 @@ func TestBundle(t *testing.T) {
 	}
 
 	// Load dummy model from assets directory
-	mdl, err := gguf.NewModel(filepath.Join("..", "assets", "dummy.gguf"))
+	b, err := builder.FromPath(filepath.Join("..", "assets", "dummy.gguf"))
 	if err != nil {
 		t.Fatalf("Failed to create model: %v", err)
 	}
+	mdl := b.Model()
 	singleGGUFID, err := mdl.ID()
 	if err != nil {
 		t.Fatalf("Failed to get model ID: %v", err)
@@ -64,10 +65,11 @@ func TestBundle(t *testing.T) {
 	}
 
 	// Load sharded dummy model from asset directory
-	shardedMdl, err := gguf.NewModel(filepath.Join("..", "assets", "dummy-00001-of-00002.gguf"))
+	shardedB, err := builder.FromPath(filepath.Join("..", "assets", "dummy-00001-of-00002.gguf"))
 	if err != nil {
 		t.Fatalf("Failed to create model: %v", err)
 	}
+	shardedMdl := shardedB.Model()
 	shardedGGUFID, err := shardedMdl.ID()
 	if err != nil {
 		t.Fatalf("Failed to get model ID: %v", err)