epic: add Florence-2 (florence2) VLM support
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
Florence-2 (`microsoft/Florence-2-large` / `-base`) is a unified multi-task vision seq2seq model. IMPORTANT: it is architecturally new for mlxcel. A DaViT vision backbone (`model_type: davit`) produces image tokens that are concatenated with the task-prompt tokens and fed to a BART-style encoder-decoder transformer (a `Florence2Encoder` plus a `Florence2Decoder` with cross-attention from decoder to encoder outputs), NOT a decoder-only stack. One text-generation interface covers captioning, OCR, detection, dense region captioning, grounding, and segmentation; task-prompt tokens select the mode, and spatial tasks emit location tokens that post-process into boxes / regions. This is the highest-risk item in the batch and should be scoped as its own epic.
Concretizes the Florence-2 entry recorded as out-of-scope-pending-seq2seq groundwork in #524.
This is the tracking epic for Florence-2 support. This repository has no `type:epic` label, so the issue keeps `type:enhancement` while serving as the epic. GitHub native sub-issue linking is not used here; the sub-issues below are tracked by the checklist in this issue.
## Sub-issues (execution order)
### Phase 1 (parallel foundations)
- [ ] #852 BART-style seq2seq encoder-decoder engine + text core
- [ ] #853 DaViT vision backbone
### Phase 2 (fusion)
- [ ] #854 Vision-language fusion + full weight loading (needs #852, #853)
### Phase 3 (modes + capstone)
- [ ] #855 Processor + task prompts + location tokens (needs #854)
- [ ] #856 End-to-end integration + real-checkpoint validation (needs #852, #853, #854, #855)
## Upstream reference
- Directory: https://github.com/Blaizzy/mlx-vlm/tree/main/mlx_vlm/models/florence2
- Top-level model (encoder-decoder wiring): https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/florence2/florence2.py
- Language (BART-style encoder + decoder): https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/florence2/language.py
- Vision (DaViT backbone): https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/florence2/vision.py
- Processor (task prompts): https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/florence2/processing_florence2.py
- Config: https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/florence2/config.py
## Public checkpoint
Confirmed available on HuggingFace:
- `microsoft/Florence-2-large`, `microsoft/Florence-2-base` (originals)
- mlx-community conversions confirmed: `mlx-community/Florence-2-large-ft-4bit` / `-8bit` / `-6bit` / `-bf16` and `mlx-community/Florence-2-base-ft-4bit` / `-8bit` / `-6bit` / `-3bit` / `-bf16` (primary validation target: a `-bf16` or `-4bit` ft variant).
## Architecture notes
- `model_type: florence2` with a `davit` vision sub-config.
- Vision: DaViT encoder with windowed attention (window_partition / window_reverse), ChannelAttention, ConvEmbed, DepthWiseConv2d, learned absolute positions.
- Language: BART-style encoder-decoder (6 encoder + 6 decoder layers, learned absolute positions). Image tokens are concatenated with prompt tokens, the encoder runs over the joint sequence, and the decoder cross-attends to encoder outputs while decoding autoregressively.
- Task-prompt tokens drive the mode (caption / OCR / detect / ground / segment); spatial tasks emit location tokens post-processed into boxes or regions.
## Implementation plan
- PROMINENT CONSTRAINT: mlxcel's generation engine is decoder-only today. Florence-2 requires a NEW encoder-decoder inference path: an encoder pass over the concatenated vision-plus-prompt sequence, then cross-attention decoding against the cached encoder outputs. The only in-tree encoder-decoder precedent is Whisper's dedicated ASR pipeline (`src/models/whisper/`); use it as the structural template for a seq2seq decode loop and encoder-output caching, but note Florence-2 needs vision-plus-text encoder input and location-token output rather than audio-to-text.
- DaViT backbone: new, but reuse existing window-attention helpers where possible (the windowed-attention patterns in the qwen2_5_vl / SAM-style encoders under `src/vision/encoders/`); ConvEmbed / DepthWiseConv2d map to the existing conv layer helpers used by the FastViT / gemma3n conv paths.
- Task-prompt handling + location-token post-processing (coordinate tokens to boxes / regions) is a dedicated processing stage; the coordinate / OCR post-processing in the OCR VLM modules is a partial reference.
- RECOMMENDED SCOPING: file this as an epic. Phase 1: seq2seq engine groundwork (encoder-decoder decode loop, cross-attention, encoder-output cache) validated on a minimal path. Phase 2: DaViT backbone. Phase 3: task-prompt modes + location-token post-processing + full multi-task validation. Land Phase 1 as reusable engine infrastructure before the model-specific work.
## Touchpoints & acceptance criteria
- [ ] NEW encoder-decoder inference path (encoder pass, cross-attention decode, encoder-output cache) landed as reusable engine infrastructure (Whisper as template).
- [ ] Text config (BART encoder-decoder) + vision config (`davit`) parse.
- [ ] DaViT vision encoder (windowed attention, window_partition / window_reverse, ChannelAttention, ConvEmbed, DepthWiseConv2d).
- [ ] Image-token concatenation with prompt tokens (Florence-2 concatenates vision + prompt for the encoder rather than scattering into placeholders).
- [ ] `from_weights` + `sanitize`: encoder-decoder weight naming, DaViT conv weights channels-last remap, mixed precision.
- [ ] Task-prompt mode handling + location-token to box / region post-processing.
- [ ] Arch-string arm `"florence2"` in `src/models/detection.rs`.
- [ ] Registration in `src/model_metadata.rs` (`for_each_model_registration!`).
- [ ] `generate_vlm` (or a dedicated seq2seq VLM path) summary + TP / distributed arch-string if applicable.
- [ ] `_tests.rs` unit tests (DaViT shapes, cross-attention, task-prompt parse).
- [ ] `docs/supported-models.md` updated.
- [ ] Real-checkpoint validation: `./target/release/mlxcel generate -m models/Florence-2-base-ft-bf16 --image -p ""` produces correct output for at least caption + OCR + detection modes. Follow `docs/adding-models.md`.
## Effort
HIGH, highest-risk item of the batch. The blocker is the seq2seq encoder-decoder engine path, which mlxcel does not have outside Whisper. Recommend an epic with the seq2seq groundwork landed first.
Contributor guide
Research direction
Start with the Phase 1 sub-issue, #852, and read the Whisper pipeline under src/models/whisper/ for the encoder-decoder structure and caching approach. Then review src/models/detection.rs, src/model_metadata.rs, and the listed _tests.rs touchpoints. Done requires reusable seq2seq groundwork before DaViT, fusion, task modes, and real-checkpoint validation can be completed.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- computer-vision, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100