Update moondream.py

fix: kv cache mismatch and decode size.

moondream.py

@@ -884,36 +884,53 @@ class MoondreamModel(nn.Module):
 
         return {"points": objects}
 
-    def _norm_size_logits(self, size_ret, B: int):
+    # moondream.py
+    def _norm_size_logits(self, size_ret: torch.Tensor | tuple, B: int):
         """
         Accepts any of:
-          • (w_logits, h_logits)
-          • Tensor (B,2,C) or (B,1,2,C) or (1,2,C) or (2,C) (B==1)
+          • tuple/list: (w_logits, h_logits)
+          • Tensor (..., 2, C)    # from batch-safe region.decode_size
+          • Tensor (B, 2*C)       # fallback
+          • Tensor (2, C) when B == 1
         Returns (w_logits, h_logits) each shaped (B, C).
         """
         if isinstance(size_ret, (tuple, list)):
             w_logits, h_logits = size_ret
         else:
             t = size_ret
-            # squeeze all singleton dims except batch & vocab
-            while t.dim() > 3:
-                t = t.squeeze(1)
-            if t.dim() == 3:            # (B,2,C)
+            # if we got (..., 2, C), squeeze a single seq dim if present
+            if t.dim() >= 3 and t.shape[-2] == 2:
+                # bring to (B, 2, C)
+                while t.dim() > 3:
+                    t = t.squeeze(1)
+                if t.dim() != 3 or t.shape[0] not in (1, B):
+                    raise RuntimeError(f"Unexpected batched size logits shape {tuple(size_ret.shape)}")
+                # expand B if needed
+                if t.shape[0] == 1 and B > 1:
+                    t = t.expand(B, -1, -1).contiguous()
                 w_logits, h_logits = t[:, 0, :], t[:, 1, :]
             elif t.dim() == 2:
-                if t.size(0) == 2 and B == 1:   # (2,C) with B==1
+                # (2, C) (B==1)  or (B, 2*C)
+                if t.shape[0] == 2 and B == 1:
                     w_logits, h_logits = t[0].unsqueeze(0), t[1].unsqueeze(0)
-                else:                            # (B,2C) fallback
-                    C2 = t.size(1); C = C2 // 2
+                else:
+                    C2 = t.shape[1]
+                    if C2 % 2 != 0:
+                        raise RuntimeError(f"Cannot split last dim {C2} into (w,h)")
+                    C = C2 // 2
                     w_logits, h_logits = t[:, :C], t[:, C:]
             else:
                 raise RuntimeError(f"Unexpected decode_size shape {tuple(t.shape)}")
-        # final squeeze if needed
+    
+        # final sanity: make sure they’re (B, C)
         if w_logits.dim() == 3: w_logits = w_logits.squeeze(1)
         if h_logits.dim() == 3: h_logits = h_logits.squeeze(1)
+        if w_logits.shape[0] != B or h_logits.shape[0] != B:
+            raise RuntimeError(f"Batched size logits mismatch: got {w_logits.shape[0]} vs B={B}")
         return w_logits.contiguous(), h_logits.contiguous()
 
 
+
     def _load_encoded_image_batched(self, encoded_image, batch_size: int):
         for b, (k, v) in zip(self.text.blocks, encoded_image.caches):
             T = k.size(2)
@@ -969,26 +986,35 @@ class MoondreamModel(nn.Module):
 
 
     def _generate_points_batched(
-        self, hidden, next_token, pos, include_size: bool = True,
-        max_objects: int = 50, lora=None, use_soft_argmax: bool = False):
+        self,
+        hidden,              # (B,1,C)
+        next_token,          # (B,1)  (unused for greedy)
+        pos,                 # int (start position in cache)
+        include_size: bool = True,
+        max_objects: int = 50,
+        lora=None,
+        use_soft_argmax: bool = True,   # reduces jitter/hallucinations
+    ):
         B = hidden.size(0)
         device = self.device
         out = [[] for _ in range(B)]
         eos_id = self.config.tokenizer.eos_id
         max_ctx = self.config.text.max_context
     
-        # 4-D mask: (B,1,1,kv_len); advance with a 1-D position vector
+        # 4-D mask: (B, 1, q_len=1, kv_len)
         mask = torch.zeros(B, 1, 1, max_ctx, device=device, dtype=torch.bool)
-        if int(pos) > 0:
-            mask[:, :, :, :int(pos)] = True
-        pos_id_vec = torch.full((1,), int(pos), device=device, dtype=torch.long)
+        p0 = int(pos)
+        if p0 > 0:
+            mask[:, :, :, :p0] = True
+        # per-row position ids (B,1)
+        pos_ids = torch.full((B, 1), p0, device=device, dtype=torch.long)
     
-        def _center01(logits):
-            # logits: (B, bins) → (B,) in [0,1]
+        # helper: (B, bins) -> (B,) in [0,1]
+        def _argmax01(logits: torch.Tensor) -> torch.Tensor:
             if use_soft_argmax:
-                p = torch.softmax(logits, dim=-1)
-                bins = torch.arange(p.size(-1), device=logits.device, dtype=torch.float32)
-                return (p * bins).sum(dim=-1) / float(p.size(-1) - 1)
+                probs = torch.softmax(logits, dim=-1)
+                bins  = torch.arange(probs.size(-1), device=logits.device, dtype=torch.float32)
+                return (probs * bins).sum(dim=-1) / float(probs.size(-1) - 1)
             idx = logits.argmax(dim=-1).to(torch.float32)
             return idx / float(logits.size(-1) - 1)
     
@@ -997,38 +1023,41 @@ class MoondreamModel(nn.Module):
     
         with torch.inference_mode():
             while alive.any() and (counts < max_objects).any():
-                # x
-                x_logits = decode_coordinate(hidden, self.region)  # (B,1,1024) or (B,1024)
-                if x_logits.dim() == 3: x_logits = x_logits.squeeze(1)
-                x_center = _center01(x_logits)
-                x_emb = encode_coordinate(x_center.to(dtype=x_logits.dtype).unsqueeze(-1), self.region).unsqueeze(1)
+                # ---------------- x ----------------
+                x_logits = decode_coordinate(hidden, self.region)   # (B,1,1024) or (B,1024)
+                if x_logits.dim() == 3:
+                    x_logits = x_logits.squeeze(1)                  # -> (B,1024)
+                x_center = _argmax01(x_logits)                      # (B,)
+                x_emb = encode_coordinate(
+                    x_center.to(dtype=x_logits.dtype).unsqueeze(-1), # (B,1)
+                    self.region
+                ).unsqueeze(1)                                       # (B,1,C)
     
-                mask[:, :, :, pos_id_vec] = True
-                logits, hidden = self._decode_one_tok(x_emb, mask, pos_id_vec, lora)
-                pos_id_vec += 1
+                # advance one token for ALIVE rows only
+                step_col = int(pos_ids[0, 0].item())
+                mask[alive, :, :, step_col] = True
+                logits, hidden = self._decode_one_tok(x_emb, mask, pos_ids, lora)
+                pos_ids[alive, 0] += 1
     
-                # y
+                # ---------------- y ----------------
                 y_logits = decode_coordinate(hidden, self.region)
-                if y_logits.dim() == 3: y_logits = y_logits.squeeze(1)
-                y_center = _center01(y_logits)
-                y_emb = encode_coordinate(y_center.to(dtype=y_logits.dtype).unsqueeze(-1), self.region).unsqueeze(1)
+                if y_logits.dim() == 3:
+                    y_logits = y_logits.squeeze(1)                  # (B,1024)
+                y_center = _argmax01(y_logits)                      # (B,)
+                y_emb = encode_coordinate(
+                    y_center.to(dtype=y_logits.dtype).unsqueeze(-1),
+                    self.region
+                ).unsqueeze(1)                                       # (B,1,C)
     
-                mask[:, :, :, pos_id_vec] = True
-                logits, hidden = self._decode_one_tok(y_emb, mask, pos_id_vec, lora)
-                pos_id_vec += 1
+                step_col = int(pos_ids[0, 0].item())
+                mask[alive, :, :, step_col] = True
+                logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
+                pos_ids[alive, 0] += 1
     
                 if include_size:
+                    # ------------- size (w,h) -------------
                     size_ret = decode_size(hidden, self.region)
-                    # Robust parse: accept (w,h) tuple OR Tensor (B,2,C)/(B,1,2,C)
-                    if isinstance(size_ret, (tuple, list)):
-                        w_logits, h_logits = size_ret
-                    else:
-                        t = size_ret
-                        if t.dim() == 4:                 # (B,1,2,C)
-                            t = t.squeeze(1)            # → (B,2,C)
-                        if t.dim() != 3 or t.size(1) != 2:
-                            raise RuntimeError(f"Unexpected decode_size shape {tuple(t.shape)}")
-                        w_logits, h_logits = t[:, 0, :], t[:, 1, :]
+                    w_logits, h_logits = self._norm_size_logits(size_ret, B)  # each (B,C)
     
                     if use_soft_argmax:
                         bins = torch.arange(w_logits.size(-1), device=device, dtype=torch.float32)
@@ -1038,14 +1067,18 @@ class MoondreamModel(nn.Module):
                         w_bin = w_logits.argmax(dim=-1).to(torch.float32)
                         h_bin = h_logits.argmax(dim=-1).to(torch.float32)
     
-                    # inverse log-scale mapping used by md2
-                    w = torch.pow(2.0, (w_bin / 1023.0) * 10.0 - 10.0)
-                    h = torch.pow(2.0, (h_bin / 1023.0) * 10.0 - 10.0)
+                    # inverse log-scale mapping used by MD2
+                    w = torch.pow(2.0, (w_bin / 1023.0) * 10.0 - 10.0)  # (B,)
+                    h = torch.pow(2.0, (h_bin / 1023.0) * 10.0 - 10.0)  # (B,)
     
-                    size_emb = encode_size(torch.stack([w, h], dim=1).to(dtype=w_logits.dtype), self.region).unsqueeze(1)
+                    size_emb = encode_size(
+                        torch.stack([w, h], dim=1).to(dtype=w_logits.dtype),  # (B,2)
+                        self.region
+                    ).unsqueeze(1)                                           # (B,1,C)
     
-                    for i in range(B):
-                        if not alive[i]: continue
+                    # record boxes only for ALIVE rows
+                    alive_idx = alive.nonzero(as_tuple=False).view(-1)
+                    for i in alive_idx.tolist():
                         xl = (x_center[i] - w[i] / 2).item()
                         xr = (x_center[i] + w[i] / 2).item()
                         yt = (y_center[i] - h[i] / 2).item()
@@ -1057,26 +1090,41 @@ class MoondreamModel(nn.Module):
                             "y_max": max(0.0, min(1.0, yb)),
                         })
     
-                    mask[:, :, :, pos_id_vec] = True
-                    logits, hidden = self._decode_one_tok(size_emb, mask, pos_id_vec, lora)
-                    pos_id_vec += 1
-                    next_tok = logits.argmax(dim=-1).squeeze(-1)
+                    step_col = int(pos_ids[0, 0].item())
+                    mask[alive, :, :, step_col] = True
+                    logits, hidden = self._decode_one_tok(size_emb, mask, pos_ids, lora)
+                    pos_ids[alive, 0] += 1
+                    next_tok = logits.argmax(dim=-1)
+                    if next_tok.dim() == 3:  # (B,1,1) possible
+                        next_tok = next_tok.squeeze(-1).squeeze(-1)
+                    elif next_tok.dim() == 2:  # (B,1)
+                        next_tok = next_tok.squeeze(1)
                 else:
-                    for i in range(B):
-                        if alive[i]:
-                            out[i].append({"x": x_center[i].item(), "y": y_center[i].item()})
-                    mask[:, :, :, pos_id_vec] = True
-                    logits, hidden = self._decode_one_tok(y_emb, mask, pos_id_vec, lora)
-                    pos_id_vec += 1
-                    next_tok = logits.argmax(dim=-1).squeeze(-1)
+                    # point mode
+                    alive_idx = alive.nonzero(as_tuple=False).view(-1)
+                    for i in alive_idx.tolist():
+                        out[i].append({"x": x_center[i].item(), "y": y_center[i].item()})
+                    step_col = int(pos_ids[0, 0].item())
+                    mask[alive, :, :, step_col] = True
+                    logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
+                    pos_ids[alive, 0] += 1
+                    next_tok = logits.argmax(dim=-1)
+                    if next_tok.dim() == 3:
+                        next_tok = next_tok.squeeze(-1).squeeze(-1)
+                    elif next_tok.dim() == 2:
+                        next_tok = next_tok.squeeze(1)
     
-                finished_now = (next_tok == eos_id) | (counts >= max_objects - 1)
-                counts = counts + ((~finished_now) & alive).to(counts.dtype)
+                # we added one object/point for all ALIVE rows this iteration
+                counts[alive] += 1
+    
+                # stop rows that hit eos OR reached max_objects
+                finished_now = (next_tok == eos_id) | (counts >= max_objects)
                 alive &= ~finished_now
     
         return out
 
 
+
     def detect_multi(self, image, objects, settings=None):
         if self.config.tokenizer.templates["detect"] is None:
             raise NotImplementedError("Model does not support object detection.")