Singularity666/editx
1
1from typing import List, Optional, Tuple, Union2 3import torch4import torch.nn as nn5import torch.nn.functional as F6from torch.nn import CrossEntropyLoss7 8from transformers import AutoConfig, AutoModelForCausalLM, \9 LlamaConfig, LlamaModel, LlamaForCausalLM, \10 CLIPVisionModel, CLIPImageProcessor11 12from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast13 14import os, diffusers15 16DEFAULT_IMAGE_TOKEN = "<image>"17DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"18DEFAULT_IM_START_TOKEN = "<im_start>"19DEFAULT_IM_END_TOKEN = "<im_end>"20 21 22class LlavaConfig(LlamaConfig):23 model_type = "llava"24 25 26class LlavaLlamaModel(LlamaModel):27 config_class = LlavaConfig28 29 def __init__(self, config: LlamaConfig):30 super(LlavaLlamaModel, self).__init__(config)31 32 if hasattr(config, "mm_vision_tower"):33 # HACK: for FSDP34 self.vision_tower = [CLIPVisionModel.from_pretrained(config.mm_vision_tower)]35 # self.vision_tower = CLIPVisionModel.from_pretrained(config.mm_vision_tower)36 37 if hasattr(config, "use_mm_proj"):38 self.mm_projector = nn.Linear(config.mm_hidden_size, config.hidden_size)39 40 def get_vision_tower(self):41 vision_tower = getattr(self, 'vision_tower', None)42 if type(vision_tower) is list:43 vision_tower = vision_tower[0]44 return vision_tower45 46 def initialize_vision_modules(self, vision_tower, mm_vision_select_layer,47 pretrain_mm_mlp_adapter=None, fsdp=None):48 self.config.mm_vision_tower = vision_tower49 50 image_processor = CLIPImageProcessor.from_pretrained(vision_tower)51 52 if not hasattr(self, 'vision_tower'):53 vision_tower = CLIPVisionModel.from_pretrained(vision_tower)54 else:55 vision_tower = self.vision_tower[0]56 vision_tower.requires_grad_(False)57 58 if fsdp is not None and len(fsdp) > 0:59 self.vision_tower = [vision_tower]60 else:61 self.vision_tower = vision_tower62 63 vision_config = vision_tower.config64 num_patches = (vision_config.image_size // vision_config.patch_size) ** 265 66 self.config.use_mm_proj = True67 self.config.mm_hidden_size = vision_config.hidden_size68 self.config.mm_vision_select_layer = mm_vision_select_layer69 70 if not hasattr(self, 'mm_projector'):71 self.mm_projector = nn.Linear(vision_config.hidden_size, self.config.hidden_size)72 73 if pretrain_mm_mlp_adapter is not None:74 mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location='cpu')75 self.mm_projector.load_state_dict({k.split('.')[-1]: v for k, v in mm_projector_weights.items()})76 77 return dict(78 image_processor=image_processor,79 image_token_len=num_patches,80 vision_config=vision_config81 )82 83 def forward(84 self,85 input_ids: torch.LongTensor = None,86 attention_mask: Optional[torch.Tensor] = None,87 past_key_values: Optional[List[torch.FloatTensor]] = None,88 inputs_embeds: Optional[torch.FloatTensor] = None,89 use_cache: Optional[bool] = None,90 output_attentions: Optional[bool] = None,91 output_hidden_states: Optional[bool] = None,92 images: Optional[torch.FloatTensor] = None,93 return_dict: Optional[bool] = None,94 ) -> Union[Tuple, BaseModelOutputWithPast]:95 96 # HACK: replace back original embeddings for LLaVA pretraining97 orig_embeds_params = getattr(self, 'orig_embeds_params', None)98 # if orig_embeds_params is not None:99 # orig_embeds_params = orig_embeds_params[0]100 # with torch.no_grad():101 # self.get_input_embeddings().weight.data[:-2] = orig_embeds_params[:-2].data102 103 if inputs_embeds is None:104 inputs_embeds = self.embed_tokens(input_ids)105 106 vision_tower = self.get_vision_tower()107 if vision_tower is not None and (input_ids.shape[1] != 1 or self.training) and images is not None:108 # TODO: this is a modified multimodal LLM -- Haotian Liu109 with torch.no_grad():110 if type(images) is list:111 # variable length images112 image_features = []113 for image in images:114 image_forward_out = vision_tower(image.unsqueeze(0), output_hidden_states=True)115 select_hidden_state_layer = getattr(self.config, "mm_vision_select_layer", -1)116 select_hidden_state = image_forward_out.hidden_states[select_hidden_state_layer]117 image_feature = select_hidden_state[:, 1:]118 image_features.append(image_feature)119 else:120 image_forward_outs = vision_tower(images.to(vision_tower.dtype), output_hidden_states=True)121 select_hidden_state_layer = getattr(self.config, "mm_vision_select_layer", -1)122 select_hidden_state = image_forward_outs.hidden_states[select_hidden_state_layer]123 image_features = select_hidden_state[:, 1:].to(images.dtype)124 if type(images) is list:125 image_features = [self.mm_projector(image_feature)[0] for image_feature in image_features]126 else:127 image_features = self.mm_projector(image_features)128 dummy_image_features = torch.zeros(256, 1024, device=inputs_embeds.device, dtype=inputs_embeds.dtype)129 dummy_image_features = self.mm_projector(dummy_image_features)130 131 new_input_embeds = []132 cur_image_idx = 0133 for cur_input_ids, cur_input_embeds in zip(input_ids, inputs_embeds):134 if (cur_input_ids == vision_tower.config.im_patch_token).sum() == 0:135 # multimodal LLM, but the current sample is not multimodal136 cur_input_embeds = cur_input_embeds + (0. * dummy_image_features).sum()137 new_input_embeds.append(cur_input_embeds)138 cur_image_idx += 1139 continue140 if vision_tower.config.use_im_start_end:141 cur_image_features = image_features[cur_image_idx]142 num_patches = cur_image_features.shape[0]143 if (cur_input_ids == vision_tower.config.im_start_token).sum() != (cur_input_ids == vision_tower.config.im_end_token).sum():144 raise ValueError("The number of image start tokens and image end tokens should be the same.")145 image_start_tokens = torch.where(cur_input_ids == vision_tower.config.im_start_token)[0]146 for image_start_token_pos in image_start_tokens:147 cur_image_features = image_features[cur_image_idx].to(device=cur_input_embeds.device)148 num_patches = cur_image_features.shape[0]149 if cur_input_ids[image_start_token_pos + num_patches + 1] != vision_tower.config.im_end_token:150 raise ValueError("The image end token should follow the image start token.")151 if orig_embeds_params is not None:152 cur_new_input_embeds = torch.cat((cur_input_embeds[:image_start_token_pos].detach(), cur_input_embeds[image_start_token_pos:image_start_token_pos+1], cur_image_features, cur_input_embeds[image_start_token_pos + num_patches + 1:image_start_token_pos + num_patches + 2], cur_input_embeds[image_start_token_pos + num_patches + 2:].detach()), dim=0)153 else:154 cur_new_input_embeds = torch.cat((cur_input_embeds[:image_start_token_pos+1], cur_image_features, cur_input_embeds[image_start_token_pos + num_patches + 1:]), dim=0)155 cur_image_idx += 1156 new_input_embeds.append(cur_new_input_embeds)157 else:158 cur_image_features = image_features[cur_image_idx]159 num_patches = cur_image_features.shape[0]160 if (cur_input_ids == vision_tower.config.im_patch_token).sum() != num_patches:161 raise ValueError("The number of image patch tokens should be the same as the number of image patches.")162 masked_indices = torch.where(cur_input_ids == vision_tower.config.im_patch_token)[0]163 mask_index_start = masked_indices[0]164 if (masked_indices != torch.arange(mask_index_start, mask_index_start+num_patches, device=masked_indices.device, dtype=masked_indices.dtype)).any():165 raise ValueError("The image patch tokens should be consecutive.")166 if orig_embeds_params is not None:167 cur_new_input_embeds = torch.cat((cur_input_embeds[:mask_index_start].detach(), cur_image_features, cur_input_embeds[mask_index_start+num_patches:].detach()), dim=0)168 else:169 cur_new_input_embeds = torch.cat((cur_input_embeds[:mask_index_start], cur_image_features, cur_input_embeds[mask_index_start+num_patches:]), dim=0)170 new_input_embeds.append(cur_new_input_embeds)171 cur_image_idx += 1172 inputs_embeds = torch.stack(new_input_embeds, dim=0)173 174 return super(LlavaLlamaModel, self).forward(175 input_ids=None, attention_mask=attention_mask, past_key_values=past_key_values,176 inputs_embeds=inputs_embeds, use_cache=use_cache,177 output_attentions=output_attentions, output_hidden_states=output_hidden_states,178 return_dict=return_dict179 )180 181class EditMapper(nn.Module):182 def __init__(self):183 super().__init__()184 185 self.llm2hid = nn.Linear(4096, 512)186 self.query = nn.Parameter(torch.randn(1, 77, 512))187 self.mapper = nn.Transformer(batch_first=True, norm_first=True,188 d_model=512, nhead=4, num_encoder_layers=4, num_decoder_layers=4,189 dim_feedforward=2048, dropout=0.0)190 self.hid2feat = nn.Linear(512, 768)191 192 def forward(self, llm, emb):193 hid = self.llm2hid(llm+emb)194 hid = self.mapper(hid, self.query.repeat(llm.shape[0], 1, 1))195 feat = self.hid2feat(hid)196 197 return feat198 199class LlavaLlamaForCausalLM(LlamaForCausalLM):200 config_class = LlavaConfig201 202 def __init__(self, config):203 super(LlamaForCausalLM, self).__init__(config)204 self.model = LlavaLlamaModel(config)205 206 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)207 208 self.edit_head = EditMapper()209 210 '''self.scheduler, self.vae, self.unet = [diffusers.DDPMScheduler.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='scheduler'),211 diffusers.AutoencoderKL.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='vae'),212 diffusers.UNet2DConditionModel.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='unet')]213 self.vae.requires_grad_(False)214 self.unet.register_to_config(in_channels=8)215 with torch.no_grad():216 conv = torch.nn.Conv2d(8, self.unet.conv_in.out_channels, self.unet.conv_in.kernel_size, self.unet.conv_in.stride, self.unet.conv_in.padding)217 conv.weight.zero_()218 conv.weight[:, :4, :, :].copy_(self.unet.conv_in.weight)219 self.unet.conv_in = conv'''220 221 # Initialize weights and apply final processing222 self.post_init()223 224 def get_model(self):225 return self.model226 227 def get_vision_tower(self):228 return self.get_model().get_vision_tower()229 230 def get_vision_tower(self):231 model = self.get_model()232 vision_tower = model.vision_tower233 if type(vision_tower) is list:234 vision_tower = vision_tower[0]235 return vision_tower236 237 def forward(238 self,239 input_ids: torch.LongTensor = None,240 attention_mask: Optional[torch.Tensor] = None,241 past_key_values: Optional[List[torch.FloatTensor]] = None,242 inputs_embeds: Optional[torch.FloatTensor] = None,243 labels: Optional[torch.LongTensor] = None,244 use_cache: Optional[bool] = None,245 output_attentions: Optional[bool] = None,246 output_hidden_states: Optional[bool] = None,247 images: Optional[torch.FloatTensor] = None,248 return_dict: Optional[bool] = None,249 p2p_inp=None, p2p_ans=None250 ) -> Union[Tuple, CausalLMOutputWithPast]:251 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions252 output_hidden_states = (253 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states254 )255 return_dict = return_dict if return_dict is not None else self.config.use_return_dict256 257 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)258 outputs = self.model(259 input_ids=input_ids,260 attention_mask=attention_mask,261 past_key_values=past_key_values,262 inputs_embeds=inputs_embeds,263 use_cache=use_cache,264 output_attentions=output_attentions,265 output_hidden_states=output_hidden_states,266 return_dict=return_dict,267 images=images268 )269 270 hidden_states = outputs[0]271 logits = self.lm_head(hidden_states)272 273 loss = None274 if labels is not None:275 # Shift so that tokens < n predict n276 shift_logits = logits[..., :-1, :].contiguous()277 shift_labels = labels[..., 1:].contiguous()278 # Flatten the tokens279 loss_fct = CrossEntropyLoss()280 shift_logits = shift_logits.view(-1, self.config.vocab_size)281 shift_labels = shift_labels.view(-1)282 # Enable model/pipeline parallelism283 shift_labels = shift_labels.to(shift_logits.device)284 loss = loss_fct(shift_logits, shift_labels)285 286 if labels is not None:287 llm = []288 for i in range(labels.shape[0]):289 try: p = labels[i].data.cpu().tolist().index(32003)-1290 except: p = len(labels[i])-9291 p = min(len(hidden_states[i])-9, p)292 llm.append(hidden_states[i][p:p+8].unsqueeze(0))293 llm = torch.cat(llm, dim=0)294 hid_edit = self.edit_head(llm, self.model.embed_tokens.weight[-8:].unsqueeze(dim=0).repeat(labels.shape[0], 1, 1))295 296 B, DROP = labels.shape[0], 0.05297 298 hid_null = self.edit_head(torch.zeros(B, 8, 4096, device=labels.device),299 self.model.embed_tokens.weight[-8:].unsqueeze(dim=0).repeat(labels.shape[0], 1, 1))300 301 with torch.no_grad():302 lat_ans, lat_inp = self.vae.encode(p2p_ans).latent_dist.sample()*self.vae.config.scaling_factor, self.vae.encode(p2p_inp).latent_dist.mode()303 lat_ans, lat_inp = [torch.from_numpy(lat_ans.data.cpu().float().numpy()).to(lat_ans.device),304 torch.from_numpy(lat_inp.data.cpu().float().numpy()).to(lat_inp.device)]305 306 noise = torch.randn_like(lat_ans)307 ts = torch.randint(0, self.scheduler.config.num_train_timesteps, (B, ), device=noise.device).long()308 lat_noise = self.scheduler.add_noise(lat_ans, noise, ts)309 310 prob = torch.rand(B, device=lat_ans.device)311 mask = (prob<(DROP*2)).reshape(B, 1, 1)312 hid_edit = torch.where(mask, hid_null, hid_edit)313 mask = (1.0-((prob>=DROP).to(lat_inp.dtype)*(prob<(DROP*3)).to(lat_inp.dtype))).reshape(B, 1, 1, 1)314 lat_inp *= mask315 316 out = self.unet(torch.cat([lat_noise, lat_inp], dim=1), ts, hid_edit).sample317 318 loss_ce, loss_edit = loss, nn.functional.mse_loss(out, noise, reduction='mean')319 if int(os.environ['LOCAL_RANK'])==0: print('loss_ce:', loss_ce, '/', 'loss_edit:', loss_edit)320 loss = loss_ce+loss_edit*0.5321 322 if not return_dict:323 output = (logits,) + outputs[1:]324 return (loss,) + output if loss is not None else output325 326 return CausalLMOutputWithPast(327 loss=loss,328 logits=logits,329 past_key_values=outputs.past_key_values,330 hidden_states=outputs.hidden_states,331 attentions=outputs.attentions,332 )333 334 def prepare_inputs_for_generation(335 self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs336 ):337 if past_key_values:338 input_ids = input_ids[:, -1:]339 340 # if `inputs_embeds` are passed, we only want to use them in the 1st generation step341 if inputs_embeds is not None and past_key_values is None:342 model_inputs = {"inputs_embeds": inputs_embeds}343 else:344 model_inputs = {"input_ids": input_ids}345 346 model_inputs.update(347 {348 "past_key_values": past_key_values,349 "use_cache": kwargs.get("use_cache"),350 "attention_mask": attention_mask,351 "images": kwargs.get("images", None),352 }353 )354 return model_inputs355 356 def initialize_vision_tokenizer(self, mm_use_im_start_end, tokenizer, device,357 tune_mm_mlp_adapter=False, pretrain_mm_mlp_adapter=None):358 vision_config = self.get_vision_tower().config359 vision_config.use_im_start_end = mm_use_im_start_end360 tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)361 self.resize_token_embeddings(len(tokenizer))362 363 if mm_use_im_start_end:364 num_new_tokens = tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)365 self.resize_token_embeddings(len(tokenizer))366 vision_config.im_start_token, vision_config.im_end_token = tokenizer.convert_tokens_to_ids([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN])367 368 if num_new_tokens > 0:369 input_embeddings = self.get_input_embeddings().weight.data370 output_embeddings = self.get_output_embeddings().weight.data371 372 input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(373 dim=0, keepdim=True)374 output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(375 dim=0, keepdim=True)376 377 input_embeddings[-num_new_tokens:] = input_embeddings_avg378 output_embeddings[-num_new_tokens:] = output_embeddings_avg379 380 if tune_mm_mlp_adapter:381 self.get_model().orig_embeds_params = [self.get_input_embeddings().weight.data.clone().to(device=device)]382 for p in self.get_input_embeddings().parameters():383 p.requires_grad = True384 for p in self.get_output_embeddings().parameters():385 p.requires_grad = False386 387 if pretrain_mm_mlp_adapter:388 mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location='cpu')389 embed_tokens_weight = mm_projector_weights['model.embed_tokens.weight']390 assert num_new_tokens == 2391 if input_embeddings.shape == embed_tokens_weight.shape:392 input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]393 elif embed_tokens_weight.shape[0] == num_new_tokens:394 input_embeddings[-num_new_tokens:] = embed_tokens_weight395 else:396 raise ValueError(f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}.")397 398 vision_config.im_patch_token = tokenizer.convert_tokens_to_ids([DEFAULT_IMAGE_PATCH_TOKEN])[0]399 400AutoConfig.register("llava", LlavaConfig)401AutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM)