Named image attachments in MLXLMCommon and the Foundation Models adapter - #535
thechriswebb wants to merge 3 commits into
Conversation
|
This seems pretty neat -- I wonder if we need to augment the |
Yes, I think this would be great. I am definitely interested for my app. |
67a8a7e to
9658a77
Compare
| /// `ImageReference.attachmentLabel` is an unconstrained `String`, so nothing | ||
| /// requires a generated label to be one of these; this is just the set an | ||
| /// app can look an image up by. | ||
| static func attachmentLabels(in entries: some Collection<Transcript.Entry>) -> [String] { |
There was a problem hiding this comment.
This is only called from tests -- is this just introspection for that purpose? It isn't public so not callable outside the repo.
| /// Qwen3-VL is deliberately not in this list. `mlx-community/Qwen3-VL-4B-Instruct-4bit` | ||
| /// cannot load: its `model.safetensors.index.json` names | ||
| /// `model-00001-of-00002.safetensors` and `model-00002-of-00002.safetensors`, | ||
| /// neither of which the repo ships, and declares about 8.9 GB while the repo | ||
| /// contains a single 3.1 GB `model.safetensors`. The 8.9 GB figure is the | ||
| /// unquantized size, so the index was carried over from the source repo and | ||
| /// never regenerated for the quantized upload. The same is true of | ||
| /// `Qwen3-VL-8B-Instruct-4bit` (names four shards, ships two) and | ||
| /// `Qwen3-VL-4B-Instruct-8bit` (names two, ships one), so it is the | ||
| /// quantization batch rather than one repo. `Qwen3-VL-2B-Instruct-4bit` is | ||
| /// packaged correctly if a Qwen3-VL model is ever wanted here. |
| /// label, because the brackets are part of what gets tokenized: Mistral-family | ||
| /// image tokens are bracketed, so a label of `IMG` renders as `[IMG]`. | ||
| @available(iOS 27.0, macOS 27.0, visionOS 27.0, *) | ||
| struct AttachmentLabelRenderer { |
There was a problem hiding this comment.
This seems like a workaround for the fact that the attachments don't carry a label. I think that is ok for this PR, but I think UserInput/LMInput probably need to carry the label so it can get proper handling in the transcript.
There was a problem hiding this comment.
The downside of delaying this is that the behavior from LLMs may change when this is fixed -- the input pattern would be different. Hopefully the behavior change would be for the best, but who knows.
There was a problem hiding this comment.
How does this work in the template? Here is mlx-community/Qwen2.5-VL-7B-Instruct-4bit:
{% set image_count = namespace(value=0) %}
{% set video_count = namespace(value=0) %}
{% for message in messages %}
{% if loop.first and message['role'] != 'system' %}
<|im_start|>system
You are a helpful assistant.<|im_end|>
{% endif %}
<|im_start|>{{ message['role'] }}
{% if message['content'] is string %}
{{ message['content'] }}<|im_end|>
{% else %}
{% for content in message['content'] %}
{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}
{% set image_count.value = image_count.value + 1 %}
{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}
<|vision_start|><|image_pad|><|vision_end|>
{% elif content['type'] == 'video' or 'video' in content %}
{% set video_count.value = video_count.value + 1 %}
{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}
<|vision_start|><|video_pad|><|vision_end|>
{% elif 'text' in content %}
{{ content['text'] }}
{% endif %}
{% endfor %}
<|im_end|>
{% endif %}
{% endfor %}
{% if add_generation_prompt %}
<|im_start|>assistant
{% endif %}
It looks like this is labeling it Picture 1.
mlx-community/gemma-4-e4b-it-4bit is much longer:
https://huggingface.co/mlx-community/gemma-4-e4b-it-4bit/blob/main/chat_template.jinja
the pertinent part looks like this:
{%- elif message['content'] is sequence -%}
{%- for item in message['content'] -%}
{%- if item['type'] == 'text' -%}
{%- if role == 'model' -%}
{{- strip_thinking(item['text']) -}}
{%- else -%}
{{- item['text'] | trim -}}
{%- endif -%}
{%- elif item['type'] == 'image' -%}
{{- '<|image|>' -}}
{%- set ns.prev_message_type = 'image' -%}
{%- elif item['type'] == 'audio' -%}
{{- '<|audio|>' -}}
{%- set ns.prev_message_type = 'audio' -%}
{%- elif item['type'] == 'video' -%}
{{- '<|video|>' -}}
{%- set ns.prev_message_type = 'video' -%}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
That doesn't seem to mark it with anything.
…re markers - Names were thrown away before the prompt was built; a name now goes directly before its own picture, matching Foundation Models, across eleven model families - In MLXLMCommon: an image takes an optional name, and one shared function assembles a message's parts, so this works with no adapter - Name refused when its bracketed form is a reserved picture token, such as [IMG] on Mistral - Two smaller parity fixes: instructions images dropped and logged, and photo orientation applied - Source break: case .url(let u) = image becomes case .url(let u) = image.source
9658a77 to
3c03f0a
Compare
thechriswebb
changed the title
- The name check asked the tokenizer whether a name held a reserved picture token, but a tokenizer flags only some of its reserved tokens as special; on GLM-OCR and FastVLM the picture token itself is not flagged, so such a name was accepted and the model then counted one more picture than it was given - A name that holds any of < > | [ ] is now refused on every model, which needs nothing from the tokenizer and so covers a marker the tokenizer cannot see at all - The tokenizer check still runs, and runs first, so the error can name the exact token; it remains the only check that catches a plain name such as IMG whose bracketed form is reserved - Names accepted before and refused now: chart [q3] (final), a <b> tag & an entity, <|, |> and <image> - New test drives the real tokenizer of all ten families and refuses each family's own picture token; without this change it fails on GLM-OCR and FastVLM
- Only the Foundation Models path checked a picture name, so a caller using MLXLMCommon directly could put a reserved picture token into the prompt; on Mistral3 that doubled the picture tokens for one picture and the process then aborted instead of raising an error - A name that holds any of < > | [ ] is now left out of the prompt wherever it is assembled, and the reason is logged; the picture itself is always kept, so the count the model checks still matches - The Foundation Models path is unchanged and still raises a typed error, because its check runs before the prompt is assembled - One list of characters now serves both checks, so the two cannot drift apart
davidkoski
left a comment
davidkoski
left a comment
There was a problem hiding this comment.
Changes look good. One issue, one potential issue:
-
os.Logging won't build on linux -- it may be time for us to adopt swift-log
- look at MLXFoundationModels too
-
we can inject image markers into the stream, e.g. Mistral3, see repro case attached below
| } else if let label = image.label { | ||
| appendText("[\(label)]") |
There was a problem hiding this comment.
I wonder if this needs to be configurable? Mistral 3's special image token is [IMG]. See labelRendersAnImagePlaceholderAsText and oneImageYieldsTwoImagePlaceholders
Have to find the right place to attach it -- it seems like it belongs with the model, though perhaps it is technically part of the template/tokens. An override in ModelConfiguration?
| import os | ||
|
|
||
| private let messageContentLogger = Logger( | ||
| subsystem: "mlx-swift-lm", category: "MessageContent") |
There was a problem hiding this comment.
This isn't platform neutral -- it won't build on linux. I see that there are some uses in the MLXFoundationModels code already and I wonder if we should change those?
Anyway, we don't have a logging solution picked but perhaps we should. I think apple/swift-log is the standard, so probably that.
| messageContentLogger.warning( | ||
| "Leaving an image name out of the prompt: it holds `\(marker, privacy: .public)`, which vision models build their image placeholders from" | ||
| ) |
There was a problem hiding this comment.
See other comment on os.Logging. swift-log doesn't support privacy markers. The standard approach is a constant string with metadata:
messageContentLogger.warning(
"Leaving an image name out of the prompt",
metadata: [
"marker": .string(marker)
]
)or interpolation :-)
There was a problem hiding this comment.
If possible please find a solution that does not require adding dependencies that can potentially increase the size of the executable. Thanks.
There was a problem hiding this comment.
Yeah, that is a good point. We have pretty minimal logging needs. Maybe something very thin in mlx-swift would do the trick.
I have something sketched out and will put up a PR in the morning.
There was a problem hiding this comment.
3 participants
The problem
Foundation Models lets an app attach labeled images to a prompt, ask about one by name, and get an answer that points back at a specific picture. None of that worked through this adapter. Labels were thrown away, so the model never saw a name and could not return one. Orientation was thrown away, so a photo from a camera arrived sideways. An image attached to the instructions was sent on the system message, which the vision models mishandle, so those requests failed or quietly misbehaved. Attachments the adapter could not handle were ignored.
What this PR does
It deliberately does not restrict what the model writes as a label. The framework's own lookup already recovers from an imprecise one, and echoing its bracket form is what lets it.
Tests
Unit tests cover each change. Integration tests run against two real vision models. One asks which of two named pictures is the blue one, so a model that simply repeats the first name it saw fails. The other does the full round trip: the model names a picture and the framework looks it back up. That passed on both models in five runs, every time on an exact name match.