RemoveBackground returns 4D MASK tensor (B x 1 x H x W) instead of standard B x H x W
- 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
Assessment
This issue has not been assessed yet.