fix: accept RGBA and grayscale images as the canvas background

Preview BG Update opened every frame as RGB. Pillow rejects that mode
for a four-channel image, a single-channel image, or a plain height-width
array, so the save stopped and the canvas never updated. Those frames are
converted to RGB first. Alpha is composited onto black, which is what the
editor draws the preview over.
This commit is contained in:
by-ae
2026-10-30 07:29:24 +10:00
parent b87635322b
commit af897e23d7
3 changed files with 97 additions and 2 deletions
+1
View File
@@ -77,6 +77,7 @@ Load shows the request error (for example `HTTP 500`) when the store cannot be r
- Verify the `guid` output from the builder is connected to Preview BG Update's `guid` input
- Check the ComfyUI console for a line starting with `[ByAE Ideogram BG]`. A failed save names the reason: no image, an empty or invalid guid, an image that could not be converted, or a disk write error
- RGBA frames are composited onto black, and single-channel frames are expanded to RGB, before the file is written. The conversion warning is for a tensor that still cannot be turned into an image
- The temp directory must be writable (usually `ComfyUI/temp`)
### Saved workflows break after updating
+44 -1
View File
@@ -167,11 +167,54 @@ def _coerce_opacity(value, default=0.5):
return default
def _uint8_frame_kind(shape):
"""How to open one HWC frame. Unknown shapes stay on the RGB path.
ComfyUI documents IMAGE as [B,H,W,C] with three channels. Pillow still
rejects mode RGB for [H,W], [H,W,1], and [H,W,4], which is what a mask
or an image that kept its alpha looks like.
"""
if (
not isinstance(shape, tuple)
or not shape
or not all(
isinstance(n, int) and not isinstance(n, bool) and n > 0 for n in shape
)
):
return "rgb"
if len(shape) == 2:
return "gray"
if len(shape) == 3 and shape[-1] == 4:
return "rgba"
if len(shape) == 3 and shape[-1] == 1:
return "gray1"
return "rgb"
def _rgb_image_from_uint8(arr):
"""RGB image for the temp PNG. Alpha is composited onto black.
The editor draws that PNG over a black canvas, so a transparent pixel
should become black rather than keep an unrelated RGB value.
"""
kind = _uint8_frame_kind(getattr(arr, "shape", None))
if kind == "rgba":
rgba = Image.fromarray(arr, mode="RGBA")
black = Image.new("RGBA", rgba.size, (0, 0, 0, 255))
return Image.alpha_composite(black, rgba).convert("RGB")
if kind == "gray":
return Image.fromarray(arr, mode="L").convert("RGB")
if kind == "gray1":
flat = arr.reshape(arr.shape[0], arr.shape[1])
return Image.fromarray(flat, mode="L").convert("RGB")
return Image.fromarray(arr, mode="RGB")
def _tensor_to_pil(image):
if image is None:
return None
arr = (image[0].detach().cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
return Image.fromarray(arr, mode="RGB")
return _rgb_image_from_uint8(arr)
def _load_bg_from_temp(guid):
+52 -1
View File
@@ -9,10 +9,20 @@ from nodes.ideogram4_nodes import (
sanitize_bg_guid, _norm_bbox, _palette, _hex_rgb, _dumps, _parse_json_list,
_caption_to_boxes, _coerce_layout_int, _coerce_opacity,
_elements_from_boxes, _preview_edges, _render_preview, ImageDraw,
save_bg_image_to_temp,
save_bg_image_to_temp, _uint8_frame_kind, _rgb_image_from_uint8, Image,
)
class _FakeFrame:
def __init__(self, shape):
self.shape = shape
self.reshaped = None
def reshape(self, *args):
self.reshaped = args
return ("flat", args)
class TestPromptHelpers(unittest.TestCase):
def test_sanitize_bg_guid_valid(self):
valid_guid = "12345678-1234-4abc-8abc-123456789012"
@@ -401,6 +411,47 @@ class TestPromptHelpers(unittest.TestCase):
self.assertIsNone(result)
self.assertTrue(any("invalid background guid" in line.lower() for line in cm.output))
def test_uint8_frame_kind_matches_channel_count(self):
self.assertEqual(_uint8_frame_kind((8, 8, 3)), "rgb")
self.assertEqual(_uint8_frame_kind((8, 8, 4)), "rgba")
self.assertEqual(_uint8_frame_kind((8, 8, 1)), "gray1")
self.assertEqual(_uint8_frame_kind((8, 8)), "gray")
# Anything else keeps the RGB attempt, which reports a conversion failure.
self.assertEqual(_uint8_frame_kind((8, 8, 2)), "rgb")
self.assertEqual(_uint8_frame_kind((0, 8, 3)), "rgb")
self.assertEqual(_uint8_frame_kind(None), "rgb")
# bool is an int, and True == 1, so a boolean channel count must not
# take the single-channel path.
self.assertEqual(_uint8_frame_kind((8, 8, True)), "rgb")
def test_rgba_frame_is_composited_onto_black(self):
arr = _FakeFrame((2, 3, 4))
Image.fromarray.reset_mock()
Image.new.reset_mock()
Image.alpha_composite.reset_mock()
_rgb_image_from_uint8(arr)
Image.fromarray.assert_called_once_with(arr, mode="RGBA")
Image.new.assert_called_once()
args, _kwargs = Image.new.call_args
self.assertEqual(args[0], "RGBA")
self.assertEqual(args[2], (0, 0, 0, 255))
Image.alpha_composite.assert_called_once()
Image.alpha_composite.return_value.convert.assert_called_once_with("RGB")
def test_single_channel_frame_is_opened_as_grayscale(self):
gray = _FakeFrame((4, 5))
Image.fromarray.reset_mock()
_rgb_image_from_uint8(gray)
Image.fromarray.assert_called_once_with(gray, mode="L")
Image.fromarray.return_value.convert.assert_called_once_with("RGB")
column = _FakeFrame((4, 5, 1))
Image.fromarray.reset_mock()
_rgb_image_from_uint8(column)
self.assertEqual(column.reshaped, (4, 5))
Image.fromarray.assert_called_once_with(("flat", (4, 5)), mode="L")
Image.fromarray.return_value.convert.assert_called_once_with("RGB")
def test_save_bg_image_warns_when_conversion_fails(self):
guid = "12345678-1234-4abc-8abc-123456789012"
image_mod = save_bg_image_to_temp.__globals__["Image"]