TEST2

modeling_florence2.py

@@ -831,8 +831,10 @@ class Florence2Attention(nn.Module):
             # reuse k, v, self_attention
             key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
             value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
-            key_states = torch.cat([past_key_value[0], key_states], dim=2)
-            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+            if past_key_value[0] is not None:
+                key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            if past_key_value[1] is not None:
+                value_states = torch.cat([past_key_value[1], value_states], dim=2)
         else:
             # self_attention
             key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
@@ -1173,6 +1175,7 @@ class Florence2SdpaAttention(Florence2Attention):
         if (
             is_cross_attention
             and past_key_value is not None
+            and past_key_value[0] is not None
             and past_key_value[0].shape[2] == key_value_states.shape[1]
         ):
             # reuse k,v, cross_attentions
@@ -1186,8 +1189,10 @@ class Florence2SdpaAttention(Florence2Attention):
             # reuse k, v, self_attention
             key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
             value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
-            key_states = torch.cat([past_key_value[0], key_states], dim=2)
-            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+            if past_key_value[0] is not None:
+                key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            if past_key_value[1] is not None:
+                value_states = torch.cat([past_key_value[1], value_states], dim=2)
         else:
             # self_attention
             key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
@@ -1801,7 +1806,7 @@ class Florence2Decoder(Florence2LanguagePreTrainedModel):
             raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
 
         # past_key_values_length
-        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values and past_key_values[0] and past_key_values[0][0] is not None else 0
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input)
@@ -2200,7 +2205,7 @@ class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel
         **kwargs,
     ):
         # cut decoder_input_ids if past_key_values is used
-        if past_key_values is not None:
+        if past_key_values is not None and past_key_values[0] is not None and past_key_values[0][0] is not None:
             past_length = past_key_values[0][0].shape[2]
 
             # Some generation methods already pass only the last input ID
@@ -3062,7 +3067,7 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel):
         **kwargs,
     ):
         # cut decoder_input_ids if past_key_values is used
-        if past_key_values is not None:
+        if past_key_values is not None and past_key_values[0] is not None and past_key_values[0][0] is not None:
             past_length = past_key_values[0][0].shape[2]
 
             # Some generation methods already pass only the last input ID