Repository navigation
feat: add shared DINOv2 feature encoder - #2549
Merged
Merged
Conversation
inureyes
marked this pull request as ready for review
October 22, 2026 18:18
inureyes
added a commit
that referenced
this pull request
Oct 22, 2026
Port standard and register-token Hugging Face DINOv2 checkpoints with shared dense-model APIs, upstream-compatible interpolation, masking, intermediate features, and giant SwiGLU support. Expose standalone image feature extraction through mlxcel predict, register architecture capabilities, and add pinned full-output reference fixtures plus operator documentation.
Load and enforce the checkpoint image processor policy, reproduce Pillow's bicubic uint8 rounding, and cover predictor output against real references. Bound checkpoint-driven allocations and harden the shared backbone APIs.
Describe the checkpoint processor policy and its Hugging Face defaults instead of the superseded patch-alignment behavior.
inureyes
force-pushed
the
feature/issue-2350-dinov2-backbone
branch
from
October 22, 2026 18:28
f084c05 to
ab61c8f
Compare
13 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Hugging Face DINOv2 checkpoints previously had no shared encoder or standalone runtime in mlxcel, which blocked the dense-model families that depend on the backbone. This adds a reusable standard/register-token encoder and exposes image feature extraction through
mlxcel predict.The encoder loads public Hugging Face standard and register checkpoints, remaps and validates their weights, supports boolean patch masks, preserves register-token ordering, returns normalized intermediate layers, and implements rectangular positional interpolation in both upstream modes. Shared consumers can construct a backbone from a scoped parent checkpoint without consuming the parent weight map. The giant preset uses its SwiGLU feed-forward layout.
The standalone predictor reads
preprocessor_config.jsonand reproduces the slow Hugging FaceBitImageProcessorpath: RGB conversion, shortest-edge resize, Pillow-compatible bicubic uint8 rounding, center crop, rescaling, and normalization. It writeslast_hidden_state,pooler_output, and spatial-onlyhidden_patch_tokens. Architecture and operator documentation describe both supported checkpoint types.Pinned full-tensor CPU f32 fixtures cover
facebook/dinov2-smallandfacebook/dinov2-with-registers-smallafter real checkpoint preprocessing. Tests also cover predictor preprocessing, positional interpolation literals,[4, 11]intermediate outputs, register exclusion, f16 masking, scoped loading, allocation limits, invalid public-layer inputs, and sanitizer behavior. Checkpoint-driven image, token, architecture, and intermediate-output sizes are bounded before expensive allocations; attention uses the shared safe dispatch.Validation:
cargo test -p mlxcel --profile test-fast --features rocm --lib dinov2 -- --nocapture --test-threads=1: 13 passed, 0 failed, 0 ignored. The explicit CPU f32 reference tests produced relative RMS values from1.197926287e-7for preprocessing through4.275870861e-6for model outputs.cargo clippy -p mlxcel --profile test-fast --features rocm --lib --tests -- -D warningscargo build -p mlxcel --profile test-fast --features rocm --bin mlxcelMLXCEL_DEVICE=gpuCLI predictions on gfx1151 for both checkpoints wrote exactlymeta.jsonandresult.safetensors. Standard-checkpoint tensor RMS values were2.844445329e-6,2.473357604e-6, and2.845811506e-6; register-checkpoint values were4.548628530e-6,2.131596900e-6, and4.596212478e-6for last hidden state, pooler output, and patch tokens respectively.mlxcel arch --jsonlistsdinov2with both model types and the predict/image/embeddings capability.cargo fmt --all -- --checkpython3 scripts/ci/check_env_var_docs.pypython3 scripts/ci/check_binary_assets.pygit diff --checkCloses #2350