environment/training-source/source-snapshot/tools/dump_reference.py

#!/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()