import math import re import torch import torch.nn as nn from transformers import CLIPVisionModel def build_vision_tower(): vision_tower = "internlm/internlm-xcomposer2d5-clip" 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.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) -> tuple[torch.Tensor, list[int]]: 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) 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 == 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.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.lora_sft_A = nn.Linear( in_features, 256, bias=False, device=device, dtype=dtype ) self.lora_sft_B = nn.Linear( 256, out_features, bias=False, device=device, dtype=dtype ) self.lora_dpo_A = nn.Linear( in_features, 256, bias=False, device=device, dtype=dtype ) self.lora_dpo_B = nn.Linear( 256, out_features, bias=False, device=device, dtype=dtype ) self.lora_web_A = nn.Linear( in_features, 512, bias=False, device=device, dtype=dtype ) self.lora_web_B = nn.Linear( 512, 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, infer_mode="base"): B, N, C = x.shape im_mask = im_mask.view(-1) x = x.reshape(-1, C) res = super().forward(x) if infer_mode == "web": res += self.lora_web_B(self.lora_web_A(x)) elif infer_mode == "write": res += self.lora_sft_B(self.lora_sft_A(x)) res += self.lora_dpo_B(self.lora_dpo_A(x)) else: pass 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)