CERBERUS/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-10-20 18:41:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

mlx: select Gemma 4 image resolution dynamically (#18603)

Browse source 
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.
This commit is contained in:
Daniel Hiltgen 2026-09-22 17:47:31 -07:00 • committed by GitHub
parent 421d74d20e
commit b8c3d1f62d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 237 additions and 76 deletions
Download patch file Download diff file Expand all files Collapse all files

3
mlxrunner/model/gemma4/gemma4.go
View file

@ -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 {

15
mlxrunner/model/gemma4/media.go
View file

@ -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
}

111
mlxrunner/model/gemma4/media_test.go
View file

@ -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")
}
})
}
}

65
mlxrunner/model/gemma4/process_image.go
View file

@ -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)

33
mlxrunner/model/gemma4/vision.go
View file

@ -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
}

86
mlxrunner/model/gemma4/vision_test.go
View file

@ -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