Update moondream.py

moondream.py

@@ -76,48 +76,54 @@ class KVCache(nn.Module):
     def update(self, pos_ids, k, v):
         """
         Supports:
-          • Prefill: k,v = (B, n_kv, q_len, d), pos_ids = (q_len,)
-          • 1-step : k,v = (B, n_kv, 1,    d), pos_ids = (B,) or scalar
+          • Prefill:  k,v  = (B, n_kv_heads, q_len, d),  pos_ids = (q_len,)
+          • 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
+        Writes into self.k_cache/self.v_cache shaped (B, n_kv_heads, T_max, d).
         """
-        kout, vout = self.k_cache, self.v_cache
+        kout, vout = self.kv_cache if hasattr(self, "kv_cache") else (self.k_cache, self.v_cache)
 
+        # normalize pos_ids
         if not torch.is_tensor(pos_ids):
             pos_ids = torch.tensor(pos_ids, device=k.device, dtype=torch.long)
         else:
-            pos_ids = pos_ids.to(k.device, dtype=torch.long)
+            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
 
-        # Expand caches’ batch dim if needed
-        if kout.size(0) != B:
-            if kout.size(0) == 1:
-                self.k_cache = kout.expand(B, -1, -1, -1).clone()
-                self.v_cache = vout.expand(B, -1, -1, -1).clone()
+        # match cache batch to B (expand-from-1 allowed)
+        if self.k_cache.size(0) != B:
+            if self.k_cache.size(0) == 1:
+                self.k_cache = self.k_cache.expand(B, -1, -1, -1).clone()
+                self.v_cache = self.v_cache.expand(B, -1, -1, -1).clone()
                 kout, vout = self.k_cache, self.v_cache
             else:
-                raise RuntimeError(f"KV cache batch mismatch: cache.B={kout.size(0)} vs k.B={B}")
+                raise RuntimeError(f"KV cache batch mismatch: cache.B={self.k_cache.size(0)} vs k.B={B}")
 
-        # Prefill (vector of positions shared across the batch)
+        # Case A: prefill — vector of length q_len (same for all rows)
         if pos_ids.dim() == 1 and pos_ids.numel() == q_len:
             for i in range(B):
                 kout[i, :, pos_ids, :] = k[i]
                 vout[i, :, pos_ids, :] = v[i]
             return kout, vout
 
-        # 1-step with per-row positions
-        if q_len == 1 and pos_ids.numel() in {1, B}:
-            if pos_ids.numel() == 1:
-                pi = int(pos_ids.item())
-                kout[:, :, pi, :] = k[:, :, 0, :]
-                vout[:, :, pi, :] = v[:, :, 0, :]
-            else:
-                pos_ids = pos_ids.view(B)
-                for i in range(B):
-                    pi = int(pos_ids[i].item())
-                    kout[i, :, pi, :] = k[i, :, 0, :]
-                    vout[i, :, pi, :] = v[i, :, 0, :]
+        # Case B: single step — q_len==1 with per-row positions
+        if q_len == 1 and pos_ids.numel() == B:
+            pos_ids = pos_ids.view(B)
+            for i in range(B):
+                pi = int(pos_ids[i].item())
+                kout[i, :, pi, :] = k[i, :, 0, :]
+                vout[i, :, pi, :] = v[i, :, 0, :]
+            return kout, vout
+
+        # Case C: scalar & q_len==1
+        if pos_ids.dim() == 0 and q_len == 1:
+            pi = int(pos_ids.item())
+            kout[:, :, pi, :] = k[:, :, 0, :]
+            vout[:, :, pi, :] = v[:, :, 0, :]
             return kout, vout
 
         raise RuntimeError(f"Unsupported KV update combo: k={tuple(k.shape)}, pos_ids={tuple(pos_ids.shape)}")
@@ -211,6 +217,7 @@ class MoondreamModel(nn.Module):
 
 
 
+
     
     
     def _setup_caches(self):
@@ -562,47 +569,46 @@ class MoondreamModel(nn.Module):
 
         return generator(next_token, pos)
 
-    def encode_image(
-        self,
-        image: Union[Image.Image, EncodedImage],
-        settings: Optional[ImageEncodingSettings] = None,
-    ) -> EncodedImage:
-        # Always start from single-row caches; avoids leftovers from batched runs
-        self._setup_caches()  # re-create caches (B=1)
-        for blk in self.text.blocks:  # make absolutely sure batch dim == 1
-            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()
+    def encode_image(self, image, settings=None) -> EncodedImage:
+        # always start from B=1 to avoid leftovers from batched runs
+        self._setup_caches()  # recreates caches with B=1
     
         if isinstance(image, EncodedImage):
             return image
         if not isinstance(image, Image.Image):
             raise ValueError("image must be a PIL Image or EncodedImage")
     
-        lora = (variant_state_dict(settings["variant"], device=self.device)
-                if settings is not None and "variant" in settings else None)
+        # hard-trim to B=1 if external code changed it
+        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) \
+                if settings and "variant" in settings else None
     
         with torch.inference_mode():
-            img_emb = self._run_vision_encoder(image)
-            bos_emb = text_encoder(torch.tensor([[self.config.tokenizer.bos_id]], device=self.device), self.text)
-            inputs_embeds = torch.cat([bos_emb, img_emb[None]], dim=1)
-            mask = self.attn_mask[:, :, :inputs_embeds.size(1), :]
-            pos_ids = torch.arange(inputs_embeds.size(1), dtype=torch.long, device=self.device)
+            img_emb = self._run_vision_encoder(image)                                   # (T_img,C)
+            bos_emb = text_encoder(torch.tensor([[self.config.tokenizer.bos_id]], device=self.device), self.text)  # (1,1,C)
+            inputs_embeds = torch.cat([bos_emb, img_emb[None]], dim=1)                  # (1,T0,C)
+    
+            mask = self.attn_mask[:, :, :inputs_embeds.size(1), :]                      # (1,1,T0,K)
+            pos_ids = torch.arange(inputs_embeds.size(1), device=self.device, dtype=torch.long)  # (T0,)
             self._prefill(inputs_embeds, mask, pos_ids, lora)
     
+        T0 = inputs_embeds.size(1)
         return EncodedImage(
-            pos=inputs_embeds.size(1),
+            pos=T0,
             caches=[
-                (
-                    b.kv_cache.k_cache[:, :, :inputs_embeds.size(1), :].clone(),
-                    b.kv_cache.v_cache[:, :, :inputs_embeds.size(1), :].clone(),
-                )
+                (b.kv_cache.k_cache[:, :, :T0, :].clone(),
+                 b.kv_cache.v_cache[:, :, :T0, :].clone())
                 for b in self.text.blocks
             ],
         )
 
 
 
+
     def query(
         self,
         image: Optional[Union[Image.Image, EncodedImage]] = None,
@@ -893,8 +899,36 @@ class MoondreamModel(nn.Module):
 
         return {"points": objects}
 
+    def _norm_size_logits(self, size_ret, 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)
+        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)
+                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
+                    w_logits, h_logits = t[0].unsqueeze(0), t[1].unsqueeze(0)
+                else:                            # (B,2C) fallback
+                    C2 = t.size(1); 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
+        if w_logits.dim() == 3: w_logits = w_logits.squeeze(1)
+        if h_logits.dim() == 3: h_logits = h_logits.squeeze(1)
+        return w_logits.contiguous(), h_logits.contiguous()
+
 
-    # -------------------- batched helpers --------------------
     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)
@@ -906,60 +940,62 @@ class MoondreamModel(nn.Module):
             b.kv_cache.k_cache[:, :, :T, :] = k.expand(batch_size, -1, -1, -1)
             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 (no detect template).")
+            raise NotImplementedError("Model does not support object detection.")
     
         rows, 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())
-        B = len(rows); T = max(lens)
+        B, T = len(rows), max(lens)
         eos = self.config.tokenizer.eos_id
     
         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)
-        torch._dynamo.mark_dynamic(prompt_emb, 1)
+        prompt_emb = text_encoder(prompt_ids, self.text)                 # (B,T,C)
+        torch._dynamo.mark_dynamic(prompt_emb, 1)                        # allow variable T
     
-        # mask: (B,1,T,kv_len)
-        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)
     
-        # IMPORTANT: for prefill pass a 1-D vector of length T (matches upstream)
         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)
+        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)
+        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)
     
-        # shared scalar end position
-        return last_hidden, next_token, int(pos + T)
+        # shared next-free position in cache (safe upper bound)
+        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) (ignored in greedy)
+        next_token,          # (B,1)  (unused in greedy)
         pos,                 # int
         include_size: bool = True,
         max_objects: int = 50,
         lora=None,
-        use_soft_argmax: bool = False,   # default OFF to match upstream numerics
+        use_soft_argmax: bool = True,   # reduces jitter
     ):
         B = hidden.size(0)
         device = self.device
@@ -967,78 +1003,67 @@ class MoondreamModel(nn.Module):
         eos_id = self.config.tokenizer.eos_id
         max_ctx = self.config.text.max_context
     
-        # 4-D mask: (B,1,1,kv_len)
+        # 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
+        if int(pos) > 0:
+            mask[:, :, :, :int(pos)] = True
+        pos_ids = torch.full((B, 1), int(pos), device=device, dtype=torch.long)
     
-        # rotary & KV path are happiest with a 1-D scalar position vector (like upstream)
-        pos_id_vec = torch.tensor([pos], device=device, dtype=torch.long)  # (1,)
+        # helper: (B, bins) -> (B,) in [0,1]
+        def _argmax01(logits):
+            if use_soft_argmax:
+                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)
     
         alive  = torch.ones(B, dtype=torch.bool, device=device)
         counts = torch.zeros(B, dtype=torch.int32, device=device)
     
-        def _center01(logits_2d):
-            # logits_2d: (B, bins)
-            if logits_2d.dim() == 3:  # (B,1,bins) -> (B,bins)
-                logits_2d = logits_2d.squeeze(1)
-            bins = logits_2d.size(-1)
-            if use_soft_argmax:
-                p = torch.softmax(logits_2d, dim=-1)
-                idx = (p * torch.arange(bins, device=logits_2d.device, dtype=torch.float32)).sum(dim=-1)
-                return idx / float(bins)  # match upstream scale
-            else:
-                return logits_2d.argmax(dim=-1).to(torch.float32) / float(bins)
-    
         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)
-                x_center = _center01(x_logits)                           # (B,) in [0,1]
-                x_emb = encode_coordinate(x_center.to(x_logits.dtype).unsqueeze(-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)
     
-                mask[alive, :, :, pos] = True
-                _, hidden = self._decode_one_tok(x_emb, mask, pos_id_vec, lora)
-                pos += 1
-                pos_id_vec[0] = pos
+                mask[alive, :, :, pos_ids[0,0]] = 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)
-                y_center = _center01(y_logits)
-                y_emb = encode_coordinate(y_center.to(y_logits.dtype).unsqueeze(-1), self.region).unsqueeze(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)
     
-                mask[alive, :, :, pos] = True
-                _, hidden = self._decode_one_tok(y_emb, mask, pos_id_vec, lora)
-                pos += 1
-                pos_id_vec[0] = pos
+                mask[alive, :, :, pos_ids[0,0]] = True
+                logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
+                pos_ids[alive, 0] += 1
     
                 if include_size:
-                    # --- size ---
                     size_ret = decode_size(hidden, self.region)
-                    # Works for tuple or stacked tensor
-                    if isinstance(size_ret, (tuple, list)):
-                        w_logits, h_logits = size_ret
+                    w_logits, h_logits = self._norm_size_logits(size_ret, B)
+    
+                    if use_soft_argmax:
+                        bins = torch.arange(w_logits.size(-1), device=device, dtype=torch.float32)
+                        w_bin = (torch.softmax(w_logits, dim=-1) * bins).sum(dim=-1)
+                        h_bin = (torch.softmax(h_logits, dim=-1) * bins).sum(dim=-1)
                     else:
-                        # expected shapes: (B,2,1024) or (B,1,2,1024)
-                        if size_ret.dim() == 3:  # (B,2,1024)
-                            w_logits, h_logits = size_ret[:, 0], size_ret[:, 1]
-                        else:                    # (B,1,2,1024)
-                            w_logits, h_logits = size_ret[:, 0, 0], size_ret[:, 0, 1]
-                    if w_logits.dim() == 3: w_logits = w_logits.squeeze(1)
-                    if h_logits.dim() == 3: h_logits = h_logits.squeeze(1)
+                        w_bin = w_logits.argmax(dim=-1).to(torch.float32)
+                        h_bin = h_logits.argmax(dim=-1).to(torch.float32)
     
-                    # bins -> size via the same inverse log2 scale as upstream
-                    w_bin = w_logits.argmax(dim=-1).to(torch.float32)
-                    h_bin = h_logits.argmax(dim=-1).to(torch.float32)
+                    # inverse log scale 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)
     
-                    size_emb = encode_size(torch.stack([w, h], dim=1).to(w_logits.dtype), self.region).unsqueeze(1)
+                    size_emb = encode_size(torch.stack([w, h], dim=1).to(dtype=w_logits.dtype), self.region).unsqueeze(1)
     
-                    # record boxes (clamped)
+                    # write boxes only for alive rows
                     for i in range(B):
-                        if not alive[i]:
-                            continue
+                        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()
@@ -1050,19 +1075,17 @@ class MoondreamModel(nn.Module):
                             "y_max": max(0.0, min(1.0, yb)),
                         })
     
-                    mask[alive, :, :, pos] = True
-                    logits, hidden = self._decode_one_tok(size_emb, mask, pos_id_vec, lora)
-                    pos += 1
-                    pos_id_vec[0] = pos
-                    next_tok = logits.argmax(dim=-1).squeeze(-1)  # (B,)
+                    mask[alive, :, :, pos_ids[0,0]] = True
+                    logits, hidden = self._decode_one_tok(size_emb, mask, pos_ids, lora)
+                    pos_ids[alive, 0] += 1
+                    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[alive, :, :, pos] = True
-                    logits, hidden = self._decode_one_tok(y_emb, mask, pos_id_vec, lora)
-                    pos += 1
-                    pos_id_vec[0] = pos
+                    mask[alive, :, :, pos_ids[0,0]] = True
+                    logits, hidden = self._decode_one_tok(y_emb, mask, pos_ids, lora)
+                    pos_ids[alive, 0] += 1
                     next_tok = logits.argmax(dim=-1).squeeze(-1)
     
                 finished_now = (next_tok == eos_id) | (counts >= max_objects - 1)
@@ -1071,10 +1094,6 @@ class MoondreamModel(nn.Module):
     
         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.")
@@ -1103,7 +1122,7 @@ class MoondreamModel(nn.Module):
                 d["label"] = lab
             res[lab] = lst
     
-        # make subsequent single-image calls stable
+        # restore B=1 so the next encode_image() starts clean
         self._reset_kv_caches(1)
         return {"objects": res}
 
@@ -1111,6 +1130,7 @@ class MoondreamModel(nn.Module):
 
 
 
+
     def _detect_gaze(
         self,
         image: EncodedImage,