mlxrunner: reject media requests the model cannot serve

MLX checkpoints that include a vision tower are already tagged with the
vision capability at import, so the server accepts image chats and ships
the image bytes with the completion request. The MLX client dropped the
bytes, and the prompt's image tags were answered as literal text.

Carry the media through to the runner and fail the request with a clear
error when the loaded model has no media support. Nothing implements the
new media interface yet, so every media request now returns the error
rather than a silently wrong answer; later changes build the image path
on top of the same interface.
This commit is contained in:
Jesse Gross 2026-08-05 17:44:34 -07:00
parent 8713570d3c
commit af5b627672
4 changed files with 132 additions and 0 deletions

View file

@ -112,6 +112,7 @@ func (c *Client) WaitUntilRunning(ctx context.Context) error {
type CompletionRequest struct {
Prompt string
Media []llm.MediaData
Options api.Options
Logprobs bool
TopLogprobs int
@ -155,6 +156,7 @@ func (c *Client) Close() error {
func (c *Client) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error {
creq := CompletionRequest{
Prompt: req.Prompt,
Media: req.Media,
Logprobs: req.Logprobs,
TopLogprobs: req.TopLogprobs,
}

View file

@ -0,0 +1,77 @@
package base
import (
// Every model's PrepareMedia decodes through image.Decode; the decoder
// set is registered once here so all models accept the same formats.
_ "image/gif"
_ "image/jpeg"
_ "image/png"
_ "golang.org/x/image/webp"
"github.com/ollama/ollama/x/mlxrunner/mlx"
)
// Segment is one run of the prompt in stream order: either a tokenized text
// run (Tokens set) or a single media item (Kind and Data set).
type Segment struct {
Tokens []int32
Kind string
Data []byte
}
// PreparedItem describes one media occurrence in the prepared stream. A
// model chooses item granularity: one per media segment, or several when
// parts encode and evaluate independently (e.g. per tile).
type PreparedItem struct {
// Range is the expansion's token range [start, end) in Tokens;
// non-empty, since cache identity enters through these positions.
Range [2]int
// Source is the index of the segment this item was prepared from; the
// item's prefix-cache identity is keyed on that segment's bytes.
Source int
// MediaData is the preprocessed encoder input with shape Dims. Dims
// enters the cache keys too: geometry changes features under
// identical bytes.
MediaData []float32
Dims []int
// Opaque carries model-private preprocessing state to EncodeMedia
// and Forward.
Opaque any
// Causal marks an expansion whose tokens attend causally, so chunks
// may split it. Unset, the first evaluation covers the whole
// expansion in one forward, as bidirectional runs require.
Causal bool
}
// PreparedRequest is the expanded input stream, every media segment's
// expansion spliced in place, with the items in stream order.
type PreparedRequest struct {
Tokens []int32
Items []PreparedItem
// Layout is an opaque request-scoped value computed in the one pass
// that sees every splice position; immutable, carried unread by the
// runner to every forward. Delivered only when Items is non-empty.
// Nil when the model derives nothing from it.
Layout any
}
// MediaModel is implemented by models that accept media inputs.
type MediaModel interface {
// PrepareMedia runs once per request on the request goroutine, CPU
// only, and returns the expanded stream. It must be deterministic for
// given segments: prefix-cache restores splice cached state with
// recomputed state.
PrepareMedia(segments []Segment) (*PreparedRequest, error)
// EncodeMedia builds one item's lazy feature graph on the MLX thread;
// it must not evaluate — the consuming forward's evaluation pulls it.
// Read the pixels from data: the runner frees the item's MediaData
// once its expansion is evaluated.
EncodeMedia(item *PreparedItem, data *mlx.Array) *mlx.Array
}

View file

@ -14,6 +14,7 @@ import (
"github.com/ollama/ollama/x/mlxrunner/batch"
"github.com/ollama/ollama/x/mlxrunner/cache"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/mlxrunner/model/base"
sampler "github.com/ollama/ollama/x/mlxrunner/sample"
"github.com/ollama/ollama/x/tokenizer"
)
@ -30,6 +31,16 @@ func (r *Runner) Prepare(request *Request) error {
return errors.New("model not loaded")
}
if len(request.Media) > 0 {
if _, ok := r.Model.(base.MediaModel); !ok {
kind := string(request.Media[0].Kind)
if kind == "" {
kind = "media"
}
return fmt.Errorf("this model does not support %s input", kind)
}
}
tokens := r.Tokenizer.Encode(request.Prompt, r.Tokenizer.AddBOS())
if len(tokens) == 0 {
return errors.New("empty prompt")

View file

@ -0,0 +1,42 @@
package mlxrunner
import (
"strings"
"testing"
"github.com/ollama/ollama/llm"
"github.com/ollama/ollama/x/mlxrunner/batch"
"github.com/ollama/ollama/x/mlxrunner/cache"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/tokenizer"
)
// textOnlyModel satisfies base.Model but not base.MediaModel.
type textOnlyModel struct{}
func (textOnlyModel) LoadWeights(map[string]*mlx.Array) error { return nil }
func (textOnlyModel) NewCaches() []cache.Cache { return nil }
func (textOnlyModel) Forward(*batch.Batch, []cache.Cache) (*mlx.Array, *mlx.Array) {
return nil, nil
}
func (textOnlyModel) Unembed(x *mlx.Array) *mlx.Array { return x }
func (textOnlyModel) Tokenizer() *tokenizer.Tokenizer { return nil }
func (textOnlyModel) MaxContextLength() int { return 0 }
func TestPrepareRejectsMediaWithoutSupport(t *testing.T) {
r := &Runner{Model: textOnlyModel{}}
req := &Request{
CompletionRequest: CompletionRequest{
Prompt: "[img-0] what is this?",
Media: []llm.MediaData{{ID: 0, Kind: llm.MediaKindImage, Data: []byte{1}}},
},
}
err := r.Prepare(req)
if err == nil {
t.Fatal("expected error for media on a text-only model")
}
if !strings.Contains(err.Error(), "does not support image input") {
t.Fatalf("unexpected error: %v", err)
}
}