Update moondream.py

fix: _generate_points_batched to advance the mask per_row & never the batch to row_0

moondream.py

@@ -993,7 +993,7 @@ class MoondreamModel(nn.Module):
         include_size: bool = True,
         max_objects: int = 50,
         lora=None,
-        use_soft_argmax: bool = True,   # reduces jitter/hallucinations
+        use_soft_argmax: bool = True,
     ):
         B = hidden.size(0)
         device = self.device
@@ -1006,6 +1006,7 @@ class MoondreamModel(nn.Module):
         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)
     
@@ -1018,41 +1019,49 @@ class MoondreamModel(nn.Module):
             idx = logits.argmax(dim=-1).to(torch.float32)
             return idx / float(logits.size(-1) - 1)
     
+        # advance-one-step for a subset of rows (alive only)
+        def _advance_rows(row_mask: torch.Tensor):
+            idx = row_mask.nonzero(as_tuple=False).flatten()
+            # set each row's next KV column true
+            for i in idx.tolist():
+                col = int(pos_ids[i, 0].item())
+                mask[i, 0, 0, col] = True
+            # decoder step (all rows run, but only alive rows’ pos_ids move)
+            return idx
+    
         alive  = torch.ones(B, dtype=torch.bool, device=device)
         counts = torch.zeros(B, dtype=torch.int32, device=device)
     
         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)                  # -> (B,1024)
                 x_center = _argmax01(x_logits)                      # (B,)
                 x_emb = encode_coordinate(
-                    x_center.to(dtype=x_logits.dtype).unsqueeze(-1), # (B,1)
+                    x_center.to(dtype=x_logits.dtype).unsqueeze(-1),  # (B,1)
                     self.region
                 ).unsqueeze(1)                                       # (B,1,C)
     
-                # advance one token for ALIVE rows only
-                step_col = int(pos_ids[0, 0].item())
-                mask[alive, :, :, step_col] = True
+                idx = _advance_rows(alive)
                 logits, hidden = self._decode_one_tok(x_emb, mask, pos_ids, lora)
-                pos_ids[alive, 0] += 1
+                pos_ids[idx, 0] += 1
     
                 # ---------------- y ----------------
                 y_logits = decode_coordinate(hidden, self.region)
                 if y_logits.dim() == 3:
-                    y_logits = y_logits.squeeze(1)                  # (B,1024)
-                y_center = _argmax01(y_logits)                      # (B,)
+                    y_logits = y_logits.squeeze(1)
+                y_center = _argmax01(y_logits)
                 y_emb = encode_coordinate(
                     y_center.to(dtype=y_logits.dtype).unsqueeze(-1),
                     self.region
-                ).unsqueeze(1)                                       # (B,1,C)
+                ).unsqueeze(1)
     
-                step_col = int(pos_ids[0, 0].item())
-                mask[alive, :, :, step_col] = True
+                idx = _advance_rows(alive)
                 logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
-                pos_ids[alive, 0] += 1
+                pos_ids[idx, 0] += 1
     
                 if include_size:
                     # ------------- size (w,h) -------------
@@ -1077,8 +1086,7 @@ class MoondreamModel(nn.Module):
                     ).unsqueeze(1)                                           # (B,1,C)
     
                     # record boxes only for ALIVE rows
-                    alive_idx = alive.nonzero(as_tuple=False).view(-1)
-                    for i in alive_idx.tolist():
+                    for i in alive.nonzero(as_tuple=False).flatten().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()
@@ -1090,41 +1098,29 @@ class MoondreamModel(nn.Module):
                             "y_max": max(0.0, min(1.0, yb)),
                         })
     
-                    step_col = int(pos_ids[0, 0].item())
-                    mask[alive, :, :, step_col] = True
+                    idx = _advance_rows(alive)
                     logits, hidden = self._decode_one_tok(size_emb, mask, pos_ids, lora)
-                    pos_ids[alive, 0] += 1
+                    pos_ids[idx, 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:
                     # point mode
-                    alive_idx = alive.nonzero(as_tuple=False).view(-1)
-                    for i in alive_idx.tolist():
+                    for i in alive.nonzero(as_tuple=False).flatten().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
+                    idx = _advance_rows(alive)
                     logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
-                    pos_ids[alive, 0] += 1
+                    pos_ids[idx, 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)
     
-                # we added one object/point for all ALIVE rows this iteration
-                counts[alive] += 1
+                # normalize next_tok to shape (B,)
+                while next_tok.dim() > 1:
+                    next_tok = next_tok.squeeze(-1)
     
-                # stop rows that hit eos OR reached max_objects
+                counts[alive] += 1
                 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.")