lablup / lablup/mlxcel

epic: add Florence-2 (florence2) VLM support

Open
#850 2 comments 0 reactions 0 assignees View on GitHub
area:inference area:models priority:low status:in-progress type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.