feat(models): Florence-2 processor, task prompts, and location tokens

[inureyes]

Summary

Completes the Florence-2 request path on top of the fused model from #854. A task marker and an image go in; a structured result with coordinates in original-image pixels comes out.

Florence-2 selects its behavior with a marker such as <OD> or <CAPTION>, and none of those markers is in the vocabulary. The processor expands each one into the English sentence the model was actually trained on (<OD> becomes "Locate the objects with category name in the image."), tokenizes that, preprocesses the image to 768x768, generates, and parses the answer's <loc_N> tokens back into boxes, quadrilaterals, and polygons against the image's pre-resize extent.

The integration entry point is Florence2Processor::run(&model, task, input, image, max_new_tokens) -> Florence2Output, which is what #856 will reach from the CLI.

What changed

New, under src/models/florence2/

• tasks.rs: the fifteen task markers, their prompt expansions copied byte for byte from the checkpoint, and routing to a post-processing kind. Seven tasks interpolate a caller-supplied input.
• coords.rs: Florence2ImageSize, Florence2BoundingBox, Florence2QuadBox, Florence2Polygon, the bin-to-pixel dequantizers, and the inverse so a caller can build the <loc_a><loc_b><loc_c><loc_d> region string the region-input tasks expect.
• scan.rs: literal string scanners (sequence-marker stripping, leading-phrase extraction, <loc_N> reading) plus the phrase-grounding stopword blacklist.
• parse.rs: the four reachable parsers over compiled-once patterns.
• postprocess.rs: Florence2TaskResult and the task dispatch. Florence2TaskResult::parse is model-free, so an answer obtained elsewhere parses without loading a checkpoint.
• processor.rs: Florence2Processor, tokenizer plus image processor plus the end-to-end run.

New elsewhere

• src/vision/processors/florence2.rs: Florence2ImageProcessor, reading preprocessor_config.json rather than hard-coding it, and refusing configurations it would silently mishandle (center cropping, a non-bicubic resample, a zero standard deviation, a zero size).
• tests/florence2_postprocess_parity.rs, tests/florence2_processor_parity.rs.
• Unit tests beside each new module.

Moved

• Florence2TextConfig from src/models/florence2/mod.rs to src/models/florence2/checkpoint.rs, next to the Florence2Config that contains it. Pure relocation, no behavior change; it keeps mod.rs under the 500-line limit now that the module list has grown.

Design notes

Florence2BoundingBox is deliberately not a reuse of rt_detr_v2::Detection. That type carries a confidence score and a label: usize index into a fixed class list; Florence-2 emits no per-box score and its labels are open-vocabulary strings decoded from the answer, so reusing it would mean two fabricated fields on every box. There is no polygon or quadrilateral type anywhere in the tree to reuse.

Florence2ImageSize has named fields rather than being a (u32, u32) tuple because upstream's own docstring disagrees with its code about the order (see corrections below), and passing the 768x768 resized extent instead of the original silently scales every coordinate.

Two upstream regexes were replaced with hand-rolled scanners, each with its equivalence argument recorded in the code. The phrase-string pattern ^\s*(.*?)(?=<od>|...|<loc_) becomes a forward scan for the first stop marker; the two behaviors that make it non-trivial (. not matching a newline, and <sep> not being in the stop list) are reproduced and tested. <loc_(\d+)> becomes a digit scan, since it is a fixed literal on the innermost path of every parser. Everything else uses fancy-regex, which matches Python's backtracking semantics.

Corrections to the issue text

• The issue suggests using the dots_ocr and paddleocr modules as a partial reference for location-token to box conversion. They contain no coordinate parsing at all. Both are plain text-emitting VLMs. The only box type in the tree is rt_detr_v2::Detection, and there is no polygon type. Nothing was reusable.
• <REGION_PROPOSAL> routes to the bboxes post-processing kind, not description_with_bboxes. Same parser, but with allow_empty_phrase set, which matters: its answers carry no category names.

Corrections to upstream

Verified against the checkpoint's own processing_florence2.py and the mlx-vlm sources.

• post_process_generation's docstring is wrong about coordinate order. It documents image_size as "height x width", but every call site unpacks it as image_width, image_height = image_size. It is (width, height). This is why the Rust type uses named fields.
• The od parse task's configured PATTERN can never match. It is r'([a-zA-Z0-9 ]+)<loc_(\\d+)>...', and the doubled backslash inside a raw string matches a literal \ followed by d. It is unreachable in practice: no task routes to 'od', and every parser except the OCR one overwrites its pattern argument on its first lines, so only the OCR pattern is ever read from config.
• The OCR area filter never runs. It is gated on area_threshold > 0 and the shipped AREA_THRESHOLD is 0.00.
• with_box_at_start is false at every call site reachable from post_process_generation, so that polygon branch is dead.
• The <ground> / <obj> prefix strip is a no-op twice over. The second assignment reads pharse_text rather than the result of the first, discarding the <ground> strip, and neither marker is in the checkpoint's vocabulary.
• A newline in a phrase silently deletes the whole chunk. ^\s*(.*?)(?=...) is not multi-line and . does not match a newline, so if a newline sits between the leading whitespace and the first location token, re.search fails and the chunk is dropped. Reproduced deliberately and pinned by od_newline.
• [^<]+ can start one byte inside a <loc_N> token. A bare location run fed to the labelled parser therefore yields the literal label loc_1> and a box shifted by one bin, rather than no result. Reproduced, pinned by od_bare_locs, and the reason <REGION_PROPOSAL> needs the empty-phrase parser.
• mlx-vlm's florence2 model applies stochastic depth at inference unless you call model.eval(). mlx's nn.Module starts in training mode and upstream's DaViT DropPath only short-circuits on not self.training. Without it the vision tower output is not reproducible. Recorded in the parity test's module docs so the next person does not lose an hour to it.
• AutoTokenizer.from_pretrained fails on this checkpoint, because it routes through AutoConfig, which does not know the custom florence2_language model type. BartTokenizerFast directly works.
• Re-confirmed from feat(models): Florence-2 vision-language fusion + full weight loading (sub of #850) #854: upstream's ModelConfig reads image_feature_source, image_pos_embed and visual_temporal_embedding from the top level of config.json, but real checkpoints store them under vision_config, and the dataclass default for image_feature_source is the reverse order.

Deliberately not implemented

Each is unreachable from post_process_generation, and each is documented in parse.rs with the reason rather than silently omitted: parse_od_from_text_and_spans and the PATTERN config indirection, the OCR area filter, the with_box_at_start polygon branch, and the <ground> / <obj> prefix strip. decode_with_spans is also not ported; its character spans exist only for score attribution, which no reachable task uses.

One intentional deviation: Florence2Task::expand rejects an input supplied to a task that takes none, and a missing input on a task that requires one. Upstream produces "What is the region ?" in the second case and lets the model answer nonsense, which is indistinguishable from a real prompt after the fact.

One knowingly accepted difference: $ in the polygon-run pattern matches only at end of string, while Python's also matches just before a trailing newline. Location runs are not emitted with trailing newlines.

Validation

Against the real checkpoint (models/Florence-2-base-ft-bf16) and the real upstream Python, in a scratch virtualenv with torch, transformers, mlx, numpy and Pillow. The post-processing reference imports the checkpoint's own processing_florence2.py unmodified, so it is upstream behavior rather than a reimplementation agreeing with itself.

• Post-processing: 32 cases across all five result shapes, all matching upstream exactly, including every edge case listed above.
• Preprocessing: tests/fixtures/test_image.png through the checkpoint's CLIPImageProcessor versus the Rust processor. Maximum absolute deviation exactly 0.0 across all 1,769,472 values; the image crate's CatmullRom and PIL's BICUBIC agree bit for bit on this 224 to 768 upscale. The test tolerance is kept loose anyway, since that agreement is a property of two independent resampling implementations.
• Tokenization: exact id match on all six tested task prompts.
• End to end: upstream mlx-vlm florence2 with weights cast bf16 to f16. Byte-identical answers on all six tasks, including <OD> producing poster<loc_0><loc_0><loc_998><loc_998> and <OCR_WITH_REGION> producing a full eight-token quadrilateral, and identical parsed coordinates for both.

No unexplained Rust-versus-Python divergence was found anywhere in this work.

Test plan

• cargo check --profile test-fast --features metal,accelerate --lib --tests
• cargo clippy --profile test-fast --features metal,accelerate --lib --tests -- -D warnings (zero warnings)
• cargo test --profile test-fast --features metal,accelerate -p mlxcel --lib florence2
• cargo test --profile test-fast --features metal,accelerate -p mlxcel --test florence2_postprocess_parity
• cargo test --profile test-fast --features metal,accelerate -p mlxcel --test florence2_processor_parity
• cargo test --profile test-fast --features metal,accelerate -p mlxcel --test florence2_fusion_parity (unchanged by the Florence2TextConfig move)
• cargo fmt --check
• python3 scripts/ci/check_cross_repo_refs.py

Closes #855

[inureyes]

Implementation Review Summary

Intent

Port the Florence-2 processor layer for #855: the fifteen task markers and their prompt expansions, 768x768 image preprocessing, tokenizer wiring, and <loc_0>..<loc_999> post-processing into boxes, quad-boxes and polygons, with Florence2Processor::run as the end-to-end entry point the capstone #856 will reach from the CLI.

Verification performed

I checked the port line by line against the checkpoint's own processing_florence2.py (the authoritative upstream for this checkpoint, published at https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/florence2/processing_florence2.py) rather than against the PR description.

• Prompt expansions: all fifteen strings byte-identical to task_prompts_without_inputs / task_prompts_with_input, including the comma in "What is the text in the image, with regions?" and the missing trailing period on "Locate {input} in the image with mask" and "What is the polygon mask of region {input}".
• Post-processing routing: matches tasks_answer_post_processing_type for all fifteen entries, including <REGION_PROPOSAL> to bboxes (not description_with_bboxes) and the allow_empty_phrase=True that the polygons parse task carries.
• Dequantization: the + 0.5 bin-center term is present in both dequantize_box and dequantize_coordinates; 1000 bins per axis; x scaled by width and y by height; the f64-quotient-narrowed-to-f32 size_per_bin reproduces torch's promotion of an int64 tensor plus a Python float. quantize reproduces floor mode with the same clamp.
• Stopword blacklist: diffed programmatically against _create_black_list_of_phrase_grounding (lines 552-578). Upstream's 100 literals with 28 repeats reduce to 72 unique entries; the Rust list is exactly those 72, sorted for binary search, with the sortedness and the count both pinned by blacklist_is_sorted_for_binary_search. Nothing truncated.
• Parsers: pure_text strips only <s>/</s> while the box parsers also strip <pad> (upstream's asymmetry, reproduced); the OCR pattern is (.+?) plus exactly eight <loc_>; the chunkers, the phrase scan's stop list, the ASCII drop-not-replace, the one-instance-per-4-loc-group fan-out, the phrase-grounding grouping and blacklist ordering, the <poly>/</poly> split gated on both delimiters being present, the ((?:<loc_\d+>)+)(?:<sep>|$) run pattern, and the odd-length tail drop all match. The two hand-rolled scanners that replace upstream regexes carry a correct equivalence argument: I re-derived the ^\s*(.*?)(?=...) backtracking behaviour (greedy \s* can consume a leading newline, . cannot cross one afterwards, <sep> is absent from the stop list) and leading_phrase reproduces it in every case I could construct.
• The <ground>/<obj> strip is not merely inert for the stated reason: the chunk pattern is [^<]+ followed only by <loc_\d+> runs, so neither marker can appear anywhere inside a chunk. The omission is safe by construction, which is stronger than the justification in parse.rs claims.
• Preprocessing: resize_exact to 768x768 with no aspect preservation, rescale 1/255, ImageNet mean/std (not CLIP, not SigLIP), NCHW [B,3,768,768], and the config read from preprocessor_config.json with hard errors on center-crop, non-bicubic resample, zero std and zero size rather than silent mishandling.
• Tokenizer: <loc_N> id equals 50269 + N contiguously through <loc_999> = 51268, verified directly against added_tokens.json (zero mismatches over all 1000). encode(prompt, add_special_tokens = true) and decode(ids, skip_special_tokens = false) are the correct polarity, and the existing MlxcelTokenizer / load_tokenizer are reused rather than reimplemented.
• Original-size threading: Florence2Processor::run takes (width, height) from preprocess_with_sizes' pre-resize original_sizes and passes it to post-processing, never the 768x768 extent. preprocess_with_sizes pushes (image.width(), image.height()) in that order, pinned by reports_the_original_size_before_resizing on a 640x480 and a 100x200 image. detection_and_ocr_parse_into_reference_coordinates additionally re-parses the same answer at a different declared size and asserts the coordinates move.
• The Florence2TextConfig relocation is byte-for-byte identical apart from the anyhow! import and dropping the now-unneeded super::Florence2TextConfig use. mod.rs re-exports it and src/models/mod.rs is unchanged for that name, so existing consumers are unaffected. mod.rs lands at 350 lines and every new file is under 500.
• Error handling on the public surface: a with-input task called with None and a no-input task called with Some are both rejected at Florence2Task::expand (a deliberate, documented divergence from upstream's "What is the region ?"); an unknown task string errors from FromStr; a truncated loc run is discarded by the 4-at-a-time grouping exactly as the regex would; an empty generation yields empty results with no panic; a non-RGB image goes through to_rgb8; a zero-size image is guarded in quantize rather than dividing by zero. No unwrap(), expect(), TODO, placeholder or mock anywhere in the new library code.
• Scope: src/models/detection.rs, src/model_metadata.rs, docs/supported-models.md, mlxcel list and the generate wiring are untouched, as feat(models): Florence-2 end-to-end integration and real-checkpoint validation (sub of #850) #856 requires. No behavioral change reaches the already-merged feat(models): Florence-2 BART-style seq2seq encoder-decoder engine + text core (sub of #850) #852/feat(models): Florence-2 DaViT vision backbone (sub of #850) #853/feat(models): Florence-2 vision-language fusion + full weight loading (sub of #850) #854 code.

Findings addressed

• src/models/florence2/scan.rs:171 and src/models/florence2/parse.rs:161 (LOW) location_bins documented that a <loc_N> digit run too wide for i32 is "accepted here rather than rejected, matching upstream's plain int()", but parse::<i32>() failed and the bin was silently dropped while the scan advanced past the token. Dropping one bin re-pairs every bin after it in the same chunk, so the four-at-a-time grouping and the eight-token OCR record would misassign coordinates rather than produce one wrong number; parse_ocr separately discarded the whole record where Python keeps it. Both now route through a shared scan::parse_bin that saturates, so the token stays counted and the grouping stays aligned. Unreachable on real output (the vocabulary carries only <loc_0>..<loc_999>), and the parity suite is unchanged. Fixed in c761645 with a regression test.

Remaining items (reported, not fixed)

• src/models/florence2/checkpoint.rs:110-124 (LOW) The relocated doc comments still carry [encoder::Florence2Encoder::from_weights], [decoder::Florence2Decoder::from_weights] and [layers::additive_causal_mask], which resolved from mod.rs but not from checkpoint.rs. They sit on private items so default cargo doc never tries to resolve them and no CI job checks intra-doc links; left alone to keep the relocation byte-identical. super::-qualifying them is the fix whenever that file is next touched.
• src/vision/processors/florence2.rs:54-63 (LOW) All Florence2ImageProcessor fields are pub, so the validation from_raw performs (non-zero size, non-zero std, bicubic-only) can be bypassed after construction. The unit tests rely on this to override height/width. Not a defect today; worth private fields plus accessors if the type ever escapes the crate.
• tests/florence2_processor_parity.rs:57 (LOW) The end-to-end fixture is square (224x224), so a width/height transposition inside run would not show up there. The ordering is covered indirectly by reports_the_original_size_before_resizing and by the non-square cases in florence2_postprocess_parity.rs, so this is a coverage gap rather than a risk; a non-square fixture would close it outright.
• src/vision/processors/florence2.rs:198-203 (LOW) An RGBA input is resized in RGBA and then flattened by to_rgb8, while PIL would convert to RGB first. Only observable on an image with a non-opaque alpha channel, which Florence-2 checkpoints are not exercised with.

Verification checklist

• All stated requirements implemented (all five acceptance criteria in feat(models): Florence-2 processor, task prompts, and location tokens (sub of #850) #855)
• No placeholder, mock or TODO code remaining
• Integrated into the project code flow (Florence2Processor::run threads task expansion, tokenize, preprocess, Florence2Model::generate_greedy, decode-with-specials and post-process against the original image size, and is exercised end to end against the real checkpoint)
• Project conventions followed (_tests.rs beside the implementation, no unwrap/expect in library code, inline format args, every file under 500 lines, upstream cited only by GitHub URL)
• Existing modules reused where applicable (MlxcelTokenizer, load_tokenizer, the ImageProcessor trait; the two cases where reuse was rejected, rt_detr_v2::Detection and a shared image processor, are argued correctly in the code)
• No unintended structural changes (the Florence2TextConfig move is a pure relocation with working re-exports)
• Tests pass: clippy zero-warning, cargo fmt --check clean, 132 --lib florence2 unit tests, florence2_postprocess_parity (6), florence2_processor_parity (4, against the real checkpoint)

[inureyes]

Security and performance review

Threat model applied: the parsed text is decoder output, which is attacker influenced in a served deployment because the image and the task input string are user supplied, so all parser input was treated as untrusted. Scope was the diff under src/models/florence2/, src/vision/processors/florence2*.rs, and the two gated integration tests. No CRITICAL or HIGH findings.

Fixed (a862e2e)

MEDIUM - unbounded allocation from preprocessor_config.json (src/vision/processors/florence2.rs:154). size is the only field in this file that becomes an allocation: preprocess_with_sizes sizes its host buffer as batch * 3 * height * width f32 before it reads a pixel. It was validated as non-zero but had no upper bound. A hostile checkpoint could make that a multi-gigabyte allocation and an out-of-memory abort rather than an error return, and a value past usize::MAX / 12 (for example 4294967296) wraps the product in a release build, leaving a short buffer for a write loop that still runs to the configured extent, so data[base + c * plane + y * width + x] panics. Added MAX_SIZE_EDGE = 8192, mirroring the MAX_LAYERS and MAX_POSITION_EMBEDDINGS bounds this PR already applies to text_config in src/models/florence2/checkpoint.rs. Real exports ship 768.

Regression guard - adversarial_answers_do_not_panic (src/models/florence2/florence2_parse_tests.rs). The scanners in scan.rs walk byte offsets rather than characters over text that can carry arbitrary Unicode, so the byte-offset slicing is the highest-risk area in the diff. The review found it correct, and this test is what keeps it correct: multi-byte characters against every boundary a scanner slices at (the leading-whitespace cut, the stop-marker lookahead, the <loc_ digit run, the bare loc_N> prefix strip), malformed location tokens with no digits or no terminator, nested <loc_ opens that exercise the scanner's advance rule, digit runs far wider than i32, unbalanced polygon and separator markers, and long repetition runs, each through all four parsers in both allow_empty_phrase modes and through post_process for all fifteen tasks.

Verified clean

Catastrophic backtracking: not reachable. fancy_regex::Regex::new (0.17.0, lib.rs::new_options) compiles to RegexImpl::Wrap, a plain regex-crate regex, whenever analyze() reports the tree is not hard. hard means lookaround, backreference, atomic group, possessive quantifier, or \K. None of the seven patterns here use any of those, so all seven run on the linear-time engine and the backtrack limit never applies. This covers the three patterns worth worrying about: (?:<loc_\d+>){4,}, ((?:<loc_\d+>)+)(?:<sep>|$), and the non-greedy (.+?) before eight location tokens in OCR_RECORD. The new adversarial test includes 4096-token runs and 2048 <loc_1><sep> repetitions; the whole 133-test florence2 suite finishes in 0.30s.

UTF-8 boundary slicing: safe at every site. leading_phrase takes its offsets from char_indices() and whitespace_end is the start of the trim_start() suffix, so both are character boundaries. location_bins only ever sets its cursor just past the ASCII <loc_ prefix or just past the ASCII >, and it advances by at least 5 bytes on the no-match branch, so it is boundary safe and cannot loop. strip_bare_loc_prefix slices past an ASCII digit run and an ASCII >. ascii_only filters over chars().

No panics in the parse path. No indexing, unwrap, or expect on attacker-controlled data. parse_bin saturates to i32::MAX instead of unwrapping. All coordinate arithmetic is f32, so no integer overflow; quantize returns NaN through clamp (which only panics on a NaN bound, not a NaN value) and the as i32 cast saturates. florence2_loc_token_id guards n < 1000 before 50269 + n. PHRASE_GROUNDING_BLACKLIST is genuinely sorted, so the binary search is sound, and blacklist_is_sorted_for_binary_search already pins it.

Complexity of the hand-rolled scanners is linear. location_bins advances monotonically, so the repeated text[index..].find calls do not overlap and total O(n). leading_phrase does at most 8 starts_with probes per character, all of which fail on the first byte inside a [^<]+ run, and it is called once per chunk over non-overlapping chunks. parse_polygons nests three iterations but each level partitions its parent, so the total stays linear in the answer length.

Resource use on the hot path. All seven regexes are LazyLock, compiled once. Florence2Processor holds one tokenizer and one image processor for its lifetime; nothing is reloaded per call. No large buffer is cloned in the request path.

Reported, not fixed (LOW)

• src/vision/processors/florence2.rs:213 - the pixel loop calls rgb.get_pixel(x, y) per pixel (about 590k bounds-checked calls) and does a divide per channel per pixel. Indexing rgb.as_raw() by row would remove the bounds checks, and hoisting 1.0 / std[c] would remove 1.8M divides. Left alone deliberately: hoisting the reciprocal changes f32 rounding, and this PR's numerical parity against the Python reference is established. Not worth reopening for what is a low single-digit percentage of a request dominated by the DaViT forward pass.
• src/models/florence2/scan.rs:121 - strip_sequence_markers chains three String::replace calls, so it copies the whole answer three times. One pass would do. Negligible at Florence-2 answer sizes.
• src/models/florence2/parse.rs:217 - parse_boxes clones cat_name once per box, so a chunk with a long phrase and many boxes is quadratic in the answer length (worst case about N^2 / 112 bytes). Bounded in practice: generate_greedy stops at max_position_embeddings, which is 1024 in the shipped config and capped at 65536 by this PR's own validation, putting the realistic worst case in the hundreds of kilobytes. Would need Arc<str> to fix, which is not worth the API churn at that bound.
• src/vision/processors/florence2.rs:201 - image.clone() on the do_resize: false path copies the whole decoded image. Unreachable with any shipped config.

For the follow-up wiring (#856)

preprocess_with_sizes accepts already-decoded DynamicImage values, which is the same boundary every other processor under src/vision/processors/ uses, so decompression-bomb defense is the caller's responsibility. When the generate and server paths are wired up, user-supplied image bytes must go through decode_request_image_with_limits / ImageInputLimits in src/server/model_worker.rs (16384x16384 and a 512 MB decode budget by default) exactly as the other VLM families do. Nothing in this PR reaches a request path yet, so there is no exposure today.

Also worth noting for that issue: Florence2Task::expand interpolates the caller's input string into the encoder prompt for the seven tasks that take one. That is the feature, not a bug, but it means the task input is prompt-controlling, and the server surface should treat it as untrusted request data rather than as a trusted internal string.