fix: filter low res masks before upsampling + rename nms and iou thresh vars

configuration_openworld_sam.py

@@ -49,8 +49,8 @@ class OpenWorldSAMConfig(PretrainedConfig):
         pixel_std=(58.395, 57.120, 57.375),
         # Inference thresholds (ADE20K instance config defaults)
         nms_on=True,
-        nms_threshold=0.5,
-        iou_threshold=0.7,
+        nms_thresh=0.5,
+        iou_thresh=0.7,
         top_k_on=False,
         detections_per_image=100,
         # Class vocabulary
@@ -70,8 +70,8 @@ class OpenWorldSAMConfig(PretrainedConfig):
         self.pixel_mean = list(pixel_mean)
         self.pixel_std = list(pixel_std)
         self.nms_on = nms_on
-        self.nms_threshold = nms_threshold
-        self.iou_threshold = iou_threshold
+        self.nms_thresh = nms_thresh
+        self.iou_thresh = iou_thresh
         self.top_k_on = top_k_on
         self.detections_per_image = detections_per_image
         self.stuff_classes = stuff_classes if stuff_classes is not None else ADE20K_150_CLASSES

modeling_openworld_sam.py

@@ -345,10 +345,14 @@ class OpenWorldSAMModel(PreTrainedModel):
                 dtype=torch.long, device=self.device,
             )
 
-            # Upsample masks to original image size
-            pred_masks = self._postprocess_masks(pred_masks, original_size_list[img_idx])
+            # Filter on low-res masks first; only the survivors get upsampled
+            # to the original image size (upsampling the full query set before
+            # filtering allocates one [num_queries, H, W] float32 tensor that
+            # can reach double-digit GiB for large vocabularies).
+            instances = self._instance_inference(
+                pred_masks, pred_logits, class_labels, original_size_list[img_idx]
+            )
 
-            instances = self._instance_inference(pred_masks, pred_logits, class_labels)
             processed_results.append({"instances": instances})
 
         return processed_results
@@ -358,7 +362,7 @@ class OpenWorldSAMModel(PreTrainedModel):
             masks.float().unsqueeze(0), orig_hw, mode="bilinear", align_corners=False
         ).squeeze(0)
 
-    def _instance_inference(self, pred_masks, iou_scores, class_labels):
+    def _instance_inference(self, pred_masks, iou_scores, class_labels, orig_hw):
         """Returns dict with keys: masks (bool), scores (float), class_ids (long)."""
         pred_masks = pred_masks.squeeze(1) if pred_masks.ndim == 4 else pred_masks
 
@@ -369,28 +373,32 @@ class OpenWorldSAMModel(PreTrainedModel):
             pred_masks, iou_scores, class_labels = pred_masks[idx], iou_scores[idx], class_labels[idx]
 
         # IoU threshold filter
-        keep = iou_scores >= self.config.iou_threshold
+        keep = iou_scores >= self.config.iou_thresh
         pred_masks, iou_scores, class_labels = pred_masks[keep], iou_scores[keep], class_labels[keep]
 
         if pred_masks.shape[0] == 0:
             empty = torch.empty(0, device=self.device)
             return {
-                "masks": torch.empty((0, pred_masks.shape[-2], pred_masks.shape[-1]),
-                                     dtype=torch.bool, device=self.device),
+                "masks": torch.empty((0, *orig_hw), dtype=torch.bool, device=self.device),
                 "scores": empty,
                 "class_ids": empty.long(),
             }
 
-        binary_masks = pred_masks > 0
 
-        # NMS
+        # NMS on low-res masks — box IoU is scale-invariant, so this doesn't
+        # need the full-res masks either.
+        low_res_binary_masks = pred_masks > 0
         if self.config.nms_on:
-            boxes = _masks_to_boxes(binary_masks)
-            nms_keep = torchvision.ops.nms(boxes, iou_scores, self.config.nms_threshold)
-            binary_masks = binary_masks[nms_keep]
+            boxes = _masks_to_boxes(low_res_binary_masks)
+            nms_keep = torchvision.ops.nms(boxes, iou_scores, self.config.nms_thresh)
+            pred_masks = pred_masks[nms_keep]
             iou_scores = iou_scores[nms_keep]
             class_labels = class_labels[nms_keep]
 
+        # Upsample only the surviving masks to the original image size
+        pred_masks = self._postprocess_masks(pred_masks, orig_hw)
+        binary_masks = pred_masks > 0
+
         return {
             "masks": binary_masks,
             "scores": iou_scores,