Update app.py

app.py

@@ -1,4 +1,3 @@
-#mgie_llava.py:
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -19,11 +18,9 @@ DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
 DEFAULT_IM_START_TOKEN = "<im_start>"
 DEFAULT_IM_END_TOKEN = "<im_end>"
 
-
 class LlavaConfig(LlamaConfig):
     model_type = "llava"
 
-
 class LlavaLlamaModel(LlamaModel):
     config_class = LlavaConfig
 
@@ -31,9 +28,7 @@ class LlavaLlamaModel(LlamaModel):
         super(LlavaLlamaModel, self).__init__(config)
 
         if hasattr(config, "mm_vision_tower"):
-            # HACK: for FSDP
             self.vision_tower = [CLIPVisionModel.from_pretrained(config.mm_vision_tower)]
-            # self.vision_tower = CLIPVisionModel.from_pretrained(config.mm_vision_tower)
 
         if hasattr(config, "use_mm_proj"):
             self.mm_projector = nn.Linear(config.mm_hidden_size, config.hidden_size)
@@ -94,22 +89,15 @@ class LlavaLlamaModel(LlamaModel):
         return_dict: Optional[bool] = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
 
-        # HACK: replace back original embeddings for LLaVA pretraining
         orig_embeds_params = getattr(self, 'orig_embeds_params', None)
-        # if orig_embeds_params is not None:
-        #     orig_embeds_params = orig_embeds_params[0]
-        #     with torch.no_grad():
-        #         self.get_input_embeddings().weight.data[:-2] = orig_embeds_params[:-2].data
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
         vision_tower = self.get_vision_tower()
         if vision_tower is not None and (input_ids.shape[1] != 1 or self.training) and images is not None:
-            # TODO: this is a modified multimodal LLM -- Haotian Liu
             with torch.no_grad():
                 if type(images) is list:
-                    # variable length images
                     image_features = []
                     for image in images:
                         image_forward_out = vision_tower(image.unsqueeze(0), output_hidden_states=True)
@@ -133,7 +121,6 @@ class LlavaLlamaModel(LlamaModel):
             cur_image_idx = 0
             for cur_input_ids, cur_input_embeds in zip(input_ids, inputs_embeds):
                 if (cur_input_ids == vision_tower.config.im_patch_token).sum() == 0:
-                    # multimodal LLM, but the current sample is not multimodal
                     cur_input_embeds = cur_input_embeds + (0. * dummy_image_features).sum()
                     new_input_embeds.append(cur_input_embeds)
                     cur_image_idx += 1
@@ -191,7 +178,7 @@ class EditMapper(nn.Module):
         self.hid2feat = nn.Linear(512, 768)
 
     def forward(self, llm, emb):
-        hid = self.llm2hid(llm+emb)
+        hid = self.llm2hid(llm + emb)
         hid = self.mapper(hid, self.query.repeat(llm.shape[0], 1, 1))
         feat = self.hid2feat(hid)
 
@@ -208,18 +195,19 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM):
 
         self.edit_head = EditMapper()
 
-        '''self.scheduler, self.vae, self.unet = [diffusers.DDPMScheduler.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='scheduler'),
-                                               diffusers.AutoencoderKL.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='vae'),
-                                               diffusers.UNet2DConditionModel.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='unet')]
+        self.scheduler, self.vae, self.unet = [
+            diffusers.DDPMScheduler.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='scheduler'),
+            diffusers.AutoencoderKL.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='vae'),
+            diffusers.UNet2DConditionModel.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='unet')
+        ]
         self.vae.requires_grad_(False)
         self.unet.register_to_config(in_channels=8)
         with torch.no_grad():
             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)
             conv.weight.zero_()
             conv.weight[:, :4, :, :].copy_(self.unet.conv_in.weight)
-            self.unet.conv_in = conv'''
+            self.unet.conv_in = conv
 
-        # Initialize weights and apply final processing
         self.post_init()
 
     def get_model(self):
@@ -228,13 +216,6 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM):
     def get_vision_tower(self):
         return self.get_model().get_vision_tower()
 
-    def get_vision_tower(self):
-        model = self.get_model()
-        vision_tower = model.vision_tower
-        if type(vision_tower) is list:
-            vision_tower = vision_tower[0]
-        return vision_tower
-
     def forward(
         self,
         input_ids: torch.LongTensor = None,
@@ -250,12 +231,9 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM):
         p2p_inp=None, p2p_ans=None
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -273,24 +251,23 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM):
 
         loss = None
         if labels is not None:
-            # Shift so that tokens < n predict n
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()
-            # Flatten the tokens
             loss_fct = CrossEntropyLoss()
             shift_logits = shift_logits.view(-1, self.config.vocab_size)
             shift_labels = shift_labels.view(-1)
-            # Enable model/pipeline parallelism
             shift_labels = shift_labels.to(shift_logits.device)
             loss = loss_fct(shift_logits, shift_labels)
 
         if labels is not None:
             llm = []
             for i in range(labels.shape[0]):
-                try: p = labels[i].data.cpu().tolist().index(32003)-1
-                except: p = len(labels[i])-9
-                p = min(len(hidden_states[i])-9, p)
-                llm.append(hidden_states[i][p:p+8].unsqueeze(0))
+                try:
+                    p = labels[i].data.cpu().tolist().index(32003) - 1
+                except:
+                    p = len(labels[i]) - 9
+                p = min(len(hidden_states[i]) - 9, p)
+                llm.append(hidden_states[i][p:p + 8].unsqueeze(0))
             llm = torch.cat(llm, dim=0)
             hid_edit = self.edit_head(llm, self.model.embed_tokens.weight[-8:].unsqueeze(dim=0).repeat(labels.shape[0], 1, 1))
 
@@ -300,25 +277,41 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM):
                                       self.model.embed_tokens.weight[-8:].unsqueeze(dim=0).repeat(labels.shape[0], 1, 1))
 
             with torch.no_grad():
-                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()
+                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()
                 lat_ans, lat_inp = [torch.from_numpy(lat_ans.data.cpu().float().numpy()).to(lat_ans.device),
                                     torch.from_numpy(lat_inp.data.cpu().float().numpy()).to(lat_inp.device)]
 
             noise = torch.randn_like(lat_ans)
-            ts = torch.randint(0, self.scheduler.config.num_train_timesteps, (B, ), device=noise.device).long()
+            ts = torch.randint(0, self.scheduler.config.num_train_timesteps, (B,), device=noise.device).long()
             lat_noise = self.scheduler.add_noise(lat_ans, noise, ts)
 
             prob = torch.rand(B, device=lat_ans.device)
-            mask = (prob<(DROP*2)).reshape(B, 1, 1)
+            mask = (prob < (DROP * 2)).reshape(B, 1, 1)
             hid_edit = torch.where(mask, hid_null, hid_edit)
-            mask = (1.0-((prob>=DROP).to(lat_inp.dtype)*(prob<(DROP*3)).to(lat_inp.dtype))).reshape(B, 1, 1, 1)
+            mask = (1.0 - ((prob >= DROP).to(lat_inp.dtype) * (prob < (DROP * 3)).to(lat_inp.dtype))).reshape(B, 1, 1, 1)
             lat_inp *= mask
 
+            # Progressive Feature Blending
+            beta_1, beta_2 = 0.7, 0.3
+            visual_features = lat_inp  # Assuming lat_inp represents the visual features
+            B_1 = beta_1 * hid_edit + (1 - beta_1) * visual_features
+            B_2 = beta_2 * hid_edit + (1 - beta_2) * visual_features
+
+            # Cross-Attention Masking
+            attention_scores = torch.matmul(hid_edit, hid_edit.transpose(-1, -2))
+            mask = torch.zeros_like(hid_edit)
+            mask[:, 3:5] = 1.0  # Emphasize central elements (e.g., "hat", "blue")
+            masked_attention_scores = attention_scores * mask
+            hid_edit = torch.matmul(F.softmax(masked_attention_scores, dim=-1), hid_edit)
+
+            # Use blended features in subsequent processing
+            hid_edit = B_1 + B_2
+
             out = self.unet(torch.cat([lat_noise, lat_inp], dim=1), ts, hid_edit).sample
 
             loss_ce, loss_edit = loss, nn.functional.mse_loss(out, noise, reduction='mean')
-            if int(os.environ['LOCAL_RANK'])==0: print('loss_ce:', loss_ce, '/', 'loss_edit:', loss_edit)
-            loss = loss_ce+loss_edit*0.5
+            if int(os.environ.get('LOCAL_RANK', 0)) == 0: print('loss_ce:', loss_ce, '/', 'loss_edit:', loss_edit)
+            loss = loss_ce + loss_edit * 0.5
 
         if not return_dict:
             output = (logits,) + outputs[1:]
@@ -338,7 +331,6 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM):
         if past_key_values:
             input_ids = input_ids[:, -1:]
 
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
         if inputs_embeds is not None and past_key_values is None:
             model_inputs = {"inputs_embeds": inputs_embeds}
         else:
@@ -370,10 +362,8 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM):
                 input_embeddings = self.get_input_embeddings().weight.data
                 output_embeddings = self.get_output_embeddings().weight.data
 
-                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
-                    dim=0, keepdim=True)
-                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
-                    dim=0, keepdim=True)
+                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
+                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
 
                 input_embeddings[-num_new_tokens:] = input_embeddings_avg
                 output_embeddings[-num_new_tokens:] = output_embeddings_avg
@@ -394,9 +384,9 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM):
                 elif embed_tokens_weight.shape[0] == num_new_tokens:
                     input_embeddings[-num_new_tokens:] = embed_tokens_weight
                 else:
-                    raise ValueError(f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}.")
+                    raise ValueError(f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Number of new tokens: {num_new_tokens}.")
 
         vision_config.im_patch_token = tokenizer.convert_tokens_to_ids([DEFAULT_IMAGE_PATCH_TOKEN])[0]
 
 AutoConfig.register("llava", LlavaConfig)
-AutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM)
+AutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM)