| import torch |
| import torch.nn as nn |
| from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel |
| from .custom_clip import _CLIPVisionModel |
| import torch.nn.functional as F |
|
|
| class CLIPVisionTower(nn.Module): |
| def __init__(self, vision_tower, args, delay_load=False): |
| super().__init__() |
|
|
| self.is_loaded = False |
| |
| self.vision_tower_name = vision_tower |
| self.select_layer = args.mm_vision_select_layer |
| self.select_feature = getattr(args, "mm_vision_select_feature", "patch") |
| self.pad_vit = getattr(args, "pad_train_clip_images", False) |
| self.resize_vision_tower = getattr(args, "resize_vision_tower", False) |
| self.resize_vision_tower_size = getattr(args, "resize_vision_tower_size", 224) |
| self.is_multipath_encoder = getattr(args,"is_multipath_encoder",False) |
| if not delay_load: |
| self.load_model() |
| else: |
| self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name) |
|
|
| def load_model(self): |
| self.image_processor = CLIPImageProcessor.from_pretrained( |
| self.vision_tower_name |
| ) |
| if self.pad_vit: |
| self.vision_tower = _CLIPVisionModel.from_pretrained( |
| self.vision_tower_name, low_cpu_mem_usage=True |
| ) |
| else: |
| self.vision_tower = CLIPVisionModel.from_pretrained( |
| self.vision_tower_name, low_cpu_mem_usage=True |
| ) |
| vision_tower = self.vision_tower |
| resize_vision_tower_size = self.resize_vision_tower_size |
| if self.resize_vision_tower: |
| origin_p_num = int(vision_tower.vision_model.embeddings.num_patches ** 0.5) |
| vision_tower_embed_dim = vision_tower.vision_model.embeddings.embed_dim |
| vision_tower.vision_model.embeddings.image_size = resize_vision_tower_size |
| vision_tower.vision_model.embeddings.num_patches = (resize_vision_tower_size // vision_tower.vision_model.embeddings.patch_size) **2 |
| vision_tower.vision_model.embeddings.num_positions = vision_tower.vision_model.embeddings.num_patches + 1 |
| vision_tower.vision_model.embeddings.register_buffer("position_ids", torch.arange(vision_tower.vision_model.embeddings.num_positions).expand((1, -1))) |
| new_p_num = int(vision_tower.vision_model.embeddings.num_patches ** 0.5) |
|
|
| origin_position_embedding_weight = vision_tower.vision_model.embeddings.position_embedding.weight |
| origin_position_embedding_weight_cls = origin_position_embedding_weight[-1:] |
| origin_position_embedding_weight = origin_position_embedding_weight[:-1].permute(1, 0).view(1, vision_tower_embed_dim, origin_p_num, origin_p_num) |
| new_position_embedding_weight = F.interpolate(origin_position_embedding_weight, (new_p_num, new_p_num), mode='bilinear', align_corners=False)[0] |
| new_position_embedding_weight = new_position_embedding_weight.flatten(-2).permute(1, 0) |
| new_position_embedding_weight = torch.cat((new_position_embedding_weight, origin_position_embedding_weight_cls), dim=0) |
| vision_tower.vision_model.embeddings.position_embedding = nn.Embedding(vision_tower.vision_model.embeddings.num_positions, vision_tower_embed_dim) |
| vision_tower.vision_model.embeddings.position_embedding.weight = torch.nn.Parameter(new_position_embedding_weight).to(origin_position_embedding_weight) |
| vision_tower.vision_model.embeddings.position_ids = vision_tower.vision_model.embeddings.position_ids.to(origin_position_embedding_weight.device) |
| self.vision_tower = vision_tower |
| self.vision_tower.requires_grad_(False) |
| self.is_loaded = True |
|
|
| def feature_select(self, image_forward_outs): |
| image_features = image_forward_outs.hidden_states[self.select_layer] |
| if self.select_feature == "patch": |
| image_features = image_features[:, 1:] |
| elif self.select_feature == "cls_patch": |
| image_features = image_features |
| else: |
| raise ValueError(f"Unexpected select feature: {self.select_feature}") |
| return image_features, [image_forward_outs.hidden_states[-11][:, 1:]] |
|
|
| @torch.no_grad() |
| def forward(self, images, attention_mask=None, output_attentions=False,output_keys=False): |
| pre_image_features = [] |
| if type(images) is list: |
| image_features = [] |
| for image in images: |
| image_forward_out = self.vision_tower( |
| image.to(device=self.device, dtype=self.dtype).unsqueeze(0), |
| output_hidden_states=True, attention_mask=attention_mask, |
| output_attentions=output_attentions, |
| output_keys=output_keys |
| ) |
| image_feature = self.feature_select(image_forward_out).to(image.dtype) |
| image_features.append(image_feature) |
| else: |
| if isinstance(self.vision_tower, _CLIPVisionModel): |
| image_forward_outs = self.vision_tower( |
| images.to(device=self.device, dtype=self.dtype), |
| output_hidden_states=True, attention_mask=attention_mask, |
| output_attentions=output_attentions, |
| output_keys=output_keys |
| ) |
| else: |
| image_forward_outs = self.vision_tower( |
| images.to(device=self.device, dtype=self.dtype), |
| output_hidden_states=True |
| ) |
| image_features, pre_image_features = self.feature_select(image_forward_outs) |
| image_features = image_features.to(images.dtype) |
| pre_image_features = [f.to(images.dtype) for f in pre_image_features] |
| torch.cuda.empty_cache() |
| attention_keys = None |
| if output_keys and hasattr(image_forward_outs, 'keys') and image_forward_outs.keys is not None: |
| attention_keys = image_forward_outs.keys[-1] |
| return image_features, pre_image_features,None, attention_keys |
|
|
| @property |
| def dummy_feature(self): |
| return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype) |
|
|
| @property |
| def dtype(self): |
| return self.vision_tower.dtype |
|
|
| @property |
| def device(self): |
| return self.vision_tower.device |
|
|
| @property |
| def config(self): |
| if self.is_loaded: |
| return self.vision_tower.config |
| else: |
| return self.cfg_only |
|
|
| @property |
| def hidden_size(self): |
| return self.config.hidden_size |
|
|
| @property |
| def num_patches(self): |
| return (self.config.image_size // self.config.patch_size) ** 2 |
|
|
