Update moondream.py

feat: adding warning to code, please do not touch the cache_reset.

moondream.py

@@ -249,49 +249,6 @@ class MoondreamModel(nn.Module):
 
         return self._vis_proj(global_features, reconstructed)
 
-
-    def encode_image(
-        self,
-        image: Union[Image.Image, EncodedImage],
-        settings: Optional[ImageEncodingSettings] = None,
-    ) -> EncodedImage:
-        # Always start from single-row caches; avoids leftovers from batched runs.
-        self._setup_caches()
-    
-        if isinstance(image, EncodedImage):
-            return image
-        elif not isinstance(image, Image.Image):
-            raise ValueError("image must be a PIL Image or EncodedImage")
-    
-        lora = (
-            variant_state_dict(settings["variant"], device=self.device)
-            if settings is not None and "variant" in settings
-            else None
-        )
-    
-        with torch.inference_mode():
-            img_emb = self._run_vision_encoder(image)
-            bos_emb = text_encoder(
-                torch.tensor([[self.config.tokenizer.bos_id]], device=self.device), self.text
-            )
-            inputs_embeds = torch.cat([bos_emb, img_emb[None]], dim=1)
-            mask = self.attn_mask[:, :, 0 : inputs_embeds.size(1), :]
-            pos_ids = torch.arange(inputs_embeds.size(1), dtype=torch.long)
-            self._prefill(inputs_embeds, mask, pos_ids, lora)
-    
-        return EncodedImage(
-            pos=inputs_embeds.size(1),
-            caches=[
-                (
-                    b.kv_cache.k_cache[:, :, : inputs_embeds.size(1), :].clone(),
-                    b.kv_cache.v_cache[:, :, : inputs_embeds.size(1), :].clone(),
-                )
-                for b in self.text.blocks
-            ],
-        )
-
-
-
     def _apply_top_p(self, probs: torch.Tensor, top_p: float):
         probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
         probs_sum = torch.cumsum(probs_sort, dim=-1)
@@ -563,6 +520,46 @@ class MoondreamModel(nn.Module):
 
         return generator(next_token, pos)
 
+    def encode_image(
+        self,
+        image: Union[Image.Image, EncodedImage],
+        settings: Optional[ImageEncodingSettings] = None,
+    ) -> EncodedImage:
+        # Always start from single-row caches; avoids leftovers from batched runs. DO NOT TOUCH THIS!!!!!!!!!
+        self._setup_caches()
+    
+        if isinstance(image, EncodedImage):
+            return image
+        elif not isinstance(image, Image.Image):
+            raise ValueError("image must be a PIL Image or EncodedImage")
+    
+        lora = (
+            variant_state_dict(settings["variant"], device=self.device)
+            if settings is not None and "variant" in settings
+            else None
+        )
+    
+        with torch.inference_mode():
+            img_emb = self._run_vision_encoder(image)
+            bos_emb = text_encoder(
+                torch.tensor([[self.config.tokenizer.bos_id]], device=self.device), self.text
+            )
+            inputs_embeds = torch.cat([bos_emb, img_emb[None]], dim=1)
+            mask = self.attn_mask[:, :, 0 : inputs_embeds.size(1), :]
+            pos_ids = torch.arange(inputs_embeds.size(1), dtype=torch.long)
+            self._prefill(inputs_embeds, mask, pos_ids, lora)
+    
+        return EncodedImage(
+            pos=inputs_embeds.size(1),
+            caches=[
+                (
+                    b.kv_cache.k_cache[:, :, : inputs_embeds.size(1), :].clone(),
+                    b.kv_cache.v_cache[:, :, : inputs_embeds.size(1), :].clone(),
+                )
+                for b in self.text.blocks
+            ],
+        )
+
     def query(
         self,
         image: Optional[Union[Image.Image, EncodedImage]] = None,