From b8c3d1f62df0e3f7da323ac10abb7b8afe66e5ef Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Tue, 22 Sep 2026 17:47:31 -0700 Subject: [PATCH] mlx: select Gemma 4 image resolution dynamically (#18603) Replace the fixed checkpoint image budget with per-image selection across the supported 70, 140, 280, 560, and 1120 budgets. Choose the publisher resize grid closest to the input resolution, accounting for aspect ratio. This preserves more detail in high-resolution documents while allowing smaller images to use fewer tokens, without adding an API parameter. Cover budget boundaries, extreme dimensions, position limits, and media expansion for both vision architectures. --- mlxrunner/model/gemma4/gemma4.go | 3 - mlxrunner/model/gemma4/media.go | 15 +--- mlxrunner/model/gemma4/media_test.go | 111 ++++++++++++++++++++---- mlxrunner/model/gemma4/process_image.go | 65 ++++++++++---- mlxrunner/model/gemma4/vision.go | 33 ++----- mlxrunner/model/gemma4/vision_test.go | 86 +++++++++++++++++- 6 files changed, 237 insertions(+), 76 deletions(-) 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