| import torch |
| import einops |
| import torch.nn as nn |
| from transformers import AutoModel |
| from torchvision import transforms |
|
|
| class DinoV3Encoder(nn.Module): |
| def __init__(self, name, feature_key, plus=False, output_dim=None, postprocess=None, n_patches=196): |
| super().__init__() |
| print("Encoder feature_key:", feature_key) |
| self.name = name |
| if plus: |
| model_name = "facebook/dinov3-vits16plus-pretrain-lvd1689m" |
| else: |
| model_name = "facebook/dinov3-vits16-pretrain-lvd1689m" |
| self.base_model = AutoModel.from_pretrained( |
| model_name, |
| |
| |
| trust_remote_code=True, |
| ) |
| self.feature_key = feature_key |
| self.emb_dim = self.base_model.config.hidden_size |
| self.output_dim = self.emb_dim |
| if feature_key == "x_norm_patchtokens": |
| self.latent_ndim = 2 |
| elif feature_key == "x_norm_clstoken": |
| self.latent_ndim = 1 |
| else: |
| raise ValueError(f"Invalid feature key: {feature_key}") |
|
|
| self.patch_size = self.base_model.config.patch_size |
|
|
| self.postprocess = postprocess |
| if postprocess is not None: |
| if postprocess == 'avg_pool': |
| self.latent_ndim = 1 |
|
|
| self.normalization = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
| |
|
|
| def forward(self, x): |
| |
| |
| |
| assert x.max() <= 1.0 and x.min() >= 0, "expect 0..1 range" |
| x = self.normalization(x) |
|
|
| prefix_shape = x.shape[:-3] |
| c, h, w = x.shape[-3:] |
|
|
| |
| prod_prefix = 1 |
| for d in prefix_shape: |
| prod_prefix *= d |
|
|
| x = x.reshape(prod_prefix, c, h, w) |
|
|
| outputs = self.base_model(pixel_values=x) |
| if self.feature_key == "x_norm_clstoken": |
| emb = outputs.last_hidden_state[:, 0, :] |
| elif self.feature_key == "x_norm_patchtokens": |
| emb = outputs.last_hidden_state[:, 5:, :] |
|
|
| emb = emb.reshape(*prefix_shape, *emb.shape[1:]) |
|
|
| if self.postprocess == 'avg_pool': |
| emb = torch.mean(emb, dim=-2) |
|
|
| if self.latent_ndim == 1: |
| emb = emb.unsqueeze(len(prefix_shape)) |
| return emb |
|
|
