Update blip3o/model/language_model/blip3o_qwen.py

blip3o/model/language_model/blip3o_qwen.py

@@ -53,167 +53,167 @@ class blip3oQwenForCausalLM(Qwen2_5_VLForConditionalGeneration, blip3oMetaForCau
         return self.model
 
 
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        ids: Optional[list] = None,
-        i_s_pos: Optional[list] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        gen_image: Optional[torch.FloatTensor] = None,
-        und_image: Optional[torch.FloatTensor] = None,
-        grid_thw: Optional[torch.FloatTensor] = None,
-        image_sizes: Optional[List[List[int]]] = None,
-        return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = 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
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+    # def forward(
+    #     self,
+    #     input_ids: torch.LongTensor = None,
+    #     attention_mask: Optional[torch.Tensor] = None,
+    #     position_ids: Optional[torch.LongTensor] = None,
+    #     past_key_values: Optional[List[torch.FloatTensor]] = None,
+    #     inputs_embeds: Optional[torch.FloatTensor] = None,
+    #     labels: Optional[torch.LongTensor] = None,
+    #     ids: Optional[list] = None,
+    #     i_s_pos: Optional[list] = None,
+    #     use_cache: Optional[bool] = None,
+    #     output_attentions: Optional[bool] = None,
+    #     output_hidden_states: Optional[bool] = None,
+    #     gen_image: Optional[torch.FloatTensor] = None,
+    #     und_image: Optional[torch.FloatTensor] = None,
+    #     grid_thw: Optional[torch.FloatTensor] = None,
+    #     image_sizes: Optional[List[List[int]]] = None,
+    #     return_dict: Optional[bool] = None,
+    #     cache_position: Optional[torch.LongTensor] = 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
+    #     )
+    #     return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         
-        if inputs_embeds is None:
-            (
-                input_ids,
-                position_ids,
-                attention_mask,
-                past_key_values,
-                inputs_embeds,
-                labels,
-                latents
-            ) = self.prepare_inputs_labels_for_multimodal(
-                input_ids,
-                position_ids,
-                attention_mask,
-                past_key_values,
-                labels,
-                gen_image,
-                und_image,
-                grid_thw,
-                i_s_pos,
-                image_sizes
-            )
-
-        outputs = self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
+    #     if inputs_embeds is None:
+    #         (
+    #             input_ids,
+    #             position_ids,
+    #             attention_mask,
+    #             past_key_values,
+    #             inputs_embeds,
+    #             labels,
+    #             latents
+    #         ) = self.prepare_inputs_labels_for_multimodal(
+    #             input_ids,
+    #             position_ids,
+    #             attention_mask,
+    #             past_key_values,
+    #             labels,
+    #             gen_image,
+    #             und_image,
+    #             grid_thw,
+    #             i_s_pos,
+    #             image_sizes
+    #         )
+
+    #     outputs = self.model(
+    #         input_ids=input_ids,
+    #         attention_mask=attention_mask,
+    #         position_ids=position_ids,
+    #         past_key_values=past_key_values,
+    #         inputs_embeds=inputs_embeds,
+    #         use_cache=use_cache,
+    #         output_attentions=output_attentions,
+    #         output_hidden_states=output_hidden_states,
+    #         return_dict=return_dict,
+    #     )
         
-        hidden_states = outputs[0]
-        logits = self.lm_head(hidden_states)
-        logits = logits.float()
+    #     hidden_states = outputs[0]
+    #     logits = self.lm_head(hidden_states)
+    #     logits = logits.float()
         
-        total_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 = torch.nn.CrossEntropyLoss()
-            shift_logits = shift_logits.view(-1, self.config.vocab_size)
-            shift_labels = shift_labels.view(-1)
-            # Enable model parallelism
-            shift_labels = shift_labels.to(shift_logits.device)
-            loss = loss_fct(shift_logits, shift_labels)
-
-
-            # compute image loss
-            # target_img_embeds = torch.clone(inputs_embeds.detach())[:,1:,:] # get target image emb
-            img_loss_funct = torch.nn.MSELoss()
-            # img_hidden_states = self.get_model().down_projector(hidden_states[:,-self.get_n_query():,:])
-            img_hidden_states = []
+    #     total_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 = torch.nn.CrossEntropyLoss()
+    #         shift_logits = shift_logits.view(-1, self.config.vocab_size)
+    #         shift_labels = shift_labels.view(-1)
+    #         # Enable model parallelism
+    #         shift_labels = shift_labels.to(shift_logits.device)
+    #         loss = loss_fct(shift_logits, shift_labels)
+
+
+    #         # compute image loss
+    #         # target_img_embeds = torch.clone(inputs_embeds.detach())[:,1:,:] # get target image emb
+    #         img_loss_funct = torch.nn.MSELoss()
+    #         # img_hidden_states = self.get_model().down_projector(hidden_states[:,-self.get_n_query():,:])
+    #         img_hidden_states = []
             
-            for b in range(hidden_states.shape[0]):
-                img_hidden_states.append(hidden_states[b,i_s_pos[b]:i_s_pos[b]+64,:])
-            img_hidden_states = torch.stack(img_hidden_states,dim=0)
-            img_hidden_states = self.get_model().down_projector(img_hidden_states)
-            # img_loss = 0.0
-            if latents is None:
-                img_loss = img_loss_funct(img_hidden_states, torch.clone(img_hidden_states.detach()))
-            else:
-                bsz = latents.shape[0]
-                # device = latents.device
-                dtype = latents.dtype
-                noise = torch.randn_like(latents, device=latents.device)
-                u = torch.rand(size=(bsz,), device="cpu")
-                indices = (u * self.get_model().noise_scheduler.config.num_train_timesteps).long()
-                timesteps = self.get_model().noise_scheduler.timesteps[indices].to(device=latents.device)
-                sigmas = self.get_sigmas(timesteps, latents.device, n_dim=latents.ndim, dtype=dtype)
-                noisy_latents = (1.0 - sigmas) * latents + sigmas * noise
-                noise_pred = self.get_model().dit(
-                    x=noisy_latents,
-                    timestep=timesteps,
-                    z_latents=self.mask_drop(img_hidden_states),
-                )
-                target = noise - latents
-                img_loss = F.mse_loss(noise_pred.float(), target.float(), reduction="mean")
-            print(f"img loss {img_loss}")
-            total_loss = img_loss
-
-        return CausalLMOutputWithPast(
-            loss=total_loss,
-            logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
+    #         for b in range(hidden_states.shape[0]):
+    #             img_hidden_states.append(hidden_states[b,i_s_pos[b]:i_s_pos[b]+64,:])
+    #         img_hidden_states = torch.stack(img_hidden_states,dim=0)
+    #         img_hidden_states = self.get_model().down_projector(img_hidden_states)
+    #         # img_loss = 0.0
+    #         if latents is None:
+    #             img_loss = img_loss_funct(img_hidden_states, torch.clone(img_hidden_states.detach()))
+    #         else:
+    #             bsz = latents.shape[0]
+    #             # device = latents.device
+    #             dtype = latents.dtype
+    #             noise = torch.randn_like(latents, device=latents.device)
+    #             u = torch.rand(size=(bsz,), device="cpu")
+    #             indices = (u * self.get_model().noise_scheduler.config.num_train_timesteps).long()
+    #             timesteps = self.get_model().noise_scheduler.timesteps[indices].to(device=latents.device)
+    #             sigmas = self.get_sigmas(timesteps, latents.device, n_dim=latents.ndim, dtype=dtype)
+    #             noisy_latents = (1.0 - sigmas) * latents + sigmas * noise
+    #             noise_pred = self.get_model().dit(
+    #                 x=noisy_latents,
+    #                 timestep=timesteps,
+    #                 z_latents=self.mask_drop(img_hidden_states),
+    #             )
+    #             target = noise - latents
+    #             img_loss = F.mse_loss(noise_pred.float(), target.float(), reduction="mean")
+    #         print(f"img loss {img_loss}")
+    #         total_loss = img_loss
+
+    #     return CausalLMOutputWithPast(
+    #         loss=total_loss,
+    #         logits=logits,
+    #         past_key_values=outputs.past_key_values,
+    #         hidden_states=outputs.hidden_states,
+    #         attentions=outputs.attentions,
+    #     )
         
 
-    @torch.no_grad()
-    def generate(
-        self,
-        inputs: Optional[torch.Tensor] = None,
-        images: Optional[torch.Tensor] = None,
-        image_sizes: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> Union[GenerateOutput, torch.LongTensor]:
-        position_ids = kwargs.pop("position_ids", None)
-        attention_mask = kwargs.pop("attention_mask", None)
-        if "inputs_embeds" in kwargs:
-            raise NotImplementedError("`inputs_embeds` is not supported")
-
-        if images is not None:
-            (
-                inputs,
-                position_ids,
-                attention_mask,
-                _,
-                inputs_embeds,
-                img_indicator,
-                _
-            ) = self.prepare_inputs_labels_for_understanding(
-                inputs,
-                position_ids,
-                attention_mask,
-                None,
-                None,
-                images,
-                image_sizes=image_sizes
-            )
-        else:
-            inputs_embeds = self.get_model().embed_tokens(inputs)
-
-        return super().generate(
-            position_ids=position_ids,
-            attention_mask=attention_mask,
-            inputs_embeds=inputs_embeds,
-            **kwargs
-        )
+    # @torch.no_grad()
+    # def generate(
+    #     self,
+    #     inputs: Optional[torch.Tensor] = None,
+    #     images: Optional[torch.Tensor] = None,
+    #     image_sizes: Optional[torch.Tensor] = None,
+    #     **kwargs,
+    # ) -> Union[GenerateOutput, torch.LongTensor]:
+    #     position_ids = kwargs.pop("position_ids", None)
+    #     attention_mask = kwargs.pop("attention_mask", None)
+    #     if "inputs_embeds" in kwargs:
+    #         raise NotImplementedError("`inputs_embeds` is not supported")
+
+    #     if images is not None:
+    #         (
+    #             inputs,
+    #             position_ids,
+    #             attention_mask,
+    #             _,
+    #             inputs_embeds,
+    #             img_indicator,
+    #             _
+    #         ) = self.prepare_inputs_labels_for_understanding(
+    #             inputs,
+    #             position_ids,
+    #             attention_mask,
+    #             None,
+    #             None,
+    #             images,
+    #             image_sizes=image_sizes
+    #         )
+    #     else:
+    #         inputs_embeds = self.get_model().embed_tokens(inputs)
+
+    #     return super().generate(
+    #         position_ids=position_ids,
+    #         attention_mask=attention_mask,
+    #         inputs_embeds=inputs_embeds,
+    #         **kwargs
+    #     )
 
     @torch.no_grad()
     def generate_image(