#!/usr/bin/env python3 """Dump a DINOv3 reference forward pass for `cargo run --example verify`. Writes three files into the output directory: pixel_values.bin f32 [3, S, S] the preprocessed input tensor features.bin f32 [tokens, 384] expected last_hidden_state model.safetensors weights, in the graph's naming Dumping the *preprocessed* pixel tensor rather than an image is deliberate: it keeps resize and normalization differences out of the comparison, so a mismatch in `verify` is a mismatch in the graph, not in the resampling filter. Saving the state dict here rather than downloading it in Rust also means `verify` needs no Hub access and no license acceptance at run time. pip install torch transformers safetensors pillow numpy huggingface-cli login # facebook/... is license-gated python tools/dump_reference.py --out ref/ The weights are gated. Accept the DINOv3 license on the canonical model page and authenticate with Hugging Face before running this script. """ import argparse import hashlib import json import pathlib import numpy as np import torch import transformers from safetensors.torch import save_file from transformers import AutoModel from transformers.models.dinov3_vit.modeling_dinov3_vit import ( apply_rotary_pos_emb, eager_attention_forward, ) DEFAULT_MODEL = "facebook/dinov3-vits16-pretrain-lvd1689m" IMAGE_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) IMAGE_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) def build_input(size: int, image: str | None) -> np.ndarray: """Return a normalized f32 CHW tensor.""" if image is not None: from PIL import Image rgb = Image.open(image).convert("RGB").resize((size, size), Image.BILINEAR) hwc = np.asarray(rgb, dtype=np.float32) / 255.0 else: # A smooth, deterministic pattern. Smooth matters: white noise makes # every patch statistically identical, which would hide a bug in the # patch ordering or in the position encoding. ys, xs = np.mgrid[0:size, 0:size].astype(np.float32) u, v = xs / size, ys / size hwc = np.stack( [ 0.5 + 0.5 * np.sin(6.0 * u + 2.0 * v), 0.5 + 0.5 * np.sin(4.0 * v - 3.0 * u * v), 0.5 + 0.5 * np.cos(5.0 * u * v + u), ], axis=-1, ).astype(np.float32) chw = ((hwc - IMAGE_MEAN) / IMAGE_STD).transpose(2, 0, 1) return np.ascontiguousarray(chw, dtype=np.float32) def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--out", default="ref", type=pathlib.Path) ap.add_argument("--model", default=DEFAULT_MODEL) ap.add_argument("--size", type=int, default=224) ap.add_argument( "--layers", type=int, default=12, help="number of leading encoder layers to execute before the final norm", ) ap.add_argument("--image", default=None, help="optional real photo instead of a test pattern") args = ap.parse_args() if args.size % 16: raise SystemExit(f"--size {args.size} is not a multiple of the patch size (16)") args.out.mkdir(parents=True, exist_ok=True) model = AutoModel.from_pretrained(args.model, dtype=torch.float32).eval() cfg = model.config print( f"{args.model}: {cfg.num_hidden_layers} layers, hidden {cfg.hidden_size}, " f"{cfg.num_register_tokens} register tokens, gated_mlp={cfg.use_gated_mlp}" ) if cfg.use_gated_mlp: raise SystemExit("the Rust graph implements the plain MLP only (ViT-S/B)") if not 0 <= args.layers <= cfg.num_hidden_layers: raise SystemExit( f"--layers must be in [0, {cfg.num_hidden_layers}], got {args.layers}" ) chw = build_input(args.size, args.image) pixels = torch.from_numpy(chw)[None] with torch.no_grad(): hidden_states = model.embeddings(pixels) embeddings = hidden_states[0].numpy().astype(np.float32) position_embeddings = model.rope_embeddings(pixels) first_layer = model.model.layer[0] first_norm1 = first_layer.norm1(hidden_states) first_q = first_layer.attention.q_proj(first_norm1) first_k = first_layer.attention.k_proj(first_norm1) first_v = first_layer.attention.v_proj(first_norm1) tokens = first_q.shape[1] heads = cfg.num_attention_heads head_dim = cfg.hidden_size // heads q_heads = first_q.view(1, tokens, heads, head_dim).transpose(1, 2) k_heads = first_k.view(1, tokens, heads, head_dim).transpose(1, 2) v_heads = first_v.view(1, tokens, heads, head_dim).transpose(1, 2) q_rope, k_rope = apply_rotary_pos_emb( q_heads, k_heads, *position_embeddings ) first_attention, _ = eager_attention_forward( first_layer.attention, q_rope, k_rope, v_heads, None, scaling=first_layer.attention.scaling, ) first_attention = first_attention.reshape(1, tokens, cfg.hidden_size) first_attention_projected = first_layer.attention.o_proj(first_attention) first_attention_scaled = first_layer.layer_scale1( first_attention_projected ) first_residual = hidden_states + first_attention_scaled first_norm2 = first_layer.norm2(first_residual) first_mlp_up = first_layer.mlp.up_proj(first_norm2) first_mlp_activated = first_layer.mlp.act_fn(first_mlp_up) first_mlp_down = first_layer.mlp.down_proj(first_mlp_activated) first_mlp_scaled = first_layer.layer_scale2(first_mlp_down) first_output = first_residual + first_mlp_scaled first_final_norm = model.norm(first_output) for layer in model.model.layer[: args.layers]: hidden_states = layer( hidden_states, position_embeddings=position_embeddings ) features = model.norm(hidden_states)[0].numpy().astype(np.float32) grid = args.size // cfg.patch_size expected_tokens = 1 + cfg.num_register_tokens + grid * grid assert features.shape == (expected_tokens, cfg.hidden_size), features.shape (args.out / "pixel_values.bin").write_bytes(chw.tobytes()) (args.out / "embeddings.bin").write_bytes( np.ascontiguousarray(embeddings).tobytes() ) for name, tensor in { "first-norm1.bin": first_norm1, "first-q.bin": first_q, "first-k.bin": first_k, "first-v.bin": first_v, "first-q-rope.bin": q_rope.transpose(1, 2).reshape( 1, tokens, cfg.hidden_size ), "first-k-rope.bin": k_rope.transpose(1, 2).reshape( 1, tokens, cfg.hidden_size ), "first-attention.bin": first_attention, "first-attention-projected.bin": first_attention_projected, "first-attention-scaled.bin": first_attention_scaled, "first-residual.bin": first_residual, "first-norm2.bin": first_norm2, "first-mlp-up.bin": first_mlp_up, "first-mlp-activated.bin": first_mlp_activated, "first-mlp-down.bin": first_mlp_down, "first-mlp-scaled.bin": first_mlp_scaled, "first-output.bin": first_output, "first-final-norm.bin": first_final_norm, }.items(): (args.out / name).write_bytes( np.ascontiguousarray(tensor[0].numpy().astype(np.float32)).tobytes() ) (args.out / "features.bin").write_bytes(np.ascontiguousarray(features).tobytes()) # `save_file` rejects shared storage, which `state_dict()` can contain. state = {k: v.contiguous().clone() for k, v in model.state_dict().items()} save_file(state, str(args.out / "model.safetensors")) def sha256(path: pathlib.Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() local_model = pathlib.Path(args.model) reference = { "schema_version": 1, "base_model": DEFAULT_MODEL if local_model.is_dir() else args.model, "model_source": str(args.model), "image_size": args.size, "encoder_layers": args.layers, "tokens": expected_tokens, "hidden_size": cfg.hidden_size, "torch_version": torch.__version__, "transformers_version": transformers.__version__, "numpy_version": np.__version__, "pixel_values_sha256": sha256(args.out / "pixel_values.bin"), "embeddings_sha256": sha256(args.out / "embeddings.bin"), "features_sha256": sha256(args.out / "features.bin"), "exported_model_sha256": sha256(args.out / "model.safetensors"), } if local_model.is_dir() and (local_model / "model.safetensors").is_file(): reference["source_model_sha256"] = sha256( local_model / "model.safetensors" ) (args.out / "reference.json").write_text( json.dumps(reference, indent=2) + "\n", encoding="utf-8" ) print( f"wrote {args.out}/ — {args.layers} layers, {args.size}x{args.size}, " f"{grid}x{grid} grid, {expected_tokens} tokens" ) print(f"features: mean {features.mean():+.4f} std {features.std():.4f}") print(f"\nnow run: cargo run --release --example verify -- {args.out}") if __name__ == "__main__": main()