| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from transformers.modeling_utils import PreTrainedModel |
| from transformers import PretrainedConfig |
|
|
|
|
| class VisualPromptEncodeConfig(PretrainedConfig): |
| model_type = 'vision_encoder' |
| _auto_class = 'AutoConfig' |
| main_input_name = "visual_prompts" |
|
|
| def __init__(self, |
| vision_hidden_size: int, |
| language_hidden_size: int, |
| patch_size: int, |
| downsample_ratio, |
| **kwargs): |
| super().__init__(**kwargs) |
|
|
|
|
| class VisualPromptEncodeModel(nn.Module): |
| def __init__(self, |
| in_channels: int, |
| vision_hidden_size: int, |
| language_hidden_size: int, |
| force_image_size: int, |
| patch_size: int, |
| downsample_ratio: int, |
| ) -> None: |
| super().__init__() |
|
|
| self.in_channels = in_channels |
| self.vision_hidden_size = vision_hidden_size |
| self.language_hidden_size = language_hidden_size |
| self.force_image_size = force_image_size |
| self.patch_size = patch_size |
| self.downsample_ratio = downsample_ratio |
|
|
| self.patch_embedding = nn.Conv2d( |
| in_channels=in_channels, out_channels=vision_hidden_size, |
| kernel_size=patch_size, stride=patch_size |
| ) |
| |
| self.mlp1 = nn.Sequential( |
| nn.LayerNorm(language_hidden_size), |
| nn.Linear(language_hidden_size, vision_hidden_size), |
| nn.GELU(), |
| nn.Linear(vision_hidden_size, vision_hidden_size) |
| ) |
|
|
| self.patch_edge_token = force_image_size// patch_size |
|
|
| def forward(self, merged_visual_prompts, visual_prompts, mark_embeddings, num_patches, num_vprompts): |
| patch_embeds = self.patch_embedding(merged_visual_prompts) |
| split_size = [npatch * nvprompt for (npatch, nvprompt) in zip(num_patches, num_vprompts)] |
| resized_visual_prompts = F.interpolate(visual_prompts.unsqueeze(1), |
| size=self.patch_edge_token, |
| mode="nearest") |
| resized_visual_prompts_per_batch = torch.split(resized_visual_prompts, split_size, dim=0) |
| split_size = [nvp for nvp in num_vprompts] |
| mark_embeddings_per_batch = torch.split(mark_embeddings, split_size, dim=0) |
| batch_vprompts_input = [] |
| for i, (per_visual_prompts, per_mark_embeddings) in enumerate(zip( |
| resized_visual_prompts_per_batch, mark_embeddings_per_batch)): |
| per_visual_prompts = per_visual_prompts.view( |
| num_vprompts[i], num_patches[i], 1, self.patch_edge_token, self.patch_edge_token) |
| per_background = torch.ones_like(per_visual_prompts) - per_visual_prompts |
| per_vprompts_input = torch.zeros( |
| (num_vprompts[i], num_patches[i], self.language_hidden_size, |
| self.patch_edge_token, self.patch_edge_token), |
| dtype=mark_embeddings.dtype).to(mark_embeddings.device) |
| per_vprompts_input = per_vprompts_input * per_background + \ |
| per_mark_embeddings[:, None, :, None, None] * per_visual_prompts |
| |
| per_vprompts_input = torch.sum(per_vprompts_input, dim=0) |
| batch_vprompts_input.append(per_vprompts_input) |
| batch_vprompts_input = torch.cat(batch_vprompts_input, dim=0) |
| batch_vprompts_input = batch_vprompts_input.permute(0, 2, 3, 1).flatten(1, 2) |
| batch_vprompts_input = self.mlp1(batch_vprompts_input).view( |
| -1, self.patch_edge_token, self.patch_edge_token, self.vision_hidden_size).permute(0, 3, 1, 2) |
| |
| patch_embeds = patch_embeds * 0.0 + batch_vprompts_input |
| return patch_embeds |
|
|
| |
|
|
| |
|
|
