torch_em.model.vit

   1import math
   2from functools import partial
   3from collections import OrderedDict
   4from typing import Tuple, List, Optional
   5
   6import torch
   7import torch.nn as nn
   8import torch.nn.functional as F
   9
  10# we catch ImportErrors here because segment_anything, micro_sam, scale_mae and timm should
  11# only be optional dependencies for torch_em
  12try:
  13    from segment_anything.modeling import ImageEncoderViT
  14    _sam_import_success = True
  15except ImportError:
  16    ImageEncoderViT = object
  17    _sam_import_success = False
  18
  19try:
  20    from timm.models.vision_transformer import VisionTransformer, PatchEmbed
  21    _timm_import_success = True
  22except ImportError:
  23    VisionTransformer = object
  24    PatchEmbed = object
  25    _timm_import_success = False
  26
  27try:
  28    from sam2.modeling.backbones.hieradet import Hiera
  29    from sam2.modeling.position_encoding import PositionEmbeddingSine
  30    from sam2.modeling.backbones.image_encoder import ImageEncoder, FpnNeck
  31    _sam2_import_success = True
  32except ImportError:
  33    ImageEncoder = object
  34    _sam2_import_success = False
  35
  36try:
  37    from dinov2.layers.attention import Attention
  38    from dinov2.layers.block import NestedTensorBlock as Block
  39    from dinov2.layers.swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
  40    from dinov2.models.vision_transformer import DinoVisionTransformer as DinoV2VisionTransformer
  41    _dinov2_import_success = True
  42except ImportError:
  43    DinoV2VisionTransformer = object
  44    _dinov2_import_success = False
  45
  46try:
  47    from dinov3.models.vision_transformer import DinoVisionTransformer as DinoV3VisionTransformer
  48    _dinov3_import_success = True
  49except ImportError:
  50    DinoV3VisionTransformer = object
  51    _dinov3_import_success = False
  52
  53
  54try:
  55    from sam3.model.vitdet import ViT as SAM3ViT, get_abs_pos
  56    _sam3_import_success = True
  57except ImportError:
  58    SAM3ViT = object
  59    _sam3_import_success = False
  60
  61try:
  62    import torchvision.models as _tv_models
  63    _torchvision_import_success = True
  64except ImportError:
  65    _tv_models = None
  66    _torchvision_import_success = False
  67
  68from .util import VIT_CONFIGS
  69
  70
  71class ViT_Sam(ImageEncoderViT):
  72    """Vision Transformer derived from the Segment Anything Codebase (https://arxiv.org/abs/2304.02643).
  73
  74    Based on:
  75    https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/image_encoder.py
  76
  77    Args:
  78        in_chans: The number of input channels.
  79        embed_dim: The embedding dimension, corresponding to the number of output channels of the vision transformer.
  80        global_attn_indexes: The global attention indices.
  81        apply_neck: Whether to apply the convolutional bottleneck after outputs of the last attention head.
  82        kwargs: Keyword arguments for the image encoder base class.
  83    """
  84    def __init__(
  85        self,
  86        in_chans: int = 3,
  87        embed_dim: int = 768,
  88        global_attn_indexes: Tuple[int, ...] = [2, 5, 8, 11],
  89        apply_neck: bool = False,
  90        **kwargs,
  91    ) -> None:
  92        if not _sam_import_success:
  93            raise RuntimeError(
  94                "The vision transformer backend can only be initialized if segment anything is installed. "
  95                "Please install segment anything from https://github.com/facebookresearch/segment-anything "
  96                "and then rerun your code."
  97            )
  98
  99        super().__init__(embed_dim=embed_dim, global_attn_indexes=global_attn_indexes, **kwargs)
 100        self.chunks_for_projection = global_attn_indexes
 101        self.in_chans = in_chans
 102        self.embed_dim = embed_dim
 103        self.apply_neck = apply_neck
 104
 105    def forward(self, x: torch.Tensor) -> torch.Tensor:
 106        """Apply the vision transformer to input data.
 107
 108        Args:
 109            x: The input data.
 110
 111        Returns:
 112            The vision transformer output.
 113        """
 114        x = self.patch_embed(x)
 115        if self.pos_embed is not None:
 116            x = x + self.pos_embed
 117
 118        list_from_encoder = []
 119        for i, blk in enumerate(self.blocks):
 120            x = blk(x)
 121            if i in self.chunks_for_projection:
 122                list_from_encoder.append(x)
 123
 124        x = x.permute(0, 3, 1, 2)
 125
 126        if self.apply_neck:
 127            x = self.neck(x)
 128
 129        list_from_encoder = [e.permute(0, 3, 1, 2) for e in list_from_encoder]
 130        return x, list_from_encoder[:3]
 131
 132
 133class ViT_CellposeSAM(nn.Module):
 134    """Vision Transformer derived from the CellposeSAM Codebase (https://doi.org/10.1038/s41592-025-02595-x).
 135
 136    This replicates CellposeSAM's actual initialization: instantiate SAM's ``ImageEncoderViT`` via
 137    ``sam_model_registry``, then modify the patch embedding, position embeddings, and set global attention.
 138    This preserves SAM's original relative position bias sizes, enabling direct checkpoint loading
 139    without any interpolation.
 140
 141    Based on: https://github.com/MouseLand/cellpose/blob/main/cellpose/vit_sam.py
 142
 143    NOTE: The pretrained CellposeSAM model uses ``vit_l`` exclusively.
 144
 145    Args:
 146        ps: The patch size (default for CellposeSAM is 8).
 147        bsize: The input image size (default for CellposeSAM is 256).
 148        apply_neck: Whether to apply the convolutional bottleneck after outputs of the last attention head.
 149    """
 150    def __init__(self, ps: int = 8, bsize: int = 256, apply_neck: bool = True) -> None:
 151        super().__init__()
 152
 153        if not _sam_import_success:
 154            raise RuntimeError(
 155                "The vision transformer backend can only be initialized if segment anything is installed. "
 156                "Please install segment anything from https://github.com/facebookresearch/segment-anything "
 157                "and then rerun your code."
 158            )
 159
 160        from segment_anything import sam_model_registry
 161
 162        # Creates the SAM vit_l encoder and applies CellposeSAM's modifications (same as cellpose.vit_sam.Transformer).
 163        encoder = sam_model_registry["vit_l"](None).image_encoder
 164
 165        w = encoder.patch_embed.proj.weight.detach()
 166        nchan = w.shape[0]
 167
 168        # CellPoseSAM changes the patch size from 16 to 'ps'.
 169        encoder.patch_embed.proj = nn.Conv2d(3, nchan, stride=ps, kernel_size=ps)
 170        encoder.patch_embed.proj.weight.data = w[:, :, ::16 // ps, ::16 // ps]
 171
 172        # Next, they subsample position embeddings for the new patch size and input resolution.
 173        ds = (1024 // 16) // (bsize // ps)
 174        encoder.pos_embed = nn.Parameter(encoder.pos_embed[:, ::ds, ::ds], requires_grad=True)
 175
 176        # Finally, they set all blocks to global attention.
 177        for blk in encoder.blocks:
 178            blk.window_size = 0
 179
 180        # Store encoder submodules directly ('state_dict' keys match CellposeSAM after prefix stripping).
 181        self.patch_embed = encoder.patch_embed
 182        self.pos_embed = encoder.pos_embed
 183        self.blocks = encoder.blocks
 184        self.neck = encoder.neck
 185
 186        # Additional attributes expected by UNETR.
 187        self.embed_dim = nchan
 188        self.img_size = bsize
 189        self.in_chans = 3
 190        self.apply_neck = apply_neck
 191
 192        # Feature extraction at evenly-spaced depths.
 193        depth = len(self.blocks)
 194        _chunks = depth // 4
 195        self.chunks_for_projection = [_chunks - 1, 2 * _chunks - 1, 3 * _chunks - 1, 4 * _chunks - 1]
 196
 197    def forward(self, x: torch.Tensor) -> torch.Tensor:
 198        """Apply the vision transformer to input data.
 199
 200        Args:
 201            x: The input data.
 202
 203        Returns:
 204            The vision transformer output.
 205        """
 206        x = self.patch_embed(x)
 207        if self.pos_embed is not None:
 208            x = x + self.pos_embed
 209
 210        list_from_encoder = []
 211        for i, blk in enumerate(self.blocks):
 212            x = blk(x)
 213            if i in self.chunks_for_projection:
 214                list_from_encoder.append(x)
 215
 216        x = x.permute(0, 3, 1, 2)
 217
 218        if self.apply_neck:
 219            x = self.neck(x)
 220
 221        list_from_encoder = [e.permute(0, 3, 1, 2) for e in list_from_encoder]
 222        return x, list_from_encoder[:3]
 223
 224
 225class ViT_MAE(VisionTransformer):
 226    """Vision Transformer derived from the Masked Auto Encoder Codebase (https://arxiv.org/abs/2111.06377).
 227
 228    Based on:
 229    https://github.com/facebookresearch/mae/blob/main/models_vit.py#L20-L53
 230
 231    Args:
 232        img_size: The size of the input for the image encoder. Input images will be resized to match this size.
 233        in_chans: The number of input channels.
 234        depth: The depth of the vision transformer.
 235        kwargs: Additional keyword arguments for the vision transformer base class.
 236    """
 237    def __init__(
 238        self,
 239        img_size: int = 1024,  # chosen to match our experiments with segment anything
 240        in_chans: int = 3,
 241        depth: int = 12,
 242        **kwargs
 243    ):
 244        if not _timm_import_success:
 245            raise RuntimeError(
 246                "The vision transformer backend can only be initialized if timm is installed. "
 247                "Please install timm (using conda/mamba) for using https://github.com/facebookresearch/mae/ "
 248                "and then rerun your code"
 249            )
 250        kwargs["num_classes"] = 0
 251        super().__init__(img_size=img_size, in_chans=in_chans, depth=depth, **kwargs)
 252        self.img_size = img_size
 253        self.in_chans = in_chans
 254        self.depth = depth
 255
 256    def convert_to_expected_dim(self, inputs_):
 257        """@private
 258        """
 259        inputs_ = inputs_[:, 1:, :]  # removing the class tokens
 260        # reshape the outputs to desired shape (N x H*W X C -> N x H x W x C)
 261        rdim = inputs_.shape[1]
 262        dshape = int(rdim ** 0.5)  # finding the square root of the outputs for obtaining the patch shape
 263        inputs_ = torch.unflatten(inputs_, 1, (dshape, dshape))
 264        inputs_ = inputs_.permute(0, 3, 1, 2)
 265        return inputs_
 266
 267    def forward_features(self, x):
 268        """@private
 269        """
 270        B = x.shape[0]
 271        x = self.patch_embed(x)
 272
 273        cls_tokens = self.cls_token.expand(B, -1, -1)
 274        x = torch.cat((cls_tokens, x), dim=1)
 275
 276        x = x + self.pos_embed
 277        x = self.pos_drop(x)
 278
 279        # chunks obtained for getting the projections for conjuctions with upsampling blocks
 280        _chunks = int(self.depth / 4)
 281        chunks_for_projection = [_chunks - 1, 2*_chunks - 1, 3*_chunks - 1, 4*_chunks - 1]
 282
 283        list_from_encoder = []
 284        for i, blk in enumerate(self.blocks):
 285            x = blk(x)
 286            if i in chunks_for_projection:
 287                list_from_encoder.append(self.convert_to_expected_dim(x))
 288
 289        x = self.convert_to_expected_dim(self.norm(x))
 290        return x, list_from_encoder[:3]
 291
 292    def forward(self, x: torch.Tensor) -> torch.Tensor:
 293        """Apply the vision transformer to input data.
 294
 295        Args:
 296            x: The input data.
 297
 298        Returns:
 299            The vision transformer output.
 300        """
 301        x, list_from_encoder = self.forward_features(x)
 302        return x, list_from_encoder
 303
 304
 305class ViT_Sam2(ImageEncoder):
 306    """Vision Transformer derived from the Segment Anything 2 Codebase (https://arxiv.org/abs/2408.00714).
 307
 308    Based on https://github.com/facebookresearch/sam2/blob/main/sam2/modeling/backbones/image_encoder.py.
 309
 310    Args:
 311        backbone_channel_list: The channels throughout the entire backbone.
 312        embed_dim: The initial embedding dimension.
 313        num_heads: The initial number of heads.
 314        stages: The number of blocks per stage.
 315        global_att_blocks: The parameter to decide which blocks have global attention.
 316        window_pos_embed_bkg_spatial_size: The spatial size of windowed positional embedding.
 317        window_spec: The window size per stage, when not using global attention.
 318        scalp: The count of lowest resolution features to discard.
 319    """
 320    def __init__(
 321        self,
 322        backbone_channel_list: List[int],
 323        img_size: int = 1024,
 324        embed_dim: int = 96,
 325        num_heads: int = 1,
 326        stages: Tuple[int, ...] = (2, 3, 16, 3),
 327        global_att_blocks: Tuple[int, ...] = (12, 16, 20),
 328        window_pos_embed_bkg_spatial_size: Tuple[int, int] = (14, 14),
 329        window_spec: Tuple[int, ...] = (8, 4, 14, 7),
 330        scalp: int = 1,
 331        **kwargs
 332    ):
 333        if not _sam2_import_success:
 334            raise RuntimeError(
 335                "The vision transformer backend can only be initialized if segment anything 2 is installed. "
 336                "Please install segment anything 2 from https://github.com/facebookresearch/sam2 "
 337                "and then rerun your code"
 338            )
 339
 340        trunk = Hiera(
 341            embed_dim=embed_dim,
 342            num_heads=num_heads,
 343            stages=stages,
 344            global_att_blocks=global_att_blocks,
 345            window_pos_embed_bkg_spatial_size=window_pos_embed_bkg_spatial_size,
 346            window_spec=window_spec,
 347        )
 348        neck = FpnNeck(
 349            position_encoding=PositionEmbeddingSine(num_pos_feats=256),
 350            d_model=256,
 351            backbone_channel_list=backbone_channel_list,
 352            fpn_top_down_levels=[2, 3],
 353            fpn_interp_model="nearest",
 354        )
 355
 356        super().__init__(trunk=trunk, neck=neck, scalp=scalp, **kwargs)
 357        self.scalp = scalp
 358        self.embed_dim = neck.d_model
 359        self.img_size = img_size
 360
 361    def forward(self, x: torch.Tensor):
 362        # The forward pass throught the backbone.
 363        features, pos = self.neck(self.trunk(x))
 364        if self.scalp > 0:  # This discard the "n" lowest resolution features.
 365            features, pos = features[:-self.scalp], pos[:-self.scalp]
 366
 367        return features[-1], features
 368
 369
 370class ViT_Sam3(SAM3ViT):
 371    """Vision Transformer derived from the Segment Anything 3 Codebase (https://arxiv.org/abs/2511.16719).
 372
 373    Based on https://github.com/facebookresearch/sam3/blob/main/sam3/model/vitdet.py.
 374
 375    Args:
 376        img_size: The input image size.
 377        embed_dim: The embedding dimension, corresponding to the number of output channels of the vision transformer.
 378        kwargs: Keyword arguments for the image encoder base class.
 379    """
 380    def __init__(self, img_size: int = 1024, embed_dim: int = 768, **kwargs):
 381        if not _sam3_import_success:
 382            raise RuntimeError(
 383                "The vision transformer backend can only be initialized if segment anything 3 is installed. "
 384                "Please install segment anything 3 from https://github.com/facebookresearch/sam3 "
 385                "and then rerun your code"
 386            )
 387
 388        super().__init__(img_size=img_size, embed_dim=embed_dim, **kwargs)
 389        self.img_size = img_size
 390        self.embed_dim = embed_dim
 391        # Native layers preserve autograd and the input dtype.
 392        for block in self.blocks:
 393            block.mlp = nn.Sequential(OrderedDict(block.mlp.named_children()))
 394
 395    def _load_from_state_dict(
 396        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs,
 397    ):
 398        # Rotary buffers depend on the input size, so keep the constructed values.
 399        for name, buffer in self.named_buffers():
 400            if name.endswith(("attn.freqs_cis", "attn.freqs_cis_real", "attn.freqs_cis_imag")):
 401                state_dict[prefix + name] = buffer
 402        super()._load_from_state_dict(
 403            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs,
 404        )
 405
 406    def forward_features(self, x):
 407        """@private
 408        """
 409        x = self.patch_embed(x)
 410        h, w = x.shape[1], x.shape[2]
 411
 412        s = 0
 413        if self.retain_cls_token:
 414            # If the 'cls_token' is retained, we don't maintain the spatial shape.
 415            x = torch.cat([self.class_embedding, x.flatten(1, 2)], dim=1)
 416            s = 1
 417
 418        if self.pos_embed is not None:
 419            x = x + get_abs_pos(
 420                self.pos_embed, self.pretrain_use_cls_token, (h, w), self.retain_cls_token, tiling=self.tile_abs_pos,
 421            )
 422
 423        x = self.ln_pre(x)
 424
 425        list_from_encoder = []
 426        for i, blk in enumerate(self.blocks):
 427            if self.use_act_checkpoint and self.training:
 428                x = torch.utils.checkpoint.checkpoint(blk, x, use_reentrant=False)
 429            else:
 430                x = blk(x)
 431
 432            x = self._convert_to_expected_dim(x, i, s)
 433
 434            if i in self.full_attn_ids:
 435                list_from_encoder.append(x)
 436
 437        return x, list_from_encoder
 438
 439    def _convert_to_expected_dim(self, x, i, s):
 440        if (i == self.full_attn_ids[-1]) or (
 441            self.return_interm_layers and i in self.full_attn_ids
 442        ):
 443            if i == self.full_attn_ids[-1]:
 444                x = self.ln_post(x)
 445
 446            feats = x[:, s:]
 447            if feats.ndim == 4:
 448                feats = feats.permute(0, 3, 1, 2)
 449            else:
 450                assert feats.ndim == 3
 451                h = w = math.sqrt(feats.shape[1])
 452                feats = feats.reshape(feats.shape[0], h, w, feats.shape[-1]).permute(0, 3, 1, 2)
 453            return feats
 454
 455        else:
 456            return x
 457
 458    def forward(self, x: torch.Tensor):
 459        """Apply the vision transformer to input data.
 460
 461        Args:
 462            x: The input data.
 463
 464        Returns:
 465            The vision transformer output.
 466        """
 467        x, list_from_encoder = self.forward_features(x)
 468        return x, list_from_encoder
 469
 470#
 471# Utilities for ScaleMAE's ViT
 472#
 473
 474
 475class CustomCompose:
 476    def __init__(self, rescale_transform, other_transforms, src_transform):
 477        self.rescale_transform = rescale_transform
 478        self.other_transforms = other_transforms
 479        self.src_transform = src_transform
 480
 481    def __call__(self, x, valid_masks=None):
 482        if valid_masks is not None:
 483            nodata = (x * (1 - valid_masks.float())).max()
 484        x_aug = self.rescale_transform(x)
 485        parms = self.rescale_transform._params
 486
 487        # sanity check, comment if this is working
 488        # valid_masks = self.rescale_transform(valid_masks.float(), params=parms)
 489        # assert (x_aug==self.rescale_transform(x, params=parms)).all() #
 490
 491        if valid_masks is not None:
 492            valid_masks = x_aug != nodata
 493            _, c, h, w = x_aug.shape
 494            zero_ratio = ((valid_masks == 0).sum((1, 2, 3)) / (h * w * c)).cpu().numpy()
 495        else:
 496            zero_ratio = -1
 497
 498        if self.other_transforms:
 499            x_aug = self.other_transforms(x_aug)
 500        x_src = self.src_transform(x_aug)
 501        dx = parms["src"][:, 1, 0] - parms["src"][:, 0, 0]
 502
 503        # dy = (parms['src'][:,2,1] - parms['src'][:,1,1])
 504        # assert (dx == dy).all()
 505
 506        h, w = x_aug.shape[-2:]
 507        # assert h == w
 508
 509        return x_aug, x_src, dx / h, zero_ratio, valid_masks
 510
 511
 512def get_2d_sincos_pos_embed_with_resolution(embed_dim, grid_size, res, cls_token=False, device="cpu"):
 513    """
 514    grid_size: int of the grid height and width
 515    res: array of size n, representing the resolution of a pixel (say, in meters),
 516    return:
 517    pos_embed: [n,grid_size*grid_size, embed_dim] or [n,1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
 518    """
 519    # res = torch.FloatTensor(res).to(device)
 520    res = res.to(device)
 521    grid_h = torch.arange(grid_size, dtype=torch.float32, device=device)
 522    grid_w = torch.arange(grid_size, dtype=torch.float32, device=device)
 523    grid = torch.meshgrid(grid_w, grid_h, indexing="xy")  # here h goes first,direction reversed for numpy
 524    grid = torch.stack(grid, dim=0)  # 2 x h x w
 525
 526    # grid = grid.reshape([2, 1, grid_size, grid_size])
 527    grid = torch.einsum("chw,n->cnhw", grid, res)  # 2 x n x h x w
 528    _, n, h, w = grid.shape
 529    pos_embed = get_2d_sincos_pos_embed_from_grid_torch(embed_dim, grid)  # (nxH*W, D/2)
 530    pos_embed = pos_embed.reshape(n, h * w, embed_dim)
 531    if cls_token:
 532        pos_embed = torch.cat(
 533            [torch.zeros([n, 1, embed_dim], dtype=torch.float32, device=pos_embed.device), pos_embed], dim=1
 534        )
 535
 536    return pos_embed
 537
 538
 539def get_2d_sincos_pos_embed_from_grid_torch(embed_dim, grid):
 540    assert embed_dim % 2 == 0
 541
 542    # use half of dimensions to encode grid_h
 543    emb_h = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid[0])  # (H*W, D/2)
 544    emb_w = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid[1])  # (H*W, D/2)
 545
 546    emb = torch.cat([emb_h, emb_w], dim=1)  # (H*W, D)
 547    return emb
 548
 549
 550def get_1d_sincos_pos_embed_from_grid_torch(embed_dim, pos):
 551    """
 552    embed_dim: output dimension for each position
 553    pos: a list of positions to be encoded: size (M,)
 554    out: (M, D)
 555    """
 556    assert embed_dim % 2 == 0
 557    # old_shape = pos
 558    omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=pos.device)
 559    omega /= embed_dim / 2.0
 560    omega = 1.0 / 10000**omega  # (D/2,)
 561
 562    pos = pos.reshape(-1)  # (M,)
 563    out = torch.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
 564
 565    emb_sin = torch.sin(out)  # (M, D/2)
 566    emb_cos = torch.cos(out)  # (M, D/2)
 567
 568    emb = torch.cat([emb_sin, emb_cos], dim=1)  # (M, D)
 569    return emb
 570
 571
 572class PatchEmbedUnSafe(PatchEmbed):
 573    """Image to Patch Embedding"""
 574
 575    def forward(self, x):
 576        B, C, H, W = x.shape
 577
 578        # NOTE: Comment code from ScaleMAE: Dropped size check in timm
 579        # assert H == self.img_size[0] and W == self.img_size[1], \
 580        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
 581
 582        x = self.proj(x).flatten(2).transpose(1, 2)
 583        return x
 584
 585
 586class ViT_ScaleMAE(VisionTransformer):
 587    """Vision Transformer dervied from the Scale Masked Auto Encoder codebase (TODO: paper and github link).
 588
 589    NOTE: For downstream tasks, the "base_resoulution" parameter needs to be adjusted manually when using
 590    the model on a different zoom factor dataset.
 591    """
 592
 593    def __init__(
 594        self, img_size=224, patch_size=16, in_chans=3, embed_dim=1024, depth=12, base_resolution=2.5, **kwargs
 595    ):
 596        if not _timm_import_success:
 597            raise RuntimeError("The ScaleMAE backend requires timm.")
 598        kwargs["num_classes"] = 0
 599        super().__init__(
 600            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, depth=depth, **kwargs
 601        )
 602        self.img_size = img_size
 603        self.in_chans = in_chans
 604        self.depth = depth
 605        self.base_resolution = base_resolution
 606        self.pos_embed = None
 607
 608        self.patch_embed = PatchEmbedUnSafe(
 609            img_size=img_size,
 610            patch_size=patch_size,
 611            in_chans=in_chans,
 612            embed_dim=embed_dim,
 613        )
 614
 615    def transform_inputs(self, x):
 616        import kornia.augmentation as K
 617        from kornia.constants import Resample
 618
 619        self._transforms = CustomCompose(
 620            rescale_transform=K.RandomResizedCrop(
 621                (2 * self.img_size, 2 * self.img_size),
 622                ratio=(1.0, 1.0),
 623                scale=(1.0, 1.0),
 624                resample=Resample.BICUBIC.name,
 625            ),
 626            other_transforms=None,
 627            src_transform=K.Resize((self.img_size, self.img_size)),
 628        )
 629        x, _, ratios, _, _ = self._transforms(x)
 630        input_res = ratios * self.base_resolution
 631        return x, input_res
 632
 633    def convert_to_expected_dim(self, x):
 634        inputs_ = x[:, 1:, :]  # removing the class tokens
 635        # reshape the outputs to desired shape (N X H*W X C -> N X H X W X C)
 636        rdim = inputs_.shape[1]
 637        dshape = int(rdim ** 0.5)  # finding square root of the outputs for obtaining the patch shape
 638        inputs_ = torch.unflatten(inputs_, 1, (dshape, dshape))
 639        inputs_ = inputs_.permute(0, 3, 1, 2)
 640        return inputs_
 641
 642    def forward_features(self, x):
 643        x, input_res = self.transform_inputs(x)
 644
 645        B, _, h, w = x.shape
 646        x = self.patch_embed(x)
 647
 648        num_patches = int((h * w) / (self.patch_embed.patch_size[0] * self.patch_embed.patch_size[1]))
 649        pos_embed = get_2d_sincos_pos_embed_with_resolution(
 650            x.shape[-1],
 651            int(num_patches ** 0.5),
 652            input_res,
 653            cls_token=True,
 654            device=x.device,
 655        )
 656
 657        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
 658        x = torch.cat((cls_tokens, x), dim=1)
 659        x = x + pos_embed
 660        x = self.pos_drop(x)
 661
 662        # chunks obtained for getting the projections for conjuctions with upsampling blocks
 663        _chunks = int(self.depth / 4)
 664        chunks_for_projection = [_chunks - 1, 2*_chunks - 1, 3*_chunks - 1, 4*_chunks - 1]
 665
 666        list_from_encoder = []
 667        for i, blk in enumerate(self.blocks):
 668            x = blk(x)
 669            if i in chunks_for_projection:
 670                list_from_encoder.append(self.convert_to_expected_dim(x))
 671
 672        x = self.convert_to_expected_dim(self.norm(x))
 673
 674        return x, list_from_encoder
 675
 676    def forward(self, x):
 677        x, list_from_encoder = self.forward_features(x)
 678        return x, list_from_encoder
 679
 680
 681def _get_dinov2_swiglu(in_features, hidden_features, **kwargs):
 682    hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
 683    return SwiGLUFFN(in_features=in_features, hidden_features=hidden_features, **kwargs)
 684
 685
 686def _get_dinov2_block(*args, ffn_layer, **kwargs):
 687    if ffn_layer is SwiGLUFFNFused:
 688        ffn_layer = _get_dinov2_swiglu
 689    return Block(*args, attn_class=Attention, ffn_layer=ffn_layer, **kwargs)
 690
 691
 692class ViT_DINOv2(DinoV2VisionTransformer):
 693    """Vision Transformer derived from the DINOv2 Codebase (https://arxiv.org/abs/2304.07193).
 694
 695    Based on:
 696    https://github.com/facebookresearch/dinov2/blob/main/dinov2/models/vision_transformer.py.
 697
 698    Args:
 699        img_size: The input image size.
 700        patch_size: The patch size.
 701        depth: The depth of the network.
 702        num_register_tokens: The number of registers added (in addition to the class tokens).
 703            It's important to know for ViTs trained with registers, to remove them at the end.
 704    """
 705    def __init__(
 706        self,
 707        img_size: int = 224,
 708        patch_size: int = 16,
 709        depth: int = 12,
 710        num_register_tokens: int = 0,
 711        **kwargs
 712    ):
 713        if not _dinov2_import_success:
 714            raise RuntimeError(
 715                "The vision transformer backend can only be initialized if DINOv2 is installed. "
 716                "Please install DINOv2 from https://github.com/facebookresearch/dinov2 "
 717                "and then rerun your code."
 718            )
 719
 720        super().__init__(
 721            img_size=img_size,
 722            depth=depth,
 723            patch_size=patch_size,
 724            num_register_tokens=num_register_tokens,
 725            **kwargs
 726        )
 727
 728        self.img_size = img_size
 729        self.num_register_tokens = num_register_tokens
 730        self.patch_size = patch_size
 731        self.attn_outs = [i for i in range(depth) if i % 3 == 2]
 732
 733    def forward(self, x, masks=None) -> torch.Tensor:
 734
 735        B = x.shape[0]
 736
 737        x = self.prepare_tokens_with_masks(x)
 738
 739        list_of_encoder = []
 740        for i, blk in enumerate(self.blocks):
 741            x = blk(x)
 742            if i in self.attn_outs:
 743                list_of_encoder.append(x)
 744
 745        x = self.norm(x)
 746        x = x[:, self.num_register_tokens + 1:].reshape(
 747            B, self.img_size // self.patch_size, self.img_size // self.patch_size, -1
 748        ).permute(0, 3, 1, 2).contiguous()
 749
 750        list_of_encoder = [
 751            o[:, self.num_register_tokens + 1:].reshape(
 752                B, self.img_size // self.patch_size, self.img_size // self.patch_size, -1
 753            ).permute(0, 3, 1, 2).contiguous() for o in list_of_encoder
 754        ]
 755
 756        return x, list_of_encoder[:3]
 757
 758
 759class ViT_DINOv3(DinoV3VisionTransformer):
 760    """Vision Transformer derived from the DINOv3 Codebase (https://arxiv.org/abs/2508.10104).
 761
 762    Based on:
 763    https://github.com/facebookresearch/dinov3/blob/main/dinov3/models/vision_transformer.py.
 764
 765    Args:
 766        img_size: The input image size.
 767        patch_size: The patch size.
 768        embed_dim: The embedding dimension.
 769        depth: The depth of the network.
 770        num_heads: The number of heads.
 771        ffn_ratio: The FFN rato.
 772        n_storage_tokens: The number of storage (class) tokens to remove.
 773        kwargs: Keyword arguments for the image encoder base class.
 774    """
 775    def __init__(
 776        self,
 777        in_chans: int = 3,
 778        img_size: int = 224,
 779        patch_size: int = 16,
 780        embed_dim: int = 768,
 781        depth: int = 12,
 782        num_heads: int = 12,
 783        ffn_ratio: float = 4.0,
 784        n_storage_tokens: int = 0,
 785        **kwargs
 786    ):
 787        if not _dinov3_import_success:
 788            raise RuntimeError(
 789                "The vision transformer backend can only be initialized if DINOv3 is installed. "
 790                "Please install DINOv3 from https://github.com/facebookresearch/dinov3 "
 791                "and then rerun your code."
 792            )
 793
 794        super().__init__(
 795            in_chans=in_chans,
 796            img_size=img_size,
 797            patch_size=patch_size,
 798            embed_dim=embed_dim,
 799            depth=depth,
 800            num_heads=num_heads,
 801            ffn_ratio=ffn_ratio,
 802            n_storage_tokens=n_storage_tokens,
 803            **kwargs
 804        )
 805
 806        self.in_chans = in_chans
 807        self.img_size = img_size
 808        self.n_storage_tokens = n_storage_tokens
 809        self.attn_outs = [i for i in range(depth) if i % 3 == 2]
 810        self.init_weights()
 811
 812    def forward(self, x) -> torch.Tensor:
 813
 814        B = x.shape[0]
 815
 816        x, hw_tuple = self.prepare_tokens_with_masks(x)
 817
 818        list_of_encoder = []
 819        for i, blk in enumerate(self.blocks):
 820            rope_sincos = self.rope_embed(H=hw_tuple[0], W=hw_tuple[1])
 821            x = blk(x, rope_sincos)
 822            if i in self.attn_outs:
 823                list_of_encoder.append(x)
 824
 825        x = self.norm(x)
 826        x = x[:, self.n_storage_tokens + 1:].reshape(
 827            B, self.img_size // self.patch_size, self.img_size // self.patch_size, -1
 828        ).permute(0, 3, 1, 2).contiguous()
 829
 830        list_of_encoder = [
 831            o[:, self.n_storage_tokens + 1:].reshape(
 832                B, self.img_size // self.patch_size, self.img_size // self.patch_size, -1
 833            ).permute(0, 3, 1, 2).contiguous() for o in list_of_encoder
 834        ]
 835
 836        return x, list_of_encoder[:3]
 837
 838
 839class ViT_Torchvision(nn.Module):
 840    """Vision Transformer from torchvision (https://arxiv.org/abs/2010.11929).
 841
 842    Wraps torchvision ViT models for use as UNETR encoders. Intermediate patch-token grids
 843    are collected at quarter-depth intervals and returned as spatial feature maps alongside
 844    the final encoder output, matching the `(x, list_from_encoder)` contract of other ViT classes.
 845
 846    Supported models: vit_b_16, vit_b_32, vit_l_16, vit_l_32, vit_h_14.
 847
 848    All five models can be used with UNETR. vit_b_16 and vit_l_16 (patch_size=16) work with
 849    any skip-connection setting. vit_b_32, vit_l_32, and vit_h_14 require use_skip_connection=False
 850    in UNETR; the decoder's internal cropping handles the spatial size difference, and postprocess_masks
 851    resizes the output back to the input resolution.
 852
 853    Args:
 854        model_name: Torchvision ViT model name (e.g. 'vit_b_16').
 855        img_size: Expected input image size used by UNETR preprocessing.
 856        in_chans: Number of input channels. If != 3, a 1x1 conv projects to 3 channels.
 857        pretrained: Whether to load ImageNet-pretrained weights.
 858    """
 859    def __init__(
 860        self,
 861        model_name: str,
 862        img_size: int = 224,
 863        in_chans: int = 3,
 864        pretrained: bool = True,
 865    ):
 866        super().__init__()
 867        if not _torchvision_import_success:
 868            raise RuntimeError(
 869                "The vision transformer backend can only be initialized if torchvision is installed. "
 870                "Please install torchvision from https://github.com/pytorch/vision and then rerun your code."
 871            )
 872
 873        fn = getattr(_tv_models, model_name)
 874        backbone = fn(weights="DEFAULT" if pretrained else None)
 875
 876        self.conv_proj = backbone.conv_proj
 877        self.class_token = backbone.class_token
 878        self.encoder = backbone.encoder  # pos_embedding, dropout, layers, ln
 879
 880        self.img_size = img_size
 881        self.in_chans = in_chans
 882        self.embed_dim = backbone.hidden_dim
 883
 884        depth = len(backbone.encoder.layers)
 885        _c = depth // 4
 886        self.chunks_for_projection = [_c - 1, 2 * _c - 1, 3 * _c - 1]
 887
 888        self.input_proj = nn.Conv2d(in_chans, 3, kernel_size=1) if in_chans != 3 else None
 889
 890    def _load_from_state_dict(
 891        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs,
 892    ):
 893        pos_embed = state_dict.get(prefix + "encoder.pos_embedding")
 894        current = self.encoder.pos_embedding
 895        if (
 896            pos_embed is not None and pos_embed.ndim == 3
 897            and pos_embed.shape[0] == current.shape[0] and pos_embed.shape[2] == current.shape[2]
 898            and pos_embed.shape[1] != current.shape[1]
 899        ):
 900            self.encoder.pos_embedding = nn.Parameter(
 901                current.new_empty(pos_embed.shape), requires_grad=current.requires_grad,
 902            )
 903        super()._load_from_state_dict(
 904            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs,
 905        )
 906
 907    def _interpolate_pos_embed(self, pos_embed: torch.Tensor, H_p: int, W_p: int) -> torch.Tensor:
 908        cls_pos, patch_pos = pos_embed[:, :1], pos_embed[:, 1:]
 909        N = patch_pos.shape[1]
 910        H_t = W_t = int(N ** 0.5)
 911        patch_pos = patch_pos.reshape(1, H_t, W_t, -1).permute(0, 3, 1, 2)
 912        patch_pos = F.interpolate(patch_pos, size=(H_p, W_p), mode="bicubic", align_corners=False)
 913        patch_pos = patch_pos.permute(0, 2, 3, 1).reshape(1, H_p * W_p, -1)
 914        return torch.cat([cls_pos, patch_pos], dim=1)
 915
 916    def forward(self, x: torch.Tensor) -> torch.Tensor:
 917        """Apply the vision transformer to input data.
 918
 919        Args:
 920            x: The input data.
 921
 922        Returns:
 923            The vision transformer output.
 924        """
 925        if self.input_proj is not None:
 926            x = self.input_proj(x)
 927
 928        x = self.conv_proj(x)  # (B, D, H_p, W_p)
 929        B, D, H_p, W_p = x.shape
 930        x = x.reshape(B, D, H_p * W_p).permute(0, 2, 1)  # (B, N, D)
 931
 932        cls = self.class_token.expand(B, -1, -1)
 933        x = torch.cat([cls, x], dim=1)  # (B, 1+N, D)
 934
 935        pos = self.encoder.pos_embedding
 936        if pos.shape[1] != x.shape[1]:
 937            pos = self._interpolate_pos_embed(pos, H_p, W_p)
 938        x = x + pos
 939        x = self.encoder.dropout(x)
 940
 941        list_from_encoder = []
 942        for i, blk in enumerate(self.encoder.layers):
 943            x = blk(x)
 944            if i in self.chunks_for_projection:
 945                feat = x[:, 1:].reshape(B, H_p, W_p, D).permute(0, 3, 1, 2).contiguous()
 946                list_from_encoder.append(feat)
 947
 948        x = self.encoder.ln(x)
 949        x = x[:, 1:].reshape(B, H_p, W_p, D).permute(0, 3, 1, 2).contiguous()
 950        return x, list_from_encoder
 951
 952
 953def get_vision_transformer(backbone: str, model: str, img_size: Optional[int] = None, **kwargs) -> nn.Module:
 954    """Get vision transformer encoder.
 955
 956    Args:
 957        backbone: The name of the vision transformer implementation.
 958            One of "sam" / "cellpose_sam" / "sam2" / "sam3" / "mae" / "scalemae" / "dinov2" / "dinov3" / "torchvision".
 959        model: The model name supported by the selected backbone.
 960        img_size: The encoder input size. If None, use the registered default for the backbone.
 961        kwargs: Additional kwargs which can be expected by the vision transformer,
 962            e.g. 'base_resolution' for `ViT_ScaleMAE`.
 963
 964    Returns:
 965        The vision transformer.
 966    """
 967    if backbone not in VIT_CONFIGS:
 968        raise ValueError(f"Unsupported backbone '{backbone}'.")
 969    if model not in VIT_CONFIGS[backbone].models:
 970        raise ValueError(f"Unsupported encoder '{model}' for backbone '{backbone}'.")
 971    if img_size is None:
 972        img_size = VIT_CONFIGS[backbone].img_size
 973
 974    if backbone == "sam":
 975        if model == "vit_b":
 976            encoder = ViT_Sam(
 977                depth=12, embed_dim=768, img_size=img_size, mlp_ratio=4,
 978                norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
 979                num_heads=12, patch_size=16, qkv_bias=True, use_rel_pos=True,
 980                global_attn_indexes=[2, 5, 8, 11],
 981                window_size=14, out_chans=256,
 982            )
 983        elif model == "vit_l":
 984            encoder = ViT_Sam(
 985                depth=24, embed_dim=1024, img_size=img_size, mlp_ratio=4,
 986                norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
 987                num_heads=16, patch_size=16, qkv_bias=True, use_rel_pos=True,
 988                global_attn_indexes=[5, 11, 17, 23],
 989                window_size=14, out_chans=256,
 990            )
 991        elif model == "vit_h":
 992            encoder = ViT_Sam(
 993                depth=32, embed_dim=1280, img_size=img_size, mlp_ratio=4,
 994                norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
 995                num_heads=16, patch_size=16, qkv_bias=True, use_rel_pos=True,
 996                global_attn_indexes=[7, 15, 23, 31],
 997                window_size=14, out_chans=256,
 998            )
 999        else:
1000            raise ValueError(f"'{model}' is not supported by SAM. Currently, 'vit_b', 'vit_l', 'vit_h' are supported.")
1001
1002    elif backbone == "cellpose_sam":
1003        if model != "vit_l":
1004            raise ValueError(f"'{model}' is not supported by CellposeSAM. Only 'vit_l' is supported.")
1005        encoder = ViT_CellposeSAM(ps=8, bsize=img_size)
1006
1007    elif backbone == "sam2":
1008        if model == "hvit_t":
1009            encoder = ViT_Sam2(
1010                img_size=img_size, embed_dim=96, num_heads=1, stages=[1, 2, 7, 2], global_att_blocks=[5, 7, 9],
1011                window_pos_embed_bkg_spatial_size=[7, 7], backbone_channel_list=[768, 384, 192, 96],
1012            )
1013        elif model == "hvit_s":
1014            encoder = ViT_Sam2(
1015                img_size=img_size, embed_dim=96, num_heads=1, stages=[1, 2, 11, 2], global_att_blocks=[7, 10, 13],
1016                window_pos_embed_bkg_spatial_size=[7, 7], backbone_channel_list=[768, 384, 192, 96],
1017            )
1018        elif model == "hvit_b":
1019            encoder = ViT_Sam2(
1020                img_size=img_size, embed_dim=112, num_heads=2, backbone_channel_list=[896, 448, 224, 112],
1021            )
1022        elif model == "hvit_l":
1023            encoder = ViT_Sam2(
1024                img_size=img_size, embed_dim=144, num_heads=2, stages=[2, 6, 36, 4], global_att_blocks=[23, 33, 43],
1025                window_spec=[8, 4, 16, 8], window_pos_embed_bkg_spatial_size=[7, 7],
1026                backbone_channel_list=[1152, 576, 288, 144],
1027            )
1028        else:
1029            raise ValueError(
1030                f"'{model}' is not supported by SAM2. Currently, 'hvit_t', 'hvit_s', 'hvit_b', 'hvit_l' are supported."
1031            )
1032
1033    elif backbone == "sam3":
1034        if model != "vit_pe":
1035            raise ValueError(
1036                "'sam3' does not have multiple model configurations. Please use 'vit_pe' as the model configuration."
1037            )
1038
1039        encoder = ViT_Sam3(
1040            img_size=img_size, pretrain_img_size=336, patch_size=14, embed_dim=1024, depth=32, num_heads=16,
1041            mlp_ratio=4.625, norm_layer="LayerNorm", drop_path_rate=0.1, qkv_bias=True, use_abs_pos=True,
1042            tile_abs_pos=True, global_att_blocks=(7, 15, 23, 31), rel_pos_blocks=(), use_rope=True,
1043            use_interp_rope=True, window_size=24, pretrain_use_cls_token=True, retain_cls_token=False, ln_pre=True,
1044            ln_post=False, return_interm_layers=False, bias_patch_embed=False, compile_mode=None,
1045        )
1046
1047    elif backbone == "mae":
1048        if model == "vit_b":
1049            encoder = ViT_MAE(
1050                img_size=img_size, patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
1051                qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6)
1052            )
1053        elif model == "vit_l":
1054            encoder = ViT_MAE(
1055                img_size=img_size, patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4,
1056                qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6)
1057            )
1058        elif model == "vit_h":
1059            encoder = ViT_MAE(
1060                img_size=img_size, patch_size=14, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4,
1061                qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6)
1062            )
1063        else:
1064            raise ValueError(f"'{model}' is not supported by MAE. Currently, 'vit_b', 'vit_l', 'vit_h' are supported.")
1065
1066    elif backbone == "scalemae":
1067        base_resolution = kwargs.get("base_resolution", 2.5)
1068
1069        if model == "vit_b":
1070            encoder = ViT_ScaleMAE(
1071                img_size=img_size, patch_size=16, embed_dim=768, depth=12, num_heads=12,
1072                mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
1073                base_resolution=base_resolution,
1074            )
1075        elif model == "vit_l":
1076            encoder = ViT_ScaleMAE(
1077                img_size=img_size, patch_size=16, embed_dim=1024, depth=24, num_heads=16,
1078                mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
1079                base_resolution=base_resolution,
1080            )
1081        elif model == "vit_h":
1082            encoder = ViT_ScaleMAE(
1083                img_size=img_size, patch_size=14, embed_dim=1280, depth=32, num_heads=16,
1084                mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
1085                base_resolution=base_resolution,
1086            )
1087        else:
1088            raise ValueError(
1089                f"'{model}' is not supported by ScaleMAE. Currently, 'vit_b', 'vit_l' and 'vit_h' are supported."
1090            )
1091
1092    elif backbone == "dinov2":
1093        block_fn = _get_dinov2_block
1094        msg = "The model name should be either 'vit_<X>' or 'vit_<X>_reg<Y>."
1095
1096        if model.startswith("vit_s"):
1097            assert model in ["vit_s", "vit_s_reg4"], msg
1098            encoder = ViT_DINOv2(
1099                img_size=img_size, patch_size=14, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
1100                block_fn=block_fn, in_chans=3, init_values=1e-5, block_chunks=0,
1101                num_register_tokens=4 if model.endswith("_reg4") else 0,
1102            )
1103        elif model.startswith("vit_b"):
1104            assert model in ["vit_b", "vit_b_reg4"], msg
1105            encoder = ViT_DINOv2(
1106                img_size=img_size, patch_size=14, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
1107                block_fn=block_fn, in_chans=3, init_values=1e-5, block_chunks=0,
1108                num_register_tokens=4 if model.endswith("_reg4") else 0,
1109            )
1110        elif model.startswith("vit_l"):
1111            assert model in ["vit_l", "vit_l_reg4"], msg
1112            encoder = ViT_DINOv2(
1113                img_size=img_size, patch_size=14, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4,
1114                block_fn=block_fn, in_chans=3, init_values=1e-5, block_chunks=0,
1115                num_register_tokens=4 if model.endswith("_reg4") else 0,
1116            )
1117        elif model.startswith("vit_g"):
1118            assert model in ["vit_g", "vit_g_reg4"], msg
1119            encoder = ViT_DINOv2(
1120                img_size=img_size, patch_size=14, embed_dim=1536, depth=40, num_heads=24, mlp_ratio=4,
1121                block_fn=block_fn, in_chans=3, init_values=1e-5, block_chunks=0,
1122                num_register_tokens=4 if model.endswith("_reg4") else 0, ffn_layer="swiglu",
1123            )
1124        else:
1125            raise ValueError(
1126                f"'{model}' is not supported by DINOv2. Currently, 'vit_s', 'vit_b', 'vit_l' and 'vit_g' are supported."
1127            )
1128
1129    elif backbone == "dinov3":
1130
1131        if model == "vit_s":
1132            encoder = ViT_DINOv3(
1133                img_size=img_size, pos_embed_rope_rescale_coords=2, pos_embed_rope_dtype="fp32", embed_dim=384,
1134                num_heads=6, layerscale_init=1.0e-05, norm_layer="layernormbf16", n_storage_tokens=4, mask_k_bias=True,
1135            )
1136        elif model == "vit_s+":
1137            encoder = ViT_DINOv3(
1138                img_size=img_size, pos_embed_rope_rescale_coords=2, pos_embed_rope_dtype="fp32", embed_dim=384,
1139                num_heads=6, ffn_ratio=6, layerscale_init=1.0e-05, norm_layer="layernormbf16",
1140                ffn_layer="swiglu", n_storage_tokens=4, mask_k_bias=True,
1141            )
1142
1143        elif model == "vit_b":
1144            encoder = ViT_DINOv3(
1145                img_size=img_size, pos_embed_rope_rescale_coords=2, pos_embed_rope_dtype="fp32",
1146                layerscale_init=1.0e-05, norm_layer="layernormbf16", n_storage_tokens=4, mask_k_bias=True,
1147            )
1148        elif model == "vit_l":
1149            encoder = ViT_DINOv3(
1150                img_size=img_size, pos_embed_rope_rescale_coords=2, pos_embed_rope_dtype="fp32", embed_dim=1024,
1151                depth=24, num_heads=16, layerscale_init=1.0e-05, norm_layer="layernormbf16",
1152                n_storage_tokens=4, mask_k_bias=True,
1153            )
1154        elif model == "vit_l+":
1155            encoder = ViT_DINOv3(
1156                img_size=img_size, pos_embed_rope_rescale_coords=2, pos_embed_rope_dtype="fp32", embed_dim=1024,
1157                depth=24, num_heads=16, ffn_ratio=6.0, layerscale_init=1.0e-05, norm_layer="layernormbf16",
1158                ffn_layer="swiglu", n_storage_tokens=4, mask_k_bias=True,
1159            )
1160        elif model == "vit_h+":
1161            encoder = ViT_DINOv3(
1162                img_size=img_size, pos_embed_rope_rescale_coords=2, pos_embed_rope_dtype="fp32", embed_dim=1280,
1163                depth=32, num_heads=20, ffn_ratio=6.0, layerscale_init=1.0e-05, norm_layer="layernormbf16",
1164                ffn_layer="swiglu", n_storage_tokens=4, mask_k_bias=True,
1165            )
1166        elif model == "vit_7b":
1167            encoder = ViT_DINOv3(
1168                img_size=img_size, pos_embed_rope_rescale_coords=2, pos_embed_rope_dtype="fp32", embed_dim=4096,
1169                depth=40, num_heads=32, ffn_ratio=3, qkv_bias=False, drop_path_rate=0.0, layerscale_init=1.0e-05,
1170                norm_layer="layernormbf16", ffn_layer="swiglu64", n_storage_tokens=4, mask_k_bias=True,
1171                untie_global_and_local_cls_norm=True,
1172            )
1173        else:
1174            raise ValueError(
1175                f"'{model}' is not supported by DINOv3. Currently, "
1176                " 'vit_s', 'vit_s+', 'vit_b', 'vit_l', 'vit_l+', 'vit_h+'. 'vit_7b' are supported."
1177            )
1178
1179    elif backbone == "torchvision":
1180        supported = ["vit_b_16", "vit_b_32", "vit_l_16", "vit_l_32", "vit_h_14"]
1181        if model not in supported:
1182            raise ValueError(f"'{model}' is not supported by the torchvision backbone. Choose from: {supported}.")
1183        encoder = ViT_Torchvision(model_name=model, img_size=img_size, **kwargs)
1184
1185    else:
1186        raise ValueError(f"The encoder constructor for '{backbone}' is not implemented.")
1187
1188    encoder.model_type = model
1189    return encoder