From bcd1483a14e2af4b5cede1c65318628388d46cb8 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Wed, 31 Jul 2024 09:51:14 -0400 Subject: [PATCH] Re-order GroundedSAMInvocation._to_numpy_masks(...) to do slightly more work on the GPU. --- invokeai/app/invocations/grounded_sam.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/invokeai/app/invocations/grounded_sam.py b/invokeai/app/invocations/grounded_sam.py index cac338c160..3a1cb9be56 100644 --- a/invokeai/app/invocations/grounded_sam.py +++ b/invokeai/app/invocations/grounded_sam.py @@ -156,11 +156,11 @@ class GroundedSAMInvocation(BaseInvocation): def _to_numpy_masks(self, masks: torch.Tensor) -> list[npt.NDArray[np.uint8]]: """Convert the tensor output from the Segment Anything model to a list of numpy masks.""" - masks = masks.cpu().float() - masks = masks.permute(0, 2, 3, 1) - masks = masks.mean(dim=-1) - masks = (masks > 0).int() - np_masks = masks.numpy().astype(np.uint8) + eps = 0.0001 + # [num_masks, channels, height, width] -> [num_masks, height, width] + masks = masks.permute(0, 2, 3, 1).float().mean(dim=-1) + masks = masks > eps + np_masks = masks.cpu().numpy().astype(np.uint8) return list(np_masks) def _apply_polygon_refinement(self, masks: list[npt.NDArray[np.uint8]]) -> list[npt.NDArray[np.uint8]]: