|
|
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 |
|
|
|
