upstream-archive/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-09-21 05:28:00 -05:00
Code Issues Packages Projects Releases Wiki Activity

gemma4: image and audio input support

Browse Source 
Safetensors gemma4 imports served by the MLX engine now answer image
and audio chats. Images run through both vision architectures: the
transformer tower (26B, 31B, e-series) and the 12B's encoder-free
unified embedder. Audio arrives through the same intake the ollama
API already accepts for gemma4 GGUFs — WAV bytes in the images field,
OpenAI input_audio parts, and /v1/audio/transcriptions uploads — with
the e2b/e4b checkpoints running clips through their conformer audio
encoder and the 12b unified checkpoint embedding the raw waveform
directly. Clips longer than 30 seconds are split evenly into chunks
of at most 30 seconds, cut at pauses, and encoded independently.

Each modality serves only checkpoints that carry it: 26B/31B have no
audio config and reject audio input, and checkpoints with an
unrecognized vision architecture still load as text-only models and
reject image requests.

The server previously hid the vision and audio capabilities for
gemma4 safetensors because the engine served neither. Both
suppressions are removed, and existing imports start advertising the
capabilities without re-importing since import already records them.
This commit is contained in:
Jesse Gross
2026-09-02 15:07:28 -07:00
parent 3ffc9a682a
commit b79067b0db
14 changed files with 2548 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/images.go
-7
View File
@@ -516,19 +516,12 @@ func (m *Model) filterUnsupportedCapabilities(capabilities []model.Capability, m
}
func suppressVisionCapability(m *Model) bool {
if isGemma4Renderer(m.Config.Renderer) && m.Config.ModelFormat == "safetensors" {
return true
}
// The current MLX Nemotron path is text-only. Do not advertise vision for
// safetensors manifests until the runner can load and serve that modality.
return isNemotron3NanoSafetensors(m)
}
func suppressAudioCapability(m *Model, arch string) bool {
if isGemma4Renderer(m.Config.Renderer) && m.Config.ModelFormat == "safetensors" {
return true
}
if m.Config.ModelFormat == "safetensors" && m.Config.Renderer == "glimmer" {
return true
}
server/images_test.go
-33
View File
@@ -589,39 +589,6 @@ func TestModelCapabilities(t *testing.T) {
},
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityThinking},
},
{
name: "gemma4 small safetensors suppresses vision and audio",
model: Model{
Config: model.ConfigV2{
ModelFormat: "safetensors",
Renderer: gemma4RendererSmall,
Capabilities: []string{"vision", "audio"},
},
Template: chatTemplate,
},
},
{
name: "gemma4 large safetensors suppresses vision and audio",
model: Model{
Config: model.ConfigV2{
ModelFormat: "safetensors",
Renderer: gemma4RendererLarge,
Capabilities: []string{"vision", "audio"},
},
Template: chatTemplate,
},
},
{
name: "default gemma4 safetensors suppresses vision and audio",
model: Model{
Config: model.ConfigV2{
ModelFormat: "safetensors",
Renderer: gemma4RendererLegacy,
Capabilities: []string{"vision", "audio"},
},
Template: chatTemplate,
},
},
}
// compare two slices of model.Capability regardless of order
server/model_list_cache.go
+1 -1
View File
@@ -397,7 +397,7 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum
}
func filterUnsupportedModelListCapabilities(capabilities []model.Capability, cfg model.ConfigV2) []model.Capability {
if cfg.ModelFormat == "safetensors" && (isGemma4Renderer(cfg.Renderer) || isNemotron3NanoSafetensorsConfig(cfg)) {
if cfg.ModelFormat == "safetensors" && isNemotron3NanoSafetensorsConfig(cfg) {
capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool {
return c == model.CapabilityVision || c == model.CapabilityAudio
})
server/renderer_resolution.go
-9
View File
@@ -99,12 +99,3 @@ func parseHumanParameterCount(s string) (uint64, bool) {
return uint64(value * multiplier), true
}
func isGemma4Renderer(renderer string) bool {
switch renderer {
case gemma4RendererLegacy, gemma4RendererSmall, gemma4RendererLarge:
return true
default:
return false
}
}
x/models/gemma4/audio.go
+517
View File
@@ -0,0 +1,517 @@
package gemma4
import (
"fmt"
"math"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/mlxrunner/model"
"github.com/ollama/ollama/x/models/nn"
)
type AudioConfig struct {
ModelType string `json:"model_type"`
HiddenSize int32 `json:"hidden_size"`
NumHiddenLayers int32 `json:"num_hidden_layers"`
NumAttentionHeads int32 `json:"num_attention_heads"`
ConvKernelSize int32 `json:"conv_kernel_size"`
ResidualWeight float32 `json:"residual_weight"`
ChunkSize int32 `json:"attention_chunk_size"`
ContextLeft int32 `json:"attention_context_left"`
ContextRight int32 `json:"attention_context_right"`
LogitCap float32 `json:"attention_logit_cap"`
InvalidLogit float32 `json:"attention_invalid_logits_value"`
RMSNormEps float32 `json:"rms_norm_eps"`
GradientClipping float32 `json:"gradient_clipping"`
// Unified-embedder field.
SamplesPerToken int32 `json:"audio_samples_per_token"`
}
func (a *AudioConfig) unified() bool {
return a.ModelType == "gemma4_unified_audio"
}
// applyAudioDefaults fills absent fields with the reference config defaults.
func applyAudioDefaults(a *AudioConfig) {
if a.HiddenSize == 0 {
a.HiddenSize = 1024
}
if a.NumHiddenLayers == 0 {
a.NumHiddenLayers = 12
}
if a.NumAttentionHeads == 0 {
a.NumAttentionHeads = 8
}
if a.ConvKernelSize == 0 {
a.ConvKernelSize = 5
}
if a.ResidualWeight == 0 {
a.ResidualWeight = 0.5
}
if a.ChunkSize == 0 {
a.ChunkSize = 12
}
if a.ContextLeft == 0 {
a.ContextLeft = 13
}
if a.LogitCap == 0 {
a.LogitCap = 50
}
if a.InvalidLogit == 0 {
a.InvalidLogit = -1e9
}
if a.RMSNormEps == 0 {
a.RMSNormEps = 1e-6
}
if a.GradientClipping == 0 {
a.GradientClipping = 1e10
}
}
// preparedAudio carries one chunk's soft-token count to softRun.
type preparedAudio struct {
numTokens int32
}
// AudioTower is the Gemma 4 audio encoder: a USM conformer over log-mel
// frames, subsampled 4x in time by two stride-2 convs.
type AudioTower struct {
SubsampleConv []*AudioSubsampleConv
InputProj nn.LinearLayer
Layers []*AudioLayer
OutputProj nn.LinearLayer
// Gradient-clipping bounds in the working dtype.
ClipMin, ClipMax *mlx.Array
}
type AudioSubsampleConv struct {
Conv *mlx.Array // [out, 3, 3, in]
Norm *mlx.Array // LayerNorm weight, no bias
}
type AudioLayer struct {
FFW1, FFW2 *AudioFeedForward
Attention *AudioAttention
LConv *AudioLightConv
PreAttnNorm *mlx.Array
PostAttnNorm *mlx.Array
OutNorm *mlx.Array
}
type AudioFeedForward struct {
PreNorm *mlx.Array
Up, Down nn.LinearLayer
PostNorm *mlx.Array
}
type AudioAttention struct {
QProj, KProj, VProj, Post nn.LinearLayer
// QScale is softplus(per_dim_scale) in f32; RelK is the sinusoidal
// relative-position table through this layer's relative_k_proj,
// [1, heads, headDim, positions] in f32.
QScale *mlx.Array
RelK *mlx.Array
}
type AudioLightConv struct {
PreNorm *mlx.Array
Start nn.LinearLayer
DWConv *mlx.Array // [hidden, kernel, 1]
ConvNorm *mlx.Array
End nn.LinearLayer
}
// audioRelPosTable builds the [positions, hidden] sinusoidal table over
// relative positions contextSize/2 down to 0, each row [sin..., cos...].
func audioRelPosTable(a *AudioConfig) *mlx.Array {
positions := int(a.ChunkSize+a.ContextLeft-1+a.ContextRight)/2 + 1
hidden := int(a.HiddenSize)
half := hidden / 2
increment := math.Log(10000) / float64(half-1)
table := make([]float32, positions*hidden)
for i := range positions {
p := float64(positions - 1 - i)
for d := range half {
angle := p * math.Exp(-float64(d)*increment)
table[i*hidden+d] = float32(math.Sin(angle))
table[i*hidden+half+d] = float32(math.Cos(angle))
}
}
return mlx.FromValues(table, positions, hidden)
}
// audioWeightRoot locates the tower's tensor prefix by probing for the
// subsample conv projection.
func audioWeightRoot(tensors map[string]*mlx.Array) (string, error) {
for _, root := range []string{"model.", ""} {
if tensors[root+"audio_tower.subsample_conv_projection.layer0.conv.weight"] != nil {
return root, nil
}
}
return "", fmt.Errorf("config declares a gemma4_audio tower but audio_tower weights are missing from the manifest")
}
func (m *Model) loadAudioWeights(tensors map[string]*mlx.Array, linears model.LinearFactory) error {
if m.Audio.unified() {
return m.loadUnifiedAudioWeights(tensors, linears)
}
a := m.Audio
root, err := audioWeightRoot(tensors)
if err != nil {
return err
}
at := root + "audio_tower."
requiredNorm := func(name string) (*mlx.Array, error) {
w := tensors[name]
if w == nil {
return nil, fmt.Errorf("missing audio weight: %s", name)
}
return w, nil
}
tower := &AudioTower{
SubsampleConv: make([]*AudioSubsampleConv, 2),
Layers: make([]*AudioLayer, a.NumHiddenLayers),
}
for i := range tower.SubsampleConv {
sp := fmt.Sprintf("%ssubsample_conv_projection.layer%d.", at, i)
conv := tensors[sp+"conv.weight"]
if conv == nil {
return fmt.Errorf("missing audio weight: %sconv.weight", sp)
}
norm, err := requiredNorm(sp + "norm.weight")
if err != nil {
return err
}
// Conv weights arrive [out, in, kH, kW]; mlx wants channels last.
tower.SubsampleConv[i] = &AudioSubsampleConv{Conv: mlx.Transpose(conv, 0, 2, 3, 1), Norm: norm}
}
if tower.InputProj, err = makeClippableLinear(linears, tensors, at+"subsample_conv_projection.input_proj_linear"); err != nil {
return err
}
table := audioRelPosTable(a)
headDim := a.HiddenSize / a.NumAttentionHeads
positions := int32(table.Dim(0))
for i := range tower.Layers {
lp := fmt.Sprintf("%slayers.%d.", at, i)
layer := &AudioLayer{
FFW1: &AudioFeedForward{},
FFW2: &AudioFeedForward{},
Attention: &AudioAttention{},
LConv: &AudioLightConv{},
}
for _, ffw := range []struct {
f *AudioFeedForward
name string
}{{layer.FFW1, "feed_forward1."}, {layer.FFW2, "feed_forward2."}} {
fp := lp + ffw.name
if ffw.f.PreNorm, err = requiredNorm(fp + "pre_layer_norm.weight"); err != nil {
return err
}
if ffw.f.Up, err = makeClippableLinear(linears, tensors, fp+"ffw_layer_1"); err != nil {
return err
}
if ffw.f.Down, err = makeClippableLinear(linears, tensors, fp+"ffw_layer_2"); err != nil {
return err
}
if ffw.f.PostNorm, err = requiredNorm(fp + "post_layer_norm.weight"); err != nil {
return err
}
}
attn := layer.Attention
if attn.QProj, err = makeClippableLinear(linears, tensors, lp+"self_attn.q_proj"); err != nil {
return err
}
if attn.KProj, err = makeClippableLinear(linears, tensors, lp+"self_attn.k_proj"); err != nil {
return err
}
if attn.VProj, err = makeClippableLinear(linears, tensors, lp+"self_attn.v_proj"); err != nil {
return err
}
if attn.Post, err = makeClippableLinear(linears, tensors, lp+"self_attn.post"); err != nil {
return err
}
perDimScale, err := requiredNorm(lp + "self_attn.per_dim_scale")
if err != nil {
return err
}
relProj, err := makeClippableLinear(linears, tensors, lp+"self_attn.relative_k_proj")
if err != nil {
return err
}
// The reference applies softplus and the rel-position projection in
// the checkpoint dtype and lifts the results to f32 with q.
attn.QScale = mlx.Softplus(perDimScale).AsType(mlx.DTypeFloat32)
relK := relProj.Forward(table.AsType(mlx.DTypeBFloat16))
relK = mlx.Reshape(relK, positions, a.NumAttentionHeads, headDim)
relK = mlx.ExpandDims(mlx.Transpose(relK, 1, 2, 0), 0)
attn.RelK = relK.AsType(mlx.DTypeFloat32)
lc := layer.LConv
if lc.PreNorm, err = requiredNorm(lp + "lconv1d.pre_layer_norm.weight"); err != nil {
return err
}
if lc.Start, err = makeClippableLinear(linears, tensors, lp+"lconv1d.linear_start"); err != nil {
return err
}
dw := tensors[lp+"lconv1d.depthwise_conv1d.weight"]
if dw == nil {
return fmt.Errorf("missing audio weight: %slconv1d.depthwise_conv1d.weight", lp)
}
lc.DWConv = mlx.Transpose(dw, 0, 2, 1)
if lc.ConvNorm, err = requiredNorm(lp + "lconv1d.conv_norm.weight"); err != nil {
return err
}
if lc.End, err = makeClippableLinear(linears, tensors, lp+"lconv1d.linear_end"); err != nil {
return err
}
if layer.PreAttnNorm, err = requiredNorm(lp + "norm_pre_attn.weight"); err != nil {
return err
}
if layer.PostAttnNorm, err = requiredNorm(lp + "norm_post_attn.weight"); err != nil {
return err
}
if layer.OutNorm, err = requiredNorm(lp + "norm_out.weight"); err != nil {
return err
}
tower.Layers[i] = layer
}
if tower.OutputProj, err = makeClippableLinear(linears, tensors, at+"output_proj"); err != nil {
return err
}
projection, err := makeClippableLinear(linears, tensors, root+"embed_audio.embedding_projection")
if err != nil {
return err
}
tower.ClipMin = mlx.FromValue(-a.GradientClipping).AsType(mlx.DTypeBFloat16)
tower.ClipMax = mlx.FromValue(a.GradientClipping).AsType(mlx.DTypeBFloat16)
m.AudioTower = tower
m.EmbedAudio = &MultimodalEmbedder{Projection: projection}
return nil
}
// loadUnifiedAudioWeights loads the encoder-free variant, whose only audio
// weight is the embedding projection.
func (m *Model) loadUnifiedAudioWeights(tensors map[string]*mlx.Array, linears model.LinearFactory) error {
for _, root := range []string{"model.", ""} {
if tensors[root+"embed_audio.embedding_projection.weight"] == nil {
continue
}
projection, err := makeClippableLinear(linears, tensors, root+"embed_audio.embedding_projection")
if err != nil {
return err
}
m.EmbedAudio = &MultimodalEmbedder{Projection: projection}
return nil
}
return fmt.Errorf("config declares a gemma4_unified_audio embedder but embed_audio weights are missing from the manifest")
}
func (m *Model) audioLoaded() bool {
return m.EmbedAudio != nil
}
func (t *AudioTower) clip(x *mlx.Array) *mlx.Array {
return mlx.Clip(x, t.ClipMin, t.ClipMax)
}
// audioAttentionMask builds the boolean [blocks, 1, chunk, context] mask
// over n frames: block b's context slot j holds frame b*chunk + j -
// (ContextLeft-1), attendable when it exists and lies within the window.
func (m *Model) audioAttentionMask(n int) *mlx.Array {
a := m.Audio
chunk, left, right := int(a.ChunkSize), int(a.ContextLeft-1), int(a.ContextRight)
ctx := chunk + left + right
blocks := (n + chunk - 1) / chunk
mask := make([]bool, blocks*chunk*ctx)
i := 0
for b := range blocks {
for q := range chunk {
gq := b*chunk + q
for j := range ctx {
gk := b*chunk + j - left
dist := gq - gk
mask[i] = gq < n && gk >= 0 && gk < n &&
((dist >= 0 && dist < left) || (dist < 0 && -dist < right))
i++
}
}
}
return mlx.FromValues(mask, blocks, 1, chunk, ctx)
}
// audioContextIndices maps block b's context slot j to index b*chunk+j of
// the padded key/value sequence.
func (m *Model) audioContextIndices(blocks int) *mlx.Array {
a := m.Audio
ctx := int(a.ChunkSize + a.ContextLeft - 1 + a.ContextRight)
idx := make([]int32, blocks*ctx)
for b := range blocks {
for j := range ctx {
idx[b*ctx+j] = int32(b*int(a.ChunkSize) + j)
}
}
return mlx.FromValues(idx, blocks*ctx)
}
// The literal 1e-6 eps here and in AudioLayer.Forward mirrors the
// reference, which passes config eps only to the lconv and subsample norms.
func (f *AudioFeedForward) Forward(t *AudioTower, h *mlx.Array, a *AudioConfig) *mlx.Array {
c := t.clip(h)
c = mlx.RMSNormFn(c, f.PreNorm, 1e-6)
c = f.Up.Forward(c)
c = mlx.SiLU(c)
c = f.Down.Forward(c)
c = t.clip(c)
c = mlx.RMSNormFn(c, f.PostNorm, 1e-6)
return mlx.Add(h, mlx.MulScalar(c, a.ResidualWeight))
}
// Forward computes chunked local attention with relative positions over the
// pre-normed [n, hidden] frames: q/k/v lift to f32, queries block into
// chunks, keys and values gather each block's context window, and the
// rel-position scores shift from position- to offset-indexed before the
// softcap, mask, and softmax.
func (at *AudioAttention) Forward(t *AudioTower, h, mask, indices *mlx.Array, a *AudioConfig) *mlx.Array {
n := h.Dim(0)
heads := a.NumAttentionHeads
headDim := a.HiddenSize / heads
chunk, left, right := int(a.ChunkSize), int(a.ContextLeft-1), int(a.ContextRight)
ctx := chunk + left + right
blocks := (n + chunk - 1) / chunk
shape := []int32{int32(n), heads, headDim}
q := mlx.Reshape(at.QProj.Forward(h).AsType(mlx.DTypeFloat32), shape...)
k := mlx.Reshape(at.KProj.Forward(h).AsType(mlx.DTypeFloat32), shape...)
v := mlx.Reshape(at.VProj.Forward(h).AsType(mlx.DTypeFloat32), shape...)
q = mlx.MulScalar(q, float32(1/(math.Sqrt(float64(headDim))*math.Ln2)))
q = mlx.Mul(q, at.QScale)
k = mlx.MulScalar(k, float32(math.Log(1+math.E)/math.Ln2))
q = mlx.PadConstant(q, []int{0}, []int{0}, []int{blocks*chunk - n})
qb := mlx.Transpose(mlx.Reshape(q, int32(blocks), int32(chunk), heads, headDim), 0, 2, 1, 3)
pad := []int{blocks*chunk - n + right}
kb := gatherContext(k, indices, left, pad, int32(blocks), int32(ctx), heads, headDim)
vb := gatherContext(v, indices, left, pad, int32(blocks), int32(ctx), heads, headDim)
ac := mlx.Matmul(qb, mlx.Transpose(kb, 0, 1, 3, 2))
bd := mlx.Matmul(qb, at.RelK)
// Rel shift: [.., chunk, positions] -> [.., chunk, context], turning
// per-position columns into per-offset columns.
positions := bd.Dim(3)
bd = mlx.PadConstant(bd, []int{3}, []int{0}, []int{ctx + 1 - positions})
bd = mlx.Reshape(bd, int32(blocks), heads, int32(chunk*(ctx+1)))
bd = bd.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, chunk*ctx))
bd = mlx.Reshape(bd, int32(blocks), heads, int32(chunk), int32(ctx))
logits := mlx.Add(ac, bd)
logits = mlx.LogitSoftcap(logits, mlx.FromValue(a.LogitCap))
logits = mlx.Where(mask, logits, mlx.FromValue(a.InvalidLogit))
weights := mlx.SoftmaxAxis(logits, -1, false)
o := mlx.Matmul(weights, vb)
o = mlx.Transpose(o, 0, 2, 1, 3)
o = mlx.Reshape(o, int32(blocks*chunk), a.HiddenSize)
o = o.Slice(mlx.Slice(0, n), mlx.Slice())
return at.Post.Forward(o.AsType(mlx.DTypeBFloat16))
}
// gatherContext pads the [n, heads, headDim] sequence so index b*chunk+j
// addresses frame b*chunk+j-left, then gathers the per-block context
// windows into [blocks, heads, context, headDim].
func gatherContext(x, indices *mlx.Array, left int, highPad []int, blocks, ctx, heads, headDim int32) *mlx.Array {
x = mlx.PadConstant(x, []int{0}, []int{left}, highPad)
x = mlx.Take(x, indices, 0)
x = mlx.Reshape(x, blocks, ctx, heads, headDim)
return mlx.Transpose(x, 0, 2, 1, 3)
}
func (lc *AudioLightConv) Forward(t *AudioTower, h *mlx.Array, a *AudioConfig) *mlx.Array {
n, hidden := int32(h.Dim(0)), a.HiddenSize
c := mlx.RMSNormFn(h, lc.PreNorm, a.RMSNormEps)
c = lc.Start.Forward(c)
c = mlx.Mul(
c.Slice(mlx.Slice(), mlx.Slice(0, int(hidden))),
mlx.Sigmoid(c.Slice(mlx.Slice(), mlx.Slice(int(hidden), int(2*hidden)))))
// Causal depthwise conv: left-pad time by kernel-1.
c = mlx.Reshape(c, 1, n, hidden)
c = mlx.PadConstant(c, []int{1}, []int{int(a.ConvKernelSize) - 1}, []int{0})
c = mlx.Conv1d(c, lc.DWConv, nil, 1, 0, 1, hidden)
c = mlx.Reshape(c, n, hidden)
c = t.clip(c)
c = mlx.RMSNormFn(c, lc.ConvNorm, a.RMSNormEps)
c = mlx.SiLU(c)
c = lc.End.Forward(c)
return mlx.Add(h, c)
}
func (l *AudioLayer) Forward(t *AudioTower, h, mask, indices *mlx.Array, a *AudioConfig) *mlx.Array {
h = l.FFW1.Forward(t, h, a)
c := t.clip(h)
c = mlx.RMSNormFn(c, l.PreAttnNorm, 1e-6)
c = l.Attention.Forward(t, c, mask, indices, a)
c = t.clip(c)
c = mlx.RMSNormFn(c, l.PostAttnNorm, 1e-6)
h = mlx.Add(h, c)
h = l.LConv.Forward(t, h, a)
h = l.FFW2.Forward(t, h, a)
h = t.clip(h)
return mlx.RMSNormFn(h, l.OutNorm, 1e-6)
}
// encodeAudio runs the conformer over one chunk's [frames, melBins] log-mel
// features, returning the lazy [numTokens, hidden] features.
func (m *Model) encodeAudio(data *mlx.Array) *mlx.Array {
a := m.Audio
t := m.AudioTower
frames := data.Dim(0)
h := mlx.Reshape(data, 1, int32(frames), audioMelBins, 1).AsType(mlx.DTypeBFloat16)
for _, l := range t.SubsampleConv {
h = mlx.Conv2d(h, l.Conv, 2, 2, 1, 1, 1, 1, 1)
h = mlx.LayerNormFn(h, l.Norm, nil, a.RMSNormEps)
h = mlx.ReLU(h)
}
n, w, c := h.Dim(1), h.Dim(2), h.Dim(3)
h = mlx.Reshape(h, int32(n), int32(w*c))
h = t.InputProj.Forward(h)
mask := m.audioAttentionMask(n)
indices := m.audioContextIndices((n + int(a.ChunkSize) - 1) / int(a.ChunkSize))
for _, layer := range t.Layers {
h = layer.Forward(t, h, mask, indices, a)
}
h = t.OutputProj.Forward(h)
h = mlx.RMSNormFn(h, nil, a.RMSNormEps)
return m.EmbedAudio.Projection.Forward(h)
}
// encodeUnifiedAudio projects raw waveform frames straight into the text
// embedding space; the unified variant has no audio tower.
func (m *Model) encodeUnifiedAudio(data *mlx.Array) *mlx.Array {
h := data.AsType(mlx.DTypeBFloat16)
h = mlx.RMSNormFn(h, nil, m.Audio.RMSNormEps)
return m.EmbedAudio.Projection.Forward(h)
}
x/models/gemma4/audio_test.go
+221
View File
@@ -0,0 +1,221 @@
package gemma4
import (
"strings"
"testing"
"github.com/ollama/ollama/x/internal/mlxtest"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/mlxrunner/model/base"
)
func newAudioTestModel() *Model {
audio := &AudioConfig{ModelType: "gemma4_audio"}
applyAudioDefaults(audio)
return &Model{
Audio: audio,
EmbedAudio: &MultimodalEmbedder{},
MM: multimodalConfig{
AudioConfig: audio,
AudioTokenID: 258881,
BOATokenID: 256000,
EOATokenID: 258883,
},
}
}
func TestPrepareAudioMedia(t *testing.T) {
m := newAudioTestModel()
prepared, err := m.PrepareMedia([]base.Segment{
{Tokens: []int32{2, 5}},
{Kind: "audio", Data: wavPCM16(16000)},
{Tokens: []int32{7}},
})
if err != nil {
t.Fatal(err)
}
// One second: 99 mel frames -> 25 soft tokens.
wantLen := 2 + 1 + 25 + 1 + 1
if len(prepared.Tokens) != wantLen {
t.Fatalf("%d tokens, want %d", len(prepared.Tokens), wantLen)
}
if prepared.Tokens[2] != m.MM.BOATokenID || prepared.Tokens[28] != m.MM.EOATokenID {
t.Fatalf("delimiters = %d, %d", prepared.Tokens[2], prepared.Tokens[28])
}
for _, tok := range prepared.Tokens[3:28] {
if tok != m.MM.AudioTokenID {
t.Fatalf("soft token = %d", tok)
}
}
if len(prepared.Items) != 1 {
t.Fatalf("%d items", len(prepared.Items))
}
item := prepared.Items[0]
if item.Range != [2]int{3, 28} || item.Source != 1 || !item.Causal {
t.Fatalf("item = %+v", item)
}
if item.Dims[0] != 99 || item.Dims[1] != audioMelBins {
t.Fatalf("dims = %v", item.Dims)
}
if p := item.Opaque.(preparedAudio); p.numTokens != 25 {
t.Fatalf("numTokens = %d", p.numTokens)
}
}
func TestPrepareAudioMediaChunks(t *testing.T) {
m := newAudioTestModel()
// 61 s: three chunks of at most 30 s, one soft-token run each, back to
// back inside one boa/eoa pair.
prepared, err := m.PrepareMedia([]base.Segment{{Kind: "audio", Data: wavPCM16(976000)}})
if err != nil {
t.Fatal(err)
}
if len(prepared.Items) != 3 {
t.Fatalf("%d items", len(prepared.Items))
}
next := 1
for i, item := range prepared.Items {
n := item.Opaque.(preparedAudio).numTokens
if item.Range != [2]int{next, next + int(n)} || item.Source != 0 || !item.Causal {
t.Fatalf("item %d range %v source %d causal %v, want [%d %d]", i, item.Range, item.Source, item.Causal, next, next+int(n))
}
if n > 750 {
t.Fatalf("item %d: %d tokens exceed one 30 s chunk", i, n)
}
next += int(n)
}
if len(prepared.Tokens) != next+1 {
t.Fatalf("%d tokens, want %d", len(prepared.Tokens), next+1)
}
if prepared.Tokens[0] != m.MM.BOATokenID || prepared.Tokens[next] != m.MM.EOATokenID {
t.Fatalf("delimiters = %d, %d", prepared.Tokens[0], prepared.Tokens[next])
}
}
func TestPrepareAudioMediaRejections(t *testing.T) {
m := newAudioTestModel()
_, err := m.PrepareMedia([]base.Segment{{Kind: "audio", Data: []byte("ID3\x04\x00junk")}})
if err == nil || !strings.Contains(err.Error(), "unrecognized audio format") {
t.Fatalf("mp3 error = %v", err)
}
_, err = m.PrepareMedia([]base.Segment{{Kind: "video", Data: []byte{1}}})
if err == nil || !strings.Contains(err.Error(), "does not support video input") {
t.Fatalf("video error = %v", err)
}
noAudio := &Model{MM: multimodalConfig{ImageTokenID: 258880}}
_, err = noAudio.PrepareMedia([]base.Segment{{Kind: "audio", Data: wavPCM16(16000)}})
if err == nil || !strings.Contains(err.Error(), "does not support audio input") {
t.Fatalf("no-audio error = %v", err)
}
}
func TestPrepareUnifiedAudioMedia(t *testing.T) {
m := newAudioTestModel()
m.Audio.ModelType = "gemma4_unified_audio"
m.Audio.SamplesPerToken = 640
// 31 s stays one item on the unified path: no chunking, one token per
// 640-sample frame, the final partial frame zero-padded.
prepared, err := m.PrepareMedia([]base.Segment{{Kind: "audio", Data: wavPCM16(496001)}})
if err != nil {
t.Fatal(err)
}
wantTokens := (496001 + 639) / 640
if want := 1 + wantTokens + 1; len(prepared.Tokens) != want {
t.Fatalf("%d tokens, want %d", len(prepared.Tokens), want)
}
if len(prepared.Items) != 1 {
t.Fatalf("%d items", len(prepared.Items))
}
item := prepared.Items[0]
if item.Range != [2]int{1, 1 + wantTokens} || !item.Causal {
t.Fatalf("item = %+v", item)
}
if item.Dims[0] != wantTokens || item.Dims[1] != 640 {
t.Fatalf("dims = %v", item.Dims)
}
if len(item.MediaData) != wantTokens*640 {
t.Fatalf("media data length = %d", len(item.MediaData))
}
}
func TestParseAudioConfig(t *testing.T) {
mm, err := parseMultimodalConfig([]byte(`{
"audio_config": {"model_type": "gemma4_unified_audio", "audio_samples_per_token": 640},
"audio_token_id": 258881, "boa_token_id": 256000, "eoa_token_index": 258883
}`))
if err != nil {
t.Fatal(err)
}
if mm.AudioConfig == nil || !mm.AudioConfig.unified() {
t.Fatalf("audio config = %+v", mm.AudioConfig)
}
if mm.EOATokenID != 258883 {
t.Fatalf("eoa = %d", mm.EOATokenID)
}
if mm.AudioConfig.RMSNormEps != 1e-6 {
t.Fatalf("eps = %v", mm.AudioConfig.RMSNormEps)
}
if _, err := parseMultimodalConfig([]byte(`{
"audio_config": {"model_type": "gemma4_unified_audio", "audio_samples_per_token": -640}}`)); err == nil {
t.Fatal("negative audio_samples_per_token accepted")
}
mm, err = parseMultimodalConfig([]byte(`{"audio_config": {"model_type": "someday_audio"}}`))
if err != nil {
t.Fatal(err)
}
if mm.AudioConfig != nil {
t.Fatalf("unknown audio model type accepted: %+v", mm.AudioConfig)
}
mm, err = parseMultimodalConfig([]byte(`{"audio_config": null}`))
if err != nil || mm.AudioConfig != nil {
t.Fatalf("null audio config: %+v, %v", mm.AudioConfig, err)
}
}
// TestAudioAttentionMask checks every (block, query, context slot) of the
// uploaded mask against the window rule, covering ragged final blocks.
func TestAudioAttentionMask(t *testing.T) {
mlxtest.Run(t, func(t *mlxtest.T) {
m := newAudioTestModel()
a := m.Audio
chunk, left := int(a.ChunkSize), int(a.ContextLeft-1)
ctx := chunk + left + int(a.ContextRight)
for _, n := range []int{5, 24, 26, 100} {
mask := m.audioAttentionMask(n).AsType(mlx.DTypeFloat32)
mlx.Eval(mask)
got := mask.Floats()
blocks := (n + chunk - 1) / chunk
if len(got) != blocks*chunk*ctx {
t.Fatalf("n=%d: %d values", n, len(got))
}
for b := range blocks {
for q := range chunk {
for j := range ctx {
gq, gk := b*chunk+q, b*chunk+j-left
dist := gq - gk
want := float32(0)
if gq < n && gk >= 0 && gk < n && dist >= 0 && dist < left {
want = 1
}
if v := got[(b*chunk+q)*ctx+j]; v != want {
t.Fatalf("n=%d block %d q %d j %d: %v, want %v", n, b, q, j, v, want)
}
}
}
}
}
})
}
x/models/gemma4/gemma4.go
+99 -19
View File
@@ -1,4 +1,4 @@
// Package gemma4 provides the Gemma 4 text model implementation for MLX.
// Package gemma4 provides the Gemma 4 model implementation for MLX.
package gemma4
import (
@@ -66,7 +66,10 @@ type TextConfig struct {
TopKExperts int32 `json:"top_k_experts"`
ExpertIntermediateSize int32 `json:"moe_intermediate_size"`
RopeParameters map[string]*RopeParams `json:"rope_parameters"`
ImageTokenIDValue int32 `json:"image_token_id"`
// UseBidirectionalAttention selects the image-span mask semantics:
// "vision" relaxes sliding layers over soft-token runs, empty means
// causal even with images; the reference's "all" is rejected at load.
UseBidirectionalAttention string `json:"use_bidirectional_attention"`
// Quantization parameters.
QuantGroupSize int `json:"-"`
@@ -75,13 +78,14 @@ type TextConfig struct {
TensorQuant map[string]*model.TensorQuantInfo `json:"-"`
// Computed fields.
SlidingScale float32 `json:"-"` // 1/sqrt(HeadDim) for sliding layers
FullScale float32 `json:"-"` // 1/sqrt(GlobalHeadDim) for full layers
SlidingRopeDims int `json:"-"` // HeadDim (full rotation for sliding)
FullRopeDims int `json:"-"` // GlobalHeadDim (partial rotation via custom freqs)
SlidingRopeBase float32 `json:"-"`
FullRopeBase float32 `json:"-"`
FullRopeFreqs *mlx.Array `json:"-"` // Precomputed proportional RoPE frequencies
BidirectionalVisionAttention bool `json:"-"` // UseBidirectionalAttention == "vision"
SlidingScale float32 `json:"-"` // 1/sqrt(HeadDim) for sliding layers
FullScale float32 `json:"-"` // 1/sqrt(GlobalHeadDim) for full layers
SlidingRopeDims int `json:"-"` // HeadDim (full rotation for sliding)
FullRopeDims int `json:"-"` // GlobalHeadDim (partial rotation via custom freqs)
SlidingRopeBase float32 `json:"-"`
FullRopeBase float32 `json:"-"`
FullRopeFreqs *mlx.Array `json:"-"` // Precomputed proportional RoPE frequencies
// Precomputed scale factors (avoid per-forward math.Sqrt/Pow).
EmbedScale float32 `json:"-"` // sqrt(hidden_size)
@@ -384,6 +388,22 @@ type Model struct {
NormScaled *mlx.Array
PerLayerProjNormWeight *mlx.Array
// Vision components; at most one of VisionTower and UnifiedEmbedder is
// set, per the checkpoint's vision_config model type.
VisionTower *VisionTower
UnifiedEmbedder *UnifiedVisionEmbedder
EmbedVision *MultimodalEmbedder
Vision *VisionConfig
// Audio components.
AudioTower *AudioTower
EmbedAudio *MultimodalEmbedder
Audio *AudioConfig
MM multimodalConfig
// Soft-token placeholder IDs of the configured media modalities.
mediaPlaceholderIDs []int32
tok *tokenizer.Tokenizer
*TextConfig
@@ -434,6 +454,8 @@ func parseTextConfig(configData []byte) (TextConfig, error) {
cfg.MaxPositionEmbeddings = 131072
}
cfg.BidirectionalVisionAttention = cfg.UseBidirectionalAttention == "vision"
// Gemma 4 uses scaling=1.0 (no 1/sqrt(head_dim) scaling); the Q/K norms
// handle magnitude control. This differs from Gemma 3 which uses
// query_pre_attn_scalar^(-0.5).
@@ -611,6 +633,12 @@ func newModel(root *model.Root) (base.Model, error) {
return nil, err
}
// The reference's "all" mode makes every layer bidirectional over the
// whole sequence, which chunked prefill and prefix reuse cannot serve.
if b := cfg.UseBidirectionalAttention; b != "" && b != "vision" {
return nil, fmt.Errorf("unsupported use_bidirectional_attention %q", b)
}
if qt := root.QuantType(); qt != "" {
cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams(qt)
if gs := root.GroupSize(); gs > 0 {
@@ -641,13 +669,32 @@ func newModel(root *model.Root) (base.Model, error) {
return nil, fmt.Errorf("parse tokenizer: %w", err)
}
mm, err := parseMultimodalConfig(configData)
if err != nil {
return nil, err
}
m := &Model{
Layers: make([]*DecoderLayer, cfg.NumHiddenLayers),
TextConfig: &cfg,
MM: mm,
Vision: mm.VisionConfig,
Audio: mm.AudioConfig,
tok: tok,
SuppressLogitBias: makeSuppressLogitBias(suppressTokens, cfg.VocabSize),
}
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 {
m.mediaPlaceholderIDs = append(m.mediaPlaceholderIDs, mm.AudioTokenID)
}
m.validateMediaTokens()
for i := range m.Layers {
donor, isShared := cfg.KVShareMap[int32(i)]
if !isShared {
@@ -976,6 +1023,17 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
return fmt.Errorf("missing precomputed final norm weight")
}
if m.Vision != nil {
if err := m.loadVisionWeights(tensors, linears); err != nil {
return err
}
}
if m.Audio != nil {
if err := m.loadAudioWeights(tensors, linears); err != nil {
return err
}
}
return nil
}
@@ -983,15 +1041,32 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden
dims := b.InputIDs.Dims()
B, L := int32(dims[0]), int32(dims[1])
positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets))
h := m.EmbedTokens.Forward(b.InputIDs)
h = mlx.MulScalar(h, m.EmbedScale)
// Compute PLE inputs if configured.
// Media placeholder rows embed the pad token and are then overwritten
// by the feature scatter.
ids := b.InputIDs
if len(b.Media) > 0 {
for _, id := range m.mediaPlaceholderIDs {
ids = mlx.Where(b.InputIDs.Equal(mlx.FromValue(int(id))), mlx.FromValue(0), ids)
}
}
h := m.EmbedTokens.Forward(ids)
h = mlx.MulScalar(h, m.EmbedScale)
if len(b.Media) > 0 {
h = m.scatterMedia(h, b)
}
// PLE's token-identity component reads the masked IDs, but its
// projection component reads the merged hidden — media rows project
// their features, not the pad embedding.
var perLayerInputs *mlx.Array
if m.HiddenSizePerLayer > 0 && m.EmbedTokensPerLayer != nil {
perLayerInputs = m.computePLEInputs(b.InputIDs, h)
perLayerInputs = m.computePLEInputs(ids, h)
}
slidingMask, fullMask := m.buildMasks(b)
// KV sharing: each donor layer stores its KVHistory here so later
// shared layers can reuse it in lieu of their own cache update.
var sharedKV map[int32]sharedHistory
@@ -1019,8 +1094,13 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden
}
}
mask := fullMask
if layer.IsSliding {
mask = slidingMask
}
var donorKV *sharedHistory
h, donorKV = layer.Forward(h, b, c, positions, B, L, m.TextConfig, pleInput, donor)
h, donorKV = layer.Forward(h, b, c, positions, B, L, m.TextConfig, pleInput, donor, mask)
// If this layer is a donor, store its cached KV for later shared layers.
if layer.IsDonor && donorKV != nil {
@@ -1150,9 +1230,9 @@ func sliceLayerDim(combined *mlx.Array, layerIdx, B, L, pleDim int32) *mlx.Array
return mlx.Squeeze(sliced, 2)
}
func (l *DecoderLayer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *TextConfig, pleInput *mlx.Array, donor *sharedHistory) (*mlx.Array, *sharedHistory) {
func (l *DecoderLayer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *TextConfig, pleInput *mlx.Array, donor *sharedHistory, mask nn.AttentionMask) (*mlx.Array, *sharedHistory) {
normed := mlx.RMSNormFn(x, l.InputNormScaled, cfg.RMSNormEps)
attnOut, kv := l.Attention.Forward(normed, b, c, positions, B, L, l.IsSliding, cfg, donor)
attnOut, kv := l.Attention.Forward(normed, b, c, positions, B, L, l.IsSliding, cfg, donor, mask)
attnOut = mlx.RMSNormFn(attnOut, l.PostAttnNormScaled, cfg.RMSNormEps)
h := mlx.Add(x, attnOut)
@@ -1200,7 +1280,7 @@ func (l *DecoderLayer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, posi
return h, kv
}
func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, isSliding bool, cfg *TextConfig, donor *sharedHistory) (*mlx.Array, *sharedHistory) {
func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, isSliding bool, cfg *TextConfig, donor *sharedHistory, baseMask nn.AttentionMask) (*mlx.Array, *sharedHistory) {
// Determine head dim and scale based on layer type.
headDim := cfg.HeadDim
scale := cfg.SlidingScale
@@ -1276,7 +1356,7 @@ func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positio
// kernel only handles L < 4 (generation). For prefill, we fall back
// to explicit matmul+softmax+matmul on CUDA.
var k, v *mlx.Array
mask := nn.CausalMask().Intersect(nn.QPaddingMask(b, q.DType()))
mask := baseMask.Intersect(nn.QPaddingMask(b, q.DType()))
if kv.history != nil {
k, v = kv.history.K(), kv.history.V()
mask = kv.history.Mask(mask)
@@ -1306,7 +1386,7 @@ func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positio
out = mlx.Reshape(out, B, cfg.NumAttentionHeads, L, headDim)
} else {
var opt nn.SDPAOption
mask := nn.CausalMask()
mask := baseMask
if kv.history != nil {
opt = nn.WithKVHistory(kv.history)
} else {
x/models/gemma4/media.go
+330
View File
@@ -0,0 +1,330 @@
package gemma4
import (
"encoding/json"
"fmt"
"log/slog"
"github.com/ollama/ollama/x/mlxrunner/batch"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/mlxrunner/model"
"github.com/ollama/ollama/x/mlxrunner/model/base"
"github.com/ollama/ollama/x/models/nn"
)
// 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"`
AudioConfig *AudioConfig `json:"audio_config"`
AudioTokenID int32 `json:"audio_token_id"`
BOATokenID int32 `json:"boa_token_id"`
EOATokenID int32 `json:"eoa_token_id"`
// Some checkpoints carry the end-of-audio token only under this key.
EOATokenIndex int32 `json:"eoa_token_index"`
}
func parseMultimodalConfig(configData []byte) (multimodalConfig, error) {
var mm multimodalConfig
if err := json.Unmarshal(configData, &mm); err != nil {
return multimodalConfig{}, fmt.Errorf("parse multimodal config: %w", err)
}
switch v := mm.VisionConfig; {
case v != nil && v.ModelType == "gemma4_vision":
if v.HeadDim == 0 && v.NumAttentionHeads > 0 {
v.HeadDim = v.HiddenSize / v.NumAttentionHeads
}
if v.PatchSize == 0 {
v.PatchSize = 16
}
if v.PoolingKernelSize == 0 {
v.PoolingKernelSize = 3
}
if v.DefaultOutputLen == 0 {
v.DefaultOutputLen = 280
}
if v.RMSNormEps == 0 {
v.RMSNormEps = 1e-6
}
v.RopeTheta = 100
if v.RopeParameters != nil && v.RopeParameters.RopeTheta > 0 {
v.RopeTheta = v.RopeParameters.RopeTheta
}
case v != nil && v.unified():
if v.PatchSize == 0 {
v.PatchSize = 16
}
if v.PoolingKernelSize == 0 {
v.PoolingKernelSize = 3
}
if v.NumSoftTokens == 0 {
v.NumSoftTokens = 280
}
if v.RMSNormEps == 0 {
v.RMSNormEps = 1e-6
}
default:
mm.VisionConfig = nil
}
switch a := mm.AudioConfig; {
case a != nil && a.ModelType == "gemma4_audio":
applyAudioDefaults(a)
case a != nil && a.unified():
if a.SamplesPerToken == 0 {
a.SamplesPerToken = 640
}
if a.SamplesPerToken < 0 {
return multimodalConfig{}, fmt.Errorf("invalid audio_samples_per_token %d", a.SamplesPerToken)
}
if a.RMSNormEps == 0 {
a.RMSNormEps = 1e-6
}
default:
mm.AudioConfig = nil
}
if mm.EOATokenID == 0 {
mm.EOATokenID = mm.EOATokenIndex
}
return mm, nil
}
// MultimodalEmbedder projects pooled features into the text embedding
// space: a scale-less RMSNorm then a linear.
type MultimodalEmbedder struct {
Projection nn.LinearLayer
}
// ClippableLinear clamps a linear's input and output to checkpoint-provided
// bounds; the loader guarantees all four are present.
type ClippableLinear struct {
Inner nn.LinearLayer
InputMin, InputMax, OutputMin, OutputMax *mlx.Array
}
func (c *ClippableLinear) Forward(x *mlx.Array) *mlx.Array {
x = mlx.Clip(x, c.InputMin, c.InputMax)
x = c.Inner.Forward(x)
return mlx.Clip(x, c.OutputMin, c.OutputMax)
}
func (c *ClippableLinear) OutputDim() int32 { return c.Inner.OutputDim() }
// makeClippableLinear builds one tower projection: the weight sits under a
// .linear suffix, with optional clamp scalars beside it.
func makeClippableLinear(linears model.LinearFactory, tensors map[string]*mlx.Array, name string) (nn.LinearLayer, error) {
inner := linears.Make(name + ".linear")
if inner == nil {
inner = linears.Make(name)
}
if inner == nil {
return nil, fmt.Errorf("missing weight: %s", name)
}
c := &ClippableLinear{
Inner: inner,
InputMin: tensors[name+".input_min"],
InputMax: tensors[name+".input_max"],
OutputMin: tensors[name+".output_min"],
OutputMax: tensors[name+".output_max"],
}
if c.InputMin == nil && c.InputMax == nil && c.OutputMin == nil && c.OutputMax == nil {
return inner, nil
}
if c.InputMin == nil || c.InputMax == nil || c.OutputMin == nil || c.OutputMax == nil {
return nil, fmt.Errorf("weight %s has a partial clamp set", name)
}
return c, nil
}
// PrepareMedia implements base.MediaModel: splice each media segment's
// placeholder expansion — the begin token, the soft-token run(s), the end
// token — into the stream, with all preprocessing on the CPU. An audio
// segment yields one item per chunk, all inside one boa/eoa pair; the runs
// attend causally, so chunked prefill may split them.
func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) {
prepared := &base.PreparedRequest{}
for s, seg := range segments {
switch seg.Kind {
case "":
prepared.Tokens = append(prepared.Tokens, seg.Tokens...)
case "image":
if !m.visionLoaded() {
return nil, fmt.Errorf("this model does not support image input")
}
pixels, positions, geom, err := m.preprocessImage(seg.Data)
if err != nil {
return nil, err
}
start := len(prepared.Tokens)
prepared.Tokens = append(prepared.Tokens, m.MM.BOITokenID)
for range geom.NumSoftTokens {
prepared.Tokens = append(prepared.Tokens, m.MM.ImageTokenID)
}
prepared.Tokens = append(prepared.Tokens, m.MM.EOITokenID)
n := int(geom.PatchesH * geom.PatchesW)
prepared.Items = append(prepared.Items, base.PreparedItem{
Range: [2]int{start, len(prepared.Tokens)},
Source: s,
MediaData: pixels,
Dims: []int{n, len(pixels) / n},
Opaque: preparedImage{positions: positions, geom: geom},
})
case "audio":
if !m.audioLoaded() {
return nil, fmt.Errorf("this model does not support audio input")
}
// A chunk's MediaData rows are mel frames for the conformer
// and one raw waveform frame per soft token for the unified
// embedder.
var chunks []audioChunk
if m.Audio.unified() {
frames, tokens, err := processUnifiedAudio(seg.Data, int(m.Audio.SamplesPerToken))
if err != nil {
return nil, err
}
chunks = []audioChunk{{data: frames, frames: tokens, numTokens: tokens}}
} else {
var err error
if chunks, err = processAudio(seg.Data); err != nil {
return nil, err
}
}
prepared.Tokens = append(prepared.Tokens, m.MM.BOATokenID)
for _, chunk := range chunks {
start := len(prepared.Tokens)
for range chunk.numTokens {
prepared.Tokens = append(prepared.Tokens, m.MM.AudioTokenID)
}
prepared.Items = append(prepared.Items, base.PreparedItem{
Range: [2]int{start, len(prepared.Tokens)},
Source: s,
MediaData: chunk.data,
Dims: []int{chunk.frames, len(chunk.data) / chunk.frames},
Opaque: preparedAudio{numTokens: int32(chunk.numTokens)},
Causal: true,
})
}
prepared.Tokens = append(prepared.Tokens, m.MM.EOATokenID)
default:
return nil, fmt.Errorf("gemma4 does not support %s input", seg.Kind)
}
}
return prepared, nil
}
// EncodeMedia implements base.MediaModel: run the matching encoder over one
// whole item, returning the lazy [soft tokens, hidden] features.
func (m *Model) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array {
switch p := item.Opaque.(type) {
case preparedAudio:
if m.AudioTower != nil {
return m.encodeAudio(data)
}
return m.encodeUnifiedAudio(data)
case preparedImage:
if m.UnifiedEmbedder != nil {
return m.encodeUnifiedImage(data, p.positions, p.geom)
}
return m.encodeImage(data, p.positions, p.geom)
}
panic("gemma4: unknown media item")
}
// softRun returns a media item's feature-bearing token range: an image
// expansion is boi + soft*N + eoi, so the run starts one past the splice;
// an audio item's range is exactly its soft run.
func softRun(item batch.MediaItem) (start, end int) {
switch p := item.Opaque.(type) {
case preparedAudio:
return item.Pos, item.Pos + int(p.numTokens)
case preparedImage:
return item.Pos + 1, item.Pos + 1 + int(p.geom.NumSoftTokens)
}
panic("gemma4: unknown media item")
}
// buildMasks returns the per-layer-type masks. Both start as unmaterialized
// causal masks; when the checkpoint enables bidirectional vision attention,
// the sliding mask relaxes each image run intersecting the chunk while the
// full-attention mask stays causal.
func (m *Model) buildMasks(b *batch.Batch) (sliding, full nn.AttentionMask) {
sliding, full = nn.CausalMask(), nn.CausalMask()
if !m.BidirectionalVisionAttention || len(b.Media) == 0 {
return
}
for _, item := range b.Media {
// Only image runs are bidirectional; audio attends causally.
if _, ok := item.Opaque.(preparedImage); !ok {
continue
}
start, end := softRun(item)
off := int(b.SeqOffsets[item.Seq])
if end <= off || start >= off+int(b.SeqQueryLens[item.Seq]) {
continue
}
sliding = sliding.Relax(item.Seq, start, end, start, end)
}
return sliding, full
}
// scatterMedia overwrites the soft-token rows this chunk covers with the
// item's projected features.
func (m *Model) scatterMedia(h *mlx.Array, b *batch.Batch) *mlx.Array {
for _, item := range b.Media {
if item.Features == nil {
continue
}
start, end := softRun(item)
off := int(b.SeqOffsets[item.Seq])
qLo := max(start, off)
qHi := min(end, off+int(b.SeqQueryLens[item.Seq]))
if qHi <= qLo {
continue
}
feat := item.Features.Slice(mlx.Slice(qLo-start, qHi-start), mlx.Slice())
feat = mlx.Reshape(feat.AsType(h.DType()), 1, int32(qHi-qLo), m.HiddenSize)
h = h.SliceUpdate(feat, mlx.Slice(item.Seq, item.Seq+1), mlx.Slice(qLo-off, qHi-off), mlx.Slice())
}
return h
}
// validateMediaTokens warns when the tokenizer's special tokens disagree
// with the config's media token IDs; the config values are authoritative.
func (m *Model) validateMediaTokens() {
type mediaToken struct {
name string
id int32
}
var tokens []mediaToken
if m.Vision != nil {
tokens = append(tokens,
mediaToken{"<|image|>", m.MM.ImageTokenID},
mediaToken{"<|image>", m.MM.BOITokenID},
mediaToken{"<image|>", m.MM.EOITokenID})
}
if m.Audio != nil {
tokens = append(tokens,
mediaToken{"<|audio|>", m.MM.AudioTokenID},
mediaToken{"<|audio>", m.MM.BOATokenID},
mediaToken{"<audio|>", m.MM.EOATokenID})
}
for _, tok := range tokens {
if id, ok := m.tok.GetSpecialToken(tok.name); ok && id != tok.id {
slog.Warn("media token mismatch", "token", tok.name, "config", tok.id, "tokenizer", id)
}
}
}
x/models/gemma4/media_test.go
+270
View File
@@ -0,0 +1,270 @@
package gemma4
import (
"bytes"
"image"
"image/png"
"testing"
"github.com/ollama/ollama/x/internal/mlxtest"
"github.com/ollama/ollama/x/mlxrunner/batch"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/mlxrunner/model/base"
)
func TestParseMultimodalConfig(t *testing.T) {
visionJSON := []byte(`{
"image_token_id": 258880, "boi_token_id": 255999, "eoi_token_id": 258882,
"vision_soft_tokens_per_image": 280,
"vision_config": {"model_type": "gemma4_vision", "hidden_size": 1152,
"num_hidden_layers": 27, "num_attention_heads": 16, "head_dim": 72,
"patch_size": 16, "pooling_kernel_size": 3, "default_output_length": 280,
"standardize": true, "rope_parameters": {"rope_theta": 100.0}},
"text_config": {"hidden_size": 2816, "use_bidirectional_attention": "vision"}}`)
mm, err := parseMultimodalConfig(visionJSON)
if err != nil {
t.Fatal(err)
}
if mm.VisionConfig == nil {
t.Fatal("vision config not retained")
}
if mm.ImageTokenID != 258880 || mm.BOITokenID != 255999 || mm.EOITokenID != 258882 {
t.Fatalf("token ids = %d/%d/%d", mm.ImageTokenID, mm.BOITokenID, mm.EOITokenID)
}
if mm.VisionConfig.RopeTheta != 100 || !mm.VisionConfig.Standardize {
t.Fatalf("vision config = %+v", mm.VisionConfig)
}
cfg, err := parseTextConfig(visionJSON)
if err != nil {
t.Fatal(err)
}
if cfg.UseBidirectionalAttention != "vision" || !cfg.BidirectionalVisionAttention {
t.Fatalf("use_bidirectional_attention = %q", cfg.UseBidirectionalAttention)
}
// The 12B unified family uses the encoder-free embedder.
unified, err := parseMultimodalConfig([]byte(`{
"image_token_id": 258880,
"vision_config": {"model_type": "gemma4_unified_vision", "mm_embed_dim": 3840,
"model_patch_size": 48, "num_soft_tokens": 280, "patch_size": 16,
"pooling_kernel_size": 3, "rms_norm_eps": 1e-6}}`))
if err != nil {
t.Fatal(err)
}
uv := unified.VisionConfig
if uv == nil || !uv.unified() {
t.Fatal("unified vision config not retained")
}
if uv.NumSoftTokens != 280 {
t.Fatalf("unified config = %+v", uv)
}
// An unknown vision architecture loads text-only.
unknown, err := parseMultimodalConfig([]byte(`{
"vision_config": {"model_type": "gemma5_vision"}}`))
if err != nil {
t.Fatal(err)
}
if unknown.VisionConfig != nil {
t.Fatal("unknown vision config should be ignored")
}
textOnly, err := parseMultimodalConfig([]byte(`{"text_config": {"hidden_size": 640}}`))
if err != nil {
t.Fatal(err)
}
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 {
return &Model{
TextConfig: &TextConfig{BidirectionalVisionAttention: true},
MM: multimodalConfig{ImageTokenID: 258880, BOITokenID: 255999, EOITokenID: 258882},
}
}
func mediaItemAt(pos int, softTokens int32) batch.MediaItem {
return batch.MediaItem{
Pos: pos,
Opaque: preparedImage{geom: ImageGeometry{NumSoftTokens: softTokens}},
}
}
func TestBuildMasks(t *testing.T) {
m := visionTestModel()
// Soft run [3, 7): boi at 2, eoi at 7.
item := mediaItemAt(2, 4)
chunk := func(off, qLen int32) *batch.Batch {
return &batch.Batch{
SeqOffsets: []int32{off},
SeqQueryLens: []int32{qLen},
Media: []batch.MediaItem{item},
}
}
sliding, full := m.buildMasks(&batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{8}})
if !sliding.IsCausal() || !full.IsCausal() {
t.Fatal("text-only batch left the causal fast path")
}
sliding, full = m.buildMasks(chunk(0, 8))
if sliding.IsCausal() {
t.Fatal("intersecting run did not relax the sliding mask")
}
if !full.IsCausal() {
t.Fatal("full-attention mask relaxed under \"vision\" semantics")
}
// A decode step past the run keeps both masks on the fast path.
sliding, full = m.buildMasks(chunk(20, 1))
if !sliding.IsCausal() || !full.IsCausal() {
t.Fatal("non-intersecting item left the causal fast path")
}
m.BidirectionalVisionAttention = false
sliding, full = m.buildMasks(chunk(0, 8))
if !sliding.IsCausal() || !full.IsCausal() {
t.Fatal("causal-only config relaxed a mask")
}
}
func TestScatterMedia(t *testing.T) {
mlxtest.Run(t, func(t *mlxtest.T) {
m := visionTestModel()
m.HiddenSize = 4
features := mlx.FromValues([]float32{
1, 1, 1, 1,
2, 2, 2, 2,
3, 3, 3, 3,
4, 4, 4, 4,
}, 4, 4)
item := mediaItemAt(1, 4) // soft run [2, 6)
item.Features = features
// Split evaluation: first chunk covers [0, 4) — soft rows 2, 3 — and the
// resumed chunk [4, 8) covers rows 4, 5 with feature rows sliced by
// overlap.
rowVal := func(h *mlx.Array, row int) float32 {
vals := h.Slice(mlx.Slice(0, 1), mlx.Slice(row, row+1), mlx.Slice()).AsType(mlx.DTypeFloat32)
mlx.Eval(vals)
return vals.Floats()[0]
}
first := m.scatterMedia(mlx.Zeros(mlx.DTypeFloat32, 1, 4, 4), &batch.Batch{
SeqOffsets: []int32{0},
SeqQueryLens: []int32{4},
Media: []batch.MediaItem{item},
})
for row, want := range []float32{0, 0, 1, 2} {
if got := rowVal(first, row); got != want {
t.Fatalf("first chunk row %d = %v, want %v", row, got, want)
}
}
resumed := m.scatterMedia(mlx.Zeros(mlx.DTypeFloat32, 1, 4, 4), &batch.Batch{
SeqOffsets: []int32{4},
SeqQueryLens: []int32{4},
Media: []batch.MediaItem{item},
})
for row, want := range []float32{3, 4, 0, 0} {
if got := rowVal(resumed, row); got != want {
t.Fatalf("resumed chunk row %d = %v, want %v", row, got, want)
}
}
// Featureless items (decode) scatter nothing.
item.Features = nil
decode := m.scatterMedia(mlx.Zeros(mlx.DTypeFloat32, 1, 1, 4), &batch.Batch{
SeqOffsets: []int32{7},
SeqQueryLens: []int32{1},
Media: []batch.MediaItem{item},
})
if got := rowVal(decode, 0); got != 0 {
t.Fatalf("featureless scatter wrote %v", got)
}
})
}
func TestPrepareMediaExpansion(t *testing.T) {
m := visionTestModel()
m.Vision = &VisionConfig{
ModelType: "gemma4_vision",
HiddenSize: 12,
PatchSize: 16,
PoolingKernelSize: 3,
DefaultOutputLen: 70,
RMSNormEps: 1e-6,
}
m.VisionTower = &VisionTower{}
var buf bytes.Buffer
if err := png.Encode(&buf, image.NewRGBA(image.Rect(0, 0, 96, 96))); err != nil {
t.Fatal(err)
}
prepared, err := m.PrepareMedia([]base.Segment{
{Tokens: []int32{9}},
{Kind: "image", Data: buf.Bytes()},
})
if err != nil {
t.Fatal(err)
}
// 96x96 at budget 70 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 {
t.Fatalf("stream length = %d, want %d", len(prepared.Tokens), 1+wantSoft+2)
}
if len(prepared.Items) != 1 {
t.Fatalf("items = %d, want 1", len(prepared.Items))
}
item := prepared.Items[0]
if item.Range != [2]int{1, 1 + wantSoft + 2} || item.Source != 1 {
t.Fatalf("item range = %v source = %d", item.Range, item.Source)
}
exp := prepared.Tokens[item.Range[0]:item.Range[1]]
if exp[0] != m.MM.BOITokenID || exp[len(exp)-1] != m.MM.EOITokenID {
t.Fatal("expansion not wrapped in boi/eoi")
}
for _, tok := range exp[1 : len(exp)-1] {
if tok != m.MM.ImageTokenID {
t.Fatalf("soft token = %d", tok)
}
}
if item.Dims[0] != 24*24 || item.Dims[1] != 16*16*3 {
t.Fatalf("dims = %v", item.Dims)
}
if len(item.MediaData) != item.Dims[0]*item.Dims[1] {
t.Fatalf("media data length = %d", len(item.MediaData))
}
p := item.Opaque.(preparedImage)
if p.geom.NumSoftTokens != int32(wantSoft) || len(p.positions) != 2*24*24 {
t.Fatalf("geometry = %+v, positions = %d", p.geom, len(p.positions))
}
if _, err := m.PrepareMedia([]base.Segment{{Kind: "audio", Data: []byte{1}}}); err == nil {
t.Fatal("audio accepted")
}
if _, err := m.PrepareMedia([]base.Segment{{Kind: "image"}}); err == nil {
t.Fatal("image with no data accepted")
}
m.VisionTower = nil
if _, err := m.PrepareMedia([]base.Segment{{Kind: "image", Data: buf.Bytes()}}); err == nil {
t.Fatal("towerless model accepted an image")
}
}
x/models/gemma4/process_audio.go
+200
View File
@@ -0,0 +1,200 @@
package gemma4
import (
"errors"
"math"
"math/cmplx"
"sync"
"github.com/ollama/ollama/x/mlxrunner/model/audio"
)
// Audio front-end constants from the reference Gemma4AudioFeatureExtractor
// and processor config; none are exposed in config.json.
const (
audioSampleRate = 16000
audioMelBins = 128
audioFrameLen = 320 // 20 ms
audioHopLen = 160 // 10 ms
audioFFTLen = 512
audioMelFloor = 1e-3
audioMaxFreq = 8000.0
// Clips longer than 30 s are encoded as independent chunks of at most
// 30 s, one soft-token run each; the reference extractor truncates at
// this length instead.
audioChunkSeconds = 30
)
// audioChunk is one span of the clip: its encoder input rows — log-mel
// frames for the conformer, one raw waveform frame per soft token for the
// unified embedder — and the soft-token count they reduce to.
type audioChunk struct {
data []float32 // [frames, len/frames]
frames int
numTokens int
}
// processAudio decodes an audio container into per-chunk log-mel features,
// resampled to 16 kHz.
func processAudio(data []byte) ([]audioChunk, error) {
samples, rate, err := audio.Decode(data)
if err != nil {
return nil, err
}
samples = audio.Resample(samples, rate, audioSampleRate)
// One frame needs frameLen+1 samples after the semicausal left pad.
if len(samples) < audioFrameLen+1-audioFrameLen/2 {
return nil, errors.New("audio too short")
}
var chunks []audioChunk
for _, chunk := range audio.Split(samples, audioSampleRate, audioChunkSeconds) {
mel, frames := melSpectrogram(chunk)
tokens := frames
for range 2 {
tokens = (tokens-1)/2 + 1
}
chunks = append(chunks, audioChunk{data: mel, frames: frames, numTokens: tokens})
}
return chunks, nil
}
// processUnifiedAudio decodes audio for the encoder-free variant: raw
// 16 kHz samples in fixed-length frames, one soft token per frame, the
// final partial frame zero-padded. No 30 s chunking; length is bounded by
// the decode duration cap and the context.
func processUnifiedAudio(data []byte, samplesPerToken int) ([]float32, int, error) {
samples, rate, err := audio.Decode(data)
if err != nil {
return nil, 0, err
}
samples = audio.Resample(samples, rate, audioSampleRate)
if len(samples) == 0 {
return nil, 0, errors.New("audio too short")
}
tokens := (len(samples) + samplesPerToken - 1) / samplesPerToken
frames := make([]float32, tokens*samplesPerToken)
copy(frames, samples)
return frames, tokens, nil
}
// audioWindow is the periodic Hann window, rounded to float32 like the
// reference before it multiplies the frame.
var audioWindow = sync.OnceValue(func() []float32 {
w := make([]float32, audioFrameLen)
for i := range w {
w[i] = float32(0.5 - 0.5*math.Cos(2*math.Pi*float64(i)/audioFrameLen))
}
return w
})
// melSpectrogram computes [frames, audioMelBins] log-mel features over the
// samples with semicausal padding: frameLen/2 zeros are prepended so the
// first frame is centered at t=0, and every frame's window covers only real
// audio past that — the reference's trailing pad-and-mask produces the same
// valid frames.
func melSpectrogram(samples []float32) ([]float32, int) {
pad := audioFrameLen / 2
if len(samples)+pad < audioFrameLen+1 {
return nil, 0
}
frames := (len(samples)+pad-(audioFrameLen+1))/audioHopLen + 1
padded := make([]float32, pad+len(samples))
copy(padded[pad:], samples)
window := audioWindow()
filters := audioMelFilters()
numBins := audioFFTLen/2 + 1
out := make([]float32, frames*audioMelBins)
fftBuf := make([]complex128, audioFFTLen)
mags := make([]float64, numBins)
for f := range frames {
frame := padded[f*audioHopLen:]
for i := range audioFrameLen {
fftBuf[i] = complex(float64(frame[i]*window[i]), 0)
}
for i := audioFrameLen; i < audioFFTLen; i++ {
fftBuf[i] = 0
}
fft(fftBuf)
for k := range numBins {
mags[k] = cmplx.Abs(fftBuf[k])
}
for m := range audioMelBins {
var mel float64
for k := range numBins {
mel += mags[k] * filters[m*numBins+k]
}
out[f*audioMelBins+m] = float32(math.Log(mel + audioMelFloor))
}
}
return out, frames
}
// audioMelFilters builds the HTK-scale triangular filterbank as [mel, bin]
// weights, mirroring the reference mel_filter_bank (no normalization).
var audioMelFilters = sync.OnceValue(func() []float64 {
hzToMel := func(f float64) float64 { return 2595 * math.Log10(1+f/700) }
melToHz := func(m float64) float64 { return 700 * (math.Pow(10, m/2595) - 1) }
melMax := hzToMel(audioMaxFreq)
corners := make([]float64, audioMelBins+2)
for i := range corners {
corners[i] = melToHz(float64(i) * melMax / float64(audioMelBins+1))
}
numBins := audioFFTLen/2 + 1
filters := make([]float64, audioMelBins*numBins)
for m := range audioMelBins {
left, center, right := corners[m], corners[m+1], corners[m+2]
for k := range numBins {
freq := float64(k) * audioSampleRate / audioFFTLen
v := min((freq-left)/(center-left), (right-freq)/(right-center))
if v > 0 {
filters[m*numBins+k] = v
}
}
}
return filters
})
// fft performs an in-place Cooley-Tukey radix-2 FFT.
func fft(x []complex128) {
n := len(x)
if n <= 1 {
return
}
j := 0
for i := 1; i < n; i++ {
bit := n >> 1
for j&bit != 0 {
j ^= bit
bit >>= 1
}
j ^= bit
if i < j {
x[i], x[j] = x[j], x[i]
}
}
for size := 2; size <= n; size <<= 1 {
half := size / 2
w := cmplx.Exp(complex(0, -2*math.Pi/float64(size)))
for start := 0; start < n; start += size {
wn := complex(1, 0)
for k := range half {
t := wn * x[start+k+half]
x[start+k+half] = x[start+k] - t
x[start+k] = x[start+k] + t
wn *= w
}
}
}
}
x/models/gemma4/process_audio_test.go
+109
View File
@@ -0,0 +1,109 @@
package gemma4
import (
"bytes"
"encoding/binary"
"math"
"strings"
"testing"
)
func wavPCM16(samples int) []byte {
var b bytes.Buffer
b.WriteString("RIFF")
binary.Write(&b, binary.LittleEndian, uint32(36+2*samples))
b.WriteString("WAVE")
b.WriteString("fmt ")
binary.Write(&b, binary.LittleEndian, uint32(16))
binary.Write(&b, binary.LittleEndian, uint16(1))
binary.Write(&b, binary.LittleEndian, uint16(1))
binary.Write(&b, binary.LittleEndian, uint32(audioSampleRate))
binary.Write(&b, binary.LittleEndian, uint32(2*audioSampleRate))
binary.Write(&b, binary.LittleEndian, uint16(2))
binary.Write(&b, binary.LittleEndian, uint16(16))
b.WriteString("data")
binary.Write(&b, binary.LittleEndian, uint32(2*samples))
b.Write(make([]byte, 2*samples))
return b.Bytes()
}
// Golden log-mel values from the reference Gemma4AudioFeatureExtractor over
// 1664 samples of 0.5*sin(2*pi*440*t/16000): bins 0, 1, 64, 127 per frame.
var melGolden = [][4]float32{
{-6.907755, 0.808556, -1.241259, -2.153034},
{-6.907755, -4.019087, -6.397553, -6.901781},
{-6.907755, -4.861755, -6.335430, -6.902662},
{-6.907755, -4.861755, -6.335430, -6.902662},
{-6.907755, -4.019087, -6.397553, -6.901781},
{-6.907755, -3.820103, -6.459379, -6.901654},
{-6.907755, -4.019087, -6.397553, -6.901781},
{-6.907755, -4.861755, -6.335430, -6.902662},
{-6.907755, -4.861755, -6.335430, -6.902662},
{-6.907755, -4.019087, -6.397553, -6.901781},
}
func TestMelSpectrogramGolden(t *testing.T) {
samples := make([]float32, 1664)
for i := range samples {
samples[i] = float32(0.5 * math.Sin(2*math.Pi*440*float64(i)/audioSampleRate))
}
mel, frames := melSpectrogram(samples)
if frames != len(melGolden) {
t.Fatalf("%d frames, want %d", frames, len(melGolden))
}
for f, want := range melGolden {
for i, bin := range []int{0, 1, 64, 127} {
got := mel[f*audioMelBins+bin]
if diff := float64(got - want[i]); math.Abs(diff) > 2e-6 {
t.Errorf("frame %d bin %d: %v, want %v", f, bin, got, want[i])
}
}
}
}
func TestProcessAudioTokenCounts(t *testing.T) {
cases := []struct {
name string
samples int
frames []int
tokens []int
}{
{"minimum", 161, []int{1}, []int{1}},
{"short", 1000, []int{6}, []int{2}},
{"thirty seconds", 480000, []int{2999}, []int{750}},
// Over the limit the clip splits evenly; each chunk must still fit.
{"just over the limit", 480001, []int{-1, -1}, []int{-1, -1}},
}
for _, tt := range cases {
t.Run(tt.name, func(t *testing.T) {
chunks, err := processAudio(wavPCM16(tt.samples))
if err != nil {
t.Fatal(err)
}
if len(chunks) != len(tt.frames) {
t.Fatalf("%d chunks, want %d", len(chunks), len(tt.frames))
}
for i, c := range chunks {
if tt.frames[i] < 0 {
if c.frames <= 0 || c.frames > 2999 || c.numTokens > 750 {
t.Errorf("chunk %d: frames %d tokens %d exceed one 30 s chunk", i, c.frames, c.numTokens)
}
} else if c.frames != tt.frames[i] || c.numTokens != tt.tokens[i] {
t.Errorf("chunk %d: frames %d tokens %d, want %d %d",
i, c.frames, c.numTokens, tt.frames[i], tt.tokens[i])
}
if len(c.data) != c.frames*audioMelBins {
t.Errorf("chunk %d: %d mel values for %d frames", i, len(c.data), c.frames)
}
}
})
}
}
func TestProcessAudioTooShort(t *testing.T) {
_, err := processAudio(wavPCM16(160))
if err == nil || !strings.Contains(err.Error(), "audio too short") {
t.Fatalf("error %v, want audio too short", err)
}
}
x/models/gemma4/process_image.go
+132
View File
@@ -0,0 +1,132 @@
package gemma4
import (
"bytes"
"fmt"
"image"
"math"
"golang.org/x/image/draw"
)
// ImageGeometry describes a preprocessed image's patch grid.
type ImageGeometry struct {
PatchesW, PatchesH int32
NumSoftTokens int32
}
// preparedImage is gemma4's model-private media state: the patch position
// grid the encoder consumes and the geometry the forward pass derives the
// soft-token run from.
type preparedImage struct {
positions []int32
geom ImageGeometry
}
// visionTargetSize ports the reference resize: sides floored to
// multiples of patchSize*poolingKernel under a patch budget, a
// zero-flooring side clamped to one multiple.
func visionTargetSize(height, width, patchSize, poolingKernel, maxPatches 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)
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 {
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)
}
if float64(targetH)*float64(targetW) > targetPx {
return 0, 0, fmt.Errorf("image %dx%d exceeds the patch budget after resize", width, height)
}
return targetH, targetW, nil
}
// preprocessImage decodes and prepares one image: aspect-preserving
// resize, rescale to [0,1], and patchify. The [-1,1] normalization stays
// in the patch embedder, so pixels here match the reference
// pixel_values.
func (m *Model) preprocessImage(data []byte) (pixels []float32, positions []int32, geom ImageGeometry, err error) {
img, _, err := image.Decode(bytes.NewReader(data))
if err != nil {
return nil, nil, ImageGeometry{}, fmt.Errorf("decode image: %w", err)
}
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)
if err != nil {
return nil, nil, ImageGeometry{}, err
}
resized := image.NewRGBA(image.Rect(0, 0, int(targetW), int(targetH)))
draw.CatmullRom.Scale(resized, resized.Bounds(), img, bounds, draw.Src, nil)
if m.Vision.unified() {
// One raster patch of pool*patchSize pixels per soft token: the
// reference merge rearranges its intermediate 16px patches back
// into this layout, with positions on the merged grid.
pixels, positions, geom = patchify(resized, targetW, targetH, patch*pool, 1)
} else {
pixels, positions, geom = patchify(resized, targetW, targetH, patch, pool)
}
return pixels, positions, geom, nil
}
// patchify converts the resized image to the tower's layout: one row per
// patchSize patch, (pixel row, pixel column, RGB) within it.
func patchify(resized *image.RGBA, targetW, targetH, patch, pool int32) ([]float32, []int32, ImageGeometry) {
pW, pH := targetW/patch, targetH/patch
numPatches := int(pW * pH)
patchLen := int(patch * patch * 3)
pixels := make([]float32, numPatches*patchLen)
positions := make([]int32, 2*numPatches)
for p := range numPatches {
gx, gy := int32(p)%pW, int32(p)/pW
positions[2*p] = gx
positions[2*p+1] = gy
writePatch(pixels[p*patchLen:], resized, int(gx*patch), int(gy*patch), int(patch))
}
return pixels, positions, ImageGeometry{PatchesW: pW, PatchesH: pH, NumSoftTokens: pW * pH / (pool * pool)}
}
func writePatch(out []float32, resized *image.RGBA, baseX, baseY, patch int) {
for py := range patch {
row := resized.PixOffset(baseX, baseY+py)
for px := range patch {
o := (py*patch + px) * 3
pix := resized.Pix[row+px*4:]
out[o] = float32(pix[0]) / 255
out[o+1] = float32(pix[1]) / 255
out[o+2] = float32(pix[2]) / 255
}
}
}
// dropAlpha flattens a non-opaque image to straight RGB: the reference
// drops alpha via RGB conversion before resizing, not by compositing.
func dropAlpha(img image.Image, bounds image.Rectangle) image.Image {
if o, ok := img.(interface{ Opaque() bool }); ok && o.Opaque() {
return img
}
flat := image.NewNRGBA(bounds)
draw.Draw(flat, bounds, img, bounds.Min, draw.Src)
for i := 3; i < len(flat.Pix); i += 4 {
flat.Pix[i] = 0xff
}
return flat
}
x/models/gemma4/vision.go
+406
View File
@@ -0,0 +1,406 @@
package gemma4
import (
"fmt"
"math"
"slices"
"strings"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/mlxrunner/model"
"github.com/ollama/ollama/x/models/nn"
)
// VisionConfig holds configuration for the Gemma 4 vision path: the
// gemma4_vision transformer tower, or the gemma4_unified_vision encoder-free
// embedder (one dense projection over merged 48px patches).
type VisionConfig struct {
ModelType string `json:"model_type"`
HiddenSize int32 `json:"hidden_size"`
NumHiddenLayers int32 `json:"num_hidden_layers"`
NumAttentionHeads int32 `json:"num_attention_heads"`
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:"-"`
}
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 {
PatchEmbedder *PatchEmbedder
Layers []*VisionLayer
StdBias *mlx.Array
StdScale *mlx.Array
}
type PatchEmbedder struct {
InputProj nn.LinearLayer
PositionEmbeddingTable *mlx.Array // [2, positions, hidden]; x and y rows summed
}
type VisionLayer struct {
InputNorm *mlx.Array
Attention *VisionAttention
PostAttnNorm *mlx.Array
PreFFNorm *mlx.Array
MLP *MLP
PostFFNorm *mlx.Array
}
type VisionAttention struct {
QProj, KProj, VProj, OProj nn.LinearLayer
QNorm, KNorm *mlx.Array // per-head; v-norm has no weight
}
// UnifiedVisionEmbedder is the encoder-free vision path: one dense
// projection over merged model patches with factorized 2D positions.
type UnifiedVisionEmbedder struct {
PatchLN1 *nn.LayerNorm
PatchDense nn.LinearLayer
PatchLN2 *nn.LayerNorm
PosEmbedding *mlx.Array // [positions, 2, mm_embed_dim]; x and y rows summed
PosNorm *nn.LayerNorm
}
// visionWeightRoot locates the tower's tensor prefix by probing for the
// patch embedder projection.
func visionWeightRoot(tensors map[string]*mlx.Array) (string, error) {
for _, root := range []string{"model.", ""} {
if tensors[root+"vision_tower.patch_embedder.input_proj.weight"] != nil {
return root, nil
}
}
return "", fmt.Errorf("config declares a gemma4_vision tower but vision_tower weights are missing from the manifest")
}
func (m *Model) loadVisionWeights(tensors map[string]*mlx.Array, linears model.LinearFactory) error {
if m.Vision.unified() {
return m.loadUnifiedVisionWeights(tensors, linears)
}
root, err := visionWeightRoot(tensors)
if err != nil {
return err
}
vt := root + "vision_tower."
v := m.Vision
requiredNorm := func(name string) (*mlx.Array, error) {
w := tensors[name]
if w == nil {
return nil, fmt.Errorf("missing vision weight: %s", name)
}
return w, nil
}
inputProj, err := makeClippableLinear(linears, tensors, vt+"patch_embedder.input_proj")
if err != nil {
return err
}
table := tensors[vt+"patch_embedder.position_embedding_table"]
if table == nil {
return fmt.Errorf("missing vision weight: %spatch_embedder.position_embedding_table", vt)
}
tower := &VisionTower{
PatchEmbedder: &PatchEmbedder{InputProj: inputProj, PositionEmbeddingTable: table},
Layers: make([]*VisionLayer, v.NumHiddenLayers),
}
for i := range tower.Layers {
lp := fmt.Sprintf("%sencoder.layers.%d.", vt, i)
layer := &VisionLayer{Attention: &VisionAttention{}, MLP: &MLP{}}
if layer.InputNorm, err = requiredNorm(lp + "input_layernorm.weight"); err != nil {
return err
}
if layer.PostAttnNorm, err = requiredNorm(lp + "post_attention_layernorm.weight"); err != nil {
return err
}
if layer.PreFFNorm, err = requiredNorm(lp + "pre_feedforward_layernorm.weight"); err != nil {
return err
}
if layer.PostFFNorm, err = requiredNorm(lp + "post_feedforward_layernorm.weight"); err != nil {
return err
}
if layer.Attention.QNorm, err = requiredNorm(lp + "self_attn.q_norm.weight"); err != nil {
return err
}
if layer.Attention.KNorm, err = requiredNorm(lp + "self_attn.k_norm.weight"); err != nil {
return err
}
if layer.Attention.QProj, err = makeClippableLinear(linears, tensors, lp+"self_attn.q_proj"); err != nil {
return err
}
if layer.Attention.KProj, err = makeClippableLinear(linears, tensors, lp+"self_attn.k_proj"); err != nil {
return err
}
if layer.Attention.VProj, err = makeClippableLinear(linears, tensors, lp+"self_attn.v_proj"); err != nil {
return err
}
if layer.Attention.OProj, err = makeClippableLinear(linears, tensors, lp+"self_attn.o_proj"); err != nil {
return err
}
if layer.MLP.GateProj, err = makeClippableLinear(linears, tensors, lp+"mlp.gate_proj"); err != nil {
return err
}
if layer.MLP.UpProj, err = makeClippableLinear(linears, tensors, lp+"mlp.up_proj"); err != nil {
return err
}
if layer.MLP.DownProj, err = makeClippableLinear(linears, tensors, lp+"mlp.down_proj"); err != nil {
return err
}
tower.Layers[i] = layer
}
if v.Standardize {
if tower.StdBias, err = requiredNorm(vt + "std_bias"); err != nil {
return err
}
if tower.StdScale, err = requiredNorm(vt + "std_scale"); err != nil {
return err
}
}
projection, err := makeClippableLinear(linears, tensors, root+"embed_vision.embedding_projection")
if err != nil {
return err
}
m.VisionTower = tower
m.EmbedVision = &MultimodalEmbedder{Projection: projection}
return nil
}
// loadUnifiedVisionWeights loads the encoder-free embedder. Its projections
// carry plain names (no .linear infix) and no clamps.
func (m *Model) loadUnifiedVisionWeights(tensors map[string]*mlx.Array, linears model.LinearFactory) error {
var root string
for _, r := range []string{"model.", ""} {
if tensors[r+"vision_embedder.patch_dense.weight"] != nil {
root = r + "vision_embedder."
break
}
}
if root == "" {
return fmt.Errorf("config declares a gemma4_unified_vision embedder but vision_embedder weights are missing from the manifest")
}
norm := func(name string) (*nn.LayerNorm, error) {
w, b := tensors[root+name+".weight"], tensors[root+name+".bias"]
if w == nil || b == nil {
return nil, fmt.Errorf("missing vision weight: %s%s", root, name)
}
// The reference uses torch LayerNorm defaults; rms_norm_eps feeds
// only the multimodal embedder's RMSNorm.
return &nn.LayerNorm{Weight: w, Bias: b, Eps: 1e-5}, nil
}
e := &UnifiedVisionEmbedder{}
var err error
if e.PatchLN1, err = norm("patch_ln1"); err != nil {
return err
}
if e.PatchDense = linears.Make(root + "patch_dense"); e.PatchDense == nil {
return fmt.Errorf("missing vision weight: %spatch_dense", root)
}
if e.PatchLN2, err = norm("patch_ln2"); err != nil {
return err
}
if e.PosEmbedding = tensors[root+"pos_embedding"]; e.PosEmbedding == nil {
return fmt.Errorf("missing vision weight: %spos_embedding", root)
}
if e.PosNorm, err = norm("pos_norm"); err != nil {
return err
}
projRoot := strings.TrimSuffix(root, "vision_embedder.")
projection, err := makeClippableLinear(linears, tensors, projRoot+"embed_vision.embedding_projection")
if err != nil {
return err
}
m.UnifiedEmbedder = e
m.EmbedVision = &MultimodalEmbedder{Projection: projection}
return nil
}
// encodeUnifiedImage embeds one image's merged patches: LN, dense, LN, factorized
// 2D positions, LN, then the shared multimodal embedder.
func (m *Model) encodeUnifiedImage(patches *mlx.Array, positions []int32, geom ImageGeometry) *mlx.Array {
e := m.UnifiedEmbedder
n := geom.PatchesH * geom.PatchesW
h := e.PatchLN1.Forward(patches.AsType(mlx.DTypeBFloat16))
h = e.PatchDense.Forward(h)
h = e.PatchLN2.Forward(h)
for axis := range 2 {
pos := make([]int32, n)
for i := range int(n) {
pos[i] = positions[2*i+axis]
}
axisTable := mlx.Squeeze(e.PosEmbedding.Slice(mlx.Slice(), mlx.Slice(axis, axis+1), mlx.Slice()), 1)
h = mlx.Add(h, mlx.Take(axisTable, mlx.FromValues(pos, int(n)), 0).AsType(h.DType()))
}
h = e.PosNorm.Forward(h)
h = mlx.RMSNormFn(h, nil, m.Vision.RMSNormEps)
return m.EmbedVision.Projection.Forward(h)
}
func (m *Model) visionLoaded() bool {
return m.VisionTower != nil || m.UnifiedEmbedder != nil
}
// visionRopeTables builds the per-axis rotation tables, shared by every
// layer: positions x inverse frequencies in f32, duplicated across the
// half's two rotation quarters, shaped to broadcast over heads.
func visionRopeTables(positions []int32, axis int, n, headDim int32, theta float32) (cos, sin *mlx.Array) {
spatial := int(headDim) / 2
nFreq := spatial / 2
pos := make([]float32, n)
for i := range int(n) {
pos[i] = float32(positions[2*i+axis])
}
invFreq := make([]float32, nFreq)
for i := range invFreq {
invFreq[i] = float32(1 / math.Pow(float64(theta), float64(2*i)/float64(spatial)))
}
f := mlx.Matmul(mlx.FromValues(pos, int(n), 1), mlx.FromValues(invFreq, 1, nFreq))
emb := mlx.Concatenate([]*mlx.Array{f, f}, -1)
cos = mlx.Reshape(mlx.Cos(emb), 1, n, 1, int32(spatial)).AsType(mlx.DTypeBFloat16)
sin = mlx.Reshape(mlx.Sin(emb), 1, n, 1, int32(spatial)).AsType(mlx.DTypeBFloat16)
return cos, sin
}
// applyVisionRoPE rotates the head dim's two halves with the x and y tables;
// rotation happens within each half independently.
func applyVisionRoPE(t, cosX, sinX, cosY, sinY *mlx.Array, headDim int32) *mlx.Array {
half := int(headDim) / 2
rot := func(p *mlx.Array) *mlx.Array {
quarter := half / 2
p1 := p.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, quarter))
p2 := p.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(quarter, half))
return mlx.Concatenate([]*mlx.Array{mlx.Neg(p2), p1}, -1)
}
tx := t.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, half))
ty := t.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(half, int(headDim)))
tx = mlx.Add(mlx.Mul(tx, cosX), mlx.Mul(rot(tx), sinX))
ty = mlx.Add(mlx.Mul(ty, cosY), mlx.Mul(rot(ty), sinY))
return mlx.Concatenate([]*mlx.Array{tx, ty}, -1)
}
func (a *VisionAttention) Forward(x *mlx.Array, cosX, sinX, cosY, sinY *mlx.Array, n int32, v *VisionConfig) *mlx.Array {
heads, headDim := v.NumAttentionHeads, v.HeadDim
q := mlx.Reshape(a.QProj.Forward(x), 1, n, heads, headDim)
k := mlx.Reshape(a.KProj.Forward(x), 1, n, heads, headDim)
val := mlx.Reshape(a.VProj.Forward(x), 1, n, heads, headDim)
q = mlx.RMSNormFn(q, a.QNorm, v.RMSNormEps)
k = mlx.RMSNormFn(k, a.KNorm, v.RMSNormEps)
val = mlx.RMSNormFn(val, nil, v.RMSNormEps)
q = applyVisionRoPE(q, cosX, sinX, cosY, sinY, headDim)
k = applyVisionRoPE(k, cosX, sinX, cosY, sinY, headDim)
q = mlx.Transpose(q, 0, 2, 1, 3)
k = mlx.Transpose(k, 0, 2, 1, 3)
val = mlx.Transpose(val, 0, 2, 1, 3)
// Fully bidirectional, and scale 1.0: the Q/K norms control magnitude.
out := mlx.FastScaledDotProductAttention(q, k, val, 1.0, "", nil)
out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), 1, n, heads*headDim)
return a.OProj.Forward(out)
}
// encodeImage runs the tower over one whole image. One image per call, no
// padding; the graph is lazy and evaluates with the consuming forward.
func (m *Model) encodeImage(pixels *mlx.Array, positions []int32, geom ImageGeometry) *mlx.Array {
v := m.Vision
n := geom.PatchesH * geom.PatchesW
// The processor rescales to [0,1]; the reference folds the [-1,1]
// normalization into the patch embedder, in f32 before the cast.
x := mlx.MulScalar(mlx.AddScalar(pixels, -0.5), 2).AsType(mlx.DTypeBFloat16)
h := m.VisionTower.PatchEmbedder.InputProj.Forward(x)
table := m.VisionTower.PatchEmbedder.PositionEmbeddingTable
for axis := range 2 {
pos := make([]int32, n)
for i := range int(n) {
pos[i] = positions[2*i+axis]
}
axisTable := mlx.Squeeze(table.Slice(mlx.Slice(axis, axis+1), mlx.Slice(), mlx.Slice()), 0)
h = mlx.Add(h, mlx.Take(axisTable, mlx.FromValues(pos, int(n)), 0).AsType(h.DType()))
}
h = mlx.Reshape(h, 1, n, v.HiddenSize)
cosX, sinX := visionRopeTables(positions, 0, n, v.HeadDim, v.RopeTheta)
cosY, sinY := visionRopeTables(positions, 1, n, v.HeadDim, v.RopeTheta)
for _, l := range m.VisionTower.Layers {
x1 := mlx.RMSNormFn(h, l.InputNorm, v.RMSNormEps)
attn := l.Attention.Forward(x1, cosX, sinX, cosY, sinY, n, v)
h = mlx.Add(h, mlx.RMSNormFn(attn, l.PostAttnNorm, v.RMSNormEps))
x2 := mlx.RMSNormFn(h, l.PreFFNorm, v.RMSNormEps)
h = mlx.Add(h, mlx.RMSNormFn(l.MLP.Forward(x2), l.PostFFNorm, v.RMSNormEps))
}
// Pool 3x3 over the patch grid, scale by sqrt(hidden), and standardize —
// in f32 with the reference's dtype round-trip (the sqrt scaling can
// exceed fp16 range, and the reference keeps this section in f32).
pool := v.PoolingKernelSize
f := mlx.Reshape(h.AsType(mlx.DTypeFloat32), geom.PatchesH/pool, pool, geom.PatchesW/pool, pool, v.HiddenSize)
f = mlx.Mean(f, 3, false)
f = mlx.Mean(f, 1, false)
f = mlx.Reshape(f, n/(pool*pool), v.HiddenSize)
f = f.AsType(mlx.DTypeBFloat16).AsType(mlx.DTypeFloat32)
f = mlx.MulScalar(f, float32(math.Sqrt(float64(v.HiddenSize))))
if m.VisionTower.StdBias != nil {
f = mlx.Mul(mlx.Sub(f, m.VisionTower.StdBias.AsType(mlx.DTypeFloat32)), m.VisionTower.StdScale.AsType(mlx.DTypeFloat32))
}
f = f.AsType(mlx.DTypeBFloat16)
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
}
x/models/gemma4/vision_test.go
+263
View File
@@ -0,0 +1,263 @@
package gemma4
import (
"image"
"math"
"slices"
"testing"
"github.com/ollama/ollama/x/internal/mlxtest"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/models/nn"
)
func TestVisionTargetSize(t *testing.T) {
cases := []struct {
name string
h, w, budget int32
wantH, wantW int32
wantSoftTokens int32
}{
{"landscape", 768, 1024, 280, 672, 912, 266},
{"square", 896, 896, 280, 768, 768, 256},
{"square small budget", 896, 896, 70, 384, 384, 64},
{"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},
}
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)
if err != nil {
t.Fatal(err)
}
if gotH != c.wantH || gotW != c.wantW {
t.Fatalf("target = %dx%d, want %dx%d", gotW, gotH, c.wantW, c.wantH)
}
soft := (gotH / 16) * (gotW / 16) / 9
if soft != c.wantSoftTokens {
t.Fatalf("soft tokens = %d, want %d", soft, c.wantSoftTokens)
}
})
}
}
// 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
// mean, sqrt(hidden) scale, RMS norm — in the reference soft-token order.
func TestEncodeImagePooling(t *testing.T) {
mlxtest.Run(t, func(t *mlxtest.T) {
const (
grid = 6
pool = 3
patchD = 12 // patch 2x2 x RGB, doubling as hidden size
patches = grid * grid
)
identity := make([]float32, patchD*patchD)
for i := range patchD {
identity[i*patchD+i] = 1
}
v := &VisionConfig{
HiddenSize: patchD,
HeadDim: 4,
PoolingKernelSize: pool,
PatchSize: 2,
RMSNormEps: 1e-6,
RopeTheta: 100,
}
m := &Model{
Vision: v,
VisionTower: &VisionTower{
PatchEmbedder: &PatchEmbedder{
InputProj: nn.NewLinear(mlx.FromValues(identity, patchD, patchD), nil),
PositionEmbeddingTable: mlx.Zeros(mlx.DTypeFloat32, 2, 8, patchD),
},
},
EmbedVision: &MultimodalEmbedder{Projection: nn.NewLinear(mlx.FromValues(identity, patchD, patchD), nil)},
}
pixels := make([]float32, patches*patchD)
positions := make([]int32, 2*patches)
for p := range patches {
positions[2*p] = int32(p % grid)
positions[2*p+1] = int32(p / grid)
for d := range patchD {
pixels[p*patchD+d] = float32((p*31+d*7)%97) / 97
}
}
geom := ImageGeometry{PatchesW: grid, PatchesH: grid, NumSoftTokens: patches / (pool * pool)}
out := m.encodeImage(mlx.FromValues(pixels, patches, patchD), positions, geom)
out = out.AsType(mlx.DTypeFloat32)
mlx.Eval(out)
got := out.Floats()
scale := float32(math.Sqrt(patchD))
for by := range grid / pool {
for bx := range grid / pool {
var block [patchD]float32
for iy := range pool {
for ix := range pool {
p := (by*pool+iy)*grid + (bx*pool + ix)
for d := range patchD {
block[d] += 2 * (pixels[p*patchD+d] - 0.5)
}
}
}
var sumsq float64
for d := range patchD {
block[d] = block[d] / (pool * pool) * scale
sumsq += float64(block[d]) * float64(block[d])
}
rms := float32(math.Sqrt(sumsq/patchD + 1e-6))
b := by*(grid/pool) + bx
for d := range patchD {
want := block[d] / rms
if diff := math.Abs(float64(got[b*patchD+d] - want)); diff > 0.03 {
t.Fatalf("soft token %d dim %d = %v, want %v", b, d, got[b*patchD+d], want)
}
}
}
}
})
}
// TestPatchifyMerged checks the unified layout against an independent index
// computation: one raster patch of pool*patchSize pixels per soft token,
// (pixel row, pixel column, RGB) across the whole merged patch.
func TestPatchifyMerged(t *testing.T) {
const patch, pool = 2, 2
const merged = patch * pool
const w, h = 8, 4 // 2x1 model patches of 4px
img := image.NewRGBA(image.Rect(0, 0, w, h))
pixel := func(x, y, c int) uint8 { return uint8(x*16 + y*3 + c) }
for y := range h {
for x := range w {
o := img.PixOffset(x, y)
for c := range 3 {
img.Pix[o+c] = pixel(x, y, c)
}
img.Pix[o+3] = 255
}
}
pixels, positions, geom := patchify(img, w, h, merged, 1)
if geom.PatchesW != 2 || geom.PatchesH != 1 || geom.NumSoftTokens != 2 {
t.Fatalf("geometry = %+v", geom)
}
if want := []int32{0, 0, 1, 0}; !slices.Equal(positions, want) {
t.Fatalf("positions = %v, want %v", positions, want)
}
patchLen := merged * merged * 3
for i, got := range pixels {
p := i / patchLen
e := i % patchLen
px, py, c := e/3%merged, e/3/merged, e%3
x := p*merged + px // single model-patch row, so no gy term
y := py
if want := float32(pixel(x, y, c)) / 255; got != want {
t.Fatalf("pixels[%d] = %v, want %v (x=%d y=%d c=%d)", i, got, want, x, y, c)
}
}
}
// TestEncodeUnifiedImage checks the encode chain end to end on an identity-sized
// config against a reference computation of LN, position add, LN, and RMS
// norm.
func TestEncodeUnifiedImage(t *testing.T) {
mlxtest.Run(t, func(t *mlxtest.T) {
const d = 12 // merged patch dim (1px teacher patches, 2x2 merge, RGB), doubling as embed dim
identity := make([]float32, d*d)
for i := range d {
identity[i*d+i] = 1
}
ones := make([]float32, d)
for i := range ones {
ones[i] = 1
}
unitLN := func() *nn.LayerNorm {
return &nn.LayerNorm{Weight: mlx.FromValues(ones, d), Bias: mlx.Zeros(mlx.DTypeFloat32, d), Eps: 1e-5}
}
posTable := make([]float32, 2*2*d)
for r := range 2 {
for a := range 2 {
for i := range d {
posTable[(r*2+a)*d+i] = float32(r+1) * float32(a+1) / 8
}
}
}
m := &Model{
Vision: &VisionConfig{ModelType: "gemma4_unified_vision", RMSNormEps: 1e-6},
UnifiedEmbedder: &UnifiedVisionEmbedder{
PatchLN1: unitLN(),
PatchDense: nn.NewLinear(mlx.FromValues(identity, d, d), nil),
PatchLN2: unitLN(),
PosEmbedding: mlx.FromValues(posTable, 2, 2, d),
PosNorm: unitLN(),
},
EmbedVision: &MultimodalEmbedder{Projection: nn.NewLinear(mlx.FromValues(identity, d, d), nil)},
}
const n = 4
patches := make([]float32, n*d)
for i := range patches {
patches[i] = float32((i*13)%29) / 29
}
positions := []int32{0, 0, 1, 0, 0, 1, 1, 1}
out := m.encodeUnifiedImage(mlx.FromValues(patches, n, d), positions, ImageGeometry{PatchesW: 2, PatchesH: 2, NumSoftTokens: n})
out = out.AsType(mlx.DTypeFloat32)
mlx.Eval(out)
got := out.Floats()
ln := func(v []float64) []float64 {
var mean, varSum float64
for _, x := range v {
mean += x
}
mean /= float64(len(v))
for _, x := range v {
varSum += (x - mean) * (x - mean)
}
varSum /= float64(len(v))
outV := make([]float64, len(v))
for i, x := range v {
outV[i] = (x - mean) / math.Sqrt(varSum+1e-5)
}
return outV
}
for p := range n {
v := make([]float64, d)
for i := range d {
v[i] = float64(patches[p*d+i])
}
v = ln(ln(v)) // LN1, identity dense, LN2
x, y := int(positions[2*p]), int(positions[2*p+1])
for i := range d {
v[i] += float64(posTable[(x*2+0)*d+i]) + float64(posTable[(y*2+1)*d+i])
}
v = ln(v)
var sumsq float64
for _, x := range v {
sumsq += x * x
}
rms := math.Sqrt(sumsq/d + 1e-6)
for i := range d {
want := v[i] / rms
if diff := math.Abs(float64(got[p*d+i]) - want); diff > 0.03 {
t.Fatalf("patch %d dim %d = %v, want %v", p, i, got[p*d+i], want)
}
}
}
})
}