| """Frozen DINOv3-B behind a trained patch embedding, and the layer projection. |
| |
| Every block's patch tokens (under the final LayerNorm) are standardized per |
| channel with fixed statistics, concatenated (12 x 768) and mapped by one linear |
| projection to the latent channels. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import TYPE_CHECKING, cast |
|
|
| import timm |
| import torch |
| from timm.layers.pos_embed_sincos import RotaryEmbeddingDinoV3 |
| from torch import Tensor, nn |
|
|
| from .config import PATCH |
|
|
| if TYPE_CHECKING: |
| from timm.models.eva import Eva |
|
|
| from .config import Dinac3Config |
|
|
| IMAGENET_MEAN = (0.485, 0.456, 0.406) |
| IMAGENET_STD = (0.229, 0.224, 0.225) |
|
|
|
|
| def _require_supported_backbone(backbone: Eva) -> None: |
| """Fail fast if timm builds a backbone the explicit forward does not cover. |
| |
| Raises: |
| ValueError: On absolute position embeddings, patch dropout, a pre-norm, |
| missing CLS/register tokens or a RoPE other than DINOv3's without |
| coordinate augmentation. |
| """ |
|
|
| supported = ( |
| backbone.pos_embed is None |
| and backbone.patch_drop is None |
| and isinstance(backbone.norm_pre, nn.Identity) |
| and backbone.cls_token is not None |
| and backbone.reg_token is not None |
| and isinstance(backbone.rope, RotaryEmbeddingDinoV3) |
| and not backbone.rope.aug_active |
| ) |
| if not supported: |
| raise ValueError("Unsupported timm DINOv3 backbone structure for dinac3") |
|
|
|
|
| class Encoder(nn.Module): |
| """Images in [-1, 1] to raw latents ``[B, C, H / 16, W / 16]``.""" |
|
|
| pixel_mean: Tensor |
| pixel_std: Tensor |
| layer_mean: Tensor |
| layer_std: Tensor |
|
|
| def __init__(self, config: Dinac3Config) -> None: |
| """Build the architecture only; every tensor comes from the artifact.""" |
|
|
| super().__init__() |
| self.backbone = cast( |
| "Eva", |
| timm.create_model( |
| config.backbone.value, |
| pretrained=False, |
| num_classes=0, |
| dynamic_img_size=True, |
| dynamic_img_pad=True, |
| ), |
| ) |
| _require_supported_backbone(self.backbone) |
| |
| |
| rope = cast("nn.Module", self.backbone.rope) |
| rope.register_buffer("periods", cast("Tensor", rope.periods), persistent=True) |
| self.prefix_tokens = self.backbone.num_prefix_tokens |
| width = self.backbone.embed_dim * len(self.backbone.blocks) |
| self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1)) |
| self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1)) |
| self.register_buffer("layer_mean", torch.zeros(width)) |
| self.register_buffer("layer_std", torch.ones(width)) |
| self.projection = nn.Linear(width, config.latent_channels) |
| self._rope_cache: dict[tuple[int, int, torch.device, torch.dtype], Tensor] = {} |
|
|
| def rope_embed(self, height: int, width: int) -> Tensor: |
| """The DINOv3 RoPE table of an image size, built once per token grid |
| (outside any compiled graph: timm builds coordinates on the CPU).""" |
|
|
| rope = cast("RotaryEmbeddingDinoV3", self.backbone.rope) |
| periods = cast("Tensor", rope.periods) |
| key = (height // PATCH, width // PATCH, periods.device, periods.dtype) |
| cached = self._rope_cache.get(key) |
| if cached is None: |
| with torch.inference_mode(False), torch.no_grad(): |
| cached = rope.get_embed(shape=[key[0], key[1]]) |
| self._rope_cache[key] = cached |
| return cached |
|
|
| def forward(self, images: Tensor, rope: Tensor) -> Tensor: |
| """Raw latents in the activation dtype; ``rope`` from :meth:`rope_embed`.""" |
|
|
| dtype = self._prefix_tokens()[0].dtype |
| x = ((images.float().add(1.0).mul(0.5) - self.pixel_mean) / self.pixel_std).to( |
| dtype=dtype |
| ) |
| features = self._block_features(x, rope) |
| tokens = torch.cat([feature.to(dtype=dtype) for feature in features], dim=-1) |
| standardized = (tokens.float() - self.layer_mean) / self.layer_std |
| latents = self.projection(standardized) |
| b, _, height, width = images.shape |
| return latents.transpose(1, 2).reshape(b, -1, height // PATCH, width // PATCH) |
|
|
| def _block_features(self, x: Tensor, rope: Tensor) -> list[Tensor]: |
| """Final-normed patch tokens of every block (timm's intermediates path).""" |
|
|
| backbone = self.backbone |
| tokens = backbone.patch_embed(x) |
| b, _, _, c = tokens.shape |
| tokens = tokens.view(b, -1, c) |
| cls_token, reg_token = self._prefix_tokens() |
| tokens = torch.cat( |
| [cls_token.expand(b, -1, -1), reg_token.expand(b, -1, -1), tokens], dim=1 |
| ) |
| features: list[Tensor] = [] |
| for block in backbone.blocks: |
| tokens = block(tokens, rope=rope) |
| features.append(backbone.norm(tokens)[:, self.prefix_tokens :]) |
| return features |
|
|
| def _prefix_tokens(self) -> tuple[Tensor, Tensor]: |
| """The CLS and register tokens (present: checked at construction).""" |
|
|
| cls_token, reg_token = self.backbone.cls_token, self.backbone.reg_token |
| if cls_token is None or reg_token is None: |
| raise RuntimeError("The DINOv3 backbone lost its CLS or register tokens") |
| return cls_token, reg_token |
|
|
