Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,607 @@
+import os
+import math
+import logging
+import numpy as np
+import kornia
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+try:
+    from .task_utils import CenterPadding, upsample_features
+except ImportError:
+    from task_utils import CenterPadding, upsample_features
+
+
+logging.getLogger().setLevel(logging.WARNING)
+
+
+class FeatureExtractor(nn.Module):
+    def __init__(self, config, device, return_class_token=False):
+        super(FeatureExtractor, self).__init__()
+        self.feature_extractor = config['pretrained']['feature_extractor']
+        self.patch_size = config['architecture']['patch_size']
+        self.backbone_ckpt_path = os.path.join(config['logging']['save_dir'], config['logging']['exp_name'], 
+                                               'dinov3_vitl16_pretrain_lvd1689m-8aa4cbdd.pth')
+        self.head_ckpt_path = os.path.join(config['logging']['save_dir'], config['logging']['exp_name'],
+                                           'dinov3_vitl16_dinotxt_vision_head_and_text_encoder-a442d8f5.pth')
+        self.return_class_token = return_class_token
+        self.device = device
+
+        if self.feature_extractor == 'dinov3_vitl16':
+            self.backbone, _ = torch.hub.load('facebookresearch/dinov3', 'dinov3_vitl16_dinotxt_tet1280d20h24l', 
+                                              backbone_weights=self.backbone_ckpt_path, weights=self.head_ckpt_path)
+            del self.backbone.text_model
+            self.model = self.backbone.visual_model.backbone.to(device)
+            self.head = self.backbone.visual_model.head.to(device)
+        else:
+            raise ValueError(f'Feature extractor {self.feature_extractor} not supported.')
+    
+    def extract_dinov3(self, images, batch_size=1024, patch_length=16, layers=[23]):
+        transform = kornia.augmentation.AugmentationSequential(
+            CenterPadding(multiple=patch_length),
+            kornia.augmentation.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        )
+        transformed_images = transform(images)
+        
+        class_tokens, text_aligned_class_tokens, patch_tokens, register_tokens, feature_maps = [], [], [], [], []
+        for i in range(0, transformed_images.shape[0], batch_size):
+            image_batch = transformed_images[i:(i + batch_size)].to(device=self.device)
+            with torch.inference_mode():
+                features_out = self.model.get_intermediate_layers(image_batch, return_class_token=True,
+                                                                  return_extra_tokens=True, n=layers)
+                class_tokens.append(features_out[-1][1])
+                patch_tokens.append(features_out[0][0])
+                register_tokens.append(features_out[0][2])
+
+                head_inputs = torch.cat([class_tokens[-1].unsqueeze(1), register_tokens[-1], patch_tokens[-1]], dim=1)
+                text_aligned_class_token = self.head(head_inputs)[0][0:1]
+                text_aligned_class_tokens.append(text_aligned_class_token)
+
+                B, _, C = patch_tokens[-1].size()
+                H, W = image_batch.shape[2], image_batch.shape[3]
+                patch_H, patch_W = math.ceil(H / patch_length), math.ceil(W / patch_length)
+                feature_maps.append(patch_tokens[-1].permute(0, 2, 1).view(B, C, patch_H, patch_W))
+
+        class_tokens = torch.cat(class_tokens, dim=0)
+        text_aligned_class_tokens = torch.cat(text_aligned_class_tokens, dim=0)
+        patch_tokens = torch.cat(patch_tokens, dim=0)
+        register_tokens = torch.cat(register_tokens, dim=0)
+        feature_maps = torch.cat(feature_maps, dim=0)
+        return class_tokens, text_aligned_class_tokens, patch_tokens, register_tokens, feature_maps
+    
+    def forward(self, images, resize=False):
+        if self.feature_extractor == 'dinov3_vitl16':
+            class_tokens, text_aligned_class_tokens, patch_tokens, register_tokens, feature_maps = self.extract_dinov3(images)
+        
+        if resize:
+            image_height, image_width = images.shape[2], images.shape[3]
+            padded_height = math.ceil(image_height / self.patch_size) * self.patch_size
+            padded_width = math.ceil(image_width / self.patch_size) * self.patch_size
+            resized_feature_maps = []
+            chunk_size = 32
+            for i in range(0, len(feature_maps), chunk_size):
+                resized_feature_maps.append(upsample_features(feature_maps[i:i + chunk_size], image_height,
+                                                              image_width, padded_height, padded_width))
+            feature_maps = torch.cat(resized_feature_maps)
+        
+        return {
+            'class_tokens': class_tokens,
+            'text_aligned_class_tokens': text_aligned_class_tokens,
+            'patch_tokens': patch_tokens,
+            'register_tokens': register_tokens,
+            'feature_maps': feature_maps
+        }
+
+
+class RegionTokenGenerator(nn.Module):
+    def __init__(self, pooling_method='average', device='cuda'):
+        super(RegionTokenGenerator, self).__init__()
+        self.model, _ = torch.hub.load('facebookresearch/dinov3', 'dinov3_vitl16_dinotxt_tet1280d20h24l')
+        self.model = self.model.visual_model.head.to(device)
+        self.pooling_method = pooling_method
+        self.device = device
+
+    def forward(self, regions, class_tokens, feature_maps, register_tokens):
+        pooled_tokens, text_aligned_tokens = [], []
+        for scale_idx in range(regions.shape[2]):
+            scale_pooled_tokens, scale_text_aligned_tokens = [], []
+            for batch_idx in range(len(regions)):
+                image_regions = regions[batch_idx, :, scale_idx]
+                image_class_tokens = class_tokens[batch_idx]
+                image_feature_maps = feature_maps[batch_idx]
+                image_register_tokens = register_tokens[batch_idx]
+                if image_regions.numel() == 0:
+                    scale_pooled_tokens.append(torch.zeros((0, image_feature_maps.shape[0]), device=self.device))
+                    scale_text_aligned_tokens.append(torch.zeros((0, image_feature_maps.shape[0]), device=self.device))
+                    continue
+
+                # Get the features that pertain to the regions
+                region_features = torch.einsum('rhw,chw->rc', image_regions.float(), image_feature_maps)
+                text_alignment_inputs = torch.cat([image_class_tokens[None], image_register_tokens, region_features], dim=0)
+                text_aligned_region_features = self.model(text_alignment_inputs[None])[0][image_register_tokens.shape[0] + 1 :]
+                
+                # Pool the region features
+                if self.pooling_method == 'average':
+                    valid_elements = image_regions.sum(dim=(1, 2), dtype=torch.float32).clamp(min=1).unsqueeze(1)
+                    region_features = region_features / valid_elements
+                    text_aligned_region_features = text_aligned_region_features / valid_elements
+                else:
+                    raise ValueError(f'Pooling method {self.pooling_method} not supported.')
+                scale_pooled_tokens.append(region_features)
+                scale_text_aligned_tokens.append(text_aligned_region_features)
+            scale_pooled_tokens = torch.stack(scale_pooled_tokens)
+            scale_text_aligned_tokens = torch.stack(scale_text_aligned_tokens)
+            pooled_tokens.append(scale_pooled_tokens.unsqueeze(2))
+            text_aligned_tokens.append(scale_text_aligned_tokens.unsqueeze(2))
+        pooled_tokens = torch.cat(pooled_tokens, dim=2)
+        text_aligned_tokens = torch.cat(text_aligned_tokens, dim=2)
+        return {
+            'pooled_tokens': pooled_tokens,
+            'text_aligned_tokens': text_aligned_tokens
+        }
+
+
+class TextEncoder(nn.Module):
+    def __init__(self, config, device, prompt='photo of a '):
+        super(TextEncoder, self).__init__()
+        self.ckpt_path = os.path.join(config['logging']['save_dir'], config['logging']['exp_name'],
+                                      'dinov3_vitl16_dinotxt_vision_head_and_text_encoder-a442d8f5.pth')
+        self.model, self.tokenizer = torch.hub.load('facebookresearch/dinov3', 'dinov3_vitl16_dinotxt_tet1280d20h24l', 
+                                                    weights=self.ckpt_path)
+        del self.model.visual_model
+        self.model = self.model.to(device)
+        self.prompt = prompt
+        self.device = device
+
+    def forward(self, texts):
+        texts = [self.prompt + t for t in texts]
+        text_tokens = self.tokenizer.tokenize(texts).to(self.device)
+        with torch.no_grad():
+            text_embeddings = self.model.encode_text(text_tokens)
+            text_embeddings = text_embeddings[:, text_embeddings.shape[1] // 2 :]
+        return text_embeddings
+
+
+class PositionalEmbedding2D(nn.Module):
+    def __init__(self, embedding_dim=64, scale=None):
+        super().__init__()
+        if scale is None or scale <= 0.0:
+            scale = 1.0
+        generator = torch.Generator()
+        generator.manual_seed(42)
+        self.register_buffer("positional_encoding_gaussian_matrix", 
+                             scale * torch.randn((2, embedding_dim // 2), generator=generator))
+
+    def _pe_encoding(self, coords):
+        coords = 2 * coords - 1
+        coords = coords @ self.positional_encoding_gaussian_matrix
+        coords = 2 * np.pi * coords
+        return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
+
+    def forward(self, size):
+        h, w = size
+        device = self.positional_encoding_gaussian_matrix.device
+        grid = torch.ones((h, w), device=device, dtype=torch.float32)
+        y_embed = grid.cumsum(dim=0) - 0.5
+        x_embed = grid.cumsum(dim=1) - 0.5
+        y_embed = y_embed / h
+        x_embed = x_embed / w
+        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
+        return pe.permute(2, 0, 1)
+
+
+class AttentionLayer(nn.Module):
+    def __init__(self, q_dim, kv_dim, hidden_dim, num_heads=8, dropout=0.1, use_bias=False, use_v_proj=True, use_out_proj=True):
+        super(AttentionLayer, self).__init__()
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        assert hidden_dim % num_heads == 0, 'Hidden dimension must be a multiple of the number of heads.'
+        self.head_dim = hidden_dim // num_heads
+        if not use_v_proj:
+            assert kv_dim == hidden_dim, 'Key and value dimensions must be the same as the hidden dimension if not using v_proj.'
+
+        self.q_proj = nn.Linear(q_dim, hidden_dim, bias=use_bias)
+        nn.init.kaiming_normal_(self.q_proj.weight, mode='fan_in', nonlinearity='linear')
+        self.k_proj = nn.Linear(kv_dim, hidden_dim, bias=use_bias)
+        nn.init.kaiming_normal_(self.k_proj.weight, mode='fan_in', nonlinearity='linear')
+        if use_v_proj:
+            self.v_proj = nn.Linear(kv_dim, hidden_dim, bias=use_bias)
+            nn.init.kaiming_normal_(self.v_proj.weight, mode='fan_in', nonlinearity='linear')
+        else:
+            self.v_proj = nn.Identity()
+        if use_bias:
+            nn.init.zeros_(self.q_proj.bias)
+            nn.init.zeros_(self.k_proj.bias)
+            if use_v_proj:
+                nn.init.zeros_(self.v_proj.bias)
+
+        self.q_norm = nn.LayerNorm(self.head_dim)
+        self.k_norm = nn.LayerNorm(self.head_dim)
+
+        self.dropout = nn.Dropout(dropout)
+        if use_out_proj:
+            self.out_proj = nn.Linear(hidden_dim, hidden_dim, bias=use_bias)
+            nn.init.kaiming_normal_(self.out_proj.weight, mode='fan_in', nonlinearity='linear')
+            if use_bias:
+                nn.init.zeros_(self.out_proj.bias)
+        else:
+            self.out_proj = nn.Identity()
+
+        self.scale = self.head_dim ** -0.5
+
+    def forward(self, q, k, v, mask=None, attn_threshold=None):
+        batch_size, q_len, _ = q.shape
+        _, kv_len, _ = k.shape
+
+        query = self.q_proj(q).view(batch_size, q_len, self.num_heads, -1).transpose(1, 2)
+        key = self.k_proj(k).view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2)
+        value = self.v_proj(v).view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2)
+
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+
+        attn_scores = torch.matmul(query, key.transpose(-2, -1)) * self.scale
+        if mask is not None:
+            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))
+        if attn_threshold is not None:
+            max_attn_scores, _ = attn_scores.max(dim=-1, keepdim=True)
+            thresholding_mask = attn_scores >= (attn_threshold * max_attn_scores)
+            attn_scores = attn_scores.masked_fill(thresholding_mask == 0, -1e9)
+        attn_weights = F.softmax(attn_scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+
+        attn_out = torch.matmul(attn_weights, value)
+        attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, q_len, self.hidden_dim)
+
+        out = self.out_proj(attn_out)
+        return out, attn_weights
+
+
+class MLPBlock(nn.Module):
+    def __init__(self, hidden_dim, intermediate_dim, dropout=0.1):
+        super(MLPBlock, self).__init__()
+        self.linear1 = nn.Linear(hidden_dim, intermediate_dim)
+        self.gelu = nn.GELU()
+        self.linear2 = nn.Linear(intermediate_dim, hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        z = self.linear1(x)
+        z = self.gelu(z)
+        z = self.dropout(z)
+        z = self.linear2(z)
+        return z
+
+
+class CrossAttentionBlock(nn.Module):
+    def __init__(self, q_dim, kv_dim, hidden_dim, mlp_dim, num_heads, dropout, use_bias):
+        super(CrossAttentionBlock, self).__init__()
+        self.query_norm = nn.LayerNorm(q_dim)
+        self.cross_attn = AttentionLayer(q_dim, kv_dim, hidden_dim, num_heads, dropout, use_bias)
+        self.dropout = nn.Dropout(dropout)
+        self.mlp_norm = nn.LayerNorm(hidden_dim)
+        self.mlp = MLPBlock(hidden_dim, mlp_dim)
+        self.out_norm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, query, context, mask=None):
+        x = self.query_norm(query)
+        x, attn_scores = self.cross_attn(q=x, k=context, v=context, mask=mask)
+        x = self.dropout(x)
+        x = x + query
+
+        y = self.mlp_norm(x)
+        y = self.mlp(y)
+        out = self.out_norm(y) + x
+        return out, attn_scores
+
+
+class TextAlignmentBlock(nn.Module):
+    def __init__(self, hidden_dim, intermediate_dim, output_dim, dropout=0.1):
+        super(TextAlignmentBlock, self).__init__()
+        self.linear1 = nn.Linear(hidden_dim, intermediate_dim)
+        self.gelu = nn.GELU()
+        self.linear2 = nn.Linear(intermediate_dim, output_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        z = self.linear1(x)
+        z = self.gelu(z)
+        z = self.dropout(z)
+        z = self.linear2(z)
+        return z
+
+
+class TokenAggregator(nn.Module):
+    def __init__(self, config):
+        super(TokenAggregator, self).__init__()
+        self.merging_iou_threshold = config['parameters']['merging_iou_threshold']
+        self.merging_similarity_threshold = config['parameters']['merging_similarity_threshold']
+        self.binarization_threshold = config['parameters'].get('binarization_threshold', 0.5)
+
+    def _compute_binary_masks(self, region_masks):
+        num_masks = region_masks.shape[0]
+        return (region_masks.reshape(num_masks, -1) > self.binarization_threshold).float()
+
+    def _compute_iou_matrix(self, binary_masks):
+        num_masks = binary_masks.shape[0]
+        if num_masks == 0:
+            return torch.zeros(0, 0, device=binary_masks.device)
+        intersection = torch.mm(binary_masks, binary_masks.t())
+        areas = binary_masks.sum(dim=1)
+        union = areas.unsqueeze(1) + areas.unsqueeze(0) - intersection
+        return intersection / torch.clamp(union, min=1.0)
+
+    def _find_connected_components(self, adjacency):
+        n = adjacency.shape[0]
+        if n == 0:
+            return []
+        
+        # Initialize labels
+        labels = torch.arange(n, device=adjacency.device)
+        
+        # Iterative label propagation
+        for _ in range(int(np.ceil(np.log2(n + 1))) + 1):
+            neighbor_labels = torch.where(adjacency, labels.unsqueeze(0).expand(n, -1), labels.unsqueeze(1).expand(-1, n))
+            new_labels = neighbor_labels.min(dim=1)[0]
+            new_labels = torch.minimum(new_labels, labels)
+            if torch.equal(new_labels, labels):
+                break
+            labels = new_labels
+        
+        # Convert to groups
+        labels_cpu = labels.cpu().tolist()
+        label_to_group = {}
+        for idx, label in enumerate(labels_cpu):
+            if label not in label_to_group:
+                label_to_group[label] = []
+            label_to_group[label].append(idx)
+        return list(label_to_group.values())
+
+    def _compute_token_similarity_matrix(self, pred_tokens):
+        num_tokens = pred_tokens.shape[0]
+        if num_tokens == 0:
+            return torch.zeros(0, 0, device=pred_tokens.device)
+        pred_tokens = F.normalize(pred_tokens, p=2, dim=-1)
+        return torch.mm(pred_tokens, pred_tokens.t())
+
+    def group_predictions(self, region_masks, pred_tokens=None):
+        num_masks = region_masks.shape[0]
+        if num_masks == 0:
+            return []
+        binary_masks = self._compute_binary_masks(region_masks)
+        iou_matrix = self._compute_iou_matrix(binary_masks)
+        mask_adjacency = iou_matrix > self.merging_iou_threshold
+        if pred_tokens is not None and pred_tokens.shape[0] == num_masks:
+            token_sim = self._compute_token_similarity_matrix(pred_tokens)
+            token_adjacency = token_sim > self.merging_similarity_threshold
+            adjacency = mask_adjacency | token_adjacency
+        else:
+            adjacency = mask_adjacency
+        return self._find_connected_components(adjacency)
+
+    def forward(self, ren_outputs, remove_singleton_groups=True):
+        pred_tokens = ren_outputs['pred_tokens']
+        region_masks = ren_outputs['region_masks']
+        text_aligned_tokens = ren_outputs['text_aligned_tokens']
+
+        pred_tokens = torch.flatten(pred_tokens, 1, 2)
+        region_masks = torch.flatten(region_masks, 1, 2)
+        text_aligned_tokens = torch.flatten(text_aligned_tokens, 1, 2)
+
+        aggregated_outputs = {'pred_tokens': [], 'region_masks': [], 'text_aligned_tokens': []}
+        for batch_idx in range(pred_tokens.shape[0]):
+            batch_pred_tokens = pred_tokens[batch_idx]
+            batch_region_masks = region_masks[batch_idx]
+            batch_text_aligned_tokens = text_aligned_tokens[batch_idx]
+
+            groups = self.group_predictions(batch_region_masks, batch_pred_tokens)
+
+            kept_groups = []
+            for local_group_idxs in groups:
+                if remove_singleton_groups and len(local_group_idxs) == 1:
+                    continue
+                global_idxs = torch.tensor(local_group_idxs, device=batch_region_masks.device)
+                group_mean_mask = batch_region_masks[global_idxs].mean(dim=0)
+                kept_groups.append({'global_idxs': global_idxs, 'mean_mask': group_mean_mask})
+
+            if len(kept_groups) == 0:
+                mask_areas = batch_region_masks.sum(dim=(-2, -1))
+                best_idx = mask_areas.argmax()
+                kept_groups.append({
+                    'global_idxs': best_idx.unsqueeze(0),
+                    'mean_mask': batch_region_masks[best_idx],
+                })
+
+            new_pred_tokens, new_region_masks, new_text_aligned_tokens = [], [], []
+            for gd in kept_groups:
+                global_idxs = gd['global_idxs']
+                new_pred_tokens.append(pred_tokens[batch_idx][global_idxs].mean(dim=0))
+                new_region_masks.append(gd['mean_mask'])
+                new_text_aligned_tokens.append(text_aligned_tokens[batch_idx][global_idxs].mean(dim=0))
+
+            aggregated_outputs['pred_tokens'].append(torch.stack(new_pred_tokens, dim=0))
+            aggregated_outputs['region_masks'].append(torch.stack(new_region_masks, dim=0))
+            aggregated_outputs['text_aligned_tokens'].append(torch.stack(new_text_aligned_tokens, dim=0))
+        return aggregated_outputs
+
+
+class RegionEncoder(nn.Module):
+    def __init__(self, config):
+        super(RegionEncoder, self).__init__()
+        hidden_dim = config['architecture']['hidden_dim']
+        text_embed_dim = config['architecture']['text_embed_dim']
+        image_resolution = config['parameters']['image_resolution']
+        patch_size = config['architecture']['patch_size']
+        feature_map_resolution = image_resolution // patch_size
+        self.feature_map_resolution = feature_map_resolution
+        self.image_resolution = image_resolution
+
+        # Create position embeddings for the prompts and feature maps
+        position_embedder = PositionalEmbedding2D(hidden_dim)
+        location_embeddings = position_embedder((image_resolution, image_resolution))
+        feature_embeddings = position_embedder((feature_map_resolution, feature_map_resolution)).flatten(-2).permute(1, 0)
+        self.register_buffer('location_embeddings', location_embeddings)
+        self.register_buffer('feature_embeddings', feature_embeddings)
+
+        # Define scale embeddings for multiscale region tokens
+        self.num_multiscale_regions = config['parameters']['num_multiscale_regions']
+        self.scale_embeddings = nn.Embedding(self.num_multiscale_regions, hidden_dim)
+        nn.init.normal_(self.scale_embeddings.weight, std=0.02)
+
+        # Instantiate the prompt and region attention layers
+        self.num_decoder_layers = config['architecture']['num_decoder_layers']
+        self.num_attention_heads = config['architecture']['num_attention_heads']
+        self.prompt_attention_layers = nn.ModuleList([
+            AttentionLayer(hidden_dim, hidden_dim, hidden_dim, num_heads=self.num_attention_heads)
+            for _ in range(self.num_decoder_layers)
+        ])
+        self.prompt_attention_norms = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(self.num_decoder_layers)])
+        self.region_attention_layers = nn.ModuleList([
+            CrossAttentionBlock(q_dim=hidden_dim, kv_dim=hidden_dim, hidden_dim=hidden_dim, mlp_dim=2 * hidden_dim,
+                                num_heads=self.num_attention_heads, dropout=0.1, use_bias=False) 
+            for _ in range(self.num_decoder_layers)
+        ])
+        self.region_attention_norms = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(self.num_decoder_layers)])
+
+        # Instantiate the region token prediction head
+        self.token_prediction_head = AttentionLayer(hidden_dim, hidden_dim, hidden_dim, num_heads=1, dropout=0.0, 
+                                                    use_v_proj=False, use_out_proj=False)
+
+        # Instantiate the text alignment head
+        self.text_alignment_block = TextAlignmentBlock(hidden_dim, 2 * hidden_dim, text_embed_dim)
+
+        # Instantiate the token aggregator
+        self.token_aggregator = TokenAggregator(config)
+
+    def load_state_dict_resolution_agnostic(self, state_dict, strict=False):
+        model_state = self.state_dict()
+        new_state = dict(state_dict)
+
+        # Interpolate location_embeddings if spatial size differs
+        if 'location_embeddings' in new_state and new_state['location_embeddings'].shape != model_state['location_embeddings'].shape:
+            old = new_state['location_embeddings']
+            target_shape = model_state['location_embeddings'].shape
+            if old.shape[0] == target_shape[0]:
+                resized = F.interpolate(old.unsqueeze(0), size=(target_shape[1], target_shape[2]), mode='bilinear', align_corners=False)
+                new_state['location_embeddings'] = resized.squeeze(0)
+            else:
+                new_state['location_embeddings'] = model_state['location_embeddings'].clone()
+
+        # Interpolate feature_embeddings if spatial size differs
+        if 'feature_embeddings' in new_state and new_state['feature_embeddings'].shape != model_state['feature_embeddings'].shape:
+            old = new_state['feature_embeddings']
+            target = model_state['feature_embeddings']
+            if old.shape[1] == target.shape[1]:
+                num_pos_old, C = old.shape
+                num_pos_new = target.shape[0]
+                h_old = int(round(num_pos_old ** 0.5))
+                w_old = num_pos_old // h_old
+                h_new = int(round(num_pos_new ** 0.5))
+                w_new = num_pos_new // h_new
+                old_2d = old.view(h_old, w_old, C).permute(2, 0, 1).unsqueeze(0)
+                resized = F.interpolate(old_2d, size=(h_new, w_new), mode='bilinear', align_corners=False)
+                new_state['feature_embeddings'] = resized.squeeze(0).permute(1, 2, 0).reshape(-1, C)
+            else:
+                new_state['feature_embeddings'] = model_state['feature_embeddings'].clone()
+
+        return self.load_state_dict(new_state, strict=strict)
+
+    def forward(self, feature_maps, grid_points, aggregate_tokens=False, remove_singleton_groups=True):
+        if isinstance(grid_points, list):
+            grid_points = torch.stack([gp.to(feature_maps.device) for gp in grid_points])
+        batch_size, num_prompts, _ = grid_points.shape
+
+        # Create scale prompt embeddings for multiscale region tokens
+        scale_prompt_embeddings = self.scale_embeddings.weight.unsqueeze(0).repeat(batch_size, 1, 1)
+        scale_prompt_embeddings = scale_prompt_embeddings.unsqueeze(1).repeat(1, num_prompts, 1, 1)
+
+        # Create spatial prompt embeddings to encode the location of the point prompts
+        spatial_prompt_embeddings = self.location_embeddings[:, grid_points[..., 0], grid_points[..., 1]]
+        spatial_prompt_embeddings = spatial_prompt_embeddings.permute(1, 2, 0).unsqueeze(2)
+        spatial_prompt_embeddings = spatial_prompt_embeddings.repeat(1, 1, self.num_multiscale_regions, 1)
+
+        # Create the query tokens
+        q = scale_prompt_embeddings
+
+        # Get the key and value tokens for the region attention layers
+        kv = feature_maps.flatten(-2).permute(0, 2, 1)
+        kv = kv + self.feature_embeddings[None]
+
+        # Apply the region attention layers and the prompt attention layers
+        for layer_idx in range(self.num_decoder_layers):
+            q += spatial_prompt_embeddings
+
+            # Apply the region attention layer
+            q = q.reshape(batch_size, num_prompts * self.num_multiscale_regions, -1)
+            q, _ = self.region_attention_layers[layer_idx](q, kv)
+            q = q.reshape(batch_size, num_prompts, self.num_multiscale_regions, -1)
+            q = self.region_attention_norms[layer_idx](q)
+
+            # Apply the prompt attention layer
+            q = q.reshape(batch_size * num_prompts, self.num_multiscale_regions, -1)
+            q, _ = self.prompt_attention_layers[layer_idx](q, q, q)
+            q = self.prompt_attention_norms[layer_idx](q)
+            q = q.reshape(batch_size, num_prompts, self.num_multiscale_regions, -1)
+        prompt_tokens = q
+
+        # Get the region tokens
+        q = prompt_tokens.reshape(batch_size, num_prompts * self.num_multiscale_regions, -1)
+        k = kv
+        v = kv - self.feature_embeddings[None]
+        pred_tokens, attn_weights = self.token_prediction_head(q, k, v)
+        pred_tokens = pred_tokens.reshape(batch_size, num_prompts, self.num_multiscale_regions, -1)
+        attn_weights = attn_weights.reshape(batch_size, num_prompts, self.num_multiscale_regions, -1)
+
+        # Get the region masks
+        region_masks = attn_weights / attn_weights.max(dim=-1, keepdim=True)[0]
+        region_masks = region_masks.reshape(batch_size, num_prompts, self.num_multiscale_regions, 
+                                            self.feature_map_resolution, self.feature_map_resolution)
+
+        # Get text aligned tokens
+        text_aligned_tokens = self.text_alignment_block(pred_tokens)
+
+        outputs = {
+            'pred_tokens': pred_tokens,
+            'region_masks': region_masks,
+            'text_aligned_tokens': text_aligned_tokens,
+        }
+        if aggregate_tokens:
+            outputs = self.token_aggregator(outputs, remove_singleton_groups=remove_singleton_groups)
+        return outputs
+
+
+if __name__ == '__main__':
+    import yaml
+    from tqdm import tqdm
+    from dataloader import COCOStuffDataset
+
+    # Load the config
+    with open('configs/train_dinov3_vitl16.yaml', 'r') as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    # Load the dataset
+    dataset = COCOStuffDataset(config, 'val')
+
+    # Generate the grid points
+    image_resolution = config['parameters']['image_resolution']
+    patch_size = config['architecture']['patch_size']
+    grid_size = image_resolution // patch_size
+    x_coords = np.linspace(patch_size // 2, image_resolution - patch_size // 2, grid_size, dtype=int)
+    y_coords = np.linspace(patch_size // 2, image_resolution - patch_size // 2, grid_size, dtype=int)
+    grid_points = np.array([(y, x) for y in y_coords for x in x_coords])
+    grid_points = torch.tensor(grid_points)[None]
+    
+    # Get the models
+    region_encoder = RegionEncoder(config).to(device)
+    feature_extractor = FeatureExtractor(config, device)
+
+    # Generate the region tokens
+    for item in tqdm(dataset):
+        image = item[0].to(device)
+        feature_maps = feature_extractor(image[None])['feature_maps']
+        ren_outputs = region_encoder(feature_maps, grid_points, aggregate_tokens=True)
+        print(ren_outputs['pred_tokens'][0].shape)
+        print(ren_outputs['region_masks'][0].shape)
+        print(ren_outputs['text_aligned_tokens'][0].shape)
+        break