| import torch |
| import torch.nn as nn |
| from transformers import CLIPVisionModel |
|
|
|
|
| class DFN5B_CLIP_ViT_H_14_378(nn.Module): |
|
|
| def __init__(self, vision_tower): |
| super().__init__() |
|
|
| self.is_loaded = False |
| self.is_resize_pos = False |
|
|
| self.vision_tower_name = vision_tower |
| self.select_layer = -1 |
| self.select_feature = 'patch' |
| self.load_model() |
|
|
| def load_model(self): |
| |
| |
| self.vision_tower = CLIPVisionModel.from_pretrained('/root/LWT/Models/DFN5B-CLIP-ViT-H-14-378') |
|
|
| 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 |
|
|
| def forward(self, images): |
|
|
| if not self.is_loaded: |
| self.load_model() |
| |
| if type(images) is list: |
| image_features = [] |
| for image in images: |
| image_forward_out = self.vision_tower( |
| image.to(device=self.device, |
| dtype=image.dtype).unsqueeze(0), |
| output_hidden_states=True) |
| image_feature = self.feature_select(image_forward_out).to( |
| image.dtype) |
| image_features.append(image_feature) |
| else: |
| |
| image_forward_outs = self.vision_tower( |
| images.to(device=self.device, dtype=images.dtype), |
| output_hidden_states=True) |
| image_features = self.feature_select(image_forward_outs).to(images.dtype) |
| return image_features |
|
|
| @property |
| def device(self): |
| return self.vision_tower.device |
| |
| @property |
| def dtype(self): |
| return self.vision_tower.dtype |
| |
