Support transformers 5: call post_init() and recompute non-persistent buffers

hf_moondream.py

@@ -6,6 +6,7 @@ from typing import Union
 
 from .config import MoondreamConfig
 from .moondream import MoondreamModel
+from .rope import precompute_freqs_cis
 
 # Files sometimes don't get loaded without these...
 from .image_crops import *
@@ -44,6 +45,29 @@ class HfMoondream(PreTrainedModel):
             MoondreamConfig.from_dict(config.config), setup_caches=False
         )
         self._is_kv_cache_setup = False
+        self.post_init()
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        # transformers>=5 builds the model on the meta device, so non-persistent
+        # buffers are not loaded from the checkpoint and must be recomputed here.
+        if module is self.model:
+            config = self.model.config
+            max_context = config.text.max_context
+            patch_w = config.vision.crop_size // config.vision.enc_patch_size
+            prefix_attn_len = 1 + patch_w**2
+            attn_mask = torch.tril(
+                torch.ones(max_context, max_context, dtype=torch.bool)
+            )
+            attn_mask[:prefix_attn_len, :prefix_attn_len] = 1
+            module.attn_mask.copy_(attn_mask)
+        elif module is self.model.text:
+            config = self.model.config.text
+            module.freqs_cis.copy_(
+                precompute_freqs_cis(
+                    config.dim // (2 * config.n_heads), config.max_context
+                )
+            )
 
     def _setup_caches(self):
         if not self._is_kv_cache_setup: