Update moondream.py

fix: EOS padding for batched gerenation causing generation duplication.

moondream.py

@@ -943,66 +943,63 @@ 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.")
     
-        rows, lens = [], []
+        # Build each row's token ids (variable length)
+        rows_ids, lens = [], []
         for lab in labels:
             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())
+            rows_ids.append(t)
+            lens.append(t.numel())
     
-        B, T = len(rows), max(lens)
-        eos = self.config.tokenizer.eos_id
+        B, T = len(rows_ids), max(lens)
     
-        # 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)
+        # Embed each row, then LEFT-pad using its own first token embedding (neutral),
+        # mirroring upstream moondream2 batching strategy.
+        embs = [text_encoder(t.unsqueeze(0), self.text)[0] for t in rows_ids]  # list of (Li, C)
+        padded = []
+        for e, L in zip(embs, lens):
+            pad = T - L
+            if pad > 0:
+                e = torch.cat([e[:1].repeat(pad, 1), e], dim=0)                # (T, C)
+            padded.append(e)
+        prompt_emb = torch.stack(padded, dim=0)                                 # (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)
-    
+        # Standard prefill over the shared image prefix [pos : pos+T]
+        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)
     
-        # Gather last real token per row
-        idx = (torch.tensor(lens, device=self.device) - 1).clamp_min(0)             # (B,)
+        hidden_BTC = self._prefill(prompt_emb, mask, pos_ids, lora)              # (B,T,C)
+        logits_BTV = lm_head(hidden_BTC, self.text)                               # (B,T,V)
+    
+        # Take the last real token of each 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,)
-    
-        return last_hidden, next_token, pos_vec
-
+        pos_end = int(pos + T)  
+        return last_hidden, next_token, pos_end
 
 
     def _generate_points_batched(
         self,
-        hidden,                 # (B,1,C)
-        next_token,             # (B,1)  (unused for greedy)
-        pos_vec,                # (B,) next-free position per row
+        hidden,              # (B,1,C)  - last token hidden state per row
+        next_token,          # (B,1)    - unused for greedy loop; kept for API
+        pos,                 # int      - first free position in cache
         include_size: bool = True,
         max_objects: int = 50,
         lora=None,
@@ -1012,18 +1009,19 @@ 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
     
-        # Build per-row masks/positions
+        # 4D mask: (B,1,1,K); we advance per-row
         mask = torch.zeros(B, 1, 1, max_ctx, device=device, dtype=torch.bool)
-        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
+        p0 = int(pos)
+        if p0 > 0:
+            mask[:, :, :, :p0] = True
+        pos_ids = torch.full((B, 1), p0, device=device, dtype=torch.long)
     
+        # helper: logits -> normalized [0..1] coordinate (soft-argmax for stability)
         def _argmax01(logits: torch.Tensor) -> torch.Tensor:
+            if logits.dim() == 3:
+                logits = logits.squeeze(1)         # (B, bins)
             if use_soft_argmax:
                 probs = torch.softmax(logits, dim=-1)
                 bins  = torch.arange(probs.size(-1), device=logits.device, dtype=torch.float32)
@@ -1031,41 +1029,36 @@ class MoondreamModel(nn.Module):
             idx = logits.argmax(dim=-1).to(torch.float32)
             return idx / float(logits.size(-1) - 1)
     
-        def _advance_rows(row_mask: torch.Tensor):
-            idx = row_mask.nonzero(as_tuple=False).flatten()
-            for i in idx.tolist():
-                col = int(pos_ids[i, 0].item())
-                mask[i, 0, 0, col] = True
-            return idx
-    
         alive  = torch.ones(B, dtype=torch.bool, device=device)
-        counts = torch.zeros(B, dtype=torch.int32, 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)
-                x_center = _argmax01(x_logits)                       # (B,)
-                x_emb = encode_coordinate(x_center.to(dtype=x_logits.dtype).unsqueeze(-1), self.region).unsqueeze(1)
+                alive_idx = alive.nonzero(as_tuple=False).squeeze(1)
+    
+                # ---------- x ----------
+                x_logits = decode_coordinate(hidden, self.region)                # (B,1,bins) or (B,bins)
+                x_center = _argmax01(x_logits)                                   # (B,)
+                x_emb = encode_coordinate(x_center.to(dtype=x_logits.dtype).unsqueeze(-1), self.region).unsqueeze(1)  # (B,1,C)
     
-                idx = _advance_rows(alive)
+                # advance one token for each alive row (per-row column)
+                mask[alive_idx, 0, 0, pos_ids[alive_idx, 0]] = True
                 logits, hidden = self._decode_one_tok(x_emb, mask, pos_ids, lora)
-                pos_ids[idx, 0] += 1
+                pos_ids[alive_idx, 0] += 1
     
-                # -------- y --------
+                # ---------- y ----------
                 y_logits = decode_coordinate(hidden, self.region)
-                if y_logits.dim() == 3: y_logits = y_logits.squeeze(1)
-                y_center = _argmax01(y_logits)
+                y_center = _argmax01(y_logits)                                   # (B,)
                 y_emb = encode_coordinate(y_center.to(dtype=y_logits.dtype).unsqueeze(-1), self.region).unsqueeze(1)
     
-                idx = _advance_rows(alive)
+                mask[alive_idx, 0, 0, pos_ids[alive_idx, 0]] = True
                 logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
-                pos_ids[idx, 0] += 1
+                pos_ids[alive_idx, 0] += 1
     
                 if include_size:
-                    size_ret = decode_size(hidden, self.region)
-                    w_logits, h_logits = self._norm_size_logits(size_ret, B)  # (B,C)
+                    # ---------- size (w,h) ----------
+                    size_ret = decode_size(hidden, self.region)                   # (...,2,bins)
+                    w_logits, h_logits = self._norm_size_logits(size_ret, B)      # each (B,bins)
     
                     if use_soft_argmax:
                         bins = torch.arange(w_logits.size(-1), device=device, dtype=torch.float32)
@@ -1075,13 +1068,14 @@ 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 (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)
     
                     size_emb = encode_size(torch.stack([w, h], dim=1).to(dtype=w_logits.dtype), self.region).unsqueeze(1)
     
-                    # record boxes only for rows still alive
-                    for i in alive.nonzero(as_tuple=False).flatten().tolist():
+                    # write outputs only for alive rows
+                    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()
@@ -1093,34 +1087,34 @@ class MoondreamModel(nn.Module):
                             "y_max": max(0.0, min(1.0, yb)),
                         })
     
-                    idx = _advance_rows(alive)
+                    mask[alive_idx, 0, 0, pos_ids[alive_idx, 0]] = True
                     logits, hidden = self._decode_one_tok(size_emb, mask, pos_ids, lora)
-                    pos_ids[idx, 0] += 1
+                    pos_ids[alive_idx, 0] += 1
+    
                     next_tok = logits.argmax(dim=-1)
+                    if next_tok.dim() == 3: next_tok = next_tok.squeeze(-1).squeeze(-1)
+                    if next_tok.dim() == 2: next_tok = next_tok.squeeze(1)
                 else:
-                    for i in alive.nonzero(as_tuple=False).flatten().tolist():
+                    # points only
+                    for i in alive_idx.tolist():
                         out[i].append({"x": x_center[i].item(), "y": y_center[i].item()})
-                    idx = _advance_rows(alive)
+                    mask[alive_idx, 0, 0, pos_ids[alive_idx, 0]] = True
                     logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
-                    pos_ids[idx, 0] += 1
+                    pos_ids[alive_idx, 0] += 1
                     next_tok = logits.argmax(dim=-1)
+                    if next_tok.dim() == 3: next_tok = next_tok.squeeze(-1).squeeze(-1)
+                    if next_tok.dim() == 2: next_tok = next_tok.squeeze(1)
     
-                # 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
+                counts[alive] += 1  # we produced one object/point for each alive row
     
-                # 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)
+                # stop rows that hit eos OR reached quota
+                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.")