Update moondream.py

fix: batched generation

moondream.py

@@ -943,7 +943,14 @@ class MoondreamModel(nn.Module):
             b.kv_cache.v_cache[:, :, :T, :] = v.expand(batch_size, -1, -1, -1)
 
 
-    def _prefill_prompt_batched(self, labels, pos: int, lora=None, temperature: float = 0.0, top_p: float = 0.0):
+    def _prefill_prompt_batched(
+        self,
+        labels,
+        pos: int,
+        lora=None,
+        temperature: float = 0.0,
+        top_p: float = 0.0,
+    ):
         tpl = self.config.tokenizer.templates["detect"]
         if tpl is None:
             raise NotImplementedError("Model does not support object detection.")
@@ -953,43 +960,49 @@ class MoondreamModel(nn.Module):
             ids = tpl["prefix"] + self.tokenizer.encode(" " + lab).ids + tpl["suffix"]
             t = torch.tensor(ids, device=self.device, dtype=torch.long)
             rows.append(t); lens.append(t.numel())
+    
         B, T = len(rows), max(lens)
         eos = self.config.tokenizer.eos_id
     
+        # Pad with EOS in the tensor, but we will still start generation per-row at its own length
         prompt_ids = torch.full((B, T), eos, device=self.device, dtype=torch.long)
         for i, ids in enumerate(rows):
             prompt_ids[i, : ids.numel()] = ids
     
-        prompt_emb = text_encoder(prompt_ids, self.text)               # (B,T,C)
+        prompt_emb = text_encoder(prompt_ids, self.text)          # (B,T,C)
         torch._dynamo.mark_dynamic(prompt_emb, 1)
     
-        base = self.attn_mask[:, :, pos:pos+T, :]                      # (1,1,T,K)
-        mask = base.expand(B, -1, -1, -1).contiguous()                 # (B,1,T,K)
+        base = self.attn_mask[:, :, pos : pos + T, :]            # (1,1,T,K)
+        mask = base.expand(B, -1, -1, -1).contiguous()           # (B,1,T,K)
     
         pos_ids = torch.arange(pos, pos + T, device=self.device, dtype=torch.long)  # (T,)
-        hidden_BTC = self._prefill(prompt_emb, mask, pos_ids, lora)    # (B,T,C)
-        logits_BTV = lm_head(hidden_BTC, self.text)                    # (B,T,V)
+        hidden_BTC = self._prefill(prompt_emb, mask, pos_ids, lora)                 # (B,T,C)
+        logits_BTV = lm_head(hidden_BTC, self.text)                                 # (B,T,V)
     
-        idx = (torch.tensor(lens, device=self.device) - 1).clamp_min(0)  # (B,)
+        # Gather last real token per row
+        idx = (torch.tensor(lens, device=self.device) - 1).clamp_min(0)             # (B,)
         last_hidden = hidden_BTC[torch.arange(B, device=self.device), idx][:, None, :]  # (B,1,C)
         last_logits = logits_BTV[torch.arange(B, device=self.device), idx]              # (B,V)
     
         if temperature == 0.0:
-            next_token = last_logits.argmax(dim=-1, keepdim=True)      # (B,1)
+            next_token = last_logits.argmax(dim=-1, keepdim=True)                   # (B,1)
         else:
             probs = torch.softmax(last_logits / temperature, dim=-1)
             probs = self._apply_top_p(probs, top_p)
-            next_token = torch.multinomial(probs, num_samples=1)       # (B,1)
+            next_token = torch.multinomial(probs, num_samples=1)                    # (B,1)
+    
+        # Per-row next positions (don’t force them all to pos+T)
+        pos_vec = (pos + torch.tensor(lens, device=self.device, dtype=torch.long))  # (B,)
     
-        pos_end = int(pos + T)  # shared next-free position
-        return last_hidden, next_token, pos_end
+        return last_hidden, next_token, pos_vec
+
 
 
     def _generate_points_batched(
         self,
-        hidden,              # (B,1,C)
-        next_token,          # (B,1)  (unused for greedy)
-        pos,                 # int (start position in cache)
+        hidden,                 # (B,1,C)
+        next_token,             # (B,1)  (unused for greedy)
+        pos_vec,                # (B,) next-free position per row
         include_size: bool = True,
         max_objects: int = 50,
         lora=None,
@@ -999,18 +1012,17 @@ class MoondreamModel(nn.Module):
         device = self.device
         out = [[] for _ in range(B)]
         eos_id = self.config.tokenizer.eos_id
+        coord_id = self.config.tokenizer.coord_id
         max_ctx = self.config.text.max_context
     
-        # 4-D mask: (B, 1, q_len=1, kv_len)
+        # Build per-row masks/positions
         mask = torch.zeros(B, 1, 1, max_ctx, device=device, dtype=torch.bool)
-        p0 = int(pos)
-        if p0 > 0:
-            mask[:, :, :, :p0] = True
+        pos_ids = pos_vec.clone().view(B, 1)  # (B,1)
+        for i in range(B):
+            p0 = int(pos_ids[i, 0].item())
+            if p0 > 0:
+                mask[i, 0, 0, :p0] = True
     
-        # per-row position ids (B,1)
-        pos_ids = torch.full((B, 1), p0, device=device, dtype=torch.long)
-    
-        # helper: (B, bins) -> (B,) in [0,1]
         def _argmax01(logits: torch.Tensor) -> torch.Tensor:
             if use_soft_argmax:
                 probs = torch.softmax(logits, dim=-1)
@@ -1019,14 +1031,11 @@ 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)
@@ -1034,39 +1043,29 @@ 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)                  # -> (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)
+                # -------- 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 = _argmax01(x_logits)                       # (B,)
+                x_emb = encode_coordinate(x_center.to(dtype=x_logits.dtype).unsqueeze(-1), self.region).unsqueeze(1)
     
                 idx = _advance_rows(alive)
                 logits, hidden = self._decode_one_tok(x_emb, mask, pos_ids, lora)
                 pos_ids[idx, 0] += 1
     
-                # ---------------- y ----------------
+                # -------- y --------
                 y_logits = decode_coordinate(hidden, self.region)
-                if y_logits.dim() == 3:
-                    y_logits = y_logits.squeeze(1)
+                if y_logits.dim() == 3: 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)
+                y_emb = encode_coordinate(y_center.to(dtype=y_logits.dtype).unsqueeze(-1), self.region).unsqueeze(1)
     
                 idx = _advance_rows(alive)
                 logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
                 pos_ids[idx, 0] += 1
     
                 if include_size:
-                    # ------------- size (w,h) -------------
                     size_ret = decode_size(hidden, self.region)
-                    w_logits, h_logits = self._norm_size_logits(size_ret, B)  # each (B,C)
+                    w_logits, h_logits = self._norm_size_logits(size_ret, B)  # (B,C)
     
                     if use_soft_argmax:
                         bins = torch.arange(w_logits.size(-1), device=device, dtype=torch.float32)
@@ -1076,16 +1075,12 @@ 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)  # (B,)
-                    h = torch.pow(2.0, (h_bin / 1023.0) * 10.0 - 10.0)  # (B,)
+                    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)
     
-                    size_emb = encode_size(
-                        torch.stack([w, h], dim=1).to(dtype=w_logits.dtype),  # (B,2)
-                        self.region
-                    ).unsqueeze(1)                                           # (B,1,C)
+                    size_emb = encode_size(torch.stack([w, h], dim=1).to(dtype=w_logits.dtype), self.region).unsqueeze(1)
     
-                    # record boxes only for ALIVE rows
+                    # record boxes only for rows still alive
                     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()
@@ -1103,7 +1098,6 @@ class MoondreamModel(nn.Module):
                     pos_ids[idx, 0] += 1
                     next_tok = logits.argmax(dim=-1)
                 else:
-                    # point mode
                     for i in alive.nonzero(as_tuple=False).flatten().tolist():
                         out[i].append({"x": x_center[i].item(), "y": y_center[i].item()})
                     idx = _advance_rows(alive)
@@ -1111,16 +1105,22 @@ class MoondreamModel(nn.Module):
                     pos_ids[idx, 0] += 1
                     next_tok = logits.argmax(dim=-1)
     
-                # normalize next_tok to shape (B,)
+                # normalize next_tok to (B,)
                 while next_tok.dim() > 1:
                     next_tok = next_tok.squeeze(-1)
     
+                # we added exactly one object/point to all alive rows
                 counts[alive] += 1
-                finished_now = (next_tok == eos_id) | (counts >= max_objects)
+    
+                # GRAMMAR STOP: only continue if the model asks to start another coord;
+                # otherwise stop row (covers EOS or any non-coord token).
+                continue_mask = (next_tok == coord_id)
+                finished_now = (~continue_mask) | (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.")
@@ -1132,17 +1132,17 @@ class MoondreamModel(nn.Module):
     
         lora = variant_state_dict(settings["variant"], device=self.device) if "variant" in settings else None
     
-        last_hidden, next_token, pos_end = self._prefill_prompt_batched(
+        last_hidden, next_token, pos_vec = self._prefill_prompt_batched(
             objects, enc.pos, lora=lora, temperature=0.0, top_p=0.0
         )
-    
+        
         det_lists = self._generate_points_batched(
-            last_hidden, next_token, pos_end,
+            last_hidden, next_token, pos_vec,
             include_size=True,
             max_objects=settings.get("max_objects", 50),
             lora=lora,
         )
-    
+
         res = {}
         for lab, lst in zip(objects, det_lists):
             for d in lst: