Add support to mmproj (docker#108)

* Add support for Multimodal projector file

* Add tests for Multimodal support

* instead of modifying Builder for testing purposes, add an implementation of Target to the test code that accomplishes the same thing

Author: ilopezluna

pkg/distribution/assets/dummy.mmproj

@@ -0,0 +1,2 @@
+This is a dummy multimodal projector file for testing purposes.
+It contains sample content to simulate a real multimodal projector file.

pkg/distribution/builder/builder.go

@@ -44,6 +44,17 @@ func (b *Builder) WithContextSize(size uint64) *Builder {
 	}
 }
 
+// WithMultimodalProjector adds a Multimodal projector file to the artifact
+func (b *Builder) WithMultimodalProjector(path string) (*Builder, error) {
+	mmprojLayer, err := partial.NewLayer(path, types.MediaTypeMultimodalProjector)
+	if err != nil {
+		return nil, fmt.Errorf("mmproj layer from %q: %w", path, err)
+	}
+	return &Builder{
+		model: mutate.AppendLayers(b.model, mmprojLayer),
+	}, nil
+}
+
 // Target represents a build target
 type Target interface {
 	Write(context.Context, types.ModelArtifact, io.Writer) error

pkg/distribution/builder/builder_test.go

@@ -0,0 +1,153 @@
+package builder_test
+
+import (
+	"context"
+	"io"
+	"path/filepath"
+	"testing"
+
+	"github.com/docker/model-distribution/builder"
+	"github.com/docker/model-distribution/types"
+)
+
+func TestWithMultimodalProjector(t *testing.T) {
+	// Create a builder from a GGUF file
+	b, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf"))
+	if err != nil {
+		t.Fatalf("Failed to create builder from GGUF: %v", err)
+	}
+
+	// Add multimodal projector
+	b2, err := b.WithMultimodalProjector(filepath.Join("..", "assets", "dummy.mmproj"))
+	if err != nil {
+		t.Fatalf("Failed to add multimodal projector: %v", err)
+	}
+
+	// Build the model
+	target := &fakeTarget{}
+	if err := b2.Build(t.Context(), target, nil); err != nil {
+		t.Fatalf("Failed to build model: %v", err)
+	}
+
+	// Verify the model has the expected layers
+	manifest, err := target.artifact.Manifest()
+	if err != nil {
+		t.Fatalf("Failed to get manifest: %v", err)
+	}
+
+	// Should have 2 layers: GGUF + multimodal projector
+	if len(manifest.Layers) != 2 {
+		t.Fatalf("Expected 2 layers, got %d", len(manifest.Layers))
+	}
+
+	// Check that one layer has the multimodal projector media type
+	foundMMProjLayer := false
+	for _, layer := range manifest.Layers {
+		if layer.MediaType == types.MediaTypeMultimodalProjector {
+			foundMMProjLayer = true
+			break
+		}
+	}
+
+	if !foundMMProjLayer {
+		t.Error("Expected to find a layer with multimodal projector media type")
+	}
+
+	// Note: We can't directly test MMPROJPath() on ModelArtifact interface
+	// but we can verify the layer was added with correct media type above
+}
+
+func TestWithMultimodalProjectorInvalidPath(t *testing.T) {
+	// Create a builder from a GGUF file
+	b, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf"))
+	if err != nil {
+		t.Fatalf("Failed to create builder from GGUF: %v", err)
+	}
+
+	// Try to add multimodal projector with invalid path
+	_, err = b.WithMultimodalProjector("nonexistent/path/to/mmproj")
+	if err == nil {
+		t.Error("Expected error when adding multimodal projector with invalid path")
+	}
+}
+
+func TestWithMultimodalProjectorChaining(t *testing.T) {
+	// Create a builder from a GGUF file
+	b, err := builder.FromGGUF(filepath.Join("..", "assets", "dummy.gguf"))
+	if err != nil {
+		t.Fatalf("Failed to create builder from GGUF: %v", err)
+	}
+
+	// Chain multiple operations: license + multimodal projector + context size
+	b, err = b.WithLicense(filepath.Join("..", "assets", "license.txt"))
+	if err != nil {
+		t.Fatalf("Failed to add license: %v", err)
+	}
+
+	b, err = b.WithMultimodalProjector(filepath.Join("..", "assets", "dummy.mmproj"))
+	if err != nil {
+		t.Fatalf("Failed to add multimodal projector: %v", err)
+	}
+
+	b = b.WithContextSize(4096)
+
+	// Build the model
+	target := &fakeTarget{}
+	if err := b.Build(t.Context(), target, nil); err != nil {
+		t.Fatalf("Failed to build model: %v", err)
+	}
+
+	// Verify the final model has all expected layers and properties
+	manifest, err := target.artifact.Manifest()
+	if err != nil {
+		t.Fatalf("Failed to get manifest: %v", err)
+	}
+
+	// Should have 3 layers: GGUF + license + multimodal projector
+	if len(manifest.Layers) != 3 {
+		t.Fatalf("Expected 3 layers, got %d", len(manifest.Layers))
+	}
+
+	// Check media types - using string comparison since we can't use types.MediaType directly
+	expectedMediaTypes := map[string]bool{
+		string(types.MediaTypeGGUF):                false,
+		string(types.MediaTypeLicense):             false,
+		string(types.MediaTypeMultimodalProjector): false,
+	}
+
+	for _, layer := range manifest.Layers {
+		if _, exists := expectedMediaTypes[string(layer.MediaType)]; exists {
+			expectedMediaTypes[string(layer.MediaType)] = true
+		}
+	}
+
+	for mediaType, found := range expectedMediaTypes {
+		if !found {
+			t.Errorf("Expected to find layer with media type %s", mediaType)
+		}
+	}
+
+	// Check context size
+	config, err := target.artifact.Config()
+	if err != nil {
+		t.Fatalf("Failed to get config: %v", err)
+	}
+
+	if config.ContextSize == nil || *config.ContextSize != 4096 {
+		t.Errorf("Expected context size 4096, got %v", config.ContextSize)
+	}
+
+	// Note: We can't directly test GGUFPath() and MMPROJPath() on ModelArtifact interface
+	// but we can verify the layers were added with correct media types above
+}
+
+var _ builder.Target = &fakeTarget{}
+
+type fakeTarget struct {
+	artifact types.ModelArtifact
+}
+
+func (ft *fakeTarget) Write(ctx context.Context, artifact types.ModelArtifact, writer io.Writer) error {
+	ft.artifact = artifact
+	return nil
+}

pkg/distribution/internal/partial/partial.go

@@ -7,7 +7,6 @@ import (
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/partial"
 	ggcr "github.com/google/go-containerregistry/pkg/v1/types"
-	"github.com/pkg/errors"
 
 	"github.com/docker/model-distribution/types"
 )
@@ -67,22 +66,31 @@ type WithLayers interface {
 }
 
 func GGUFPath(i WithLayers) (string, error) {
+	return layerPathByMediaType(i, types.MediaTypeGGUF)
+}
+
+func MMPROJPath(i WithLayers) (string, error) {
+	return layerPathByMediaType(i, types.MediaTypeMultimodalProjector)
+}
+
+// layerPathByMediaType is a generic helper function that finds a layer by media type and returns its path
+func layerPathByMediaType(i WithLayers, mediaType ggcr.MediaType) (string, error) {
 	layers, err := i.Layers()
 	if err != nil {
 		return "", fmt.Errorf("get layers: %w", err)
 	}
 	for _, l := range layers {
 		mt, err := l.MediaType()
-		if err != nil || mt != types.MediaTypeGGUF {
+		if err != nil || mt != mediaType {
 			continue
 		}
-		ggufLayer, ok := l.(*Layer)
+		layer, ok := l.(*Layer)
 		if !ok {
-			return "", errors.New("gguf Layer is not available locally")
+			return "", fmt.Errorf("%s Layer is not available locally", mediaType)
 		}
-		return ggufLayer.Path, nil
+		return layer.Path, nil
 	}
-	return "", errors.New("model does not contain a GGUF layer")
+	return "", fmt.Errorf("model does not contain a %s layer", mediaType)
 }
 
 func ManifestForLayers(i WithLayers) (*v1.Manifest, error) {

pkg/distribution/internal/partial/partial_test.go

@@ -0,0 +1,116 @@
+package partial_test
+
+import (
+	"path/filepath"
+	"testing"
+
+	"github.com/docker/model-distribution/internal/gguf"
+	"github.com/docker/model-distribution/internal/mutate"
+	"github.com/docker/model-distribution/internal/partial"
+	"github.com/docker/model-distribution/types"
+)
+
+func TestMMPROJPath(t *testing.T) {
+	// Create a model from GGUF file
+	mdl, err := gguf.NewModel(filepath.Join("..", "..", "assets", "dummy.gguf"))
+	if err != nil {
+		t.Fatalf("Failed to create model from GGUF: %v", err)
+	}
+
+	// Add multimodal projector layer
+	mmprojLayer, err := partial.NewLayer(filepath.Join("..", "..", "assets", "dummy.mmproj"), types.MediaTypeMultimodalProjector)
+	if err != nil {
+		t.Fatalf("Failed to create multimodal projector layer: %v", err)
+	}
+
+	mdlWithMMProj := mutate.AppendLayers(mdl, mmprojLayer)
+
+	// Test MMPROJPath function
+	mmprojPath, err := partial.MMPROJPath(mdlWithMMProj)
+	if err != nil {
+		t.Fatalf("Failed to get multimodal projector path: %v", err)
+	}
+
+	expectedPath := filepath.Join("..", "..", "assets", "dummy.mmproj")
+	if mmprojPath != expectedPath {
+		t.Errorf("Expected multimodal projector path %s, got %s", expectedPath, mmprojPath)
+	}
+}
+
+func TestMMPROJPathNotFound(t *testing.T) {
+	// Create a model from a GGUF file without a Multimodal projector
+	mdl, err := gguf.NewModel(filepath.Join("..", "..", "assets", "dummy.gguf"))
+	if err != nil {
+		t.Fatalf("Failed to create model from GGUF: %v", err)
+	}
+
+	// Test MMPROJPath function should return error
+	_, err = partial.MMPROJPath(mdl)
+	if err == nil {
+		t.Error("Expected error when getting multimodal projector path from model without multimodal projector layer")
+	}
+
+	expectedErrorMsg := "model does not contain a application/vnd.docker.ai.mmproj layer"
+	if err.Error() != expectedErrorMsg {
+		t.Errorf("Expected error message %q, got %q", expectedErrorMsg, err.Error())
+	}
+}
+
+func TestGGUFPath(t *testing.T) {
+	// Create a model from GGUF file
+	mdl, err := gguf.NewModel(filepath.Join("..", "..", "assets", "dummy.gguf"))
+	if err != nil {
+		t.Fatalf("Failed to create model from GGUF: %v", err)
+	}
+
+	// Test GGUFPath function
+	ggufPath, err := partial.GGUFPath(mdl)
+	if err != nil {
+		t.Fatalf("Failed to get GGUF path: %v", err)
+	}
+
+	expectedPath := filepath.Join("..", "..", "assets", "dummy.gguf")
+	if ggufPath != expectedPath {
+		t.Errorf("Expected GGUF path %s, got %s", expectedPath, ggufPath)
+	}
+}
+
+func TestLayerPathByMediaType(t *testing.T) {
+	// Create a model from GGUF file
+	mdl, err := gguf.NewModel(filepath.Join("..", "..", "assets", "dummy.gguf"))
+	if err != nil {
+		t.Fatalf("Failed to create model from GGUF: %v", err)
+	}
+
+	// Add license layer
+	licenseLayer, err := partial.NewLayer(filepath.Join("..", "..", "assets", "license.txt"), types.MediaTypeLicense)
+	if err != nil {
+		t.Fatalf("Failed to create license layer: %v", err)
+	}
+
+	// Add a Multimodal projector layer
+	mmprojLayer, err := partial.NewLayer(filepath.Join("..", "..", "assets", "dummy.mmproj"), types.MediaTypeMultimodalProjector)
+	if err != nil {
+		t.Fatalf("Failed to create multimodal projector layer: %v", err)
+	}
+
+	mdlWithLayers := mutate.AppendLayers(mdl, licenseLayer, mmprojLayer)
+
+	// Test that we can find each layer type
+	ggufPath, err := partial.GGUFPath(mdlWithLayers)
+	if err != nil {
+		t.Fatalf("Failed to get GGUF path: %v", err)
+	}
+	if ggufPath != filepath.Join("..", "..", "assets", "dummy.gguf") {
+		t.Errorf("Expected GGUF path to be: %s, got: %s", filepath.Join("..", "..", "assets", "dummy.gguf"), ggufPath)
+	}
+
+	mmprojPath, err := partial.MMPROJPath(mdlWithLayers)
+	if err != nil {
+		t.Fatalf("Failed to get multimodal projector path: %v", err)
+	}
+	if mmprojPath != filepath.Join("..", "..", "assets", "dummy.mmproj") {
+		t.Errorf("Expected multimodal projector path to be: %s, got: %s", filepath.Join("..", "..", "assets", "dummy.mmproj"), mmprojPath)
+	}
+
+}

pkg/distribution/internal/store/model.go

@@ -114,6 +114,10 @@ func (m *Model) GGUFPath() (string, error) {
 	return mdpartial.GGUFPath(m)
 }
 
+func (m *Model) MMPROJPath() (string, error) {
+	return mdpartial.MMPROJPath(m)
+}
+
 func (m *Model) Tags() []string {
 	return m.tags
 }