Fix mask redundant dims

2024-09-25 03:40:34 +05:00
parent e3a7de4b8b
commit 03ff6c2da0
1 changed files with 4 additions and 1 deletions
@@ -199,11 +199,14 @@ class RembgByBiRefNet:
            image = apply_mask_to_image(image.cpu(), mask.cpu())

            _images.append(image)
-            _masks.append(mask)
+            _masks.append(mask.squeeze(0))

        out_images = torch.cat(_images, dim=0)
        out_masks = torch.cat(_masks, dim=0)

+        if out_masks.shape[0] == 1:
+            out_masks = out_masks.squeeze(0)
+
        return out_images, out_masks