From 24513ed9035cb45d8089439d18e3882bfe0139f1 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 7 Aug 2026 17:36:45 +0900 Subject: [PATCH] feat(models): wire Florence-2 into detection, loading, and the CLI generate path Capstone of the Florence-2 epic (#850): the engine (#852), DaViT tower (#853), fusion (#854), and processor (#855) all landed as libraries with no runtime reachability. This change registers the family end to end so a real checkpoint loads and answers task prompts through `mlxcel generate`. Florence-2 is an encoder-decoder (seq2seq) VLM: the BART decoder cross-attends to cached encoder output over the fused image+prompt sequence, using its own dual-cache (`Florence2SeqCache`) rather than the decoder-only `KVCache` list. It therefore cannot ride the shared autoregressive loop. The integration follows the DiffusionGemma precedent: a `LoadedModel::Florence2VLM` variant whose CLI route exits into a dedicated task pipeline before the standard loop, plus a startup refusal on the server until a seq2seq worker path exists (letting a decoder-only worker pick it up would serve garbage). What changed: - `src/models/florence2/runtime.rs` (new): `Florence2VlmModel` bundling the fused model with its processor, `run_task` (prompt expand -> preprocess -> greedy decode -> coordinate parse, also returning the token count for stats), `parse_task_prompt` for the CLI `-p` syntax (``, bare `od`, ` text`, `...`), and an honest-minimal `LanguageModel` impl (teacher-forced BART forward, `supports_batching = false`) for trait completeness. - `src/models/detection.rs`: `"florence2"` arm -> `ModelType::Florence2VLM`; registered in `src/model_metadata.rs` (`kind: Vlm, directory: Vlm`, adapter loading refused with a named message); `ModelType` variant, `ALL_MODEL_TYPES`, `metadata()` ("Other VLM" family, bf16/f16-only note), and the exhaustiveness test list updated in `src/models/mod.rs`. - `src/loading/vlm_florence2.rs` (new) plus the `try_load_vlm_model_from_dir` arm; `LoadedModel::Florence2VLM` variant and `delegate_language_model!` arm. - `src/commands/generate_florence2.rs` (new): the CLI driver. Rejects `--audio`/`--video`, requires exactly one `--image`, parses the task prompt, and routes image decoding through `decode_image_payloads_with_limits` (decompression-bomb defense handed over from #855). Renders parsed boxes/quads/polygons one instance per line in original-image pixels, with a raw-text fallback that distinguishes "model found nothing" from "parser rejected the answer". Early exit wired into `run_generate_once` after the diffusion exits. - `src/server/startup.rs`: `start_server` bails for `ModelType::Florence2VLM` with a message pointing at the CLI, before any worker spawns. - `src/distributed/tensor_parallel/inference.rs`: placeholder `fallback_architecture` arm keeps the dispatch table total; TP never serves this family. - `docs/supported-models.md`: Florence-2 entry documenting the seq2seq pipeline, the fifteen task markers, and the bf16/f16-only constraint (quantized mlx-community conversions are rejected at load; no quantized code path exists for the BART stack or the DaViT tower). Validation: unit tests beside the code (`florence2_runtime_tests.rs` prompt-parse matrix, render tests in `generate_florence2.rs`, a detection test with the real config shape), plus real-checkpoint CLI runs on `models/Florence-2-base-ft-bf16` for caption, OCR, and detection reproducing the byte-identical parity answers from #855 (recorded in the PR). Refs #850. --- docs/supported-models.md | 1 + src/commands/generate.rs | 12 + src/commands/generate_florence2.rs | 286 ++++++++++++++++++ src/commands/mod.rs | 1 + src/distributed/tensor_parallel/inference.rs | 5 + src/loaded_model.rs | 6 + src/loading/mod.rs | 1 + src/loading/vlm.rs | 3 + src/loading/vlm_florence2.rs | 34 +++ src/model_metadata.rs | 1 + src/models/detection.rs | 2 + src/models/detection_tests.rs | 36 +++ .../florence2/florence2_runtime_tests.rs | 103 +++++++ src/models/florence2/mod.rs | 2 + src/models/florence2/runtime.rs | 262 ++++++++++++++++ src/models/mod.rs | 20 +- src/server/startup.rs | 17 ++ 17 files changed, 787 insertions(+), 5 deletions(-) create mode 100644 src/commands/generate_florence2.rs create mode 100644 src/loading/vlm_florence2.rs create mode 100644 src/models/florence2/florence2_runtime_tests.rs create mode 100644 src/models/florence2/runtime.rs diff --git a/docs/supported-models.md b/docs/supported-models.md index 59ff869eb..cafcec55e 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -118,6 +118,7 @@ Implemented VLM variants include: - MiniCPM-O - Moondream 3 - Moondream 2 (`moondream2` / `moondream1`): reuses Moondream3's linear-patch ViT vision tower and overlap-crop preprocessor, paired with a Phi-1.5-style dense text decoder (fused QKV, partial rotary embedding, parallel attention/MLP, tanh-GELU) instead of Moondream3's sparse-MoE decoder. Images are split into a resized global crop and a grid of overlapping local crops; the local crop features are trimmed of their overlap margins, stitched, and adaptively average-pooled back to the 27x27 encoder grid before being concatenated with the global features and projected to the text hidden size. The BOS token and the 729 projected image tokens form a bidirectional prefix ahead of the causal text prompt. Checkpoint revisions from 2025-06-21 onwards are trained against the `moondream/starmie-v1` tokenizer with Moondream3-style control-token templates (bos = eos = 0), while the official repository still ships the older GPT-2 `tokenizer.json` next to them; mlxcel detects the revision from the bundled `moondream.py`, resolves the starmie tokenizer from the Hub (cached after the first fetch, or place starmie's `tokenizer.json` in the model directory when offline), and keeps the GPT-2 tokenizer with `Question:`/`Answer:` framing for the 2025-01-09 .. 2025-04-14 revisions where that contract is the correct one. +- Florence-2 (`florence2`): Microsoft's task-prompted vision foundation model, an encoder-decoder (seq2seq) VLM. A DaViT tower (4 stages, window + grid attention) encodes the 768x768 image into a `[1, 576, 1024]` feature grid; a learned 2D position embedding, a cosine temporal embedding, and the `image_feature_source` pooling recipe produce 577 projected tokens that are concatenated in front of the task-prompt embeddings (no image placeholder token). The fused sequence runs through a BART encoder, and a causal BART decoder with encoder cross-attention generates the answer greedily from `decoder_start_token_id`. The `-p/--prompt` string selects one of fifteen task modes, for example ``, ``, ``, ``, ``, ``, ``, or ` some text` for the seven tasks that take input text; markers are case-insensitive and the angle brackets are optional. Spatial answers are decoded through the checkpoint's 1000-bin `` tokens and parsed into pixel-space boxes, quad boxes, or polygons against the original image size. Because generation is seq2seq (cross-attention against cached encoder output), the family runs through its own CLI pipeline (`mlxcel generate -m --image -p ''`) rather than the shared autoregressive loop, and `mlxcel-server` refuses the checkpoint at startup for now. bf16/f16 exports only, one image per request: use `mlx-community/Florence-2-base-ft-bf16` (or `-large-ft-bf16`). The 4-bit / 8-bit / 6-bit / 3-bit mlx-community conversions are rejected at load with a named error, because neither the BART stack nor the DaViT tower has a quantized code path yet. - Phi-3 Vision, Phi4MM, Phi4 SigLIP VLM - Molmo2 and Molmo-Point - Nemotron-H Nano Omni: ships a Conformer/Parakeet audio encoder and accepts spoken audio from the CLI with `--audio `. Input audio is resampled to 16 kHz before the encoder. Server-side `input_audio` in `POST /v1/chat/completions` is also supported: the audio block is spliced inside the last user turn, before the `<|im_end|>` end-of-turn marker that the ChatML template uses (id 151 in the released checkpoint). diff --git a/src/commands/generate.rs b/src/commands/generate.rs index 8890c274a..9b4cdd42f 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -2077,6 +2077,18 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> { &user_prompt, ); } + // Florence-2 is an encoder-decoder (seq2seq) VLM: the decoder + // cross-attends to cached encoder output over the fused image+prompt + // sequence, so route it to its task pipeline before the + // autoregressive loop (issue #856). The raw `-p` string is the task + // prompt; the tokenized chat-template form above does not apply. + if let mlxcel::LoadedModel::Florence2VLM(florence2_model) = &model { + return super::generate_florence2::run_florence2_generation( + florence2_model, + &args, + &user_prompt, + ); + } // Reject an off-ladder `--image-soft-tokens` before loading any image: // the budget drives the resize target, so an unsupported value is a // user error, not something to clamp silently. diff --git a/src/commands/generate_florence2.rs b/src/commands/generate_florence2.rs new file mode 100644 index 000000000..090da003d --- /dev/null +++ b/src/commands/generate_florence2.rs @@ -0,0 +1,286 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! CLI driver for Florence-2 task generation (issue #856). +//! +//! Florence-2 is an encoder-decoder (seq2seq) VLM: the encoder consumes the +//! fused image-plus-prompt sequence and the decoder cross-attends to it, so +//! it cannot run on the autoregressive loop `run_generation_mode` drives. +//! `run_generate_once` routes `LoadedModel::Florence2VLM` here before that +//! loop, mirroring the DiffusionGemma / LLaDA-2 early exits. +//! +//! The `-p/--prompt` string selects one of the fifteen task modes +//! (``, ``, ``, ...), optionally followed by the input +//! text the task interpolates. Image decoding routes through the shared +//! [`mlxcel::ImageInputLimits`] admission bounds so an oversized or +//! decompression-bomb payload is rejected before any pixel work. + +use std::time::Instant; + +use anyhow::{Context, Result, anyhow, ensure}; + +use mlxcel::models::florence2::{Florence2TaskResult, Florence2VlmModel, parse_task_prompt}; + +use super::generate::print_generation_preamble; +use crate::GenerateArgs; + +/// Run one Florence-2 task from the CLI flag surface and print the parsed +/// answer plus a generation-stats line. +pub(crate) fn run_florence2_generation( + model: &Florence2VlmModel, + args: &GenerateArgs, + user_prompt: &str, +) -> Result<()> { + ensure!( + args.generation.audio.is_none(), + "Florence-2 does not take --audio input" + ); + ensure!( + args.generation.video.is_empty(), + "Florence-2 does not take --video input" + ); + ensure!( + !args.generation.image.is_empty(), + "Florence-2 is an image-task model: pass --image together with a task prompt \ + such as -p '', -p '', or -p ''" + ); + ensure!( + args.generation.image.len() == 1, + "Florence-2 processes one image per request; got {} --image paths", + args.generation.image.len() + ); + + let (task, input) = parse_task_prompt(user_prompt).map_err(|e| anyhow!("-p/--prompt: {e}"))?; + + // Decode through the shared admission limits (decompression-bomb + // defense, issue #855 handoff): `preprocess_with_sizes` takes an + // already-decoded image, so the bound has to hold at this boundary. + let image_path = &args.generation.image[0]; + let bytes = std::fs::read(image_path) + .with_context(|| format!("Failed to read image {image_path:?}"))?; + let mut images = + mlxcel::decode_image_payloads_with_limits(&[bytes], mlxcel::current_image_input_limits()) + .with_context(|| format!("Failed to decode image {image_path:?}"))?; + let image = images + .pop() + .ok_or_else(|| anyhow!("image decoding returned no image for {image_path:?}"))?; + + print_generation_preamble(user_prompt)?; + println!(); + + let started = Instant::now(); + let run = model.run_task(task, input.as_deref(), &image, args.generation.max_tokens)?; + let elapsed = started.elapsed().as_secs_f64(); + + println!( + "{}", + render_task_result(&run.output.result, &run.output.raw_text) + ); + println!(); + let tps = if elapsed > 0.0 { + run.generated_tokens as f64 / elapsed + } else { + 0.0 + }; + println!( + "[Generated {} tokens in {:.2}s = {:.2} tok/s]", + run.generated_tokens, elapsed, tps + ); + if args.generation.profile { + println!("[Raw answer] {}", run.output.raw_text); + } + + mlxcel_core::clear_memory_cache(); + Ok(()) +} + +/// Render a parsed task result for the terminal. +/// +/// Spatial results are printed one instance per line as +/// `label: [coordinates]` in original-image pixels. An answer that parsed to +/// nothing falls back to the raw decoded text, which is the only way to tell +/// "the model found nothing" from "the parser rejected the answer". +fn render_task_result(result: &Florence2TaskResult, raw_text: &str) -> String { + let rendered = match result { + Florence2TaskResult::Text(text) => text.trim().to_string(), + Florence2TaskResult::Boxes { boxes, labels } => boxes + .iter() + .zip(labels) + .map(|(bbox, label)| { + format!( + "{}: [{:.1}, {:.1}, {:.1}, {:.1}]", + display_label(label), + bbox.xmin, + bbox.ymin, + bbox.xmax, + bbox.ymax + ) + }) + .collect::>() + .join("\n"), + Florence2TaskResult::QuadBoxes { quad_boxes, labels } => quad_boxes + .iter() + .zip(labels) + .map(|(quad, label)| { + format!("{}: [{}]", display_label(label), join_points(&quad.points)) + }) + .collect::>() + .join("\n"), + Florence2TaskResult::Polygons { polygons, labels } => { + render_polygon_instances(polygons, labels) + } + Florence2TaskResult::BoxesOrPolygons { + boxes, + box_labels, + polygons, + polygon_labels, + } => { + let mut lines: Vec = boxes + .iter() + .zip(box_labels) + .map(|(bbox, label)| { + format!( + "{}: [{:.1}, {:.1}, {:.1}, {:.1}]", + display_label(label), + bbox.xmin, + bbox.ymin, + bbox.xmax, + bbox.ymax + ) + }) + .collect(); + let polygon_block = render_polygon_instances(polygons, polygon_labels); + if !polygon_block.is_empty() { + lines.push(polygon_block); + } + lines.join("\n") + } + // `Florence2TaskResult` is `#[non_exhaustive]`: render future + // variants through the raw-text fallback below. + _ => String::new(), + }; + if rendered.is_empty() { + format!("[no parsed instances] raw answer: {}", raw_text.trim()) + } else { + rendered + } +} + +fn render_polygon_instances( + polygons: &[Vec], + labels: &[String], +) -> String { + polygons + .iter() + .zip(labels) + .map(|(outlines, label)| { + let rendered_outlines = outlines + .iter() + .map(|polygon| format!("[{}]", join_points(&polygon.points))) + .collect::>() + .join(", "); + format!("{}: {}", display_label(label), rendered_outlines) + }) + .collect::>() + .join("\n") +} + +fn join_points(points: &[f32]) -> String { + points + .iter() + .map(|point| format!("{point:.1}")) + .collect::>() + .join(", ") +} + +/// `` predicts unnamed regions; keep the line shape stable. +fn display_label(label: &str) -> &str { + if label.is_empty() { "(region)" } else { label } +} + +#[cfg(test)] +mod tests { + use super::*; + use mlxcel::models::{Florence2BoundingBox, Florence2Polygon, Florence2QuadBox}; + + #[test] + fn renders_text_result() { + let result = Florence2TaskResult::Text("A car parked on the street.".to_string()); + assert_eq!( + render_task_result(&result, ""), + "A car parked on the street." + ); + } + + #[test] + fn renders_boxes_one_instance_per_line() { + let result = Florence2TaskResult::Boxes { + boxes: vec![ + Florence2BoundingBox { + xmin: 1.0, + ymin: 2.0, + xmax: 3.0, + ymax: 4.0, + }, + Florence2BoundingBox { + xmin: 5.0, + ymin: 6.0, + xmax: 7.0, + ymax: 8.0, + }, + ], + labels: vec!["car".to_string(), String::new()], + }; + let rendered = render_task_result(&result, ""); + assert_eq!( + rendered, + "car: [1.0, 2.0, 3.0, 4.0]\n(region): [5.0, 6.0, 7.0, 8.0]" + ); + } + + #[test] + fn renders_quad_boxes_and_polygons() { + let quads = Florence2TaskResult::QuadBoxes { + quad_boxes: vec![Florence2QuadBox { + points: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + }], + labels: vec!["HELLO".to_string()], + }; + assert_eq!( + render_task_result(&quads, ""), + "HELLO: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]" + ); + + let polygons = Florence2TaskResult::Polygons { + polygons: vec![vec![Florence2Polygon { + points: vec![1.0, 2.0, 3.0, 4.0], + }]], + labels: vec!["dog".to_string()], + }; + assert_eq!( + render_task_result(&polygons, ""), + "dog: [1.0, 2.0, 3.0, 4.0]" + ); + } + + #[test] + fn empty_result_falls_back_to_raw_text() { + let result = Florence2TaskResult::Boxes { + boxes: vec![], + labels: vec![], + }; + let rendered = render_task_result(&result, ""); + assert!(rendered.contains("raw answer"), "got: {rendered}"); + } +} diff --git a/src/commands/mod.rs b/src/commands/mod.rs index 0bacf23b0..b19d0f671 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -23,6 +23,7 @@ pub(crate) mod detect; pub(crate) mod download; pub(crate) mod generate; mod generate_diffusion; +mod generate_florence2; mod generate_llada2; mod generate_vlm; pub(crate) mod inspect; diff --git a/src/distributed/tensor_parallel/inference.rs b/src/distributed/tensor_parallel/inference.rs index d3982e1be..5c5fd0a1e 100644 --- a/src/distributed/tensor_parallel/inference.rs +++ b/src/distributed/tensor_parallel/inference.rs @@ -400,6 +400,11 @@ fn fallback_architecture(model_type: ModelType) -> &'static str { // worker loop, never routed to tensor-parallel text inference; the // planner's supported-architecture validation rejects this string. ModelType::Llada2Moe => "llada2_moe", + // Florence-2 is an encoder-decoder (seq2seq) VLM served through its + // own CLI task pipeline, never routed to tensor-parallel text + // inference; the planner's supported-architecture validation rejects + // this string. Placeholder keeps the dispatch table total. + ModelType::Florence2VLM => "florence2", // Whisper is an ASR model served through the audio endpoints, never // routed to tensor-parallel text inference; the loader rejects it // earlier. Return a placeholder so the dispatch table stays total. diff --git a/src/loaded_model.rs b/src/loaded_model.rs index fa7ae978b..8e300b797 100644 --- a/src/loaded_model.rs +++ b/src/loaded_model.rs @@ -112,6 +112,11 @@ pub enum LoadedModel { MiniCPMV46VLM(vision::MiniCPMV46VLModel), Moondream3VLM(vision::Moondream3VLModel), Moondream2VLM(vision::Moondream2VLModel), + /// Florence-2 (issue #856): DaViT tower + BART encoder-decoder text + /// stack. Seq2seq, so the CLI routes it to the Florence-2 task pipeline + /// before the autoregressive loop and the server refuses it at startup; + /// the `LanguageModel` delegation below exists for trait completeness. + Florence2VLM(models::Florence2VlmModel), Gemma3n(models::Gemma3nModel), Gemma3nVLM(vision::Gemma3nVLModel), Phi(models::PhiModel), @@ -271,6 +276,7 @@ macro_rules! delegate_language_model { LoadedModel::MiniCPMV46VLM(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::Moondream3VLM(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::Moondream2VLM(inner) => LanguageModel::$method(inner, $($arg),*), + LoadedModel::Florence2VLM(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::Gemma3n(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::Gemma3nVLM(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::Phi(inner) => LanguageModel::$method(inner, $($arg),*), diff --git a/src/loading/mod.rs b/src/loading/mod.rs index b4bdc23c0..a1686cd8a 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -217,6 +217,7 @@ fn try_load_vlm_model_from_dir( ModelType::MiniCPMV46VLM => Some(load_minicpmv4_6_vlm(model_path)?), ModelType::Moondream3VLM => Some(load_moondream3_vlm(model_path)?), ModelType::Moondream2VLM => Some(load_moondream2_vlm(model_path)?), + ModelType::Florence2VLM => Some(load_florence2_vlm(model_path)?), ModelType::Gemma3nVLM => Some(load_gemma3n_vlm(model_path)?), ModelType::Phi4MMVLM => Some(load_phi4mm_vlm(model_path)?), ModelType::Phi4SigLipVLM => Some(load_phi4_siglip_vlm(model_path)?), diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index b87c6c0a2..468dd32c2 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -47,6 +47,8 @@ mod dots_ocr; mod ernie4_5_vl; #[path = "vlm_fastvlm.rs"] mod fastvlm; +#[path = "vlm_florence2.rs"] +mod florence2; #[path = "vlm_gemma.rs"] mod gemma; #[path = "vlm_gemma_unified.rs"] @@ -97,6 +99,7 @@ pub(crate) use deepseekocr::{ pub(crate) use dots_ocr::load_dots_ocr_vl; pub(crate) use ernie4_5_vl::load_ernie4_5_moe_vlm; pub(crate) use fastvlm::load_fastvlm_vlm; +pub(crate) use florence2::load_florence2_vlm; pub(crate) use gemma::{load_gemma3_vlm, load_gemma3n_vlm, load_gemma4_vlm}; pub(crate) use gemma_unified::load_gemma4_unified; pub(crate) use granite_vision::load_granite_vision_vlm; diff --git a/src/loading/vlm_florence2.rs b/src/loading/vlm_florence2.rs new file mode 100644 index 000000000..63bbb70f8 --- /dev/null +++ b/src/loading/vlm_florence2.rs @@ -0,0 +1,34 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Florence-2 (`florence2`) VLM loader. +//! +//! Loads the whole runtime unit through +//! [`models::Florence2VlmModel::load`]: the DaViT tower plus BART seq2seq +//! text stack from the checkpoint safetensors, and the processor (BART +//! tokenizer plus 768x768 image preprocessor) from the same directory. The +//! model rejects quantized checkpoints with a named error before any weight +//! is loaded; only bf16 / f16 exports such as +//! `mlx-community/Florence-2-base-ft-bf16` are supported. + +use anyhow::Result; +use std::path::Path; + +use crate::LoadedModel; +use crate::models; + +pub(crate) fn load_florence2_vlm(model_path: &Path) -> Result { + let model = models::Florence2VlmModel::load(model_path)?; + Ok(LoadedModel::Florence2VLM(model)) +} diff --git a/src/model_metadata.rs b/src/model_metadata.rs index ac57cdf23..4adcfb463 100644 --- a/src/model_metadata.rs +++ b/src/model_metadata.rs @@ -150,6 +150,7 @@ macro_rules! for_each_model_registration { MiniCPMV46VLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("MiniCPM-V 4.6 VLM does not support adapter loading; use load_model() instead") }; Moondream3VLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("Moondream3 VLM does not support adapter loading; use load_model() instead") }; Moondream2VLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("Moondream2 VLM does not support adapter loading; use load_model() instead") }; + Florence2VLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("Florence-2 does not support adapter loading; use load_model() instead") }; GptOss => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::GptOssModel::load, args: models::gpt_oss::ModelArgs, weight_builder: models::GptOssModel::from_weights, wrap: |m| LoadedModel::GptOss(models::GptOssWrapper::new(m)) } }; Qwen2Moe => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::Qwen2MoeModel::load, args: models::qwen2_moe::ModelArgs, weight_builder: models::Qwen2MoeModel::from_weights, wrap: LoadedModel::Qwen2Moe } }; Gemma3n => { kind: Text, directory: Nonstandard, weight: Some(WeightLoadRoute::Special), adapter: None }; diff --git a/src/models/detection.rs b/src/models/detection.rs index e87b8e74e..a0cc2a405 100644 --- a/src/models/detection.rs +++ b/src/models/detection.rs @@ -323,6 +323,8 @@ pub fn get_model_type(model_path: &Path) -> Result { "molmo" => Ok(ModelType::MolmoVLM), "molmo2" => Ok(ModelType::Molmo2VLM), "molmo_point" => Ok(ModelType::MolmoPointVLM), + // Florence-2 (DaViT tower + BART encoder-decoder text stack). + "florence2" => Ok(ModelType::Florence2VLM), // Speech-to-text (encoder-decoder ASR). "whisper" => Ok(ModelType::Whisper), _ => Err(anyhow::anyhow!( diff --git a/src/models/detection_tests.rs b/src/models/detection_tests.rs index 7985867c8..0867ddc82 100644 --- a/src/models/detection_tests.rs +++ b/src/models/detection_tests.rs @@ -87,6 +87,42 @@ fn whisper_model_type_is_detected() { fs::remove_dir_all(model_dir).unwrap(); } +#[test] +fn florence2_model_type_is_detected() { + // Florence-2 declares `model_type: "florence2"` at the top level. The + // real checkpoint's `vision_config.model_type` is an empty string, so + // detection must key off the top-level value alone (a `vision_config` + // is present but never consulted for this family). + let model_dir = temp_path("florence2_vlm"); + fs::create_dir_all(&model_dir).unwrap(); + fs::write( + model_dir.join("config.json"), + r#"{ + "model_type": "florence2", + "architectures": ["Florence2ForConditionalGeneration"], + "is_encoder_decoder": true, + "projection_dim": 768, + "text_config": { + "model_type": "florence2_language", + "d_model": 768, + "encoder_layers": 6, + "decoder_layers": 6, + "vocab_size": 51289 + }, + "vision_config": { + "model_type": "", + "dim_embed": [128, 256, 512, 1024] + } + }"#, + ) + .unwrap(); + + let detected = super::detection::get_model_type(&model_dir).unwrap(); + assert_eq!(detected, ModelType::Florence2VLM); + + fs::remove_dir_all(model_dir).unwrap(); +} + #[test] fn gpt2_model_type_is_detected() { // GPT-2 configs use the original OpenAI field names (`n_embd` / `n_head` / diff --git a/src/models/florence2/florence2_runtime_tests.rs b/src/models/florence2/florence2_runtime_tests.rs new file mode 100644 index 000000000..58ee4eed9 --- /dev/null +++ b/src/models/florence2/florence2_runtime_tests.rs @@ -0,0 +1,103 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Unit tests for the Florence-2 runtime surface: CLI task-prompt parsing. +//! Model-backed behavior (loading, LanguageModel forward) needs the real +//! checkpoint and is covered by `tests/florence2_fusion_parity.rs` and the +//! CLI validation in the integration issue. + +use super::parse_task_prompt; +use crate::models::florence2::Florence2Task; + +#[test] +fn parses_bare_marker_forms() { + for prompt in ["", "OD", "od", " "] { + let (task, input) = parse_task_prompt(prompt).unwrap(); + assert_eq!(task, Florence2Task::ObjectDetection, "prompt {prompt:?}"); + assert_eq!(input, None, "prompt {prompt:?}"); + } +} + +#[test] +fn parses_every_task_marker_round_trip() { + for task in Florence2Task::ALL { + let (parsed, input) = parse_task_prompt(task.token()).unwrap(); + assert_eq!(parsed, task); + assert_eq!(input, None); + } +} + +#[test] +fn parses_marker_with_input_text() { + let (task, input) = parse_task_prompt(" a green car").unwrap(); + assert_eq!(task, Florence2Task::CaptionToPhraseGrounding); + assert_eq!(input.as_deref(), Some("a green car")); + + // Bare marker followed by input text. + let (task, input) = parse_task_prompt("caption_to_phrase_grounding a green car").unwrap(); + assert_eq!(task, Florence2Task::CaptionToPhraseGrounding); + assert_eq!(input.as_deref(), Some("a green car")); +} + +#[test] +fn parses_region_input_without_separating_space() { + let (task, input) = + parse_task_prompt("").unwrap(); + assert_eq!(task, Florence2Task::RegionToCategory); + assert_eq!( + input.as_deref(), + Some("") + ); +} + +#[test] +fn rejects_unknown_and_malformed_prompts() { + // Unknown marker. + let err = parse_task_prompt("").unwrap_err(); + assert!( + err.contains(""), + "error should list valid markers: {err}" + ); + + // Free-form text with no marker. + let err = parse_task_prompt("describe this image").unwrap_err(); + assert!( + err.contains(""), + "error should list valid markers: {err}" + ); + + // Unclosed marker. + let err = parse_task_prompt("` takes no input; the parser still splits the syntax and `expand` + // rejects it, keeping the strict boundary in one place. + let (task, input) = parse_task_prompt(" spurious text").unwrap(); + assert_eq!(task, Florence2Task::ObjectDetection); + assert_eq!(input.as_deref(), Some("spurious text")); + assert!(task.expand(input.as_deref()).is_err()); + + // A task that requires input still errors in `expand` when it is absent. + let (task, input) = parse_task_prompt("").unwrap(); + assert_eq!(task, Florence2Task::OpenVocabularyDetection); + assert_eq!(input, None); + assert!(task.expand(None).is_err()); +} diff --git a/src/models/florence2/mod.rs b/src/models/florence2/mod.rs index 823c5493f..dc6d17f07 100644 --- a/src/models/florence2/mod.rs +++ b/src/models/florence2/mod.rs @@ -44,6 +44,7 @@ mod model; mod parse; mod postprocess; mod processor; +mod runtime; mod scan; mod tasks; @@ -61,6 +62,7 @@ pub use coords::{ pub use model::Florence2Model; pub use postprocess::{Florence2PostProcessingType, Florence2TaskResult}; pub use processor::{Florence2Output, Florence2Processor}; +pub use runtime::{Florence2RunOutput, Florence2VlmModel, parse_task_prompt}; pub use tasks::Florence2Task; use std::path::Path; diff --git a/src/models/florence2/runtime.rs b/src/models/florence2/runtime.rs new file mode 100644 index 000000000..0735a1c90 --- /dev/null +++ b/src/models/florence2/runtime.rs @@ -0,0 +1,262 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Florence-2 runtime registration: the loadable model unit and the CLI +//! task-prompt surface. +//! +//! [`Florence2VlmModel`] bundles the fused [`Florence2Model`] with its +//! [`Florence2Processor`] so the runtime holds one unit that can take a task +//! prompt and an image and return a parsed answer. It is what +//! `LoadedModel::Florence2VLM` stores. +//! +//! Florence-2 is an encoder-decoder (seq2seq) model: the answer is decoded +//! against cached encoder output through cross-attention, with its own +//! [`super::Florence2SeqCache`] rather than the decoder-only `KVCache` list. +//! It therefore cannot run on the autoregressive loop every other +//! `LoadedModel` family uses. The CLI routes this family to +//! [`Florence2VlmModel::run_task`] before the standard generation loop +//! (mirroring the DiffusionGemma early exit in `commands/generate.rs`), and +//! `mlxcel-server` refuses the checkpoint at startup until a seq2seq worker +//! path exists. The [`LanguageModel`] impl below exists for trait +//! completeness (warmup, tooling), not as a generation path. + +use std::path::Path; + +use anyhow::{Result, anyhow}; +use image::DynamicImage; + +use mlxcel_core::generate::LanguageModel; +use mlxcel_core::layers::KVCache; +use mlxcel_core::{MlxArray, UniquePtr}; + +use super::coords::Florence2ImageSize; +use super::model::Florence2Model; +use super::processor::{Florence2Output, Florence2Processor}; +use super::tasks::Florence2Task; + +/// The loadable Florence-2 runtime unit: fused model plus processor. +/// +/// Holds MLX weight handles, so the owning provider serializes access. +pub struct Florence2VlmModel { + model: Florence2Model, + processor: Florence2Processor, +} + +/// One [`Florence2VlmModel::run_task`] call: the parsed answer plus the +/// generated-token count the CLI needs for its throughput line. +pub struct Florence2RunOutput { + /// Raw decoded answer and its parsed form. + pub output: Florence2Output, + /// Number of decoder tokens generated (EOS excluded). + pub generated_tokens: usize, +} + +impl Florence2VlmModel { + /// Load model and processor from a checkpoint directory. Quantized + /// checkpoints are rejected by [`Florence2Model::load`] with a named + /// error; use a bf16 or f16 export. + pub fn load(model_path: &Path) -> Result { + let model = Florence2Model::load(model_path)?; + let processor = Florence2Processor::from_pretrained(model_path)?; + Ok(Self { model, processor }) + } + + /// The fused vision-language model. + pub fn model(&self) -> &Florence2Model { + &self.model + } + + /// The task-prompt and image processor. + pub fn processor(&self) -> &Florence2Processor { + &self.processor + } + + /// Run one task end to end: expand and tokenize the task prompt, + /// preprocess the image, decode greedily, and parse the answer against + /// the original image size. + /// + /// Same pipeline as [`Florence2Processor::run`], kept separate so the + /// caller also gets the generated-token count for stats. + pub fn run_task( + &self, + task: Florence2Task, + input: Option<&str>, + image: &DynamicImage, + max_new_tokens: usize, + ) -> Result { + let prompt_ids = self.processor.encode_prompt(task, input)?; + let processed = self + .processor + .image_processor() + .preprocess_with_sizes(std::slice::from_ref(image)); + let (width, height) = *processed + .original_sizes + .first() + .ok_or_else(|| anyhow!("Florence-2: image preprocessing returned no images"))?; + + let generated = + self.model + .generate_greedy(&processed.pixel_values, &prompt_ids, max_new_tokens)?; + let generated_tokens = generated.len(); + let raw_text = self.processor.decode_answer(&generated)?; + let result = + self.processor + .post_process(&raw_text, task, Florence2ImageSize::new(width, height)); + Ok(Florence2RunOutput { + output: Florence2Output { raw_text, result }, + generated_tokens, + }) + } +} + +/// Parse a CLI `-p/--prompt` string into a Florence-2 task and its optional +/// input text. +/// +/// Accepted forms: +/// - `""` or `"od"`: a bare task marker, angle brackets optional, +/// case-insensitive (per [`Florence2Task::from_str`]). +/// - `" a green car"` or +/// `"CAPTION_TO_PHRASE_GROUNDING a green car"`: marker followed by the +/// input text the task interpolates. +/// - `""`: a region +/// input needs no separating space, everything after the first closing +/// `>` is the input. +/// +/// Whether the task actually takes input text is validated later by +/// [`Florence2Task::expand`], which rejects a missing or superfluous input; +/// this function only splits the syntax. Anything that does not start with a +/// recognized marker is an error listing the valid markers, matching the +/// deliberate strictness of `expand` (upstream silently misparses sloppy +/// prompts into nonsense questions). +pub fn parse_task_prompt(prompt: &str) -> Result<(Florence2Task, Option), String> { + let trimmed = prompt.trim(); + if trimmed.is_empty() { + return Err(format!( + "empty Florence-2 prompt; pass a task marker such as {}", + valid_task_markers() + )); + } + + if let Some(rest) = trimmed.strip_prefix('<') { + let Some(end) = rest.find('>') else { + return Err(format!( + "Florence-2 prompt {trimmed:?} opens a task marker without a closing '>'" + )); + }; + let task: Florence2Task = rest[..end].parse().map_err(|_| { + format!( + "unknown Florence-2 task marker <{}>; valid markers: {}", + &rest[..end], + valid_task_markers() + ) + })?; + let input = rest[end + 1..].trim(); + let input = (!input.is_empty()).then(|| input.to_string()); + return Ok((task, input)); + } + + // Bare form: the whole prompt is a task name, or a task name followed by + // input text. + if let Ok(task) = trimmed.parse::() { + return Ok((task, None)); + } + if let Some((head, rest)) = trimmed.split_once(char::is_whitespace) + && let Ok(task) = head.parse::() + { + let input = rest.trim(); + let input = (!input.is_empty()).then(|| input.to_string()); + return Ok((task, input)); + } + + Err(format!( + "prompt {trimmed:?} does not start with a Florence-2 task; valid markers: {}", + valid_task_markers() + )) +} + +fn valid_task_markers() -> String { + Florence2Task::ALL + .iter() + .map(|task| task.token()) + .collect::>() + .join(", ") +} + +impl LanguageModel for Florence2VlmModel { + /// Honest minimal trait forward: a text-only teacher-forced BART pass. + /// The prompt is encoded by the bidirectional encoder, and the decoder + /// consumes the same ids shifted right behind `decoder_start_token_id` + /// (HuggingFace `shift_tokens_right`), returning per-position logits. + /// + /// A fresh seq2seq cache is built per call and the passed decoder-only + /// `caches` are ignored: cross-attention K/V cannot live in `KVCache`, + /// so an incremental decode driven through this trait would silently + /// re-encode each step. The CLI routes Florence-2 to + /// [`Florence2VlmModel::run_task`] before the autoregressive loop and the + /// server refuses the checkpoint at startup, so this exists for trait + /// completeness (warmup, tooling) rather than as a generation path. + fn forward( + &self, + input_ids: &MlxArray, + _caches: &mut [KVCache], + _mask: Option<&MlxArray>, + ) -> UniquePtr { + let text = self.model.text_model(); + let encoder_hidden = text.encode_tokens(input_ids); + + let shape = mlxcel_core::array_shape(input_ids); + let (batch, seq) = (shape[0], shape[1]); + let start = text.config().decoder_start_token_id; + let start_col = + mlxcel_core::from_slice_i32(&vec![start; batch.max(1) as usize], &[batch, 1]); + let decoder_input = if seq > 1 { + let prefix = mlxcel_core::slice(input_ids, &[0, 0], &[batch, seq - 1]); + mlxcel_core::concatenate(&start_col, &prefix, 1) + } else { + start_col + }; + + let mut cache = self.model.make_cache(); + text.decode(&decoder_input, &encoder_hidden, &mut cache) + } + + /// Florence-2 keeps its own dual (self + cross attention) cache, built by + /// [`Florence2Model::make_cache`]; there is nothing to store in the + /// decoder-only `KVCache` list. + fn make_caches(&self) -> Vec { + Vec::new() + } + + fn num_layers(&self) -> usize { + self.model.config().text.decoder_layers as usize + } + + fn eos_token_ids(&self) -> Vec { + vec![self.model.config().text.eos_token_id] + } + + /// Seq2seq generation is a model-owned loop over a single sequence; the + /// batched/paged scheduler must never pick this model up. + fn supports_batching(&self) -> bool { + false + } + + fn supports_padded_prefill(&self) -> bool { + false + } +} + +#[cfg(test)] +#[path = "florence2_runtime_tests.rs"] +mod florence2_runtime_tests; diff --git a/src/models/mod.rs b/src/models/mod.rs index a89328f36..add106c87 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -179,9 +179,9 @@ pub use falcon_h1::FalconH1Model; pub use florence2::{ FLORENCE2_LOC_TOKEN_BASE, FLORENCE2_VISION_PREFIX, Florence2BoundingBox, Florence2Config, Florence2DaViT, Florence2ImageSize, Florence2Model, Florence2Output, Florence2Polygon, - Florence2PostProcessingType, Florence2Processor, Florence2QuadBox, Florence2SeqCache, - Florence2Task, Florence2TaskResult, Florence2TextConfig, Florence2TextModel, - Florence2VisionConfig, florence2_loc_token_id, + Florence2PostProcessingType, Florence2Processor, Florence2QuadBox, Florence2RunOutput, + Florence2SeqCache, Florence2Task, Florence2TaskResult, Florence2TextConfig, Florence2TextModel, + Florence2VisionConfig, Florence2VlmModel, florence2_loc_token_id, }; pub use gemma::GemmaModel; pub use gemma2::Gemma2Model; @@ -335,9 +335,13 @@ pub enum ModelType { MiniCPMV46VLM, // MiniCPM-V 4.6 (SigLIP + VitMerger + Merger + Qwen3.5 text) Moondream3VLM, // Moondream3 (custom ViT + custom text decoder, query/caption image path) Moondream2VLM, // Moondream2 (SigLIP-style ViT + Phi text decoder + crop tiling) - Gemma3n, // Gemma 3n (text-only) + /// Florence-2 (`florence2`): DaViT vision tower + BART encoder-decoder + /// text stack. Encoder-decoder (seq2seq), so it is served through its own + /// task pipeline (CLI early exit), not the autoregressive decode loop. + Florence2VLM, + Gemma3n, // Gemma 3n (text-only) Gemma3nVLM, // Gemma 3n VLM (MobileNetV5 + Gemma3n) - Phi, // Phi 1/2 + Phi, // Phi 1/2 /// Phixtral (`phi-msft` with `num_local_experts`): a Mixtral-style /// sparse MoE on the Phi-2 parallel-residual backbone. Shares the /// `phi-msft` model_type with dense Phi and is told apart by @@ -563,6 +567,7 @@ pub const ALL_MODEL_TYPES: &[ModelType] = &[ ModelType::MiniCPMV46VLM, ModelType::Moondream3VLM, ModelType::Moondream2VLM, + ModelType::Florence2VLM, ModelType::Gemma3n, ModelType::Gemma3nVLM, ModelType::Phi, @@ -969,6 +974,10 @@ impl ModelType { } ModelType::Moondream3VLM => ("Moondream 3 (custom ViT + custom decoder)", "Other VLM"), ModelType::Moondream2VLM => ("Moondream 2 (SigLIP-style ViT + Phi text)", "Other VLM"), + ModelType::Florence2VLM => ( + "Florence-2 (DaViT + BART seq2seq, task prompts, bf16/f16 only)", + "Other VLM", + ), ModelType::MiniCPMOVLM => ( "MiniCPM-o (dynamic SigLIP + resampler + Qwen3-VL text)", "Other VLM", @@ -1082,6 +1091,7 @@ mod metadata_tests { MiniCPMV46VLM, Moondream3VLM, Moondream2VLM, + Florence2VLM, Gemma3n, Gemma3nVLM, Phi, diff --git a/src/server/startup.rs b/src/server/startup.rs index 296efb23d..526c3dc54 100644 --- a/src/server/startup.rs +++ b/src/server/startup.rs @@ -1572,6 +1572,23 @@ fn install_surgery_pipeline_for_server(startup: &ServerStartupConfig) -> Result< pub async fn start_server(mut startup: ServerStartupConfig) -> Result<()> { initialize_server_logging(&startup)?; + // Florence-2 (issue #856): the model loads through `load_model` for the + // CLI task pipeline, but its encoder-decoder (seq2seq) generation cannot + // run on the decoder-only worker loops this server spawns; letting a + // worker pick it up would serve garbage from the trait-completeness + // forward. Refuse the checkpoint here, before any worker starts, until a + // seq2seq worker path exists. + if matches!( + crate::models::get_model_type(&startup.model_path), + Ok(crate::models::ModelType::Florence2VLM) + ) { + anyhow::bail!( + "Florence-2 is an encoder-decoder (seq2seq) VLM that mlxcel-server cannot serve \ + yet. Run it through the CLI instead: mlxcel generate -m --image \ + -p '' (or another task marker such as or )." + ); + } + // Issue #688 (M1/M2 hardening): disable CUDA graph capture for hazard-family // models (Gemma 4) here, on the main startup thread, before any generation or // pipeline worker is spawned and before the first GPU eval latches MLX's