feat(unpack): enhance unpacking logic to support V0.2 layer-per-file packaging

Author: ilopezluna

pkg/distribution/builder/from_directory.go

@@ -190,7 +190,7 @@ func FromDirectory(dirPath string, opts ...DirectoryOption) (*Builder, error) {
 	// TODO: Extract additional metadata from weight files if needed
 	// For safetensors, we might want to read config.json from the directory
 
-	// Build the model
+	// Build the model with V0.2 config (layer-per-file with annotations)
 	created := time.Now()
 	mdl := &partial.BaseModel{
 		ModelConfigFile: types.ConfigFile{
@@ -203,7 +203,8 @@ func FromDirectory(dirPath string, opts ...DirectoryOption) (*Builder, error) {
 				DiffIDs: diffIDs,
 			},
 		},
-		LayerList: layers,
+		LayerList:       layers,
+		ConfigMediaType: types.MediaTypeModelConfigV02, // V0.2: layer-per-file with filepath annotations
 	}
 
 	return &Builder{

pkg/distribution/internal/bundle/unpack.go

@@ -14,7 +14,33 @@ import (
 )
 
 // Unpack creates and return a Bundle by unpacking files and config from model into dir.
+// It auto-detects the packaging version:
+//   - V0.2 (layer-per-file with annotations): Uses UnpackFromLayers for full path preservation
+//   - V0.1 (legacy): Uses the original unpacking logic based on GGUFPaths(), SafetensorsPaths(), etc.
 func Unpack(dir string, model types.Model) (*Bundle, error) {
+	// Check if the model uses V0.2 packaging (layer-per-file with annotations)
+	if artifact, ok := model.(types.ModelArtifact); ok {
+		if isV02Model(artifact) {
+			return UnpackFromLayers(dir, artifact)
+		}
+	}
+
+	// V0.1 legacy unpacking
+	return unpackLegacy(dir, model)
+}
+
+// isV02Model checks if the model was packaged using V0.2 format (layer-per-file with annotations).
+// It does this by checking the config media type in the manifest.
+func isV02Model(model types.ModelArtifact) bool {
+	manifest, err := model.Manifest()
+	if err != nil {
+		return false
+	}
+	return manifest.Config.MediaType == types.MediaTypeModelConfigV02
+}
+
+// unpackLegacy is the original V0.1 unpacking logic that uses model.GGUFPaths(), model.SafetensorsPaths(), etc.
+func unpackLegacy(dir string, model types.Model) (*Bundle, error) {
 	bundle := &Bundle{
 		dir: dir,
 	}

pkg/distribution/internal/partial/model.go

@@ -14,6 +14,10 @@ import (
 type BaseModel struct {
 	ModelConfigFile types.ConfigFile
 	LayerList       []oci.Layer
+	// ConfigMediaType specifies the media type for the config descriptor.
+	// If empty, defaults to MediaTypeModelConfigV01 for backward compatibility.
+	// Set to MediaTypeModelConfigV02 for layer-per-file packaging (FromDirectory).
+	ConfigMediaType oci.MediaType
 }
 
 var _ types.ModelArtifact = &BaseModel{}
@@ -125,3 +129,9 @@ func (m *BaseModel) Config() (types.ModelConfig, error) {
 func (m *BaseModel) Descriptor() (types.Descriptor, error) {
 	return Descriptor(m)
 }
+
+// GetConfigMediaType returns the config media type for the model.
+// If not set, returns empty string and ManifestForLayers will default to V0.1.
+func (m *BaseModel) GetConfigMediaType() oci.MediaType {
+	return m.ConfigMediaType
+}

pkg/distribution/internal/partial/partial.go

@@ -191,6 +191,11 @@ func matchesMediaType(layerMT, targetMT oci.MediaType) bool {
 	}
 }
 
+// WithConfigMediaType provides access to the config media type version.
+type WithConfigMediaType interface {
+	GetConfigMediaType() oci.MediaType
+}
+
 func ManifestForLayers(i WithLayers) (*oci.Manifest, error) {
 	raw, err := i.RawConfigFile()
 	if err != nil {
@@ -200,8 +205,17 @@ func ManifestForLayers(i WithLayers) (*oci.Manifest, error) {
 	if err != nil {
 		return nil, fmt.Errorf("compute config hash: %w", err)
 	}
+
+	// Use the config media type from the model if available, otherwise default to V0.1
+	configMediaType := types.MediaTypeModelConfigV01
+	if cmt, ok := i.(WithConfigMediaType); ok {
+		if mt := cmt.GetConfigMediaType(); mt != "" {
+			configMediaType = mt
+		}
+	}
+
 	cfgDsc := oci.Descriptor{
-		MediaType: types.MediaTypeModelConfigV01,
+		MediaType: configMediaType,
 		Size:      int64(len(raw)),
 		Digest:    cfgHash,
 	}

pkg/distribution/types/config.go

@@ -10,9 +10,16 @@ import (
 type MediaType = oci.MediaType
 
 const (
-	// MediaTypeModelConfigV01 is the media type for the model config json.
+	// MediaTypeModelConfigV01 is the media type for the model config json (legacy format).
+	// V0.1 uses model.GGUFPaths(), model.SafetensorsPaths(), etc. for unpacking.
 	MediaTypeModelConfigV01 MediaType = "application/vnd.docker.ai.model.config.v0.1+json"
 
+	// MediaTypeModelConfigV02 is the media type for models using the layer-per-file approach.
+	// V0.2 packages each file as an individual layer with filepath annotations.
+	// This format preserves nested directory structure (e.g., text_encoder/model.safetensors).
+	// Used by builder.FromDirectory.
+	MediaTypeModelConfigV02 MediaType = "application/vnd.docker.ai.model.config.v0.2+json"
+
 	// MediaTypeGGUF indicates a file in GGUF version 3 format, containing a tensor model.
 	MediaTypeGGUF MediaType = "application/vnd.docker.ai.gguf.v3"