From b46df68f545ff97a1301c76270273c4e1d9ce01b Mon Sep 17 00:00:00 2001 From: Dinuka Perera Date: Thu, 21 May 2026 13:10:52 +0530 Subject: [PATCH 1/4] Fix is_last off-by-one in MaskGenerationPipeline for partial batches --- src/transformers/pipelines/mask_generation.py | 2 +- .../test_pipelines_mask_generation.py | 109 ++++++++++++++++++ 2 files changed, 110 insertions(+), 1 deletion(-) diff --git a/src/transformers/pipelines/mask_generation.py b/src/transformers/pipelines/mask_generation.py index 920ce9d0a39c..ffbe5868f7a7 100644 --- a/src/transformers/pipelines/mask_generation.py +++ b/src/transformers/pipelines/mask_generation.py @@ -231,7 +231,7 @@ def preprocess( for i in range(0, n_points, points_per_batch): batched_points = grid_points[:, i : i + points_per_batch, :, :] labels = input_labels[:, i : i + points_per_batch] - is_last = i == n_points - points_per_batch + is_last = i + points_per_batch >= n_points yield { "input_points": batched_points, "input_labels": labels, diff --git a/tests/pipelines/test_pipelines_mask_generation.py b/tests/pipelines/test_pipelines_mask_generation.py index 07ce9c5d91ca..550318dfa154 100644 --- a/tests/pipelines/test_pipelines_mask_generation.py +++ b/tests/pipelines/test_pipelines_mask_generation.py @@ -13,6 +13,7 @@ # limitations under the License. import unittest +from unittest.mock import MagicMock, patch import numpy as np from huggingface_hub.utils import insecure_hashlib @@ -93,6 +94,114 @@ def get_test_pipeline( def run_pipeline_test(self, mask_generator, examples): pass + @require_torch + def test_preprocess_is_last_partial_batch(self): + """ + Regression test for the is_last flag in MaskGenerationPipeline.preprocess. + + When n_points is not a multiple of points_per_batch, the old formula + ``i == n_points - points_per_batch`` never evaluates to True on the final + iteration, so PipelinePackIterator's ``while not is_last`` loop can never + terminate cleanly and the accumulated results for the last batch are lost. + + The fix ``i + points_per_batch >= n_points`` always marks the last batch + correctly regardless of whether n_points is an exact multiple. + """ + import torch + + n_points = 100 # not a multiple of points_per_batch=64 + points_per_batch = 64 + + # Build a lightweight mock pipeline that skips model/image-processor I/O + # and jumps straight to the batching loop. + # No spec= so dynamically-set Pipeline attributes like image_processor are accessible. + mock_pipeline = MagicMock() + mock_pipeline.dtype = torch.float32 + mock_pipeline.device = torch.device("cpu") + + # Fake outputs from image_processor.generate_crop_boxes + crop_boxes = torch.zeros(1, 4) + grid_points = torch.zeros(1, n_points, 1, 2) + input_labels = torch.zeros(1, n_points) + cropped_images = [MagicMock()] + + mock_pipeline.image_processor.size = {"longest_edge": 1024} + mock_pipeline.image_processor.generate_crop_boxes.return_value = ( + crop_boxes, + grid_points, + cropped_images, + input_labels, + ) + + # image_processor(images=...) → dict with pixel_values tensor + fake_model_inputs = {"pixel_values": torch.zeros(1, 3, 16, 16)} + mock_pipeline.image_processor.return_value.to.return_value = fake_model_inputs + + # model.get_image_embeddings → plain embedding tensor (SAM-style single output) + mock_pipeline.model.get_image_embeddings.return_value = torch.zeros(1, 256, 16, 16) + + # device_placement / inference_context are no-op context managers + mock_pipeline.device_placement.return_value.__enter__ = MagicMock(return_value=None) + mock_pipeline.device_placement.return_value.__exit__ = MagicMock(return_value=False) + mock_pipeline.get_inference_context.return_value.return_value.__enter__ = MagicMock(return_value=None) + mock_pipeline.get_inference_context.return_value.return_value.__exit__ = MagicMock(return_value=False) + mock_pipeline._ensure_tensor_on_device.side_effect = lambda x, device: x + + with patch("transformers.pipelines.mask_generation.load_image", return_value=MagicMock()): + gen = MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch) + yielded = list(gen) + + # With n_points=100, points_per_batch=64 we expect 2 batches (i=0 and i=64) + self.assertEqual(len(yielded), 2, "Expected exactly 2 batches for 100 points with batch_size=64") + + # The first batch must NOT be flagged as last + self.assertFalse(yielded[0]["is_last"], "First batch should not be is_last") + + # The second (final, partial) batch MUST be flagged as last + self.assertTrue(yielded[1]["is_last"], "Final partial batch must have is_last=True") + + @require_torch + def test_preprocess_is_last_exact_multiple(self): + """is_last is also correct when n_points is an exact multiple of points_per_batch.""" + import torch + + n_points = 128 # exact multiple: 128 / 64 == 2 + points_per_batch = 64 + + mock_pipeline = MagicMock() + mock_pipeline.dtype = torch.float32 + mock_pipeline.device = torch.device("cpu") + + crop_boxes = torch.zeros(1, 4) + grid_points = torch.zeros(1, n_points, 1, 2) + input_labels = torch.zeros(1, n_points) + cropped_images = [MagicMock()] + + mock_pipeline.image_processor.size = {"longest_edge": 1024} + mock_pipeline.image_processor.generate_crop_boxes.return_value = ( + crop_boxes, + grid_points, + cropped_images, + input_labels, + ) + + fake_model_inputs = {"pixel_values": torch.zeros(1, 3, 16, 16)} + mock_pipeline.image_processor.return_value.to.return_value = fake_model_inputs + mock_pipeline.model.get_image_embeddings.return_value = torch.zeros(1, 256, 16, 16) + mock_pipeline.device_placement.return_value.__enter__ = MagicMock(return_value=None) + mock_pipeline.device_placement.return_value.__exit__ = MagicMock(return_value=False) + mock_pipeline.get_inference_context.return_value.return_value.__enter__ = MagicMock(return_value=None) + mock_pipeline.get_inference_context.return_value.return_value.__exit__ = MagicMock(return_value=False) + mock_pipeline._ensure_tensor_on_device.side_effect = lambda x, device: x + + with patch("transformers.pipelines.mask_generation.load_image", return_value=MagicMock()): + gen = MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch) + yielded = list(gen) + + self.assertEqual(len(yielded), 2) + self.assertFalse(yielded[0]["is_last"]) + self.assertTrue(yielded[1]["is_last"]) + @slow @require_torch def test_small_model_pt(self): From b05aff2e72f098935eba6adb521b1ee351874ab3 Mon Sep 17 00:00:00 2001 From: Dinuka Perera Date: Thu, 21 May 2026 16:35:22 +0530 Subject: [PATCH 2/4] Code cleaning up --- .../test_pipelines_mask_generation.py | 48 ++++--------------- 1 file changed, 9 insertions(+), 39 deletions(-) diff --git a/tests/pipelines/test_pipelines_mask_generation.py b/tests/pipelines/test_pipelines_mask_generation.py index 550318dfa154..2ab1d11b288a 100644 --- a/tests/pipelines/test_pipelines_mask_generation.py +++ b/tests/pipelines/test_pipelines_mask_generation.py @@ -96,30 +96,16 @@ def run_pipeline_test(self, mask_generator, examples): @require_torch def test_preprocess_is_last_partial_batch(self): - """ - Regression test for the is_last flag in MaskGenerationPipeline.preprocess. - - When n_points is not a multiple of points_per_batch, the old formula - ``i == n_points - points_per_batch`` never evaluates to True on the final - iteration, so PipelinePackIterator's ``while not is_last`` loop can never - terminate cleanly and the accumulated results for the last batch are lost. - - The fix ``i + points_per_batch >= n_points`` always marks the last batch - correctly regardless of whether n_points is an exact multiple. - """ + """is_last must be True on the final batch when n_points is not a multiple of points_per_batch.""" import torch n_points = 100 # not a multiple of points_per_batch=64 points_per_batch = 64 - # Build a lightweight mock pipeline that skips model/image-processor I/O - # and jumps straight to the batching loop. - # No spec= so dynamically-set Pipeline attributes like image_processor are accessible. mock_pipeline = MagicMock() mock_pipeline.dtype = torch.float32 mock_pipeline.device = torch.device("cpu") - # Fake outputs from image_processor.generate_crop_boxes crop_boxes = torch.zeros(1, 4) grid_points = torch.zeros(1, n_points, 1, 2) input_labels = torch.zeros(1, n_points) @@ -132,15 +118,8 @@ def test_preprocess_is_last_partial_batch(self): cropped_images, input_labels, ) - - # image_processor(images=...) → dict with pixel_values tensor - fake_model_inputs = {"pixel_values": torch.zeros(1, 3, 16, 16)} - mock_pipeline.image_processor.return_value.to.return_value = fake_model_inputs - - # model.get_image_embeddings → plain embedding tensor (SAM-style single output) + mock_pipeline.image_processor.return_value.to.return_value = {"pixel_values": torch.zeros(1, 3, 16, 16)} mock_pipeline.model.get_image_embeddings.return_value = torch.zeros(1, 256, 16, 16) - - # device_placement / inference_context are no-op context managers mock_pipeline.device_placement.return_value.__enter__ = MagicMock(return_value=None) mock_pipeline.device_placement.return_value.__exit__ = MagicMock(return_value=False) mock_pipeline.get_inference_context.return_value.return_value.__enter__ = MagicMock(return_value=None) @@ -148,24 +127,18 @@ def test_preprocess_is_last_partial_batch(self): mock_pipeline._ensure_tensor_on_device.side_effect = lambda x, device: x with patch("transformers.pipelines.mask_generation.load_image", return_value=MagicMock()): - gen = MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch) - yielded = list(gen) - - # With n_points=100, points_per_batch=64 we expect 2 batches (i=0 and i=64) - self.assertEqual(len(yielded), 2, "Expected exactly 2 batches for 100 points with batch_size=64") + yielded = list(MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch)) - # The first batch must NOT be flagged as last - self.assertFalse(yielded[0]["is_last"], "First batch should not be is_last") - - # The second (final, partial) batch MUST be flagged as last - self.assertTrue(yielded[1]["is_last"], "Final partial batch must have is_last=True") + self.assertEqual(len(yielded), 2) + self.assertFalse(yielded[0]["is_last"]) + self.assertTrue(yielded[1]["is_last"]) @require_torch def test_preprocess_is_last_exact_multiple(self): """is_last is also correct when n_points is an exact multiple of points_per_batch.""" import torch - n_points = 128 # exact multiple: 128 / 64 == 2 + n_points = 128 # exact multiple of points_per_batch=64 points_per_batch = 64 mock_pipeline = MagicMock() @@ -184,9 +157,7 @@ def test_preprocess_is_last_exact_multiple(self): cropped_images, input_labels, ) - - fake_model_inputs = {"pixel_values": torch.zeros(1, 3, 16, 16)} - mock_pipeline.image_processor.return_value.to.return_value = fake_model_inputs + mock_pipeline.image_processor.return_value.to.return_value = {"pixel_values": torch.zeros(1, 3, 16, 16)} mock_pipeline.model.get_image_embeddings.return_value = torch.zeros(1, 256, 16, 16) mock_pipeline.device_placement.return_value.__enter__ = MagicMock(return_value=None) mock_pipeline.device_placement.return_value.__exit__ = MagicMock(return_value=False) @@ -195,8 +166,7 @@ def test_preprocess_is_last_exact_multiple(self): mock_pipeline._ensure_tensor_on_device.side_effect = lambda x, device: x with patch("transformers.pipelines.mask_generation.load_image", return_value=MagicMock()): - gen = MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch) - yielded = list(gen) + yielded = list(MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch)) self.assertEqual(len(yielded), 2) self.assertFalse(yielded[0]["is_last"]) From 71dc1bf81ee5c5e17e4eff8ffd3c191ae0b50165 Mon Sep 17 00:00:00 2001 From: Dinuka Perera Date: Thu, 21 May 2026 16:59:50 +0530 Subject: [PATCH 3/4] build fix --- tests/pipelines/test_pipelines_mask_generation.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/pipelines/test_pipelines_mask_generation.py b/tests/pipelines/test_pipelines_mask_generation.py index 2ab1d11b288a..45d1ad755bac 100644 --- a/tests/pipelines/test_pipelines_mask_generation.py +++ b/tests/pipelines/test_pipelines_mask_generation.py @@ -127,7 +127,8 @@ def test_preprocess_is_last_partial_batch(self): mock_pipeline._ensure_tensor_on_device.side_effect = lambda x, device: x with patch("transformers.pipelines.mask_generation.load_image", return_value=MagicMock()): - yielded = list(MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch)) + gen = MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch) + yielded = list(gen) self.assertEqual(len(yielded), 2) self.assertFalse(yielded[0]["is_last"]) @@ -166,7 +167,8 @@ def test_preprocess_is_last_exact_multiple(self): mock_pipeline._ensure_tensor_on_device.side_effect = lambda x, device: x with patch("transformers.pipelines.mask_generation.load_image", return_value=MagicMock()): - yielded = list(MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch)) + gen = MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch) + yielded = list(gen) self.assertEqual(len(yielded), 2) self.assertFalse(yielded[0]["is_last"]) From c4cbf01ce346d878a6798de0c034dbe835bd17dc Mon Sep 17 00:00:00 2001 From: Matt Date: Thu, 21 May 2026 14:14:29 +0100 Subject: [PATCH 4/4] Simplify is_last test using a real tiny pipeline Replace the two heavily-mocked test methods with a single subTest-driven check that calls preprocess on a real MaskGenerationPipeline backed by hf-internal-testing/tiny-random-SamModel. Covers both the partial-batch (100 points) and exact-multiple (64 points) cases without mocking. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../test_pipelines_mask_generation.py | 89 ++----------------- 1 file changed, 9 insertions(+), 80 deletions(-) diff --git a/tests/pipelines/test_pipelines_mask_generation.py b/tests/pipelines/test_pipelines_mask_generation.py index 45d1ad755bac..4d25bb9c0cc1 100644 --- a/tests/pipelines/test_pipelines_mask_generation.py +++ b/tests/pipelines/test_pipelines_mask_generation.py @@ -13,7 +13,6 @@ # limitations under the License. import unittest -from unittest.mock import MagicMock, patch import numpy as np from huggingface_hub.utils import insecure_hashlib @@ -94,85 +93,15 @@ def get_test_pipeline( def run_pipeline_test(self, mask_generator, examples): pass - @require_torch - def test_preprocess_is_last_partial_batch(self): - """is_last must be True on the final batch when n_points is not a multiple of points_per_batch.""" - import torch - - n_points = 100 # not a multiple of points_per_batch=64 - points_per_batch = 64 - - mock_pipeline = MagicMock() - mock_pipeline.dtype = torch.float32 - mock_pipeline.device = torch.device("cpu") - - crop_boxes = torch.zeros(1, 4) - grid_points = torch.zeros(1, n_points, 1, 2) - input_labels = torch.zeros(1, n_points) - cropped_images = [MagicMock()] - - mock_pipeline.image_processor.size = {"longest_edge": 1024} - mock_pipeline.image_processor.generate_crop_boxes.return_value = ( - crop_boxes, - grid_points, - cropped_images, - input_labels, - ) - mock_pipeline.image_processor.return_value.to.return_value = {"pixel_values": torch.zeros(1, 3, 16, 16)} - mock_pipeline.model.get_image_embeddings.return_value = torch.zeros(1, 256, 16, 16) - mock_pipeline.device_placement.return_value.__enter__ = MagicMock(return_value=None) - mock_pipeline.device_placement.return_value.__exit__ = MagicMock(return_value=False) - mock_pipeline.get_inference_context.return_value.return_value.__enter__ = MagicMock(return_value=None) - mock_pipeline.get_inference_context.return_value.return_value.__exit__ = MagicMock(return_value=False) - mock_pipeline._ensure_tensor_on_device.side_effect = lambda x, device: x - - with patch("transformers.pipelines.mask_generation.load_image", return_value=MagicMock()): - gen = MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch) - yielded = list(gen) - - self.assertEqual(len(yielded), 2) - self.assertFalse(yielded[0]["is_last"]) - self.assertTrue(yielded[1]["is_last"]) - - @require_torch - def test_preprocess_is_last_exact_multiple(self): - """is_last is also correct when n_points is an exact multiple of points_per_batch.""" - import torch - - n_points = 128 # exact multiple of points_per_batch=64 - points_per_batch = 64 - - mock_pipeline = MagicMock() - mock_pipeline.dtype = torch.float32 - mock_pipeline.device = torch.device("cpu") - - crop_boxes = torch.zeros(1, 4) - grid_points = torch.zeros(1, n_points, 1, 2) - input_labels = torch.zeros(1, n_points) - cropped_images = [MagicMock()] - - mock_pipeline.image_processor.size = {"longest_edge": 1024} - mock_pipeline.image_processor.generate_crop_boxes.return_value = ( - crop_boxes, - grid_points, - cropped_images, - input_labels, - ) - mock_pipeline.image_processor.return_value.to.return_value = {"pixel_values": torch.zeros(1, 3, 16, 16)} - mock_pipeline.model.get_image_embeddings.return_value = torch.zeros(1, 256, 16, 16) - mock_pipeline.device_placement.return_value.__enter__ = MagicMock(return_value=None) - mock_pipeline.device_placement.return_value.__exit__ = MagicMock(return_value=False) - mock_pipeline.get_inference_context.return_value.return_value.__enter__ = MagicMock(return_value=None) - mock_pipeline.get_inference_context.return_value.return_value.__exit__ = MagicMock(return_value=False) - mock_pipeline._ensure_tensor_on_device.side_effect = lambda x, device: x - - with patch("transformers.pipelines.mask_generation.load_image", return_value=MagicMock()): - gen = MaskGenerationPipeline.preprocess(mock_pipeline, "fake_image.png", points_per_batch=points_per_batch) - yielded = list(gen) - - self.assertEqual(len(yielded), 2) - self.assertFalse(yielded[0]["is_last"]) - self.assertTrue(yielded[1]["is_last"]) + def test_preprocess_is_last(self): + mask_generator = pipeline("mask-generation", model="hf-internal-testing/tiny-random-SamModel") + mask_generator.image_processor.pad_size = {"height": 24, "width": 24} + image = "./tests/fixtures/tests_samples/COCO/000000039769.png" + for points_per_batch in (100, 64): + with self.subTest(points_per_batch=points_per_batch): + batches = list(mask_generator.preprocess(image, points_per_batch=points_per_batch)) + self.assertTrue(batches[-1]["is_last"]) + self.assertFalse(any(b["is_last"] for b in batches[:-1])) @slow @require_torch