Skip to content

feat(models): Florence-2 DaViT vision backbone (sub of #850) #853

Description

@inureyes

Summary

The DaViT hierarchical vision backbone for Florence-2: dual attention (spatial windowed plus channel) over 4 stages, producing the image features that the fusion stage concatenates with the task prompt.

Part of epic

#850

Dependencies

none

Upstream reference

Architecture notes

  • ConvEmbed patch embedding per stage: patch_size = [7,3,3,3], patch_stride = [4,2,2,2], patch_padding = [3,1,1,1], with patch_prenorm.
  • DepthWiseConv2d positional convs: conv_at_attn, conv_at_ffn.
  • SpatialBlock: windowed attention (window_partition / window_reverse, window_size = 12, WindowAttention).
  • ChannelBlock: ChannelAttention.
  • 4 stages: depths = [1,1,9,1], num_heads = [4,8,16,32], qkv_bias, mlp_ratio = 4.
  • Output image features with hidden_size = 768.
  • Parse VisionConfig (model_type: davit).

Implementation plan

  • Add the DaViT encoder following the tree's vision-encoder layout under src/vision/encoders/ (and/or the Florence-2 family directory src/models/florence2/).
  • Reuse windowed-attention helpers under src/vision/encoders/: the window_partition / window_reverse patterns from qwen2_5_vl and kimi_vl.
  • Reuse the conv layer helpers used by the FastViT and gemma3n conv paths for ConvEmbed and DepthWiseConv2d.
  • DaViT conv weights need a channels-last remap in the sanitize step. Mirror the conv channels-last remap already used by the FastViT / gemma3n loaders.

Acceptance criteria

  • VisionConfig (model_type: davit) parses.
  • ConvEmbed, DepthWiseConv2d, SpatialBlock (WindowAttention), and ChannelBlock (ChannelAttention) implemented across the 4 stages.
  • DaViT conv weights channels-last remap handled in sanitize.
  • DaViT forward on a real image produces features matching the mlx-vlm reference (per-stage shapes plus numerical parity).
  • Integration: the backbone is exposed through the API that the fusion sub-issue (#S3) calls, so it plugs into the model forward path rather than living as a standalone script. Full model-path wiring is completed by #S3, which depends on this.

Effort

MEDIUM-HIGH.

Note: this sub-issue is tracked by the checklist in #850.

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:modelsModel architectures, weights, loading, metadatapriority:mediumMedium prioritystatus:doneCompletedtype:enhancementNew features, capabilities, or significant additions

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions