feat: configurable mask_decoder_chunk_size

modeling_openworld_sam.py

@@ -319,16 +319,35 @@ class OpenWorldSAMModel(PreTrainedModel):
                 f[img_idx].unsqueeze(0) for f in _features["high_res_feats"]
             ]
 
+            # The decoder's repeat_image path duplicates image_embeddings and
+            # image_pe once per candidate query, so scoring all num_classes *
+            # num_tokens candidates in one call can reach double-digit GiB
+            # for large vocabularies. `chunk_size` bounds peak memory
+            # rather than total candidate count.
+            image_embed_img = _features["image_embed"][img_idx].unsqueeze(0)
+            image_pe = self.evf_sam2.visual_model.sam_prompt_encoder.get_dense_pe()
+            chunk_size = self.config.mask_decoder_chunk_size
+            num_total_tokens = sparse_embeddings.shape[0]
+
+            low_res_masks_chunks = []
+            iou_pred_chunks = []
+
             with torch.no_grad():
-                low_res_masks, iou_pred, _, _ = self.evf_sam2.visual_model.sam_mask_decoder(
-                    image_embeddings=_features["image_embed"][img_idx].unsqueeze(0),
-                    image_pe=self.evf_sam2.visual_model.sam_prompt_encoder.get_dense_pe(),
-                    sparse_prompt_embeddings=sparse_embeddings,
-                    dense_prompt_embeddings=dense_embeddings,
-                    multimask_output=False,
-                    repeat_image=True,
-                    high_res_features=high_res_features,
-                )
+                for start in range(0, num_total_tokens, chunk_size):
+                    end = min(start + chunk_size, num_total_tokens)
+                    chunk_low_res_masks, chunk_iou_pred, _, _ = self.evf_sam2.visual_model.sam_mask_decoder(
+                        image_embeddings=image_embed_img,
+                        image_pe=image_pe,
+                        sparse_prompt_embeddings=sparse_embeddings[start:end],
+                        dense_prompt_embeddings=dense_embeddings[start:end],
+                        multimask_output=False,
+                        repeat_image=True,
+                        high_res_features=high_res_features,
+                    )
+                    low_res_masks_chunks.append(chunk_low_res_masks)
+                    iou_pred_chunks.append(chunk_iou_pred)
+            low_res_masks = torch.cat(low_res_masks_chunks, dim=0)
+            iou_pred = torch.cat(iou_pred_chunks, dim=0)
 
             pred_masks = low_res_masks.squeeze(1)  # [total_tokens, H_low, W_low]
             pred_logits = iou_pred.squeeze(1)       # [total_tokens]