From af897e23d7d439b7ce83de62422cdcf924af43a4 Mon Sep 17 00:00:00 2001 From: by-ae <25704492+by-ae@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:29:24 +1000 Subject: [PATCH] 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. --- README.md | 1 + nodes/ideogram4_nodes.py | 45 +++++++++++++++++++++++++++++++++- tests/test_nodes.py | 53 +++++++++++++++++++++++++++++++++++++++- 3 files changed, 97 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index a60a9db..e1452a5 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/nodes/ideogram4_nodes.py b/nodes/ideogram4_nodes.py index ca3b647..7d14886 100644 --- a/nodes/ideogram4_nodes.py +++ b/nodes/ideogram4_nodes.py @@ -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): diff --git a/tests/test_nodes.py b/tests/test_nodes.py index e7169d8..2723845 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -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"]