diff --git a/create/nemotron_h.go b/create/nemotron_h.go index c811577f6..0fc60edcf 100644 --- a/create/nemotron_h.go +++ b/create/nemotron_h.go @@ -27,10 +27,14 @@ func newNemotronHImportTransform(rawConfig json.RawMessage) (quantizePolicy, err return nemotronHImportTransform{numLayers: numLayers}, nil } -func nemotronHIsUnsupportedModalityTensor(name string) bool { +// Nemotron's modality tower names do not match the shared predicates. +func nemotronHIsVisionTower(name string) bool { return strings.HasPrefix(name, "vision_model.") || - strings.HasPrefix(name, "mlp1.") || - strings.HasPrefix(name, "sound_encoder.") || + strings.HasPrefix(name, "mlp1.") +} + +func nemotronHIsAudioTower(name string) bool { + return strings.HasPrefix(name, "sound_encoder.") || strings.HasPrefix(name, "sound_projection.") } @@ -65,7 +69,7 @@ func (t nemotronHImportTransform) promoteSensitive(name string) bool { } func (t nemotronHImportTransform) quantizationType(name string, shape []int32, quantize string) string { - if nemotronHIsUnsupportedModalityTensor(name) || nemotronHShouldKeepBF16ForDirectNonAffine(name) { + if nemotronHIsVisionTower(name) || nemotronHIsAudioTower(name) || nemotronHShouldKeepBF16ForDirectNonAffine(name) { return "" } diff --git a/mlxrunner/model/nemotron_h/nemotron_h.go b/mlxrunner/model/nemotron_h/nemotron_h.go index 221a7d9bf..1e8cacab9 100644 --- a/mlxrunner/model/nemotron_h/nemotron_h.go +++ b/mlxrunner/model/nemotron_h/nemotron_h.go @@ -27,6 +27,7 @@ var ( _ model.Model = (*Model)(nil) _ model.SelfDraft = (*Model)(nil) _ model.DraftModel = (*mtpDraft)(nil) + _ model.MediaModel = (*Model)(nil) ) type Config struct { @@ -82,6 +83,15 @@ type Model struct { MTP *MTPHead + VisionEncoder *RadioVisionEncoder + Projector *VisionProjector + VisionConfig *VisionConfig + + imageStartTokenID int32 + imageTokenID int32 + imageEndTokenID int32 + visionErr error + tok *tokenizer.Tokenizer *Config @@ -322,10 +332,31 @@ func newModel(root *model.Root) (model.Model, error) { return nil, fmt.Errorf("parse tokenizer: %w", err) } + var preprocessorData []byte + var visionErr error + if _, ok := root.Manifest.ConfigLayer("preprocessor_config.json"); ok { + data, err := root.Manifest.ReadConfig("preprocessor_config.json") + if err != nil { + visionErr = fmt.Errorf("load preprocessor_config.json: %w", err) + } else { + preprocessorData = data + } + } + var visionConfig *VisionConfig + if visionErr == nil { + visionConfig, visionErr = parseVisionConfig(configData, preprocessorData) + } + m := &Model{ - Layers: make([]*Layer, cfg.NumHiddenLayers), - Config: &cfg, - tok: tok, + Layers: make([]*Layer, cfg.NumHiddenLayers), + Config: &cfg, + tok: tok, + visionErr: visionErr, + } + if visionConfig != nil { + if err := m.configureVision(visionConfig); err != nil { + m.visionErr = err + } } for i, typ := range cfg.LayerTypes { m.Layers[i] = &Layer{Type: typ} @@ -1078,6 +1109,11 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { cfg := m.Config linears := model.NewLinearFactory(tensors, cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant) + if m.VisionConfig != nil { + if err := m.loadVisionWeights(tensors, linears); err != nil { + return err + } + } useQuantizedExperts := supportsGatherQMM(cfg.QuantMode, cfg.QuantBits) if !useQuantizedExperts && cfg.TensorQuant != nil { for _, tq := range cfg.TensorQuant { @@ -1453,6 +1489,9 @@ func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden B, L := int32(dims[0]), int32(dims[1]) h := m.EmbedTokens.Forward(tokens) + if len(b.Media) > 0 { + h = m.scatterMedia(h, b, 0) + } for i, layer := range m.Layers { var c cache.Cache if caches != nil && i < len(caches) { @@ -1507,7 +1546,13 @@ func (m *mtpDraft) Forward(b *batch.Batch, _, draftCaches []cache.Cache) (hidden dims := b.InputIDs.Dims() B, L := int32(dims[0]), int32(dims[1]) - emb := m.MTP.Enorm.Forward(m.EmbedTokens.Forward(b.InputIDs), m.LayerNormEpsilon) + raw := m.EmbedTokens.Forward(b.InputIDs) + if len(b.Media) > 0 { + // The pair at slot S embeds the look-ahead token S+1, so each row's + // column 0 holds the prompt token one past its offset. + raw = (*Model)(m).scatterMedia(raw, b, 1) + } + emb := m.MTP.Enorm.Forward(raw, m.LayerNormEpsilon) h := m.MTP.Hnorm.Forward(b.Hidden, m.LayerNormEpsilon) fused := m.MTP.FC.Forward(emb.Concatenate(-1, h)) diff --git a/mlxrunner/model/nemotron_h/nemotron_h_test.go b/mlxrunner/model/nemotron_h/nemotron_h_test.go index 2b0e46627..7a8ce3a7f 100644 --- a/mlxrunner/model/nemotron_h/nemotron_h_test.go +++ b/mlxrunner/model/nemotron_h/nemotron_h_test.go @@ -1,6 +1,7 @@ package nemotron_h import ( + "fmt" "math" "slices" "strings" @@ -9,6 +10,7 @@ import ( "github.com/ollama/ollama/mlx" "github.com/ollama/ollama/mlx/mlxtest" "github.com/ollama/ollama/mlxrunner/cache" + "github.com/ollama/ollama/mlxrunner/model" "github.com/ollama/ollama/mlxrunner/nn" ) @@ -537,3 +539,282 @@ func TestFoldSharedExpertsExtendsGlobalScales(t *testing.T) { } }) } + +func TestParseVisionConfigOmni(t *testing.T) { + cfg, err := parseVisionConfig([]byte(`{ + "ps_version": "v2", + "vision_config": { + "version": "radio_v2.5-h", + "patch_size": 16, + "min_num_patches": 1024, + "max_num_patches": 13312, + "args": {"model": "vit_huge_patch16_224"} + }, + "patch_size": 16, + "downsample_ratio": 0.5, + "img_context_token_id": 18, + "img_context_token": "", + "img_start_token": "", + "img_end_token": "", + "vit_hidden_size": 1280, + "projector_hidden_size": 20480 + }`), []byte(`{ + "patch_size": 16, + "downsample_ratio": 0.5, + "norm_mean": [0.48145466, 0.4578275, 0.40821073], + "norm_std": [0.26862954, 0.26130258, 0.27577711] + }`)) + if err != nil { + t.Fatalf("parseVisionConfig returned error: %v", err) + } + if cfg == nil { + t.Fatal("parseVisionConfig returned nil config") + } + if got, want := cfg.ModelName, "vit_huge_patch16_224"; got != want { + t.Fatalf("ModelName = %q, want %q", got, want) + } + if got, want := cfg.PSVersion, "v2"; got != want { + t.Fatalf("PSVersion = %q, want %q", got, want) + } + if got, want := cfg.NumAttentionHeads, int32(16); got != want { + t.Fatalf("NumAttentionHeads = %d, want %d", got, want) + } + if cfg.PatchSize != 0 || cfg.HiddenSize != 0 || cfg.NumHiddenLayers != 0 || cfg.HeadDim != 0 { + t.Fatalf("weight-derived dimensions initialized before weights: %+v", cfg) + } + if got, want := cfg.DownsampleFactor, int32(2); got != want { + t.Fatalf("DownsampleFactor = %d, want %d", got, want) + } + if got, want := cfg.ImageTokenID, int32(18); got != want { + t.Fatalf("ImageTokenID = %d, want %d", got, want) + } + if got, want := cfg.MinNumPatches, 1024; got != want { + t.Fatalf("MinNumPatches = %d, want %d", got, want) + } + if got, want := cfg.MaxNumPatches, 13312; got != want { + t.Fatalf("MaxNumPatches = %d, want %d", got, want) + } + if got, want := cfg.MaxModelLen, 16384; got != want { + t.Fatalf("MaxModelLen = %d, want %d", got, want) + } +} + +func TestParseVisionConfigRejectsMalformedNormalization(t *testing.T) { + config := []byte(`{ + "ps_version": "v2", + "vision_config": {"version": "c-radio_v4-h", "args": {"model": "vit_huge_patch16_224"}}, + "norm_mean": [0.1, 0.2] + }`) + if _, err := parseVisionConfig(config, nil); err == nil { + t.Fatal("expected malformed config norm_mean error") + } + + config = []byte(`{"ps_version":"v2","vision_config":{"version":"c-radio_v4-h","args":{"model":"vit_huge_patch16_224"}}}`) + preprocessor := []byte(`{"norm_std": [0.1, 0.2, 0.3, 0.4]}`) + if _, err := parseVisionConfig(config, preprocessor); err == nil { + t.Fatal("expected malformed preprocessor norm_std error") + } +} + +func TestParseVisionConfigRequiresDynamicResolutionBounds(t *testing.T) { + for _, config := range []string{ + `{"ps_version":"v2","vision_config":{"version":"c-radio_v4-h","min_num_patches":1024,"args":{"model":"vit_huge_patch16_224"}}}`, + `{"ps_version":"v2","vision_config":{"version":"c-radio_v4-h","max_num_patches":13312,"args":{"model":"vit_huge_patch16_224"}}}`, + `{"ps_version":"v2","vision_config":{"version":"c-radio_v4-h","args":{"model":"vit_huge_patch16_224"}}}`, + } { + if _, err := parseVisionConfig([]byte(config), nil); err == nil || !strings.Contains(err.Error(), "requires min_num_patches and max_num_patches") { + t.Fatalf("parseVisionConfig error = %v, want missing dynamic-resolution bounds", err) + } + } + + _, err := parseVisionConfig([]byte(`{ + "ps_version": "v2", + "vision_config": { + "version": "c-radio_v4-h", + "min_num_patches": 2048, + "max_num_patches": 1024, + "args": {"model": "vit_huge_patch16_224"} + } + }`), nil) + if err == nil || !strings.Contains(err.Error(), "min_num_patches (2048) exceeds max_num_patches (1024)") { + t.Fatalf("parseVisionConfig error = %v, want reversed bounds error", err) + } +} + +func TestParseVisionConfigRejectsUnknownArchitecture(t *testing.T) { + for _, tt := range []struct { + name string + config string + want string + }{ + { + name: "missing model", + config: `{"ps_version":"v2","vision_config":{"min_num_patches":1024,"max_num_patches":13312}}`, + want: `unsupported RADIO model ""`, + }, + { + name: "unknown model", + config: `{"ps_version":"v2","vision_config":{"min_num_patches":1024,"max_num_patches":13312,"args":{"model":"vit_custom_patch16_224"}}}`, + want: `unsupported RADIO model "vit_custom_patch16_224"`, + }, + { + name: "missing ps_version", + config: `{"vision_config":{"min_num_patches":1024,"max_num_patches":13312,"args":{"model":"vit_huge_patch16_224"}}}`, + want: `unsupported RADIO ps_version ""`, + }, + { + name: "unsupported ps_version", + config: `{"ps_version":"v1","vision_config":{"min_num_patches":1024,"max_num_patches":13312,"args":{"model":"vit_huge_patch16_224"}}}`, + want: `unsupported RADIO ps_version "v1"`, + }, + { + name: "mismatched head count", + config: `{"ps_version":"v2","vision_config":{"num_attention_heads":8,"min_num_patches":1024,"max_num_patches":13312,"args":{"model":"vit_huge_patch16_224"}}}`, + want: `RADIO model "vit_huge_patch16_224" has num_attention_heads=16, config has 8`, + }, + } { + t.Run(tt.name, func(t *testing.T) { + _, err := parseVisionConfig([]byte(tt.config), nil) + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("parseVisionConfig error = %v, want %q", err, tt.want) + } + }) + } +} + +func TestRadioPatchDimensions(t *testing.T) { + hiddenSize, patchSize, err := radioPatchDimensions([]int{1280, 3 * 16 * 16}) + if err != nil { + t.Fatal(err) + } + if hiddenSize != 1280 || patchSize != 16 { + t.Fatalf("dimensions = (%d, %d), want (1280, 16)", hiddenSize, patchSize) + } + + for _, dims := range [][]int{{1280}, {1280, 767}, {1280, 3 * 15}} { + if _, _, err := radioPatchDimensions(dims); err == nil { + t.Fatalf("radioPatchDimensions(%v) succeeded, want error", dims) + } + } +} + +func TestRadioLayerCount(t *testing.T) { + const prefix = "vision_model.radio_model.model." + tensors := map[string]*mlx.Array{ + prefix + "blocks.0.norm1.weight": nil, + prefix + "blocks.1.norm1.weight": nil, + prefix + "blocks.1.mlp.fc1.weight": nil, + "unrelated.blocks.2.weight": nil, + } + if got, err := radioLayerCount(tensors, prefix); err != nil || got != 2 { + t.Fatalf("radioLayerCount = %d, %v, want 2, nil", got, err) + } + + delete(tensors, prefix+"blocks.1.norm1.weight") + delete(tensors, prefix+"blocks.1.mlp.fc1.weight") + tensors[prefix+"blocks.2.norm1.weight"] = nil + if _, err := radioLayerCount(tensors, prefix); err == nil || !strings.Contains(err.Error(), "missing RADIO transformer block 1") { + t.Fatalf("radioLayerCount error = %v, want missing block 1", err) + } +} + +func TestVisionConfigLoadsArchitectureFromWeights(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + const prefix = "vision_model.radio_model.model." + weight := mlx.Zeros(mlx.DTypeBFloat16, 384, 3*16*16) + tensors := map[string]*mlx.Array{prefix + "patch_generator.embedder.weight": weight} + for i := range 12 { + tensors[fmt.Sprintf("%sblocks.%d.norm1.weight", prefix, i)] = weight + } + + cfg := &VisionConfig{ModelName: "vit_small_patch16_224", NumAttentionHeads: 6} + if err := cfg.loadArchitecture(tensors, prefix); err != nil { + t.Fatal(err) + } + if cfg.PatchSize != 16 || cfg.HiddenSize != 384 || cfg.NumHiddenLayers != 12 || cfg.HeadDim != 64 { + t.Fatalf("derived architecture = %+v", cfg) + } + }) +} + +func TestLoadVisionWeightsRejectsNonSquarePositionGrid(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + const prefix = "vision_model.radio_model.model." + tensors := map[string]*mlx.Array{ + prefix + "patch_generator.embedder.weight": mlx.Zeros(mlx.DTypeBFloat16, 4, 3), + prefix + "patch_generator.cls_token.token": mlx.Zeros(mlx.DTypeBFloat16, 1, 4), + prefix + "patch_generator.pos_embed": mlx.Zeros(mlx.DTypeBFloat16, 1, 6, 4), + } + m := &Model{ + VisionConfig: &VisionConfig{}, + VisionEncoder: &RadioVisionEncoder{}, + Projector: &VisionProjector{}, + } + linears := model.NewLinearFactory(tensors, 0, 0, "", nil) + if err := m.loadVisionWeights(tensors, linears); err == nil || !strings.Contains(err.Error(), "6 positions, want a square grid") { + t.Fatalf("loadVisionWeights error = %v, want non-square position-grid error", err) + } + }) +} + +func TestResizeGrid2DPreservesSourceDType(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + x := mlx.FromValues([]float32{0, 1, 2, 3}, 1, 2, 2, 1).AsType(mlx.DTypeBFloat16) + got := resizeGrid2D(x, 3, 3) + mlx.Eval(got) + if got.DType() != mlx.DTypeBFloat16 { + t.Fatalf("resizeGrid2D dtype = %s, want bfloat16", got.DType()) + } + }) +} + +func TestNemotronImagePatchBudget(t *testing.T) { + cfg := &VisionConfig{ + DownsampleFactor: 2, + MinNumPatches: 1024, + MaxNumPatches: 13312, + MaxModelLen: 16384, + } + + // Context-bound, not memory-bound: model.MediaModel requires PrepareMedia + // to be deterministic for given segments. + if got, want := nemotronImagePatchBudget(cfg), 13312; got != want { + t.Fatalf("budget = %d, want %d", got, want) + } + if got, want := nemotronImagePatchBudget(&VisionConfig{DownsampleFactor: 2, MinNumPatches: 1024, MaxNumPatches: 13312, MaxModelLen: 512}), 2032; got != want { + t.Fatalf("context-limited budget = %d, want %d", got, want) + } +} + +func TestNemotronImagePatchGrid(t *testing.T) { + cfg := &VisionConfig{ + PatchSize: 16, + DownsampleFactor: 2, + MinNumPatches: 1024, + MaxNumPatches: 13312, + } + + for _, tt := range []struct { + name string + height int + width int + budget int + wantMax int + wantMin int + wantFactor int + }{ + {name: "square minimum", height: 512, width: 512, budget: 1024, wantMax: 1024, wantMin: 1024, wantFactor: 2}, + {name: "large square capped", height: 2048, width: 2048, budget: 13312, wantMax: 13312, wantMin: 1024, wantFactor: 2}, + {name: "wide image capped", height: 512, width: 2048, budget: 4096, wantMax: 4096, wantMin: 1024, wantFactor: 2}, + } { + t.Run(tt.name, func(t *testing.T) { + patchH, patchW := nemotronImagePatchGrid(tt.height, tt.width, tt.budget, cfg) + if got := patchH * patchW; got > tt.wantMax || got < tt.wantMin { + t.Fatalf("patches = %d (%dx%d), want within [%d,%d]", got, patchH, patchW, tt.wantMin, tt.wantMax) + } + if patchH%tt.wantFactor != 0 || patchW%tt.wantFactor != 0 { + t.Fatalf("patch grid = %dx%d, want both divisible by %d", patchH, patchW, tt.wantFactor) + } + }) + } +} diff --git a/mlxrunner/model/nemotron_h/vision.go b/mlxrunner/model/nemotron_h/vision.go new file mode 100644 index 000000000..8c60a1e1a --- /dev/null +++ b/mlxrunner/model/nemotron_h/vision.go @@ -0,0 +1,653 @@ +package nemotron_h + +import ( + "encoding/json" + "fmt" + "math" + "strconv" + "strings" + + "github.com/ollama/ollama/mlx" + "github.com/ollama/ollama/mlxrunner/model" + "github.com/ollama/ollama/mlxrunner/nn" +) + +const ( + nemotronVisionDefaultLayerNormEpsilon = float32(1e-6) + nemotronVisionDefaultProjectorNormEps = float32(1e-5) + nemotronVisionDefaultMaxModelLen = 16384 + nemotronVisionReservedTokens = 4 +) + +type radioModelSpec struct { + patchSize int32 + hiddenSize int32 + numHiddenLayers int32 + numAttentionHeads int32 +} + +// radioModelSpecs is the architecture set used by RADIO's timm model names. +var radioModelSpecs = map[string]radioModelSpec{ + "vit_small_patch16_224": {patchSize: 16, hiddenSize: 384, numHiddenLayers: 12, numAttentionHeads: 6}, + "vit_base_patch16_224": {patchSize: 16, hiddenSize: 768, numHiddenLayers: 12, numAttentionHeads: 12}, + "vit_large_patch16_224": {patchSize: 16, hiddenSize: 1024, numHiddenLayers: 24, numAttentionHeads: 16}, + "vit_huge_patch16_224": {patchSize: 16, hiddenSize: 1280, numHiddenLayers: 32, numAttentionHeads: 16}, +} + +type VisionConfig struct { + ModelName string + PSVersion string + Version string + PatchSize int32 + HiddenSize int32 + NumHiddenLayers int32 + NumAttentionHeads int32 + HeadDim int32 + LayerNormEps float32 + ProjectorNormEps float32 + DownsampleFactor int32 + MinNumPatches int + MaxNumPatches int + MaxModelLen int + Mean [3]float32 + Std [3]float32 + ImageTokenID int32 + ImageToken string + ImageStartToken string + ImageEndToken string + + configuredPatchSize int32 + configuredHiddenSize int32 + configuredNumHiddenLayers int32 +} + +type RadioVisionEncoder struct { + PatchEmbed nn.LinearLayer + ClassToken *mlx.Array + Position *mlx.Array + PositionGridSize int32 + Layers []*RadioVisionLayer +} + +type RadioVisionLayer struct { + Norm1 *nn.LayerNorm + QKV nn.LinearLayer + Proj nn.LinearLayer + Norm2 *nn.LayerNorm + FC1 nn.LinearLayer + FC2 nn.LinearLayer +} + +type VisionProjector struct { + Norm *nn.RMSNorm + FC1 nn.LinearLayer + FC2 nn.LinearLayer +} + +func parseVisionConfig(configData, preprocessorData []byte) (*VisionConfig, error) { + var env struct { + PSVersion string `json:"ps_version"` + VisionConfig *struct { + Version string `json:"version"` + PatchSize int32 `json:"patch_size"` + MinNumPatches int `json:"min_num_patches"` + MaxNumPatches int `json:"max_num_patches"` + HiddenSize int32 `json:"hidden_size"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + NumHeads int32 `json:"num_attention_heads"` + Args struct { + Model string `json:"model"` + MinNumPatches int `json:"min_num_patches"` + MaxNumPatches int `json:"max_num_patches"` + } `json:"args"` + } `json:"vision_config"` + PatchSize int32 `json:"patch_size"` + DownsampleRatio float32 `json:"downsample_ratio"` + ImgContextTokenID int32 `json:"img_context_token_id"` + ImgContextToken string `json:"img_context_token"` + ImgStartToken string `json:"img_start_token"` + ImgEndToken string `json:"img_end_token"` + VitHiddenSize int32 `json:"vit_hidden_size"` + NormMean []float32 `json:"norm_mean"` + NormStd []float32 `json:"norm_std"` + } + if err := json.Unmarshal(configData, &env); err != nil { + return nil, fmt.Errorf("parse vision config: %w", err) + } + if env.VisionConfig == nil { + return nil, nil + } + modelName := strings.TrimSpace(env.VisionConfig.Args.Model) + spec, ok := radioModelSpecs[modelName] + if !ok { + return nil, fmt.Errorf("unsupported RADIO model %q", modelName) + } + psVersion := strings.TrimSpace(env.PSVersion) + if psVersion != "v2" { + return nil, fmt.Errorf("unsupported RADIO ps_version %q", psVersion) + } + if heads := env.VisionConfig.NumHeads; heads > 0 && heads != spec.numAttentionHeads { + return nil, fmt.Errorf("RADIO model %q has num_attention_heads=%d, config has %d", modelName, spec.numAttentionHeads, heads) + } + + cfg := &VisionConfig{ + ModelName: modelName, + PSVersion: psVersion, + Version: strings.TrimSpace(env.VisionConfig.Version), + NumAttentionHeads: spec.numAttentionHeads, + LayerNormEps: nemotronVisionDefaultLayerNormEpsilon, + ProjectorNormEps: nemotronVisionDefaultProjectorNormEps, + MinNumPatches: firstPositiveInt(env.VisionConfig.MinNumPatches, env.VisionConfig.Args.MinNumPatches), + MaxNumPatches: firstPositiveInt(env.VisionConfig.MaxNumPatches, env.VisionConfig.Args.MaxNumPatches), + MaxModelLen: nemotronVisionDefaultMaxModelLen, + Mean: [3]float32{0.48145466, 0.4578275, 0.40821073}, + Std: [3]float32{0.26862954, 0.26130258, 0.27577711}, + ImageTokenID: env.ImgContextTokenID, + ImageToken: firstNonEmpty(env.ImgContextToken, ""), + ImageStartToken: firstNonEmpty(env.ImgStartToken, ""), + ImageEndToken: firstNonEmpty(env.ImgEndToken, ""), + configuredPatchSize: firstPositiveInt32(env.VisionConfig.PatchSize, env.PatchSize), + configuredHiddenSize: firstPositiveInt32(env.VisionConfig.HiddenSize, env.VitHiddenSize), + configuredNumHiddenLayers: env.VisionConfig.NumHiddenLayers, + } + if env.DownsampleRatio > 0 { + cfg.DownsampleFactor = int32(math.Round(float64(1 / env.DownsampleRatio))) + } + if cfg.DownsampleFactor <= 0 { + cfg.DownsampleFactor = 2 + } + if err := setTriplet(cfg.Mean[:], env.NormMean, "norm_mean"); err != nil { + return nil, fmt.Errorf("vision config: %w", err) + } + if err := setTriplet(cfg.Std[:], env.NormStd, "norm_std"); err != nil { + return nil, fmt.Errorf("vision config: %w", err) + } + + if len(preprocessorData) > 0 { + if err := cfg.applyPreprocessorConfig(preprocessorData); err != nil { + return nil, err + } + } + + switch cfg.Version { + case "", "radio_v2.5-h", "c-radio_v4-h": + default: + return nil, fmt.Errorf("unsupported RADIO version %q", cfg.Version) + } + if cfg.DownsampleFactor != 2 { + return nil, fmt.Errorf("unsupported RADIO downsample factor=%d", cfg.DownsampleFactor) + } + if cfg.MinNumPatches <= 0 || cfg.MaxNumPatches <= 0 { + return nil, fmt.Errorf("RADIO dynamic resolution requires min_num_patches and max_num_patches") + } + if cfg.MinNumPatches > cfg.MaxNumPatches { + return nil, fmt.Errorf("RADIO min_num_patches (%d) exceeds max_num_patches (%d)", cfg.MinNumPatches, cfg.MaxNumPatches) + } + return cfg, nil +} + +func (cfg *VisionConfig) applyPreprocessorConfig(data []byte) error { + var pre struct { + PatchSize int32 `json:"patch_size"` + DownsampleRatio float32 `json:"downsample_ratio"` + NormMean []float32 `json:"norm_mean"` + NormStd []float32 `json:"norm_std"` + MinNumPatches int `json:"min_num_patches"` + MaxNumPatches int `json:"max_num_patches"` + MaxModelLen int `json:"max_model_len"` + } + if err := json.Unmarshal(data, &pre); err != nil { + return fmt.Errorf("parse preprocessor_config.json: %w", err) + } + if pre.PatchSize > 0 { + if cfg.configuredPatchSize > 0 && pre.PatchSize != cfg.configuredPatchSize { + return fmt.Errorf("parse preprocessor_config.json: patch_size=%d, config has %d", pre.PatchSize, cfg.configuredPatchSize) + } + cfg.configuredPatchSize = pre.PatchSize + } + if pre.DownsampleRatio > 0 { + cfg.DownsampleFactor = int32(math.Round(float64(1 / pre.DownsampleRatio))) + } + cfg.MinNumPatches = firstPositiveInt(pre.MinNumPatches, cfg.MinNumPatches) + cfg.MaxNumPatches = firstPositiveInt(pre.MaxNumPatches, cfg.MaxNumPatches) + cfg.MaxModelLen = firstPositiveInt(pre.MaxModelLen, cfg.MaxModelLen) + if err := setTriplet(cfg.Mean[:], pre.NormMean, "norm_mean"); err != nil { + return fmt.Errorf("parse preprocessor_config.json: %w", err) + } + if err := setTriplet(cfg.Std[:], pre.NormStd, "norm_std"); err != nil { + return fmt.Errorf("parse preprocessor_config.json: %w", err) + } + return nil +} + +func firstPositiveInt32(values ...int32) int32 { + for _, value := range values { + if value > 0 { + return value + } + } + return 0 +} + +func firstPositiveInt(values ...int) int { + for _, value := range values { + if value > 0 { + return value + } + } + return 0 +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if strings.TrimSpace(value) != "" { + return value + } + } + return "" +} + +func setTriplet(dst []float32, src []float32, name string) error { + if len(src) == 0 { + return nil + } + if len(src) != len(dst) { + return fmt.Errorf("%s must contain %d values, got %d", name, len(dst), len(src)) + } + copy(dst, src) + return nil +} + +func (m *Model) configureVision(cfg *VisionConfig) error { + imageStartTokenID, ok := m.tok.GetSpecialToken(cfg.ImageStartToken) + if !ok { + return fmt.Errorf("tokenizer is missing %s", cfg.ImageStartToken) + } + imageTokenID, ok := m.tok.GetSpecialToken(cfg.ImageToken) + if !ok { + return fmt.Errorf("tokenizer is missing %s", cfg.ImageToken) + } + if cfg.ImageTokenID > 0 && imageTokenID != cfg.ImageTokenID { + return fmt.Errorf("tokenizer %s id = %d, config has %d", cfg.ImageToken, imageTokenID, cfg.ImageTokenID) + } + imageEndTokenID, ok := m.tok.GetSpecialToken(cfg.ImageEndToken) + if !ok { + return fmt.Errorf("tokenizer is missing %s", cfg.ImageEndToken) + } + + m.VisionConfig = cfg + m.VisionEncoder = &RadioVisionEncoder{} + m.Projector = &VisionProjector{} + m.imageStartTokenID = imageStartTokenID + m.imageTokenID = imageTokenID + m.imageEndTokenID = imageEndTokenID + return nil +} + +func (m *Model) loadVisionWeights(tensors map[string]*mlx.Array, linears model.LinearFactory) error { + visionPrefix := resolveVisionPrefix(tensors) + if visionPrefix == "" { + return fmt.Errorf("missing RADIO vision weights") + } + + ve := m.VisionEncoder + ve.PatchEmbed = linears.Make(visionPrefix + "patch_generator.embedder") + ve.ClassToken = tensors[visionPrefix+"patch_generator.cls_token.token"] + ve.Position = tensors[visionPrefix+"patch_generator.pos_embed"] + if ve.PatchEmbed == nil || ve.ClassToken == nil || ve.Position == nil { + return fmt.Errorf("missing RADIO patch generator weights") + } + positionGridSize, err := radioPositionGridSize(ve.Position) + if err != nil { + return err + } + ve.PositionGridSize = positionGridSize + + cfg := m.VisionConfig + if err := cfg.loadArchitecture(tensors, visionPrefix); err != nil { + return err + } + ve.Layers = make([]*RadioVisionLayer, cfg.NumHiddenLayers) + for i := range cfg.NumHiddenLayers { + prefix := fmt.Sprintf("%sblocks.%d", visionPrefix, i) + layer := &RadioVisionLayer{ + Norm1: newLayerNorm(tensors, prefix+".norm1", cfg.LayerNormEps), + QKV: linears.Make(prefix + ".attn.qkv"), + Proj: linears.Make(prefix + ".attn.proj"), + Norm2: newLayerNorm(tensors, prefix+".norm2", cfg.LayerNormEps), + FC1: linears.Make(prefix + ".mlp.fc1"), + FC2: linears.Make(prefix + ".mlp.fc2"), + } + if layer.Norm1 == nil || layer.QKV == nil || layer.Proj == nil || + layer.Norm2 == nil || layer.FC1 == nil || layer.FC2 == nil { + return fmt.Errorf("vision layer %d: missing RADIO transformer weights", i) + } + ve.Layers[i] = layer + } + + projector := m.Projector + normWeight := tensors["mlp1.0.weight"] + projector.FC1 = linears.Make("mlp1.1") + projector.FC2 = linears.Make("mlp1.3") + if normWeight == nil || projector.FC1 == nil || projector.FC2 == nil { + return fmt.Errorf("missing Nemotron vision projector weights") + } + projector.Norm = nn.NewRMSNorm(normWeight, cfg.ProjectorNormEps) + return nil +} + +func (cfg *VisionConfig) loadArchitecture(tensors map[string]*mlx.Array, visionPrefix string) error { + weight := tensors[visionPrefix+"patch_generator.embedder.weight"] + if weight == nil { + return fmt.Errorf("missing RADIO patch embedder weight") + } + hiddenSize, patchSize, err := radioPatchDimensions(weight.Dims()) + if err != nil { + return fmt.Errorf("RADIO patch embedder: %w", err) + } + numHiddenLayers, err := radioLayerCount(tensors, visionPrefix) + if err != nil { + return err + } + + spec, ok := radioModelSpecs[cfg.ModelName] + if !ok { + return fmt.Errorf("unsupported RADIO model %q", cfg.ModelName) + } + cfg.NumAttentionHeads = spec.numAttentionHeads + for _, check := range []struct { + name string + got int32 + want int32 + }{ + {name: "patch_size", got: patchSize, want: spec.patchSize}, + {name: "hidden_size", got: hiddenSize, want: spec.hiddenSize}, + {name: "num_hidden_layers", got: numHiddenLayers, want: spec.numHiddenLayers}, + } { + if check.got != check.want { + return fmt.Errorf("RADIO model %q has %s=%d, weights have %d", cfg.ModelName, check.name, check.want, check.got) + } + } + for _, check := range []struct { + name string + configured int32 + derived int32 + }{ + {name: "patch_size", configured: cfg.configuredPatchSize, derived: patchSize}, + {name: "hidden_size", configured: cfg.configuredHiddenSize, derived: hiddenSize}, + {name: "num_hidden_layers", configured: cfg.configuredNumHiddenLayers, derived: numHiddenLayers}, + } { + if check.configured > 0 && check.configured != check.derived { + return fmt.Errorf("RADIO %s=%d, weights have %d", check.name, check.configured, check.derived) + } + } + if hiddenSize%cfg.NumAttentionHeads != 0 { + return fmt.Errorf("vision hidden_size (%d) must be divisible by num_attention_heads (%d)", hiddenSize, cfg.NumAttentionHeads) + } + + cfg.PatchSize = patchSize + cfg.HiddenSize = hiddenSize + cfg.NumHiddenLayers = numHiddenLayers + cfg.HeadDim = hiddenSize / cfg.NumAttentionHeads + return nil +} + +func radioPatchDimensions(dims []int) (hiddenSize, patchSize int32, err error) { + if len(dims) != 2 || dims[0] <= 0 || dims[1] <= 0 { + return 0, 0, fmt.Errorf("weight shape %v, want [hidden_size, 3*patch_size*patch_size]", dims) + } + if dims[1]%3 != 0 { + return 0, 0, fmt.Errorf("input dimension %d is not divisible by 3 RGB channels", dims[1]) + } + patchArea := dims[1] / 3 + patch := int(math.Sqrt(float64(patchArea))) + if patch*patch != patchArea { + return 0, 0, fmt.Errorf("input dimension %d does not describe square RGB patches", dims[1]) + } + return int32(dims[0]), int32(patch), nil +} + +func radioLayerCount(tensors map[string]*mlx.Array, visionPrefix string) (int32, error) { + blockPrefix := visionPrefix + "blocks." + indices := make(map[int]struct{}) + for name := range tensors { + rest, ok := strings.CutPrefix(name, blockPrefix) + if !ok { + continue + } + indexText, _, ok := strings.Cut(rest, ".") + if !ok { + continue + } + index, err := strconv.Atoi(indexText) + if err != nil || index < 0 { + continue + } + indices[index] = struct{}{} + } + if len(indices) == 0 { + return 0, fmt.Errorf("missing RADIO transformer block weights") + } + for i := range len(indices) { + if _, ok := indices[i]; !ok { + return 0, fmt.Errorf("missing RADIO transformer block %d", i) + } + } + return int32(len(indices)), nil +} + +func resolveVisionPrefix(tensors map[string]*mlx.Array) string { + for _, prefix := range []string{ + "vision_model.radio_model.model.", + "model.vision_model.radio_model.model.", + } { + if tensors[prefix+"patch_generator.embedder.weight"] != nil { + return prefix + } + } + return "" +} + +func newLayerNorm(tensors map[string]*mlx.Array, prefix string, eps float32) *nn.LayerNorm { + weight := tensors[prefix+".weight"] + bias := tensors[prefix+".bias"] + if weight == nil || bias == nil { + return nil + } + return &nn.LayerNorm{Weight: weight, Bias: bias, Eps: eps} +} + +func (m *Model) visionTokenCount(height, width int) int { + cfg := m.VisionConfig + if cfg == nil || height <= 0 || width <= 0 { + return 0 + } + patchH := height / int(cfg.PatchSize) + patchW := width / int(cfg.PatchSize) + return (patchH / int(cfg.DownsampleFactor)) * (patchW / int(cfg.DownsampleFactor)) +} + +func (m *Model) forwardVision(pixelValues *mlx.Array) *mlx.Array { + cfg := m.VisionConfig + ve := m.VisionEncoder + dims := pixelValues.Dims() + B := int32(dims[0]) + H, W := int32(dims[2]), int32(dims[3]) + patchH := H / cfg.PatchSize + patchW := W / cfg.PatchSize + + h := ve.PatchEmbed.Forward(radioPatchify(pixelValues, cfg.PatchSize)) + h = mlx.Add(h, radioPositionEmbedding(ve.Position, ve.PositionGridSize, patchH, patchW)) + cls := mlx.ExpandDims(ve.ClassToken, 0) + if B > 1 { + cls = mlx.Tile(cls, []int32{B, 1, 1}) + } + h = mlx.Concatenate([]*mlx.Array{cls, h}, 1) + + for _, layer := range ve.Layers { + h = radioVisionLayerForward(layer, h, B, cfg) + } + + clsTokens := int32(ve.ClassToken.Dim(0)) + h = mlx.SliceStartStop( + h, + []int32{0, clsTokens, 0}, + []int32{B, clsTokens + patchH*patchW, cfg.HiddenSize}, + ) + h = mlx.Reshape(h, B, patchH, patchW, cfg.HiddenSize) + h = radioPixelShuffle(h, cfg.DownsampleFactor) + h = mlx.Reshape(h, B, (patchH/cfg.DownsampleFactor)*(patchW/cfg.DownsampleFactor), cfg.HiddenSize*cfg.DownsampleFactor*cfg.DownsampleFactor) + return m.Projector.Forward(h) +} + +func (p *VisionProjector) Forward(x *mlx.Array) *mlx.Array { + x = p.Norm.Forward(x, 0) + x = p.FC1.Forward(x) + x = mlx.ReLUSquared(x) + return p.FC2.Forward(x) +} + +func radioVisionLayerForward(layer *RadioVisionLayer, x *mlx.Array, B int32, cfg *VisionConfig) *mlx.Array { + S := int32(x.Dim(1)) + normed := layer.Norm1.Forward(x) + attn := radioAttentionForward(layer, normed, B, S, cfg) + h := mlx.Add(x, attn) + mlp := layer.FC1.Forward(layer.Norm2.Forward(h)) + mlpDType := mlp.DType() + mlp = mlx.GELU(mlp.AsType(mlx.DTypeFloat32)).AsType(mlpDType) + mlp = layer.FC2.Forward(mlp) + return mlx.Add(h, mlp) +} + +func radioAttentionForward(layer *RadioVisionLayer, x *mlx.Array, B, S int32, cfg *VisionConfig) *mlx.Array { + qkv := layer.QKV.Forward(x) + q := sliceAxis(qkv, 2, 0, cfg.HiddenSize) + k := sliceAxis(qkv, 2, cfg.HiddenSize, 2*cfg.HiddenSize) + v := sliceAxis(qkv, 2, 2*cfg.HiddenSize, 3*cfg.HiddenSize) + + q = mlx.Transpose(mlx.Reshape(q, B, S, cfg.NumAttentionHeads, cfg.HeadDim), 0, 2, 1, 3) + k = mlx.Transpose(mlx.Reshape(k, B, S, cfg.NumAttentionHeads, cfg.HeadDim), 0, 2, 1, 3) + v = mlx.Transpose(mlx.Reshape(v, B, S, cfg.NumAttentionHeads, cfg.HeadDim), 0, 2, 1, 3) + + out := mlx.FastScaledDotProductAttention(q, k, v, float32(1/math.Sqrt(float64(cfg.HeadDim))), "", nil) + out = mlx.Transpose(out, 0, 2, 1, 3) + out = mlx.Reshape(out, B, S, cfg.HiddenSize) + return layer.Proj.Forward(out) +} + +func radioPatchify(pixelValues *mlx.Array, patchSize int32) *mlx.Array { + dims := pixelValues.Dims() + B := int32(dims[0]) + C := int32(dims[1]) + H, W := int32(dims[2]), int32(dims[3]) + patchH := H / patchSize + patchW := W / patchSize + + x := mlx.Reshape(pixelValues, B, C, patchH, patchSize, patchW, patchSize) + x = mlx.Transpose(x, 0, 2, 4, 1, 3, 5) + return mlx.Reshape(x, B, patchH*patchW, C*patchSize*patchSize) +} + +func radioPixelShuffle(x *mlx.Array, factor int32) *mlx.Array { + dims := x.Dims() + B := int32(dims[0]) + H := int32(dims[1]) + W := int32(dims[2]) + C := int32(dims[3]) + + x = mlx.Reshape(x, B, H, W/factor, factor*C) + x = mlx.Transpose(x, 0, 2, 1, 3) + x = mlx.Reshape(x, B, W/factor, H/factor, factor*factor*C) + return mlx.Transpose(x, 0, 2, 1, 3) +} + +func radioPositionGridSize(pos *mlx.Array) (int32, error) { + dims := pos.Dims() + if len(dims) < 2 { + return 0, fmt.Errorf("RADIO position embedding has %d dimensions, want at least 2", len(dims)) + } + numPositions := dims[len(dims)-2] + if numPositions <= 0 { + return 0, fmt.Errorf("RADIO position embedding has no positions") + } + src := int32(math.Round(math.Sqrt(float64(numPositions)))) + if int(src*src) != numPositions { + return 0, fmt.Errorf("RADIO position embedding has %d positions, want a square grid", numPositions) + } + return src, nil +} + +func radioPositionEmbedding(pos *mlx.Array, src, patchH, patchW int32) *mlx.Array { + dims := pos.Dims() + hidden := int32(dims[len(dims)-1]) + grid := mlx.Reshape(pos, 1, src, src, hidden) + target := max(patchH, patchW) + if target != src { + grid = resizeGrid2D(grid, target, target) + } + if patchH != target || patchW != target { + grid = mlx.SliceStartStop( + grid, + []int32{0, 0, 0, 0}, + []int32{1, patchH, patchW, hidden}, + ) + } + return mlx.Reshape(grid, 1, patchH*patchW, hidden) +} + +func resizeGrid2D(x *mlx.Array, outH, outW int32) *mlx.Array { + if int32(x.Dim(1)) == outH && int32(x.Dim(2)) == outW { + return x + } + dtype := x.DType() + x = x.AsType(mlx.DTypeFloat32) + if inH := int32(x.Dim(1)); inH != outH { + lo, hi, loWeight, hiWeight := linearResizeWeights(inH, outH) + loX := mlx.Take(x, mlx.NewArrayInt32(lo, []int32{outH}), 1) + hiX := mlx.Take(x, mlx.NewArrayInt32(hi, []int32{outH}), 1) + loW := mlx.FromValues(loWeight, 1, int(outH), 1, 1) + hiW := mlx.FromValues(hiWeight, 1, int(outH), 1, 1) + x = mlx.Add(mlx.Mul(loX, loW), mlx.Mul(hiX, hiW)) + } + if inW := int32(x.Dim(2)); inW != outW { + lo, hi, loWeight, hiWeight := linearResizeWeights(inW, outW) + loX := mlx.Take(x, mlx.NewArrayInt32(lo, []int32{outW}), 2) + hiX := mlx.Take(x, mlx.NewArrayInt32(hi, []int32{outW}), 2) + loW := mlx.FromValues(loWeight, 1, 1, int(outW), 1) + hiW := mlx.FromValues(hiWeight, 1, 1, int(outW), 1) + x = mlx.Add(mlx.Mul(loX, loW), mlx.Mul(hiX, hiW)) + } + return x.AsType(dtype) +} + +func linearResizeWeights(in, out int32) ([]int32, []int32, []float32, []float32) { + lo := make([]int32, out) + hi := make([]int32, out) + loWeight := make([]float32, out) + hiWeight := make([]float32, out) + scale := float64(in) / float64(out) + for i := range out { + src := (float64(i)+0.5)*scale - 0.5 + floor := math.Floor(src) + l := int32(floor) + h := l + 1 + wHigh := float32(src - floor) + if l < 0 { + l = 0 + } + if h < 0 { + h = 0 + } + if l >= in { + l = in - 1 + } + if h >= in { + h = in - 1 + } + lo[i] = l + hi[i] = h + loWeight[i] = 1 - wHigh + hiWeight[i] = wHigh + } + return lo, hi, loWeight, hiWeight +} diff --git a/mlxrunner/model/nemotron_h/vision_prompt.go b/mlxrunner/model/nemotron_h/vision_prompt.go new file mode 100644 index 000000000..29c13457c --- /dev/null +++ b/mlxrunner/model/nemotron_h/vision_prompt.go @@ -0,0 +1,179 @@ +package nemotron_h + +import ( + "bytes" + "fmt" + "image" + "math" + + "golang.org/x/image/draw" + + "github.com/ollama/ollama/mlx" + "github.com/ollama/ollama/mlxrunner/batch" + "github.com/ollama/ollama/mlxrunner/model" +) + +// preparedImage is the model-private state travelling with a prepared item. +type preparedImage struct { + height int + width int +} + +// PrepareMedia implements model.MediaModel: splice each image segment's +// placeholder expansion — image start, the feature-token run, image end — +// into the stream, decoding and resizing on the CPU. +func (m *Model) PrepareMedia(segments []model.Segment) (*model.PreparedRequest, error) { + prepared := &model.PreparedRequest{} + for s, seg := range segments { + if seg.Data == nil { + prepared.Tokens = append(prepared.Tokens, seg.Tokens...) + continue + } + if seg.Kind != "image" { + return nil, fmt.Errorf("nemotron_h does not support %s input", seg.Kind) + } + if m.VisionEncoder == nil || m.Projector == nil || m.VisionConfig == nil { + if m.visionErr != nil { + return nil, fmt.Errorf("nemotron_h vision is unavailable: %w", m.visionErr) + } + return nil, fmt.Errorf("this model does not support %s input", seg.Kind) + } + + pixels, height, width, err := m.preprocessImage(seg.Data, nemotronImagePatchBudget(m.VisionConfig)) + if err != nil { + return nil, err + } + + start := len(prepared.Tokens) + prepared.Tokens = append(prepared.Tokens, m.imageStartTokenID) + for range m.visionTokenCount(height, width) { + prepared.Tokens = append(prepared.Tokens, m.imageTokenID) + } + prepared.Tokens = append(prepared.Tokens, m.imageEndTokenID) + + prepared.Items = append(prepared.Items, model.PreparedItem{ + Range: [2]int{start, len(prepared.Tokens)}, + Source: s, + MediaData: pixels, + Dims: []int{1, 3, height, width}, + Opaque: preparedImage{height: height, width: width}, + // Image rows use the text stack's causal mask, so chunked prefill + // may split the expansion after the vision features are encoded. + Causal: true, + }) + } + return prepared, nil +} + +// EncodeMedia implements model.MediaModel: run the RADIO tower and projector +// over one image, returning the lazy [1, tokens, hidden] features. +func (m *Model) EncodeMedia(_ *model.PreparedItem, data *mlx.Array) *mlx.Array { + return mlx.Squeeze(m.forwardVision(data.AsType(m.VisionEncoder.Position.DType())), 0) +} + +// scatterMedia overwrites each expansion's feature rows with the encoded +// features, for whatever part of the expansion this forward covers. The +// expansion is imageStart + features + imageEnd, so the run starts one past +// the splice. delta shifts each row's offset to the sequence position of +// batch column 0, which the MTP draft needs: its slots hold look-ahead +// tokens rather than the tokens at their own offset. +func (m *Model) scatterMedia(h *mlx.Array, b *batch.Batch, delta int) *mlx.Array { + for _, item := range b.Media { + if item.Features == nil { + continue + } + img := item.Opaque.(preparedImage) + start := item.Pos + 1 + end := start + m.visionTokenCount(img.height, img.width) + + off := int(b.SeqOffsets[item.Seq]) + delta + 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 +} + +func (m *Model) preprocessImage(data []byte, maxPatches int) (pixels []float32, height, width int, err error) { + cfg := m.VisionConfig + if cfg == nil { + return nil, 0, 0, fmt.Errorf("model has no vision config") + } + + src, _, err := image.Decode(bytes.NewReader(data)) + if err != nil { + return nil, 0, 0, fmt.Errorf("decode image: %w", err) + } + bounds := src.Bounds() + srcW, srcH := bounds.Dx(), bounds.Dy() + if srcW <= 0 || srcH <= 0 { + return nil, 0, 0, fmt.Errorf("invalid image dimensions %dx%d", srcW, srcH) + } + + patchH, patchW := nemotronImagePatchGrid(srcH, srcW, maxPatches, cfg) + targetH := patchH * int(cfg.PatchSize) + targetW := patchW * int(cfg.PatchSize) + dst := image.NewRGBA(image.Rect(0, 0, targetW, targetH)) + draw.CatmullRom.Scale(dst, dst.Bounds(), src, bounds, draw.Src, nil) + + pixels = make([]float32, 3*targetH*targetW) + plane := targetH * targetW + for y := range targetH { + for x := range targetW { + offset := y*dst.Stride + x*4 + i := y*targetW + x + pixels[i] = (float32(dst.Pix[offset])/255 - cfg.Mean[0]) / cfg.Std[0] + pixels[plane+i] = (float32(dst.Pix[offset+1])/255 - cfg.Mean[1]) / cfg.Std[1] + pixels[2*plane+i] = (float32(dst.Pix[offset+2])/255 - cfg.Mean[2]) / cfg.Std[2] + } + } + + return pixels, targetH, targetW, nil +} + +func nemotronImagePatchBudget(cfg *VisionConfig) int { + return max(cfg.MinNumPatches, min(cfg.MaxNumPatches, (cfg.MaxModelLen-nemotronVisionReservedTokens)*int(cfg.DownsampleFactor)*int(cfg.DownsampleFactor))) +} + +func nemotronImagePatchGrid(srcH, srcW, maxPatches int, cfg *VisionConfig) (int, int) { + patchSize := float64(cfg.PatchSize) + closestH := max(1, int(math.RoundToEven(float64(srcH)/patchSize+0.5))) + closestW := max(1, int(math.RoundToEven(float64(srcW)/patchSize+0.5))) + sourcePatches := max(1, closestH*closestW) + budget := max(cfg.MinNumPatches, min(maxPatches, cfg.MaxNumPatches)) + + scale := math.Min(math.Sqrt(float64(budget)/float64(sourcePatches)), 1) + targetH := max(1, int(math.Floor(scale*float64(closestH)))) + targetW := max(1, int(math.Floor(scale*float64(closestW)))) + if budget > cfg.MinNumPatches && targetH*targetW < cfg.MinNumPatches { + scaleUp := math.Sqrt(float64(cfg.MinNumPatches) / float64(targetH*targetW)) + targetH = max(1, int(math.Ceil(scaleUp*float64(targetH)))) + targetW = max(1, int(math.Ceil(scaleUp*float64(targetW)))) + } + + divisor := int(cfg.DownsampleFactor) + if rem := targetH % divisor; rem != 0 { + inc := divisor - rem + if (targetH+inc)*targetW <= budget { + targetH += inc + } else { + targetH = max(divisor, targetH-rem) + } + } + if rem := targetW % divisor; rem != 0 { + inc := divisor - rem + if targetH*(targetW+inc) <= budget { + targetW += inc + } else { + targetW = max(divisor, targetW-rem) + } + } + + return targetH, targetW +} diff --git a/mlxrunner/model/nemotron_h/vision_prompt_test.go b/mlxrunner/model/nemotron_h/vision_prompt_test.go new file mode 100644 index 000000000..7d587400e --- /dev/null +++ b/mlxrunner/model/nemotron_h/vision_prompt_test.go @@ -0,0 +1,164 @@ +package nemotron_h + +import ( + "bytes" + "errors" + "image" + "image/png" + "testing" + + "github.com/ollama/ollama/mlx" + "github.com/ollama/ollama/mlx/mlxtest" + "github.com/ollama/ollama/mlxrunner/batch" + "github.com/ollama/ollama/mlxrunner/model" +) + +func TestPrepareMediaMarksImageExpansionCausal(t *testing.T) { + var imageData bytes.Buffer + if err := png.Encode(&imageData, image.NewNRGBA(image.Rect(0, 0, 1, 1))); err != nil { + t.Fatal(err) + } + + m := &Model{ + VisionEncoder: &RadioVisionEncoder{}, + Projector: &VisionProjector{}, + VisionConfig: &VisionConfig{ + PatchSize: 1, + DownsampleFactor: 1, + MinNumPatches: 1, + MaxNumPatches: 1, + MaxModelLen: 16, + Std: [3]float32{1, 1, 1}, + }, + imageStartTokenID: 10, + imageTokenID: 11, + imageEndTokenID: 12, + } + + prepared, err := m.PrepareMedia([]model.Segment{ + {Tokens: []int32{1}}, + {Kind: "image", Data: imageData.Bytes()}, + {Tokens: []int32{2}}, + }) + if err != nil { + t.Fatal(err) + } + if len(prepared.Items) != 1 { + t.Fatalf("items = %d, want 1", len(prepared.Items)) + } + if got, want := prepared.Items[0].Range, [2]int{1, 4}; got != want { + t.Fatalf("range = %v, want %v", got, want) + } + if !prepared.Items[0].Causal { + t.Fatal("image expansion must be causal in the text stack") + } +} + +func TestPrepareMediaReportsDisabledVision(t *testing.T) { + m := &Model{visionErr: errors.New("unsupported RADIO version \"radio-v5\"")} + + prepared, err := m.PrepareMedia([]model.Segment{{Tokens: []int32{1, 2}}}) + if err != nil { + t.Fatal(err) + } + if got, want := prepared.Tokens, []int32{1, 2}; len(got) != len(want) || got[0] != want[0] || got[1] != want[1] { + t.Fatalf("text tokens = %v, want %v", got, want) + } + + _, err = m.PrepareMedia([]model.Segment{{Kind: "image", Data: []byte{1}}}) + if err == nil { + t.Fatal("PrepareMedia unexpectedly accepted image input") + } + if got, want := err.Error(), "nemotron_h vision is unavailable: unsupported RADIO version \"radio-v5\""; got != want { + t.Fatalf("PrepareMedia error = %q, want %q", got, want) + } +} + +// scatterMediaFixture builds a model and batch whose expansion covers +// positions 1..4: one item spliced at position 0, so imageStart sits at 0 and +// four feature rows follow. PatchSize and DownsampleFactor of 1 make +// visionTokenCount simply height*width. +func scatterMediaFixture() (*Model, *batch.Batch, *mlx.Array) { + m := &Model{ + Config: &Config{HiddenSize: 2}, + VisionConfig: &VisionConfig{PatchSize: 1, DownsampleFactor: 1}, + } + + features := mlx.FromValues([]float32{ + 1, 2, + 3, 4, + 5, 6, + 7, 8, + }, 4, 2) + + b := &batch.Batch{ + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{6}, + Media: []batch.MediaItem{{ + Seq: 0, + Pos: 0, + Features: features, + Opaque: preparedImage{height: 1, width: 4}, + }}, + } + + return m, b, mlx.Zeros(mlx.DTypeFloat32, 1, 6, 2) +} + +// The target forward's column 0 is the sequence position in SeqOffsets, so +// the feature rows land on the expansion's own positions 1..4. +func TestScatterMediaTargetForward(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + m, b, h := scatterMediaFixture() + got := m.scatterMedia(h, b, 0) + mlx.Eval(got) + + assertAllClose(t, "target-forward scatter", got.Floats(), []float32{ + 0, 0, // position 0: imageStart, untouched + 1, 2, // positions 1..4: feature rows + 3, 4, + 5, 6, + 7, 8, + 0, 0, // position 5: past the expansion + }, 1e-5) + }) +} + +// The MTP draft's slot S embeds the look-ahead token S+1, so its column 0 +// holds the token one past its offset and every feature row shifts down by +// one column. Getting this wrong embeds raw placeholder tokens into the +// draft, which degrades acceptance silently rather than failing outright. +func TestScatterMediaDraftForwardShiftsByOne(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + m, b, h := scatterMediaFixture() + got := m.scatterMedia(h, b, 1) + mlx.Eval(got) + + assertAllClose(t, "draft-forward scatter", got.Floats(), []float32{ + 1, 2, // shifted one column earlier than the target forward + 3, 4, + 5, 6, + 7, 8, + 0, 0, + 0, 0, + }, 1e-5) + }) +} + +// A forward whose query range misses the expansion writes nothing, so a +// decode step past the image leaves the hidden rows alone. +func TestScatterMediaSkipsNonOverlappingQuery(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + m, b, h := scatterMediaFixture() + b.SeqOffsets = []int32{8} + b.SeqQueryLens = []int32{1} + + got := m.scatterMedia(h, b, 0) + mlx.Eval(got) + + assertAllClose(t, "non-overlapping scatter", got.Floats(), []float32{ + 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, + }, 1e-5) + }) +} diff --git a/model/parsers/nemotron3nano.go b/model/parsers/nemotron3nano.go index ac6033f64..b4f206fb5 100644 --- a/model/parsers/nemotron3nano.go +++ b/model/parsers/nemotron3nano.go @@ -97,6 +97,12 @@ func (p *Nemotron3NanoParser) Add(s string, done bool) (content string, thinking p.buffer.Reset() p.buffer.WriteString(trimmed) } + if done { + thinking = p.buffer.String() + p.buffer.Reset() + p.maybeThinkingOpenAtBOL = false + return "", thinking, nil, nil + } return "", "", nil, nil } } @@ -132,6 +138,10 @@ func (p *Nemotron3NanoParser) Add(s string, done bool) (content string, thinking // No end marker - emit unambiguous thinking thinking = p.emitThinking(bufStr) + if done { + thinking += p.buffer.String() + p.buffer.Reset() + } return "", thinking, nil, nil } diff --git a/model/parsers/nemotron3nano_test.go b/model/parsers/nemotron3nano_test.go index 5232c0bcc..da621a648 100644 --- a/model/parsers/nemotron3nano_test.go +++ b/model/parsers/nemotron3nano_test.go @@ -276,6 +276,99 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) { } } +func TestNemotron3NanoParser_DoneFlushesBufferedText(t *testing.T) { + tests := []struct { + name string + thinkingEnabled bool + chunks []string + expectedContent string + expectedThinking string + }{ + { + name: "partial think close", + thinkingEnabled: true, + chunks: []string{"reasoning ", "", false) + if err != nil { + t.Fatal(err) + } + finalContent, finalThinking, finalCalls, err := p.Add("", true) + if err != nil { + t.Fatal(err) + } + if got, want := content+finalContent, ""; got != want { + t.Fatalf("content = %q, want %q", got, want) + } + if got, want := thinking+finalThinking, "reasoning"; got != want { + t.Fatalf("thinking = %q, want %q", got, want) + } + if got := append(calls, finalCalls...); len(got) != 0 { + t.Fatalf("calls = %v, want none", got) + } +} + func TestNemotron3NanoParser_HasToolSupport(t *testing.T) { p := &Nemotron3NanoParser{} if !p.HasToolSupport() { diff --git a/model/parsers/qwen38_test.go b/model/parsers/qwen38_test.go index 9870767d6..7feb3674f 100644 --- a/model/parsers/qwen38_test.go +++ b/model/parsers/qwen38_test.go @@ -113,7 +113,7 @@ func TestQwen38ParserMalformedAndControlLikeContent(t *testing.T) { { name: "truncated tool call", continuation: "PlanBefore", - wantContent: "Before", + wantContent: "Before", wantThinking: "Plan", }, } diff --git a/model/parsers/qwen3coder.go b/model/parsers/qwen3coder.go index 1a8e3cf12..c34f03de7 100644 --- a/model/parsers/qwen3coder.go +++ b/model/parsers/qwen3coder.go @@ -60,6 +60,18 @@ func (p *Qwen3CoderParser) Add(s string, done bool) (content string, thinking st p.acc.WriteString(s) events := p.parseEvents() + if done { + switch p.state { + case qwenParserState_LookingForToolStart: + if p.acc.Len() > 0 { + events = append(events, qwenEventContent{content: p.acc.String()}) + } + case qwenParserState_CollectingToolContent: + events = append(events, qwenEventContent{content: toolOpenTag + p.acc.String()}) + } + p.acc.Reset() + p.state = qwenParserState_LookingForToolStart + } var toolCalls []api.ToolCall var sb strings.Builder diff --git a/model/parsers/qwen3coder_test.go b/model/parsers/qwen3coder_test.go index 7142567ed..726f99a4b 100644 --- a/model/parsers/qwen3coder_test.go +++ b/model/parsers/qwen3coder_test.go @@ -1062,6 +1062,55 @@ func TestQwen3CoderParserToolCallIndexing(t *testing.T) { } } +func TestQwen3CoderParserDoneFlushesBufferedContent(t *testing.T) { + for _, tt := range []struct { + name string + chunks []string + want string + }{ + {name: "trailing whitespace", chunks: []string{"answer", " \n"}, want: "answer \n"}, + {name: "partial opener", chunks: []string{"answer ", ""}, want: ""}, + {name: "truncated tool call", chunks: []string{"", ""}, want: ""}, + } { + t.Run(tt.name, func(t *testing.T) { + parser := Qwen3CoderParser{} + parser.Init(nil, nil, nil) + + var got string + for _, chunk := range tt.chunks { + content, _, calls, err := parser.Add(chunk, false) + if err != nil { + t.Fatal(err) + } + if len(calls) != 0 { + t.Fatalf("calls = %v, want none", calls) + } + got += content + } + content, _, calls, err := parser.Add("", true) + if err != nil { + t.Fatal(err) + } + got += content + if len(calls) != 0 { + t.Fatalf("calls = %v, want none", calls) + } + if got != tt.want { + t.Fatalf("content = %q, want %q", got, tt.want) + } + + content, _, calls, err = parser.Add("next", true) + if err != nil { + t.Fatal(err) + } + if content != "next" || len(calls) != 0 { + t.Fatalf("parser was not reset after done: content = %q, calls = %v", content, calls) + } + }) + } +} + func TestQwen3CoderParserToolCallIndexingStreaming(t *testing.T) { parser := Qwen3CoderParser{} parser.Init(nil, nil, nil) diff --git a/server/images.go b/server/images.go index 41376ae0d..6415c1c0b 100644 --- a/server/images.go +++ b/server/images.go @@ -458,26 +458,15 @@ func (m *Model) filterUnsupportedCapabilities(capabilities []model.Capability, m return c == model.CapabilityAudio }) } - if suppressVisionCapability(m) { - capabilities = slices.DeleteFunc(capabilities, func(c model.Capability) bool { - return c == model.CapabilityVision - }) - } return capabilities } -func suppressVisionCapability(m *Model) bool { - // 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 m.Config.ModelFormat == "safetensors" && m.Config.Renderer == "glimmer" { return true } - if isNemotron3NanoSafetensors(m) { + if isNemotronSafetensors(m) { return true } @@ -491,14 +480,16 @@ func suppressAudioCapability(m *Model, arch string) bool { return false } -func isNemotron3NanoSafetensors(m *Model) bool { - return isNemotron3NanoSafetensorsConfig(m.Config) +func isNemotronSafetensors(m *Model) bool { + return isNemotronSafetensorsConfig(m.Config) } -func isNemotron3NanoSafetensorsConfig(cfg model.ConfigV2) bool { +func isNemotronSafetensorsConfig(cfg model.ConfigV2) bool { return cfg.ModelFormat == "safetensors" && (cfg.Parser == "nemotron-3-nano" || cfg.Renderer == "nemotron-3-nano" || + cfg.Parser == "nemotron-3.5-nano" || + cfg.Renderer == "nemotron-3.5-nano" || cfg.ModelFamily == "nemotron_h_omni" || slices.Contains(cfg.ModelFamilies, "nemotron_h_omni")) } diff --git a/server/images_test.go b/server/images_test.go index 6f24daee9..7852f1cac 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -598,7 +598,7 @@ func TestModelCapabilities(t *testing.T) { expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision}, }, { - name: "nemotron3 safetensors suppresses vision and audio but keeps thinking", + name: "nemotron3 safetensors exposes vision and suppresses audio", model: Model{ Config: model.ConfigV2{ ModelFormat: "safetensors", @@ -608,7 +608,20 @@ func TestModelCapabilities(t *testing.T) { }, Template: chatTemplate, }, - expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityThinking}, + expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityTools, model.CapabilityThinking}, + }, + { + name: "nemotron3.5 safetensors exposes vision and suppresses audio", + model: Model{ + Config: model.ConfigV2{ + ModelFormat: "safetensors", + Parser: "nemotron-3.5-nano", + Renderer: "nemotron-3.5-nano", + Capabilities: []string{"completion", "vision", "audio"}, + }, + Template: chatTemplate, + }, + expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityTools, model.CapabilityThinking}, }, } diff --git a/server/model_list_test.go b/server/model_list_test.go index 44fd83036..5bbcdbd97 100644 --- a/server/model_list_test.go +++ b/server/model_list_test.go @@ -93,7 +93,7 @@ func TestListModelsFollowsManifestChanges(t *testing.T) { } } -func TestCapabilitiesSuppressNemotronSafetensorsMedia(t *testing.T) { +func TestCapabilitiesExposeNemotronSafetensorsVision(t *testing.T) { caps := []model.Capability{ model.CapabilityCompletion, model.CapabilityTools, @@ -111,15 +111,14 @@ func TestCapabilitiesSuppressNemotronSafetensorsMedia(t *testing.T) { model.CapabilityCompletion, model.CapabilityTools, model.CapabilityThinking, + model.CapabilityVision, } { if !slices.Contains(got, capability) { t.Errorf("capabilities = %v, want %s", got, capability) } } - for _, capability := range []model.Capability{model.CapabilityVision, model.CapabilityAudio} { - if slices.Contains(got, capability) { - t.Errorf("capabilities = %v, did not expect %s", got, capability) - } + if slices.Contains(got, model.CapabilityAudio) { + t.Errorf("capabilities = %v, did not expect audio", got) } }