diff --git a/mlxrunner/model/gemma4/gemma4.go b/mlxrunner/model/gemma4/gemma4.go index 8349ab4b1..48fa0a8c9 100644 --- a/mlxrunner/model/gemma4/gemma4.go +++ b/mlxrunner/model/gemma4/gemma4.go @@ -764,9 +764,6 @@ func newModel(root *model.Root) (model.Model, error) { } if m.Vision != nil { - if err := validateVisionSoftTokenBudget(m.visionSoftTokenBudget()); err != nil { - return nil, err - } m.mediaPlaceholderIDs = append(m.mediaPlaceholderIDs, mm.ImageTokenID) } if m.Audio != nil { diff --git a/mlxrunner/model/gemma4/media.go b/mlxrunner/model/gemma4/media.go index 6b1af4cae..d50612149 100644 --- a/mlxrunner/model/gemma4/media.go +++ b/mlxrunner/model/gemma4/media.go @@ -14,11 +14,10 @@ import ( // multimodalConfig carries the top-level config.json fields the text_config // unwrap in parseTextConfig discards. type multimodalConfig struct { - VisionConfig *VisionConfig `json:"vision_config"` - ImageTokenID int32 `json:"image_token_id"` - BOITokenID int32 `json:"boi_token_id"` - EOITokenID int32 `json:"eoi_token_id"` - VisionSoftTokensPerImage int32 `json:"vision_soft_tokens_per_image"` + VisionConfig *VisionConfig `json:"vision_config"` + ImageTokenID int32 `json:"image_token_id"` + BOITokenID int32 `json:"boi_token_id"` + EOITokenID int32 `json:"eoi_token_id"` AudioConfig *AudioConfig `json:"audio_config"` AudioTokenID int32 `json:"audio_token_id"` @@ -45,9 +44,6 @@ func parseMultimodalConfig(configData []byte) (multimodalConfig, error) { if v.PoolingKernelSize == 0 { v.PoolingKernelSize = 3 } - if v.DefaultOutputLen == 0 { - v.DefaultOutputLen = 280 - } if v.RMSNormEps == 0 { v.RMSNormEps = 1e-6 } @@ -62,9 +58,6 @@ func parseMultimodalConfig(configData []byte) (multimodalConfig, error) { if v.PoolingKernelSize == 0 { v.PoolingKernelSize = 3 } - if v.NumSoftTokens == 0 { - v.NumSoftTokens = 280 - } if v.RMSNormEps == 0 { v.RMSNormEps = 1e-6 } diff --git a/mlxrunner/model/gemma4/media_test.go b/mlxrunner/model/gemma4/media_test.go index cafd8c76a..1bcb3332a 100644 --- a/mlxrunner/model/gemma4/media_test.go +++ b/mlxrunner/model/gemma4/media_test.go @@ -4,6 +4,7 @@ import ( "bytes" "image" "image/png" + "slices" "testing" "github.com/ollama/ollama/mlx/mlxtest" @@ -58,7 +59,7 @@ func TestParseMultimodalConfig(t *testing.T) { if uv == nil || !uv.unified() { t.Fatal("unified vision config not retained") } - if uv.NumSoftTokens != 280 { + if uv.PatchSize != 16 || uv.PoolingKernelSize != 3 { t.Fatalf("unified config = %+v", uv) } @@ -79,13 +80,6 @@ func TestParseMultimodalConfig(t *testing.T) { if textOnly.VisionConfig != nil { t.Fatal("text-only checkpoint grew a vision config") } - - if err := validateVisionSoftTokenBudget(280); err != nil { - t.Fatal(err) - } - if err := validateVisionSoftTokenBudget(300); err == nil { - t.Fatal("budget 300 accepted") - } } func visionTestModel() *Model { @@ -202,12 +196,12 @@ func TestScatterMedia(t *testing.T) { func TestPrepareMediaExpansion(t *testing.T) { m := visionTestModel() m.Vision = &VisionConfig{ - ModelType: "gemma4_vision", - HiddenSize: 12, - PatchSize: 16, - PoolingKernelSize: 3, - DefaultOutputLen: 70, - RMSNormEps: 1e-6, + ModelType: "gemma4_vision", + HiddenSize: 12, + PatchSize: 16, + PoolingKernelSize: 3, + RMSNormEps: 1e-6, + positionEmbeddingSize: 10240, } m.VisionTower = &VisionTower{} @@ -224,7 +218,7 @@ func TestPrepareMediaExpansion(t *testing.T) { t.Fatal(err) } - // 96x96 at budget 70 resizes to 384x384: 24x24 patches, 64 soft tokens, + // 96x96 selects 70 and resizes to 384x384: 24x24 patches, 64 soft tokens, // spliced after the one-token text run. wantSoft := 64 if len(prepared.Tokens) != 1+wantSoft+2 || prepared.Tokens[0] != 9 { @@ -268,3 +262,90 @@ func TestPrepareMediaExpansion(t *testing.T) { t.Fatal("towerless model accepted an image") } } + +func TestPrepareMediaDynamicImages(t *testing.T) { + var small, document bytes.Buffer + if err := png.Encode(&small, image.NewRGBA(image.Rect(0, 0, 320, 240))); err != nil { + t.Fatal(err) + } + if err := png.Encode(&document, image.NewRGBA(image.Rect(0, 0, 1800, 1400))); err != nil { + t.Fatal(err) + } + for _, kind := range []string{"gemma4_vision", "gemma4_unified_vision"} { + t.Run(kind, func(t *testing.T) { + // Old fixed-budget metadata must not pin either image to 280. + mm, err := parseMultimodalConfig([]byte(`{ + "vision_soft_tokens_per_image":280, + "vision_config":{"model_type":"` + kind + `", + "default_output_length":280,"num_soft_tokens":280}}`)) + if err != nil { + t.Fatal(err) + } + m := visionTestModel() + m.Vision = mm.VisionConfig + m.Vision.positionEmbeddingSize = 10240 + patch, pool := 16, 3 + if m.Vision.unified() { + m.UnifiedEmbedder = &UnifiedVisionEmbedder{} + patch, pool = 48, 1 + } else { + m.VisionTower = &VisionTower{} + } + + segments := []model.Segment{ + {Tokens: []int32{9}}, + {Kind: "image", Data: small.Bytes()}, + {Tokens: []int32{10, 11}}, + {Kind: "image", Data: document.Bytes()}, + } + prepared, err := m.PrepareMedia(segments) + if err != nil { + t.Fatal(err) + } + if len(prepared.Items) != 2 || len(prepared.Tokens) != 1+65+2+1075 { + t.Fatalf("items = %d, tokens = %d", len(prepared.Items), len(prepared.Tokens)) + } + for i, want := range []struct{ start, end, w, h, soft int }{ + {1, 66, 432, 336, 63}, + {68, 1143, 1776, 1392, 1073}, + } { + item := prepared.Items[i] + if item.Range != [2]int{want.start, want.end} || item.Source != 2*i+1 { + t.Fatalf("item %d: range = %v, source = %d", i, item.Range, item.Source) + } + p := item.Opaque.(preparedImage) + wantGeom := ImageGeometry{PatchesW: int32(want.w / patch), PatchesH: int32(want.h / patch), NumSoftTokens: int32(want.soft)} + if p.geom != wantGeom || len(p.positions) != 2*want.soft*pool*pool { + t.Fatalf("item %d: geometry = %+v, positions = %d", i, p.geom, len(p.positions)) + } + if !slices.Equal(item.Dims, []int{want.soft * pool * pool, patch * patch * 3}) || len(item.MediaData) != want.w*want.h*3 { + t.Fatalf("item %d: dims = %v, pixels = %d", i, item.Dims, len(item.MediaData)) + } + if prepared.Tokens[want.start] != m.MM.BOITokenID || prepared.Tokens[want.end-1] != m.MM.EOITokenID { + t.Fatalf("item %d: missing image delimiters", i) + } + for _, token := range prepared.Tokens[want.start+1 : want.end-1] { + if token != m.MM.ImageTokenID { + t.Fatalf("item %d: soft token = %d", i, token) + } + } + } + + // Preparing the same image again in a continuation retains its + // representation, independently of the other image or text. + repeated, err := m.PrepareMedia(append(segments[:2:2], model.Segment{Tokens: []int32{12}})) + if err != nil { + t.Fatal(err) + } + a, b := prepared.Items[0], repeated.Items[0] + pa, pb := a.Opaque.(preparedImage), b.Opaque.(preparedImage) + if pa.geom != pb.geom || !slices.Equal(pa.positions, pb.positions) || !slices.Equal(a.Dims, b.Dims) || !slices.Equal(a.MediaData, b.MediaData) { + t.Fatal("repeated image changed representation") + } + m.Vision.positionEmbeddingSize = 1 + if _, err := m.PrepareMedia(segments); err == nil { + t.Fatal("out-of-bounds position grid accepted") + } + }) + } +} diff --git a/mlxrunner/model/gemma4/process_image.go b/mlxrunner/model/gemma4/process_image.go index fa8a8341e..f7d10b4a2 100644 --- a/mlxrunner/model/gemma4/process_image.go +++ b/mlxrunner/model/gemma4/process_image.go @@ -24,33 +24,57 @@ type preparedImage struct { } // visionTargetSize ports the reference resize: sides floored to -// multiples of patchSize*poolingKernel under a patch budget, a +// multiples of patchSize*poolingKernel under a soft-token budget, a // zero-flooring side clamped to one multiple. -func visionTargetSize(height, width, patchSize, poolingKernel, maxPatches int32) (targetH, targetW int32, err error) { +func visionTargetSize(height, width int, patchSize, poolingKernel, budget int32) (targetH, targetW int32, err error) { if height <= 0 || width <= 0 { return 0, 0, fmt.Errorf("invalid image size %dx%d", width, height) } - sideMult := patchSize * poolingKernel - targetPx := float64(maxPatches) * float64(patchSize) * float64(patchSize) + side := int64(patchSize) * int64(poolingKernel) + if patchSize <= 0 || poolingKernel <= 0 || budget <= 0 || + side > math.MaxInt32/int64(budget) || int64(poolingKernel)*int64(poolingKernel) > math.MaxInt32/int64(budget) { + return 0, 0, fmt.Errorf("invalid vision geometry: patch %d, pool %d, budget %d", patchSize, poolingKernel, budget) + } + sideMult := float64(side) + targetPx := float64(budget) * sideMult * sideMult factor := math.Sqrt(targetPx / (float64(height) * float64(width))) - targetH = int32(math.Floor(factor*float64(height)/float64(sideMult))) * sideMult - targetW = int32(math.Floor(factor*float64(width)/float64(sideMult))) * sideMult - if targetH == 0 && targetW == 0 { + h := math.Floor(factor*float64(height)/sideMult) * sideMult + w := math.Floor(factor*float64(width)/sideMult) * sideMult + if h == 0 && w == 0 { return 0, 0, fmt.Errorf("image %dx%d is too small to process", width, height) } - maxSide := (maxPatches / (poolingKernel * poolingKernel)) * sideMult - if targetH == 0 { - targetH = sideMult - targetW = min(int32(math.Floor(float64(width)/float64(height)))*sideMult, maxSide) - } else if targetW == 0 { - targetW = sideMult - targetH = min(int32(math.Floor(float64(height)/float64(width)))*sideMult, maxSide) + maxSide := float64(budget) * sideMult + if h == 0 { + h = sideMult + w = min(math.Floor(float64(width)/float64(height))*sideMult, maxSide) + } else if w == 0 { + w = sideMult + h = min(math.Floor(float64(height)/float64(width))*sideMult, maxSide) } - if float64(targetH)*float64(targetW) > targetPx { + if h*w > targetPx || h > maxSide || w > maxSide { return 0, 0, fmt.Errorf("image %dx%d exceeds the patch budget after resize", width, height) } + return int32(h), int32(w), nil +} + +// dynamicVisionTargetSize chooses the smallest publisher-supported mode that +// preserves both source dimensions, or the largest mode when none fits. +func dynamicVisionTargetSize(height, width int, patch, pool int32) (int32, int32, error) { + var targetH, targetW int32 + for _, budget := range [...]int32{70, 140, 280, 560, 1120} { + var err error + targetH, targetW, err = visionTargetSize(height, width, patch, pool, budget) + if err != nil { + return 0, 0, err + } + // Compare the rounded grid: choosing the nearest area can discard + // detail even when a larger supported mode preserves it. + if int(targetH) >= height && int(targetW) >= width { + return targetH, targetW, nil + } + } return targetH, targetW, nil } @@ -66,12 +90,19 @@ func (m *Model) preprocessImage(data []byte) (pixels []float32, positions []int3 patch, pool := m.Vision.PatchSize, m.Vision.PoolingKernelSize bounds := img.Bounds() - img = dropAlpha(img, bounds) - targetH, targetW, err := visionTargetSize(int32(bounds.Dy()), int32(bounds.Dx()), patch, pool, m.visionSoftTokenBudget()*pool*pool) + targetH, targetW, err := dynamicVisionTargetSize(bounds.Dy(), bounds.Dx(), patch, pool) if err != nil { return nil, nil, ImageGeometry{}, err } + positionPatch := patch + if m.Vision.unified() { + positionPatch *= pool + } + if int(max(targetH, targetW)/positionPatch) > m.Vision.positionEmbeddingSize { + return nil, nil, ImageGeometry{}, fmt.Errorf("image patch grid exceeds vision position embedding size %d", m.Vision.positionEmbeddingSize) + } + img = dropAlpha(img, bounds) resized := image.NewRGBA(image.Rect(0, 0, int(targetW), int(targetH))) draw.CatmullRom.Scale(resized, resized.Bounds(), img, bounds, draw.Src, nil) diff --git a/mlxrunner/model/gemma4/vision.go b/mlxrunner/model/gemma4/vision.go index aee4400e5..872069778 100644 --- a/mlxrunner/model/gemma4/vision.go +++ b/mlxrunner/model/gemma4/vision.go @@ -3,7 +3,6 @@ package gemma4 import ( "fmt" "math" - "slices" "strings" "github.com/ollama/ollama/mlx" @@ -22,37 +21,20 @@ type VisionConfig struct { HeadDim int32 `json:"head_dim"` PatchSize int32 `json:"patch_size"` PoolingKernelSize int32 `json:"pooling_kernel_size"` - DefaultOutputLen int32 `json:"default_output_length"` RMSNormEps float32 `json:"rms_norm_eps"` Standardize bool `json:"standardize"` RopeParameters *RopeParams `json:"rope_parameters"` - // Unified-embedder field. - NumSoftTokens int32 `json:"num_soft_tokens"` - RopeTheta float32 `json:"-"` + + // Retained at load time so CPU preprocessing need not access MLX arrays. + positionEmbeddingSize int } func (v *VisionConfig) unified() bool { return v.ModelType == "gemma4_unified_vision" } -// visionSoftTokenBudgets are the soft-token counts the reference processor -// supports; the budget fixes the patch budget as budget*pooling². -var visionSoftTokenBudgets = []int32{70, 140, 280, 560, 1120} - -// visionSoftTokenBudget returns the per-image soft-token budget the checkpoint -// requests. -func (m *Model) visionSoftTokenBudget() int32 { - if m.MM.VisionSoftTokensPerImage > 0 { - return m.MM.VisionSoftTokensPerImage - } - if m.Vision.unified() { - return m.Vision.NumSoftTokens - } - return m.Vision.DefaultOutputLen -} - // VisionTower is the Gemma 4 image encoder: a bidirectional transformer over // patch embeddings with 2D positions, pooled 3x3 into soft tokens. type VisionTower struct { @@ -196,6 +178,7 @@ func (m *Model) loadVisionWeights(tensors map[string]*mlx.Array, linears model.L m.VisionTower = tower m.EmbedVision = &MultimodalEmbedder{Projection: projection} + v.positionEmbeddingSize = table.Dim(1) return nil } @@ -249,6 +232,7 @@ func (m *Model) loadUnifiedVisionWeights(tensors map[string]*mlx.Array, linears m.UnifiedEmbedder = e m.EmbedVision = &MultimodalEmbedder{Projection: projection} + m.Vision.positionEmbeddingSize = e.PosEmbedding.Dim(0) return nil } @@ -397,10 +381,3 @@ func (m *Model) encodeImage(pixels *mlx.Array, positions []int32, geom ImageGeom f = mlx.RMSNormFn(f, nil, v.RMSNormEps) return m.EmbedVision.Projection.Forward(f) } - -func validateVisionSoftTokenBudget(budget int32) error { - if !slices.Contains(visionSoftTokenBudgets, budget) { - return fmt.Errorf("unsupported vision soft token budget %d", budget) - } - return nil -} diff --git a/mlxrunner/model/gemma4/vision_test.go b/mlxrunner/model/gemma4/vision_test.go index 23bce1146..d0a7b78d0 100644 --- a/mlxrunner/model/gemma4/vision_test.go +++ b/mlxrunner/model/gemma4/vision_test.go @@ -15,7 +15,8 @@ import ( func TestVisionTargetSize(t *testing.T) { cases := []struct { name string - h, w, budget int32 + h, w int + budget int32 wantH, wantW int32 wantSoftTokens int32 }{ @@ -25,10 +26,19 @@ func TestVisionTargetSize(t *testing.T) { {"square large budget", 896, 896, 1120, 1584, 1584, 1089}, {"extreme aspect clamps", 10, 10000, 280, 48, 13440, 280}, {"tiny upscales", 20, 20, 280, 768, 768, 256}, + // Reference: Transformers 2c4914fb939fe9de0d8e7a798af4684d552f18b4, + // models/gemma4/image_processing_gemma4.py:get_aspect_ratio_preserving_size, + // called with (height=1400, width=1800, patch_size=16, + // max_patches=budget*9, pooling_kernel_size=3). + {"document 70", 1400, 1800, 70, 336, 432, 63}, + {"document 140", 1400, 1800, 140, 480, 624, 130}, + {"document 280", 1400, 1800, 280, 672, 864, 252}, + {"document 560", 1400, 1800, 560, 960, 1248, 520}, + {"document 1120", 1400, 1800, 1120, 1392, 1776, 1073}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { - gotH, gotW, err := visionTargetSize(c.h, c.w, 16, 3, c.budget*9) + gotH, gotW, err := visionTargetSize(c.h, c.w, 16, 3, c.budget) if err != nil { t.Fatal(err) } @@ -43,6 +53,78 @@ func TestVisionTargetSize(t *testing.T) { } } +func TestDynamicVisionTargetSize(t *testing.T) { + cases := []struct { + name string + h, w int + wantH, wantW int32 + }{ + {"tiny", 1, 1, 384, 384}, + {"small counting image", 120, 210, 288, 528}, + {"small landscape", 240, 320, 336, 432}, + {"small text preserves resolution", 400, 500, 480, 624}, + {"portrait text preserves resolution", 500, 400, 624, 480}, + {"area alone does not preserve height", 49, 858, 96, 2352}, + {"small square", 512, 512, 528, 528}, + {"native square", 768, 768, 768, 768}, + {"large square", 1024, 1024, 1104, 1104}, + {"landscape", 720, 1280, 816, 1488}, + {"document", 1400, 1800, 1392, 1776}, + {"portrait document", 1800, 1400, 1776, 1392}, + {"4K", 2160, 3840, 1200, 2112}, + {"thin", 10, 10000, 48, 13440}, + {"tall", 10000, 10, 13440, 48}, + {"huge width", 1, math.MaxInt, 48, 53760}, + {"huge height", math.MaxInt, 1, 53760, 48}, + {"huge square", math.MaxInt, math.MaxInt, 1584, 1584}, + {"70 at boundary", 384, 384, 384, 384}, + {"140 above boundary", 385, 385, 528, 528}, + {"140 at boundary", 528, 528, 528, 528}, + {"280 above boundary", 529, 529, 768, 768}, + {"280 at boundary", 768, 768, 768, 768}, + {"560 above boundary", 769, 769, 1104, 1104}, + {"560 at boundary", 1104, 1104, 1104, 1104}, + {"1120 above boundary", 1105, 1105, 1584, 1584}, + {"1120 at boundary", 1584, 1584, 1584, 1584}, + {"1120 caps resolution", 1585, 1585, 1584, 1584}, + {"wide preserves resolution", 300, 2688, 336, 3360}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + h, w, err := dynamicVisionTargetSize(c.h, c.w, 16, 3) + if err != nil { + t.Fatal(err) + } + if h != c.wantH || w != c.wantW { + t.Fatalf("target = %dx%d, want %dx%d", w, h, c.wantW, c.wantH) + } + }) + } +} + +func TestVisionTargetSizeInvalid(t *testing.T) { + cases := []struct { + name string + h, w int + patch, pool, budget int32 + }{ + {"empty", 0, 100, 16, 3, 1120}, + {"negative", 100, -1, 16, 3, 1120}, + {"zero patch", 100, 100, 0, 3, 1120}, + {"negative pool", 100, 100, 16, -1, 1120}, + {"zero budget", 100, 100, 16, 3, 0}, + {"dimension overflow", 100, 100, math.MaxInt32, 3, 1120}, + {"patch count overflow", 100, 100, 1, 2000, 1120}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if _, _, err := visionTargetSize(c.h, c.w, c.patch, c.pool, c.budget); err == nil { + t.Fatal("invalid geometry accepted") + } + }) + } +} + // TestEncodeImagePooling checks the encode chain end to end on a towerless // config (identity projections, zero position table, no encoder layers): // the output must equal the reference computation — normalize, 3x3 grid