update generate function

modeling_florence2.py

@@ -2655,7 +2655,7 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
             return image_features, image_attention_mask
 
         task_prefix_embeds = inputs_embeds
-
+        
         if task_prefix_attention_mask is None:
             task_prefix_attention_mask = torch.ones(batch_size, task_prefix_embeds.size(1), device=device)
 
@@ -2721,12 +2721,14 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
         >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
         "A green car parked in front of a yellow building."
         ```"""
+        print("asdasdasdasdasdasdasdasda")
+
         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
-
+        print("asdasdasdasdasdasdasdasda")
         image_features = None
         if inputs_embeds is None:
             # 1. Extra the input embeddings
@@ -2736,8 +2738,10 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
             if pixel_values is not None:
                 # (batch_size, num_image_tokens, hidden_size)
                 image_features = self._encode_image(pixel_values)
-                inputs_embeds, attention_mask = self._merge_input_ids_with_image_features(image_features, inputs_embeds, task_prefix_attention_mask=attention_mask)
-
+                inputs_embeds, attention_mask = self._merge_input_ids_with_image_features(image_features, inputs_embeds)
+        
+        print(attention_mask)
+        
         if inputs_embeds is not None:
             attention_mask = attention_mask.to(inputs_embeds.dtype)
         outputs = self.language_model(
@@ -2783,6 +2787,7 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
         input_ids, 
         inputs_embeds=None,
         pixel_values=None,
+        attention_mask=None,
         **kwargs
         ):
 
@@ -2793,11 +2798,12 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
             # 2. Merge text and images
             if pixel_values is not None:
                 image_features = self._encode_image(pixel_values)
-                inputs_embeds, attention_mask = self._merge_input_ids_with_image_features(image_features, inputs_embeds)
+                inputs_embeds, attention_mask = self._merge_input_ids_with_image_features(image_features, inputs_embeds, task_prefix_attention_mask=attention_mask)
         
         return self.language_model.generate(
             input_ids=None,
             inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
             **kwargs
         )