Overview

This pull request addresses a batch size mismatch issue that arises when users set the parameter crops_n_layers > 0 in the Segment Anything Model (SAM) powered mask-generation pipeline. Previously, attempting to run the pipeline caused a runtime error, as the model would receive mismatched batch dimensions for image embeddings and input points. fixes issue #35530 . @Rocketknight1 kindly review this at the earliest .

Fixes #35530

The Problem

When segmenting images in "segment everything" mode with cropping enabled (crops_n_layers=1 or more), the pipeline splits the image into multiple crops. It generates a grid of points for each crop to compute masks in parallel. In some scenarios, particularly where the input is not batched before the cropping procedure, the pipeline would produce additional crops as separate batch entries. This led to:
• A mismatch in the batch dimension between the generated image embeddings and the corresponding point embeddings.
• A ValueError reporting different batch sizes, preventing the pipeline from completing the segmentation process.

Proposed Solution

We introduced internal improvements that unify the handling of image crops and points so they remain synchronized:
• Ensured that all crops from a single original image preserve the same batch index.
• Correctly merged the batched embeddings back together, with consistent shapes, so each batch item has the correct number of crops and associated point embeddings.

Approach

1. Revisited the internal crop generation logic to ensure that each cropped region is tracked as part of the same batch, instead of being treated as separate entries.
2. Adjusted how point embeddings are grouped per crop, clarifying that they should share the original batch dimension.
3. Reworked the post-processing step to handle the final combination of crops so the final output remains consistent with the user’s single input image.

Screenshots

Before

After

Prior to this fix, the Before screenshot would visibly show an error stack trace indicating the batch mismatch. The After screenshot demonstrates a smooth run, producing valid masks across cropped regions.