Sync upload modeling_florence2.py

models/LLM/Florence-2-large-PromptGen-v2.0/modeling_florence2.py

@@ -1796,6 +1796,12 @@ class Florence2Decoder(Florence2LanguagePreTrainedModel):
             raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
 
         # past_key_values_length
+        if past_key_values is not None:
+            if hasattr(past_key_values, "get_seq_length") and past_key_values.get_seq_length() == 0:
+                past_key_values = None
+            elif not isinstance(past_key_values, tuple):
+                past_key_values = tuple(past_key_values)
+                
         past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
 
         if inputs_embeds is None:
@@ -1939,7 +1945,7 @@ class Florence2Decoder(Florence2LanguagePreTrainedModel):
 
 
 class Florence2LanguageModel(Florence2LanguagePreTrainedModel):
-    _tied_weights_keys = {"decoder.embed_tokens.weight": "encoder.embed_tokens.weight"}
+    _tied_weights_keys = None
 
     def __init__(self, config: Florence2LanguageConfig):
         super().__init__(config)
@@ -1955,8 +1961,8 @@ class Florence2LanguageModel(Florence2LanguagePreTrainedModel):
 
     def _tie_weights(self):
         if self.config.tie_word_embeddings:
-            self._tie_or_clone_weights(self.encoder.embed_tokens, self.shared)
-            self._tie_or_clone_weights(self.decoder.embed_tokens, self.shared)
+            self.encoder.embed_tokens.weight = self.shared.weight
+            self.decoder.embed_tokens.weight = self.shared.weight
 
     def get_input_embeddings(self):
         return self.shared
@@ -2062,10 +2068,7 @@ class Florence2LanguageModel(Florence2LanguagePreTrainedModel):
 
 class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel, GenerationMixin):
     base_model_prefix = "model"
-    _tied_weights_keys = {
-        "model.decoder.embed_tokens.weight": "model.encoder.embed_tokens.weight",
-        "lm_head.weight": "model.encoder.embed_tokens.weight"
-    }
+    _tied_weights_keys = None
     _keys_to_ignore_on_load_missing = ["final_logits_bias"]
 
     def __init__(self, config: Florence2LanguageConfig):
@@ -2077,6 +2080,10 @@ class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel
         # Initialize weights and apply final processing
         self.post_init()
 
+    def _tie_weights(self):
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.model.shared.weight
+
     def get_encoder(self):
         return self.model.get_encoder()
 
@@ -2199,7 +2206,10 @@ class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel
     ):
         # cut decoder_input_ids if past_key_values is used
         if past_key_values is not None:
-            past_length = past_key_values[0][0].shape[2]
+            if hasattr(past_key_values, "get_seq_length"):
+                past_length = past_key_values.get_seq_length()
+            else:
+                past_length = past_key_values[0][0].shape[2]
 
             # Some generation methods already pass only the last input ID
             if decoder_input_ids.shape[1] > past_length:
@@ -2228,6 +2238,9 @@ class Florence2LanguageForConditionalGeneration(Florence2LanguagePreTrainedModel
 
     @staticmethod
     def _reorder_cache(past_key_values, beam_idx):
+        if hasattr(past_key_values, "reorder_cache"):
+            past_key_values.reorder_cache(beam_idx)
+            return past_key_values
         reordered_past = ()
         for layer_past in past_key_values:
             # cached cross_attention states don't have to be reordered -> they are always the same
@@ -2336,6 +2349,17 @@ class Florence2PreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _skip_keys_device_placement = "past_key_values"
 
+    @property
+    def device(self) -> torch.device:
+        try:
+            return next(self.parameters()).device
+        except StopIteration:
+            return torch.device("cpu")
+
+    @device.setter
+    def device(self, value):
+        pass
+
     @property
     def _supports_flash_attn_2(self):
         """
@@ -2491,8 +2515,12 @@ class Florence2VisionModelWithProjection(Florence2PreTrainedModel):
         if self.image_pos_embed is not None:
             x = x.view(batch_size * T, -1, x.shape[-1])
             num_tokens = x.shape[-2]
-            h, w = int(num_tokens ** 0.5), int(num_tokens ** 0.5)
-            assert h * w == num_tokens, 'only support square feature maps for now'
+            h, w = H // 32, W // 32
+            if h * w != num_tokens:
+                aspect_ratio = H / W
+                w = int(round((num_tokens / aspect_ratio) ** 0.5))
+                h = num_tokens // w
+            assert h * w == num_tokens, f'failed to resolve feature map shape: num_tokens={num_tokens}, H={H}, W={W}'
             x = x.view(batch_size * T, h, w, x.shape[-1])
             pos_embed = self.image_pos_embed(x)
             x = x + pos_embed
@@ -2549,10 +2577,7 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel, GenerationMixi
         language_model = Florence2LanguageForConditionalGeneration(config=config.text_config)
 
         if language_model._tied_weights_keys is not None:
-            if isinstance(language_model._tied_weights_keys, dict):
-                self._tied_weights_keys = {f"language_model.{k}": f"language_model.{v}" for k, v in language_model._tied_weights_keys.items()}
-            else:
-                self._tied_weights_keys = [f"language_model.{k}" for k in language_model._tied_weights_keys]
+            self._tied_weights_keys = None
         self.language_model = language_model
 
         self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
@@ -2614,8 +2639,12 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel, GenerationMixi
         if self.image_pos_embed is not None:
             x = x.view(batch_size * T, -1, x.shape[-1])
             num_tokens = x.shape[-2]
-            h, w = int(num_tokens ** 0.5), int(num_tokens ** 0.5)
-            assert h * w == num_tokens, 'only support square feature maps for now'
+            h, w = H // 32, W // 32
+            if h * w != num_tokens:
+                aspect_ratio = H / W
+                w = int(round((num_tokens / aspect_ratio) ** 0.5))
+                h = num_tokens // w
+            assert h * w == num_tokens, f'failed to resolve feature map shape: num_tokens={num_tokens}, H={H}, W={W}'
             x = x.view(batch_size * T, h, w, x.shape[-1])
             pos_embed = self.image_pos_embed(x)
             x = x + pos_embed
@@ -2822,7 +2851,10 @@ class Florence2ForConditionalGeneration(Florence2PreTrainedModel, GenerationMixi
     ):
         # cut decoder_input_ids if past_key_values is used
         if past_key_values is not None:
-            past_length = past_key_values[0][0].shape[2]
+            if hasattr(past_key_values, "get_seq_length"):
+                past_length = past_key_values.get_seq_length()
+            else:
+                past_length = past_key_values[0][0].shape[2]
 
             # Some generation methods already pass only the last input ID
             if decoder_input_ids.shape[1] > past_length: