Fix runtime buffers after load

hf_moondream.py

@@ -45,6 +45,13 @@ class HfMoondream(PreTrainedModel):
         self._is_kv_cache_setup = False
         self.post_init()
 
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        output = super().from_pretrained(*args, **kwargs)
+        model = output[0] if isinstance(output, tuple) else output
+        model.model._refresh_runtime_buffers()
+        return output
+
     def _setup_caches(self):
         if not self._is_kv_cache_setup:
             self.model._setup_caches()

moondream.py

@@ -22,6 +22,7 @@ from .region import (
 )
 from .layers import QuantizedLinear
 from .lora import load_adapter, normalize_adapter_id
+from .rope import precompute_freqs_cis
 from .utils import remove_outlier_points
 
 ImageEncodingSettings = TypedDict(
@@ -131,15 +132,7 @@ class MoondreamModel(nn.Module):
             torch.empty(config.region.size_feat_dim // 2, 2, dtype=dtype).T
         )
 
-        attn_mask = torch.tril(
-            torch.ones(
-                1, 1, config.text.max_context, config.text.max_context, dtype=torch.bool
-            )
-        )
-        patch_w = config.vision.crop_size // config.vision.enc_patch_size
-        prefix_attn_len = 1 + patch_w**2
-        attn_mask[..., :prefix_attn_len, :prefix_attn_len] = 1
-        self.register_buffer("attn_mask", attn_mask, persistent=False)
+        self.register_buffer("attn_mask", self._build_attn_mask(), persistent=False)
 
         self.use_flex_decoding = True
         self._causal_block_mask = None
@@ -171,6 +164,28 @@ class MoondreamModel(nn.Module):
             )
         return self._point_gen_indices
 
+    def _build_attn_mask(self):
+        attn_mask = torch.tril(
+            torch.ones(
+                1,
+                1,
+                self.config.text.max_context,
+                self.config.text.max_context,
+                dtype=torch.bool,
+            )
+        )
+        patch_w = self.config.vision.crop_size // self.config.vision.enc_patch_size
+        prefix_attn_len = 1 + patch_w**2
+        attn_mask[..., :prefix_attn_len, :prefix_attn_len] = 1
+        return attn_mask
+
+    def _refresh_runtime_buffers(self):
+        self.attn_mask = self._build_attn_mask().to(device=self.device)
+        self.text.freqs_cis = precompute_freqs_cis(
+            self.config.text.dim // (2 * self.config.text.n_heads),
+            self.config.text.max_context,
+        ).to(device=self.device)
+
     def _setup_caches(self):
         c = self.config.text
         for b in self.text.blocks: