import torch import torch.nn as nn import re import math from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig def build_vision_tower(): vision_tower = 'internlm-xcomposer2d5-ol-7b/base/IXC2d5_clip_l_560' return CLIPVisionTower(vision_tower) def build_vision_projector(input_dim=4096): projector_type = 'mlp2x_gelu' mm_hidden_size = input_dim mid_hidden_size = 4096 hidden_size = 4096 mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type) if mlp_gelu_match: mlp_depth = int(mlp_gelu_match.group(1)) modules = [nn.Linear(mm_hidden_size, mid_hidden_size)] for _ in range(1, mlp_depth): modules.append(nn.GELU()) modules.append(nn.Linear(mid_hidden_size, mid_hidden_size)) return nn.Sequential(*modules) if projector_type == 'identity': return IdentityMap() raise ValueError(f'Unknown projector type: {projector_type}') class IdentityMap(nn.Module): def __init__(self): super().__init__() def forward(self, x, *args, **kwargs): return x @property def config(self): return {"mm_projector_type": 'identity'} class CLIPVisionTower(nn.Module): def __init__(self, vision_tower): super().__init__() self.is_loaded = False self.vision_tower_name = vision_tower # self.conv_dim = 8192 # self.conv = torch.nn.Conv2d(1024, self.conv_dim,3,2,1) self.select_layer = -1 self.select_feature = 'patch' self.load_model() def load_model(self): self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name) self.vision_tower.requires_grad_(False) self.is_loaded = True def resize_pos(self): print('Dummy Resized') 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, glb_GN, sub_GN): if not self.is_loaded: self.load_model() assert type(images) is list shapes = [] input_imgs = [] for img in images: _, C, H, W = img.shape shapes.append([H // 560, W // 560]) sub_img = img.reshape(1, 3, H // 560, 560, W // 560, 560).permute(0, 2, 4, 1, 3, 5).reshape(-1, 3, 560, 560).contiguous() glb_img = torch.nn.functional.interpolate(img.float(), size=(560, 560), mode='bicubic', ).to(sub_img.dtype) input_imgs.append(glb_img) input_imgs.append(sub_img) input_imgs = torch.cat(input_imgs, dim=0) ''' if input_imgs.shape[0] > 50: image_f_1 = self.vision_tower(input_imgs[:50].to(device=self.device, dtype=self.dtype), output_hidden_states=True).hidden_states[self.select_layer][:, 1:] with torch.no_grad(): image_f_2 = self.vision_tower(input_imgs[50:].to(device=self.device, dtype=self.dtype), output_hidden_states=True).hidden_states[self.select_layer][:, 1:] image_features = torch.cat([image_f_1, image_f_2], dim=0).to(input_imgs.dtype) else: image_features = self.vision_tower(input_imgs.to(device=self.device, dtype=self.dtype), output_hidden_states=True).hidden_states[self.select_layer][:, 1:].to(input_imgs.dtype) ''' image_features = \ self.vision_tower(input_imgs.to(device=self.device, dtype=self.dtype), output_hidden_states=True).hidden_states[ self.select_layer][:, 1:].to(input_imgs.dtype) _, N, C = image_features.shape H = int(math.sqrt(N)) assert N == 40 ** 2 output_imgs = [] output_len = [] for [h, w] in shapes: B_ = h * w glb_img = image_features[:1] ### 1, N, C glb_img = glb_img.reshape(1, H, H, C).reshape(1, H // 2, 2, H // 2, 2, C).contiguous().permute(0, 1, 3, 2, 4, 5).reshape(1, H // 2, H // 2, 4 * C).contiguous() temp_glb_GN = sub_GN.repeat(1, H // 2, 1, 1) glb_img = torch.cat([glb_img, temp_glb_GN], dim=2).reshape(1, -1, 4 * C) sub_img = image_features[1:1 + B_] ### ?, N, C sub_img = sub_img.reshape(B_, H, H, C).reshape(B_, H // 2, 2, H // 2, 2, C).contiguous().permute(0, 1, 3, 2, 4, 5).reshape( B_, -1, 4 * C).contiguous() sub_img = sub_img.reshape(1, h, w, 20, 20, -1).permute(0, 1, 3, 2, 4, 5).reshape(1, h * 20, w * 20, 4 * C) temp_sub_GN = sub_GN.repeat(1, h * 20, 1, 1) sub_img = torch.cat([sub_img, temp_sub_GN], dim=2).reshape(1, -1, 4 * C) output_imgs.append(torch.cat([glb_img, glb_GN, sub_img], dim=1)) temp_len = int((h * w + 1) * 400 + 1 + (h + 1) * 20) assert temp_len == output_imgs[-1].shape[1] output_len.append(temp_len) image_features = image_features[1 + h * w:] output_imgs = torch.cat(output_imgs, dim=1) return output_imgs, output_len @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 class PLoRA(nn.Linear): def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None, lora_r=8, lora_alpha=16, lora_dropout=0.05, lora_len=0, **kwargs) -> None: super().__init__(in_features, out_features, bias, device, dtype) self.lora_r = lora_r self.lora_alpha = lora_alpha self.lora_len = lora_len if lora_dropout > 0.: self.lora_dropout = nn.Dropout(p=lora_dropout) else: self.lora_dropout = lambda x: x self.lora_scaling = self.lora_alpha / self.lora_r self.Plora_A = nn.Linear(in_features, self.lora_r, bias=False, device=device, dtype=dtype) self.Plora_B = nn.Linear(self.lora_r, out_features, bias=False, device=device, dtype=dtype) self.reset_parameters() def reset_parameters(self): if hasattr(self, 'lora_A'): # initialize A the same way as the default for nn.Linear and B to zero nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) nn.init.zeros_(self.lora_B.weight) # print ("lora weight init {} {}".format(torch.mean(self.lora_A.weight), torch.mean(self.lora_B.weight))) def forward(self, x, im_mask=None): B, N, C = x.shape im_mask = im_mask.view(-1) x = x.reshape(-1, C) res = super().forward(x) if im_mask is not None: if torch.sum(im_mask) > 0: part_x = x[im_mask] res[im_mask] += self.Plora_B(self.Plora_A( self.lora_dropout(part_x))) * self.lora_scaling else: part_x = x[:1] res[:1] += self.Plora_B(self.Plora_A( self.lora_dropout(part_x))) * 0 return res.reshape(B, N, -1)