qev_fast vision: fix KeyError when a full image cache evicts an image the same call uses

SHA256SUMS

@@ -21,7 +21,7 @@ ed65a44a68aefcc6397068db52274433cd372a2c5fff6b1f7f901ef543429941  qev_fast/engin
 e9a6ef607a94a11f659f4d57fa1588c8b99b5288535ed9fb7db0a144fd817c55  qev_fast/prefix.py
 c849a299772c7186b5b30feb9db777ef12673271e739051ce0968f6bbc741b27  qev_fast/runner.py
 9a7e7a159b4b6c9c1755ccbaf6c11b433bac78f8ec67474833e61821d77999ef  qev_fast/systemone.py
-53c44b0101b889464a78cf7b8c8bfe518765610f39478df1b7a3e8f680cb02e4  qev_fast/vision.py
+3713ce5af6d548cd5db324feae5d31f99379ddb21c544d3cc0bad98939ad159a  qev_fast/vision.py
 e8d061404d312ca2cfbec5db8d5adbacb9fd2888e20997833f8080ed317e4184  qev_image.py
 1cfc5fbdc1b4317bd2df8f2e965ee466f95a33af5db8c949d157f86c93996d4e  qev_image_server.py
 01239024d6bce903ae16a8bda08a56b6ecbf43d630b9bd590fb1a73e9c11e54c  qev_server.py

qev_fast/vision.py

@@ -105,19 +105,27 @@ class Vision:
                 seen.add(p.key)
         self.stats["images"] += len(preps)
         self.stats["cache_hits"] += len(preps) - len(todo)
+        fresh = {}
         if todo:
             pix = torch.cat([p.pix for p in todo]).to(self.device, self.visual.dtype)
             thw = torch.cat([p.thw for p in todo]).to(self.device)
             out = self.visual(pix, grid_thw=thw, return_dict=True).pooler_output
             self.stats["tower_calls"] += 1
             for p, e in zip(todo, torch.split(out, [p.n for p in todo]), strict=True):
-                self.cache[p.key] = e
-            while len(self.cache) > self.cache_size:
-                self.cache.popitem(last=False)
+                fresh[p.key] = e
+        # Read every embedding of this call before the cache is trimmed: trimming can evict a cached image this same
+        # call uses (a full cache, one cached and one new image), and a call can hold more images than the cache.
         res = []
         for p in preps:
-            self.cache.move_to_end(p.key)
-            res.append(self.cache[p.key])
+            e = fresh.get(p.key)
+            if e is None:
+                e = self.cache[p.key]
+                self.cache.move_to_end(p.key)
+            res.append(e)
+        for k, e in fresh.items():
+            self.cache[k] = e
+        while len(self.cache) > self.cache_size:
+            self.cache.popitem(last=False)
         return res
 
     def positions(self, ids: list, preps: list) -> np.ndarray: