Comfy-Org / Comfy-Org/ComfyUI

RemoveBackground returns 4D MASK tensor (B x 1 x H x W) instead of standard B x H x W

Open
#14,008 1 comment 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

## Title

RemoveBackground returns 4D MASK tensor (`B x 1 x H x W`) instead of standard `B x H x W`

## Body

### Summary

The new BiRefNet `RemoveBackground` node declares a `MASK` output, but currently returns a 4D tensor shaped `B x 1 x H x W`.

ComfyUI `MASK` tensors are conventionally `B x H x W`. Several existing mask producers follow this shape, and some consumers validate/unpack masks as 3D tensors. As a result, piping `RemoveBackground` into strict mask consumers can fail even when the image and mask spatial sizes match.

### Affected code

Introduced by PR #12747 / commit `d3c18c163`.

`comfy_extras/nodes_bg_removal.py` declares a `MASK` output:

```python
outputs=[
IO.Mask.Output("mask", tooltip="Generated foreground mask")
]
```

`comfy/bg_removal_model.py` preserves the singleton channel dimension and returns it:

```python
out = torch.nn.functional.interpolate(out, size=(H, W), mode="bicubic", antialias=False)

mask = out.sigmoid().to(device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype())
if mask.ndim == 3:
mask = mask.unsqueeze(0)
if mask.shape[1] != 1:
mask = mask.movedim(-1, 1)

return mask
```

For BiRefNet, the model output is `B x 1 x h x w`, so after resize this returns `B x 1 x H x W`.

### Why this is a problem

Existing mask producers return `B x H x W`, for example:

- `ImageToMask`: `image[:, :, :, channel]`
- `SolidMask`: `(1, height, width)`
- `LoadImage` / `LoadImageMask`: batched alpha mask as `B x H x W`

Strict consumers also expect 3D masks. For example, OpenAI image edit paths in `comfy_api_nodes/nodes_openai.py` do:

```python
if mask.shape[1:] != image.shape[1:-1]:
raise Exception("Mask and Image must be the same size")
_, height, width = mask.shape
```

For an image shaped `B x H x W x C`, the current `RemoveBackground` output gives:

```text
image.shape[1:-1] == (H, W)
mask.shape == (B, 1, H, W)
mask.shape[1:] == (1, H, W)
```

So the size check fails even though the spatial dimensions are correct.

### Expected behavior

`RemoveBackground` should return a standard ComfyUI `MASK` tensor shaped `B x H x W`.

### Suggested fix

Squeeze only the singleton channel dimension before returning:

```python
if mask.ndim == 4 and mask.shape[1] == 1:
mask = mask[:, 0, :, :]

return mask
```

This preserves the batch dimension, including `B == 1`, and avoids the unsafe behavior of a broad `mask.squeeze()`.

### Notes

The BiRefNet model-internal `B x 1 x H x W` prediction is fine. The issue is only that the public `IO.Mask.Output` should follow the normal ComfyUI `MASK` shape contract.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.