nodes/1038lab-ComfyUI-RMBG
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
AI Lab
2025-11-25 13:52:43 -08:00
committed by GitHub
parent ce0d2c57b6
commit 148cffa236
1 changed files with 11 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sam3/model/decoder.py
+11 -8
View File
@@ -277,10 +277,8 @@ class TransformerDecoder(nn.Module):
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 @@ class TransformerDecoder(nn.Module):
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 @@ class TransformerDecoder(nn.Module):
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,)