Update moondream.py

moondream.py

@@ -77,29 +77,28 @@ class KVCache(nn.Module):
             "v_cache", torch.zeros(*cache_shape, device=device, dtype=dtype)
         )
 
+    # In class KVCache, REPLACE the whole update() with this:
     def update(self, pos_ids, k, v):
         """
         Supports:
           • Prefill:  k,v  = (B, n_kv_heads, q_len, d),  pos_ids = (q_len,)
-          • 1-step:   k,v  = (B, n_kv_heads, 1, d),      pos_ids = (B,) or (B,1)
-          • Legacy:   k,v  = (1, n_kv_heads, q_len, d),  pos_ids = scalar int
+          • 1-step:   k,v  = (B, n_kv_heads, 1, d),      pos_ids = (B,1) or (B,)
+          • Legacy:   k,v  = (1, n_kv_heads, q_len, d),  pos_ids = scalar
         Writes into self.k_cache/self.v_cache shaped (B, n_kv_heads, T_max, d).
         """
         kout, vout = self.k_cache, self.v_cache
     
+        # Normalize pos_ids
         if not torch.is_tensor(pos_ids):
-            # Scalar legacy path
-            kout[:, :, pos_ids, :] = k
-            vout[:, :, pos_ids, :] = v
-            return kout, vout
-    
-        pos_ids = pos_ids.to(dtype=torch.long, device=k.device)
+            pos_ids = torch.tensor(pos_ids, device=k.device, dtype=torch.long)
+        else:
+            pos_ids = pos_ids.to(device=k.device, dtype=torch.long)
     
         if k.dim() != 4 or v.dim() != 4:
             raise RuntimeError(f"KV update expects k,v 4D. Got k={tuple(k.shape)} v={tuple(v.shape)}")
         B, Hkv, q_len, D = k.shape
     
-        # Make sure cache batch matches B (expand-from-1 is ok, otherwise error)
+        # Ensure cache batch matches B (expand-from-1 allowed)
         if kout.size(0) != B:
             if kout.size(0) == 1:
                 self.k_cache = kout.expand(B, -1, -1, -1).clone()
@@ -108,23 +107,23 @@ class KVCache(nn.Module):
             else:
                 raise RuntimeError(f"KV cache batch mismatch: cache.B={kout.size(0)} vs k.B={B}")
     
-        # Case A: PREFILL — pos_ids indexes a contiguous range per row
+        # Case A: PREFILL — vector of length q_len (same for all B rows)
         if pos_ids.dim() == 1 and pos_ids.numel() == q_len:
             for i in range(B):
-                kout[i, :, pos_ids, :] = k[i]            # (Hkv, q_len, D)
+                kout[i, :, pos_ids, :] = k[i]   # (Hkv, q_len, D)
                 vout[i, :, pos_ids, :] = v[i]
             return kout, vout
     
-        # Case B: STEP — q_len == 1 and one position per row
+        # Case B: 1-STEP — q_len == 1 with (B,) or (B,1) per-row positions
         if q_len == 1 and pos_ids.numel() == B:
-            pos_ids_flat = pos_ids.view(-1)              # handle (B,1) or (B,)
+            pos_ids = pos_ids.view(B)
             for i in range(B):
-                pi = int(pos_ids_flat[i].item())
+                pi = int(pos_ids[i].item())
                 kout[i, :, pi, :] = k[i, :, 0, :]
                 vout[i, :, pi, :] = v[i, :, 0, :]
             return kout, vout
     
-        # Case C: scalar for everyone
+        # Case C: scalar for everyone & q_len == 1
         if pos_ids.dim() == 0 and q_len == 1:
             pi = int(pos_ids.item())
             kout[:, :, pi, :] = k[:, :, 0, :]
@@ -140,6 +139,7 @@ class KVCache(nn.Module):
 
 
 
+
 class MoondreamModel(nn.Module):
 
     def __init__(
@@ -589,10 +589,12 @@ class MoondreamModel(nn.Module):
         elif not isinstance(image, Image.Image):
             raise ValueError("image must be a PIL Image or EncodedImage")
 
+        # At the VERY TOP of encode_image(), right after the type checks:
         for blk in self.text.blocks:
             if blk.kv_cache.k_cache.size(0) != 1:
                 blk.kv_cache.k_cache = blk.kv_cache.k_cache[:1].contiguous()
                 blk.kv_cache.v_cache = blk.kv_cache.v_cache[:1].contiguous()
+
     
         lora = (
             variant_state_dict(settings["variant"], device=self.device)
@@ -971,107 +973,135 @@ class MoondreamModel(nn.Module):
         # CRITICAL: per-row next position
         pos_vec = torch.tensor(lens, device=self.device, dtype=torch.long) + pos      # (B,)
     
-        return last_hidden, next_token, pos_vec
+        # At the end of _prefill_prompt_batched(), return a Python int:
+        pos_end = int((pos + T))
+        return last_hidden, next_token, pos_end
+
 
-    # In class MoondreamModel, replace the whole method:
     def _generate_points_batched(
         self,
         hidden,              # (B,1,C)
-        next_token,          # (B,1)  (not used when temperature=0, but ok)
-        pos: int,            # shared scalar next position
+        next_token,          # (B,1) (unused in greedy, but OK)
+        pos,                 # int or Tensor; normalized below
         include_size: bool = True,
         max_objects: int = 50,
         lora=None,
+        use_soft_argmax: bool = True,    # NEW: reduces jitter/hallucinations
     ):
-        """
-        Vectorized version of _generate_points() that decodes x -> y -> size -> next-token
-        for all rows in the batch simultaneously. Returns list-of-lists of dicts, len B.
-        """
         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
     
+        # Normalize pos to a scalar int (supports int, (1,), (B,), (B,1))
+        if torch.is_tensor(pos):
+            pos = int(pos.max().item())  # safe upper bound; we manage per-row with pos_ids/alive
+    
         # 4-D mask: (B, 1, q_len=1, kv_len)
         mask = torch.zeros(B, 1, 1, max_ctx, device=device, dtype=torch.bool)
         if pos > 0:
             mask[:, :, :, :pos] = True
-        # IMPORTANT: position_ids must be (B, 1) for rotary; KVCache.update accepts (B,1) too
+    
+        # position_ids must be (B,1)
         pos_ids = torch.full((B, 1), pos, device=device, dtype=torch.long)
     
         alive = torch.ones(B, dtype=torch.bool, device=device)
         counts = torch.zeros(B, dtype=torch.int32, device=device)
     
+        # helpers ---------------------------------------------------------
+        def _argmax01(logits):
+            # logits: (B, bins)
+            if use_soft_argmax:
+                probs = torch.softmax(logits, dim=-1)
+                bins = torch.arange(probs.size(-1), device=logits.device, dtype=torch.float32)
+                idx = (probs * bins).sum(dim=-1) / (probs.size(-1) - 1)
+                return idx  # 0..1
+            else:
+                idx = logits.argmax(dim=-1).to(torch.float32)
+                return idx / float(logits.size(-1) - 1)
+    
         with torch.inference_mode():
             while alive.any() and (counts < max_objects).any():
-                # --- x coordinate ---
-                x_logits = decode_coordinate(hidden, self.region)       # (B,1,1024) or (B,1024)
+                # --- 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 = x_logits.argmax(dim=-1).to(torch.float32) / float(x_logits.size(-1))  # (B,)
-                x_in = x_center.to(dtype=x_logits.dtype).unsqueeze(-1)  # (B,1)
-                x_emb = encode_coordinate(x_in, self.region).unsqueeze(1)   # (B,1,C)
+                x_center = _argmax01(x_logits)                    # (B,) in [0,1]
+                x_in = x_center.to(dtype=x_logits.dtype).unsqueeze(-1)   # (B,1)
+                x_emb = encode_coordinate(x_in, self.region).unsqueeze(1)  # (B,1,C)
     
-                # advance attention one step
-                mask[:, :, :, pos] = True
+                # advance attention one step FOR ALIVE ROWS ONLY
+                mask[alive, :, :, pos] = True
                 logits, hidden = self._decode_one_tok(x_emb, mask, pos_ids, lora)
-                pos += 1
-                pos_ids[:, 0] = pos
+                pos_ids[alive, 0] += 1
+                pos += 1  # scalar next free slot
     
-                # --- y coordinate ---
+                # --- y ---------------------------------------------------
                 y_logits = decode_coordinate(hidden, self.region)
                 if y_logits.dim() == 3:
                     y_logits = y_logits.squeeze(1)
-                y_center = y_logits.argmax(dim=-1).to(torch.float32) / float(y_logits.size(-1))
-                y_in = y_center.to(dtype=y_logits.dtype).unsqueeze(-1)  # (B,1)
+                y_center = _argmax01(y_logits)                    # (B,)
+                y_in = y_center.to(dtype=y_logits.dtype).unsqueeze(-1)
                 y_emb = encode_coordinate(y_in, self.region).unsqueeze(1)
     
-                mask[:, :, :, pos] = True
+                mask[alive, :, :, pos] = True
                 logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
+                pos_ids[alive, 0] += 1
                 pos += 1
-                pos_ids[:, 0] = pos
     
                 if include_size:
-                    # --- size ---
-                    size_logits = decode_size(hidden, self.region)  # tuple/list [w_logits, h_logits]
-                    # Support both (B,1,1024) and (B,1024)
+                    # --- size --------------------------------------------
+                    size_logits = decode_size(hidden, self.region)
                     w_logits = size_logits[0].squeeze(1)
                     h_logits = size_logits[1].squeeze(1)
-                    w_bin = w_logits.argmax(dim=-1).to(torch.float32)
-                    h_bin = h_logits.argmax(dim=-1).to(torch.float32)
+                    if use_soft_argmax:
+                        # convert expected-bin -> size (same mapping as paper/code)
+                        w_bin = (torch.softmax(w_logits, dim=-1) * torch.arange(w_logits.size(-1), device=device)).sum(dim=-1)
+                        h_bin = (torch.softmax(h_logits, dim=-1) * torch.arange(h_logits.size(-1), device=device)).sum(dim=-1)
+                    else:
+                        w_bin = w_logits.argmax(dim=-1).to(torch.float32)
+                        h_bin = h_logits.argmax(dim=-1).to(torch.float32)
                     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_in = torch.stack([w, h], dim=1).to(dtype=w_logits.dtype)  # (B,2)
-                    size_emb = encode_size(size_in, self.region).unsqueeze(1)       # (B,1,C)
+                    size_emb = encode_size(size_in, self.region).unsqueeze(1)      # (B,1,C)
     
-                    # record boxes
+                    # record boxes only for alive rows
                     for i in range(B):
-                        if alive[i]:
-                            out[i].append({
-                                "x_min": (x_center[i] - w[i] / 2).item(),
-                                "y_min": (y_center[i] - h[i] / 2).item(),
-                                "x_max": (x_center[i] + w[i] / 2).item(),
-                                "y_max": (y_center[i] + h[i] / 2).item(),
-                            })
+                        if not alive[i]:
+                            continue
+                        xl = (x_center[i] - w[i] / 2).item()
+                        xr = (x_center[i] + w[i] / 2).item()
+                        yt = (y_center[i] - h[i] / 2).item()
+                        yb = (y_center[i] + h[i] / 2).item()
+                        # clamp for safety
+                        out[i].append({
+                            "x_min": max(0.0, min(1.0, xl)),
+                            "y_min": max(0.0, min(1.0, yt)),
+                            "x_max": max(0.0, min(1.0, xr)),
+                            "y_max": max(0.0, min(1.0, yb)),
+                        })
     
-                    mask[:, :, :, pos] = True
+                    mask[alive, :, :, pos] = True
                     logits, hidden = self._decode_one_tok(size_emb, mask, pos_ids, lora)
+                    pos_ids[alive, 0] += 1
                     pos += 1
-                    pos_ids[:, 0] = pos
                     next_tok = logits.argmax(dim=-1).squeeze(-1)  # (B,)
                 else:
                     for i in range(B):
                         if alive[i]:
                             out[i].append({"x": x_center[i].item(), "y": y_center[i].item()})
-                    mask[:, :, :, pos] = True
+                    mask[alive, :, :, pos] = True
                     logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
+                    pos_ids[alive, 0] += 1
                     pos += 1
-                    pos_ids[:, 0] = pos
                     next_tok = logits.argmax(dim=-1).squeeze(-1)
     
+                # stop only rows that hit eos (or reached max objects)
                 finished_now = (next_tok == eos_id) | (counts >= max_objects - 1)
-                counts = counts + (~finished_now & alive).to(counts.dtype)
+                counts = counts + ((~finished_now) & alive).to(counts.dtype)
                 alive &= ~finished_now
     
         return out
@@ -1080,6 +1110,7 @@ class MoondreamModel(nn.Module):
 
 
 
+
     def detect_multi(self, image, objects, settings=None):
         """
         Parallel multi-label detection.