| import torch |
| import torch.nn as nn |
| import timm |
|
|
|
|
| class TimmVisionEncoder(nn.Module): |
| def __init__(self, pretrained_encoder_name: str, load_pretrained: bool = True): |
| super().__init__() |
| self.model = timm.create_model( |
| pretrained_encoder_name, |
| pretrained=load_pretrained, |
| dynamic_img_size=True, |
| ) |
| self.hidden_size = self.model.embed_dim |
| self.model.norm = nn.LayerNorm(self.hidden_size, elementwise_affine=False) |
| self.num_prefix_tokens = self.model.num_prefix_tokens |
| patch_size = self.model.patch_embed.patch_size |
| self.patch_size = int( |
| patch_size[0] if isinstance(patch_size, tuple) else patch_size |
| ) |
| if self.patch_size == 14: |
| input_size = 224 |
| else: |
| input_size = 256 |
| self.model.set_input_size(input_size, patch_size) |
|
|
| data_config = timm.data.resolve_model_data_config(self.model) |
| self.register_buffer( |
| "pixel_mean", |
| torch.tensor(data_config["mean"])[None, :, None, None], |
| persistent=True, |
| ) |
| self.register_buffer( |
| "pixel_std", |
| torch.tensor(data_config["std"])[None, :, None, None], |
| persistent=True, |
| ) |
|
|
| def forward( |
| self, |
| pixel_values: torch.Tensor, |
| return_prefix_tokens: bool = False, |
| ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| x = pixel_values.float() / 255.0 |
| weight_dtype = self.model.patch_embed.proj.weight.dtype |
| x = ((x - self.pixel_mean) / self.pixel_std).to(dtype=weight_dtype) |
| output = self.model.forward_features(x) |
| if isinstance(output, dict): |
| output = output["x"] |
| prefix_tokens = output[:, : self.num_prefix_tokens, :] |
| patch_tokens = output[:, self.num_prefix_tokens :, :] |
| if return_prefix_tokens: |
| return patch_tokens, prefix_tokens |
| return patch_tokens |
|
|
|
|
| class DinoV2Encoder(TimmVisionEncoder): |
| pass |
|
|
|
|
| class DinoV3Encoder(TimmVisionEncoder): |
| pass |
|
|
|
|
| ALL_ENCODERS = { |
| "dinov3_small": (DinoV3Encoder, "vit_small_patch16_dinov3.lvd1689m"), |
| "dinov3_base": (DinoV3Encoder, "vit_base_patch16_dinov3.lvd1689m"), |
| "dinov3_large": (DinoV3Encoder, "vit_large_patch16_dinov3.lvd1689m"), |
| "dinov2_small_reg": (DinoV2Encoder, "vit_small_patch14_reg4_dinov2.lvd142m"), |
| "dinov2_base_reg": (DinoV2Encoder, "vit_base_patch14_reg4_dinov2.lvd142m"), |
| "dinov2_large_reg": (DinoV2Encoder, "vit_large_patch14_reg4_dinov2.lvd142m"), |
| } |
|
|
|
|
| def build_encoder(encoder_name: str, pretrained: bool = False): |
| if encoder_name not in ALL_ENCODERS: |
| raise ValueError( |
| f"Unknown encoder {encoder_name!r}. Available: {list(ALL_ENCODERS)}" |
| ) |
| model_cls, model_id = ALL_ENCODERS[encoder_name] |
| return model_cls(model_id, load_pretrained=pretrained) |
|
|
|
|
| __all__ = [ |
| "ALL_ENCODERS", |
| "DinoV2Encoder", |
| "DinoV3Encoder", |
| "TimmVisionEncoder", |
| "build_encoder", |
| ] |
|
|
