Upload HfMoondream

config.json

@@ -9,5 +9,5 @@
   "config": {},
   "model_type": "moondream1",
   "torch_dtype": "float16",
-  "transformers_version": "4.48.0"
+  "transformers_version": "4.44.0"
 }

generation_config.json

@@ -1,4 +1,4 @@
 {
   "_from_model_config": true,
-  "transformers_version": "4.48.0"
+  "transformers_version": "4.44.0"
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fadcffea8c17fe8a20ea68af3a013cf3184a63787ee4453cc9eb75206c7c1f9b
+oid sha256:96dce588e4a319fde7af3c70fbf27e726f4850e22522d0fdc4b165d5e6003ad5
 size 3854538376

moondream.py

@@ -15,13 +15,26 @@ from .region import decode_coordinate, encode_coordinate, decode_size, encode_si
 from .utils import remove_outlier_points
 
 
-SamplingSettings = TypedDict(
-    "SamplingSettings",
-    {"max_tokens": int},
+TextSamplingSettings = TypedDict(
+    "TextSamplingSettings",
+    {
+        "max_tokens": int,
+        "temperature": float,
+        "top_p": float,
+    },
+    total=False,
+)
+
+ObjectSamplingSettings = TypedDict(
+    "ObjectSamplingSettings",
+    {"max_objects": int},
     total=False,
 )
 
 DEFAULT_MAX_TOKENS = 768
+DEFAULT_TEMPERATURE = 0.5
+DEFAULT_TOP_P = 0.3
+DEFAULT_MAX_OBJECTS = 50
 
 
 @dataclass(frozen=True)
@@ -144,7 +157,7 @@ class MoondreamModel(nn.Module):
     def _decode_one_tok(
         self, x: torch.Tensor, attn_mask: torch.Tensor, pos_ids: torch.Tensor
     ):
-        hidden = text_decoder(x[None], self.text, attn_mask, pos_ids, self.config.text)
+        hidden = text_decoder(x, self.text, attn_mask, pos_ids, self.config.text)
         logits = lm_head(hidden, self.text)
         return logits, hidden
 
@@ -209,7 +222,19 @@ class MoondreamModel(nn.Module):
             ],
         )
 
-    def _prefill_prompt(self, prompt_tokens: torch.Tensor, pos: int):
+    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)
+        mask = probs_sum - probs_sort > top_p
+        probs_sort[mask] = 0.0
+        probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+        next_probs = torch.zeros_like(probs)
+        next_probs.scatter_(dim=-1, index=probs_idx, src=probs_sort)
+        return next_probs
+
+    def _prefill_prompt(
+        self, prompt_tokens: torch.Tensor, pos: int, temperature: float, top_p: float
+    ):
         with torch.inference_mode():
             prompt_emb = text_encoder(prompt_tokens, self.text)
             torch._dynamo.mark_dynamic(prompt_emb, 1)
@@ -217,7 +242,14 @@ class MoondreamModel(nn.Module):
             pos_ids = torch.arange(pos, pos + prompt_emb.size(1), dtype=torch.long)
             hidden = self._prefill(prompt_emb, mask, pos_ids)
             logits = lm_head(hidden, self.text)
-            next_token = torch.argmax(logits, dim=-1)
+
+            if temperature == 0:
+                next_token = torch.argmax(logits, dim=-1).unsqueeze(1)
+            else:
+                probs = torch.softmax(logits / temperature, dim=-1)
+                probs = self._apply_top_p(probs, top_p)
+                next_token = torch.multinomial(probs, num_samples=1)
+
         pos = pos + prompt_emb.size(1)
         return logits, hidden, next_token, pos
 
@@ -225,9 +257,23 @@ class MoondreamModel(nn.Module):
         self,
         prompt_tokens: torch.Tensor,
         pos: int,
-        max_tokens: int,
+        settings: Optional[TextSamplingSettings] = None,
     ):
-        _, _, next_token, pos = self._prefill_prompt(prompt_tokens, pos)
+        max_tokens = (
+            settings.get("max_tokens", DEFAULT_MAX_TOKENS)
+            if settings
+            else DEFAULT_MAX_TOKENS
+        )
+        temperature = (
+            settings.get("temperature", DEFAULT_TEMPERATURE)
+            if settings
+            else DEFAULT_TEMPERATURE
+        )
+        top_p = settings.get("top_p", DEFAULT_TOP_P) if settings else DEFAULT_TOP_P
+
+        _, _, next_token, pos = self._prefill_prompt(
+            prompt_tokens, pos, temperature, top_p
+        )
 
         def generator(next_token, pos):
             mask = torch.zeros(1, 1, 2048, device=self.device, dtype=torch.bool)
@@ -275,7 +321,14 @@ class MoondreamModel(nn.Module):
                     mask[:, :, pos], pos_ids[0] = 1, pos
                     logits, _ = self._decode_one_tok(next_emb, mask, pos_ids)
                     pos += 1
-                    next_token = torch.argmax(logits, dim=-1)
+
+                    if temperature == 0:
+                        next_token = torch.argmax(logits, dim=-1).unsqueeze(1)  # (1, 1)
+                    else:
+                        probs = torch.softmax(logits / temperature, dim=-1)  # (1, V)
+                        probs = self._apply_top_p(probs, top_p)
+                        next_token = torch.multinomial(probs, num_samples=1)  # (1, 1)
+
                     generated_tokens += 1
 
             # Flush any remaining text in the cache
@@ -292,7 +345,7 @@ class MoondreamModel(nn.Module):
         image: Union[Image.Image, EncodedImage],
         question: str,
         stream: bool = False,
-        settings: Optional[SamplingSettings] = None,
+        settings: Optional[TextSamplingSettings] = None,
     ):
         if self.config.tokenizer.templates["query"] is None:
             raise NotImplementedError("Model does not support querying.")
@@ -309,12 +362,8 @@ class MoondreamModel(nn.Module):
             device=self.device,
         )
 
-        max_tokens = DEFAULT_MAX_TOKENS
-        if settings:
-            max_tokens = settings.get("max_tokens", DEFAULT_MAX_TOKENS)
-
         def generator():
-            for token in self._generate_text(prompt_tokens, image.pos, max_tokens):
+            for token in self._generate_text(prompt_tokens, image.pos, settings):
                 yield token
 
         if stream:
@@ -332,7 +381,7 @@ class MoondreamModel(nn.Module):
         image: Union[Image.Image, EncodedImage],
         length: Literal["normal", "short", "long"] = "normal",
         stream: bool = False,
-        settings: Optional[SamplingSettings] = None,
+        settings: Optional[TextSamplingSettings] = None,
     ):
         if self.config.tokenizer.templates["caption"] is None:
             raise NotImplementedError("Model does not support captioning.")
@@ -346,12 +395,8 @@ class MoondreamModel(nn.Module):
             [self.config.tokenizer.templates["caption"][length]], device=self.device
         )
 
-        max_tokens = DEFAULT_MAX_TOKENS
-        if settings:
-            max_tokens = settings.get("max_tokens", DEFAULT_MAX_TOKENS)
-
         def generator():
-            for token in self._generate_text(prompt_tokens, image.pos, max_tokens):
+            for token in self._generate_text(prompt_tokens, image.pos, settings):
                 yield token
 
         if stream:
@@ -365,7 +410,7 @@ class MoondreamModel(nn.Module):
         next_token: torch.Tensor,
         pos: int,
         include_size: bool = True,
-        max_points: int = 50,
+        max_objects: int = DEFAULT_MAX_OBJECTS,
     ):
         out = []
         mask = torch.zeros(1, 1, 2048, device=self.device, dtype=torch.bool)
@@ -375,13 +420,13 @@ class MoondreamModel(nn.Module):
         with torch.inference_mode():
             while (
                 next_token.item() != self.config.tokenizer.eos_id
-                and len(out) < max_points
+                and len(out) < max_objects
             ):
                 x_logits = decode_coordinate(hidden, self.region)
                 x_center = torch.argmax(x_logits, dim=-1) / x_logits.size(-1)
                 next_emb = encode_coordinate(
                     x_center.to(dtype=x_logits.dtype), self.region
-                )
+                ).unsqueeze(0)
 
                 # Decode y-coordinate
                 mask[:, :, pos], pos_ids[0] = 1, pos
@@ -391,7 +436,7 @@ class MoondreamModel(nn.Module):
                 y_center = torch.argmax(y_logits, dim=-1) / y_logits.size(-1)
                 next_emb = encode_coordinate(
                     y_center.to(dtype=y_logits.dtype), self.region
-                )
+                ).unsqueeze(0)
 
                 # Decode size
                 if include_size:
@@ -409,12 +454,16 @@ class MoondreamModel(nn.Module):
                     w = torch.pow(2.0, (w_bin.float() / 1023.0) * 10.0 - 10.0)
                     h = torch.pow(2.0, (h_bin.float() / 1023.0) * 10.0 - 10.0)
 
-                    next_emb = encode_size(
-                        torch.tensor(
-                            [w, h], device=self.device, dtype=size_logits.dtype
-                        ),
-                        self.region,
-                    )[None]
+                    next_emb = (
+                        encode_size(
+                            torch.tensor(
+                                [w, h], device=self.device, dtype=size_logits.dtype
+                            ),
+                            self.region,
+                        )
+                        .unsqueeze(0)
+                        .unsqueeze(0)
+                    )
 
                     # Add object
                     out.append(
@@ -440,7 +489,7 @@ class MoondreamModel(nn.Module):
         self,
         image: Union[Image.Image, EncodedImage],
         object: str,
-        settings: Optional[SamplingSettings] = None,
+        settings: Optional[ObjectSamplingSettings] = None,
     ):
         if self.config.tokenizer.templates["detect"] is None:
             raise NotImplementedError("Model does not support object detection.")
@@ -457,11 +506,18 @@ class MoondreamModel(nn.Module):
             device=self.device,
         )
 
-        _, hidden, next_token, pos = self._prefill_prompt(prompt_tokens, image.pos)
+        _, hidden, next_token, pos = self._prefill_prompt(
+            prompt_tokens, image.pos, temperature=0, top_p=0
+        )
         hidden = hidden[:, -1:, :]
 
+        max_objects = (
+            settings.get("max_objects", DEFAULT_MAX_OBJECTS)
+            if settings
+            else DEFAULT_MAX_OBJECTS
+        )
         objects = self._generate_points(
-            hidden, next_token, pos, include_size=True, max_points=50
+            hidden, next_token, pos, include_size=True, max_objects=max_objects
         )
 
         return {"objects": objects}
@@ -470,7 +526,7 @@ class MoondreamModel(nn.Module):
         self,
         image: Union[Image.Image, EncodedImage],
         object: str,
-        settings: Optional[SamplingSettings] = None,
+        settings: Optional[ObjectSamplingSettings] = None,
     ):
         if self.config.tokenizer.templates["point"] is None:
             raise NotImplementedError("Model does not support pointing.")
@@ -487,11 +543,18 @@ class MoondreamModel(nn.Module):
             device=self.device,
         )
 
-        _, hidden, next_token, pos = self._prefill_prompt(prompt_tokens, image.pos)
+        _, hidden, next_token, pos = self._prefill_prompt(
+            prompt_tokens, image.pos, temperature=0, top_p=0
+        )
         hidden = hidden[:, -1:, :]
 
+        max_objects = (
+            settings.get("max_objects", DEFAULT_MAX_OBJECTS)
+            if settings
+            else DEFAULT_MAX_OBJECTS
+        )
         objects = self._generate_points(
-            hidden, next_token, pos, include_size=False, max_points=50
+            hidden, next_token, pos, include_size=False, max_objects=max_objects
         )
 
         return {"points": objects}
@@ -545,7 +608,7 @@ class MoondreamModel(nn.Module):
                 return None
 
             gaze = self._generate_points(
-                hidden, next_token, pos, include_size=False, max_points=1
+                hidden, next_token, pos, include_size=False, max_objects=1
             )
             return gaze[0]