Add files via upload

Author: 1038lab

sam3/model/decoder.py

@@ -277,10 +277,8 @@ def __init__(
 
             if resolution is not None and stride is not None:
                 feat_size = resolution // stride
-                coords_h, coords_w = self._get_coords(
-                    feat_size, feat_size, device="cuda"
-                )
-                self.compilable_cord_cache = (coords_h, coords_w)
+                # Build cache lazily on the actual reference_boxes device to avoid
+                # mixing CPU/GPU tensors when users run on CPU.
                 self.compilable_stored_size = (feat_size, feat_size)
 
         self.roi_pooler = (
@@ -332,7 +330,11 @@ def _get_rpb_matrix(self, reference_boxes, feat_size):
         H, W = feat_size
         boxes_xyxy = box_cxcywh_to_xyxy(reference_boxes).transpose(0, 1)
         bs, num_queries, _ = boxes_xyxy.shape
-        if self.compilable_cord_cache is None:
+        if (
+            self.compilable_cord_cache is None
+            or any(c.device != reference_boxes.device for c in self.compilable_cord_cache)
+            or self.compilable_stored_size != (H, W)
+        ):
             self.compilable_cord_cache = self._get_coords(H, W, reference_boxes.device)
             self.compilable_stored_size = (H, W)
 
@@ -345,11 +347,12 @@ def _get_rpb_matrix(self, reference_boxes, feat_size):
         else:
             # cache miss, will create compilation issue
             # In case we're not compiling, we'll still rely on the dict-based cache
-            if feat_size not in self.coord_cache:
-                self.coord_cache[feat_size] = self._get_coords(
+            cache_key = (feat_size, reference_boxes.device)
+            if cache_key not in self.coord_cache:
+                self.coord_cache[cache_key] = self._get_coords(
                     H, W, reference_boxes.device
                 )
-            coords_h, coords_w = self.coord_cache[feat_size]
+            coords_h, coords_w = self.coord_cache[cache_key]
 
             assert coords_h.shape == (H,)
             assert coords_w.shape == (W,)