lklimkiewicz's picture
Upload InternLMXComposer2ForCausalLM
d4f316d verified
import torch
import torch.nn as nn
import re
import math
from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig
def build_vision_tower():
vision_tower = 'openai/clip-vit-large-patch14-336'
return CLIPVisionTower(vision_tower)
def build_vision_projector():
projector_type = 'mlp2x_gelu'
mm_hidden_size = 4096
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.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//336, W//336])
sub_img = img.reshape(1,3,H//336,336,W//336,336).permute(0,2,4,1,3,5).reshape(-1,3,336,336).contiguous()
glb_img = torch.nn.functional.interpolate(img.float(), size=(336,336), mode='bicubic',).to(sub_img.dtype)
input_imgs.append(glb_img)
input_imgs.append(sub_img)
input_imgs = torch.cat(input_imgs, dim=0)
image_forward_outs = self.vision_tower(input_imgs.to(device=self.device, dtype=self.dtype), output_hidden_states=True)
image_features = self.feature_select(image_forward_outs).to(input_imgs.dtype) ### B*?, N, C
_, N, C = image_features.shape
H = int(math.sqrt(N))
assert N == 24 ** 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, 12, 12, -1).permute(0,1,3,2,4,5).reshape(1,h*12,w*12,4*C)
temp_sub_GN = sub_GN.repeat(1, h*12, 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)*144 + 1 + (h+1)*12)
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
x = x.reshape(-1, C)
im_mask = im_mask.view(-1)
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)