From a9ab2b62dac17743c62fc1cab8468340496d302e Mon Sep 17 00:00:00 2001 From: Christian Byrne Date: Sat, 15 Aug 2026 10:27:24 -0700 Subject: [PATCH] [Partner Nodes] Stop adding an opaque alpha channel to API node images (#15369) * Stop adding an opaque alpha channel to API node images bytesio_to_image_tensor converted every downloaded image to RGBA, so nodes whose API returns no transparency still emitted a 4 channel IMAGE. Keep the alpha when the decoded image has one, stay RGB when it does not. --------- Signed-off-by: bigcat88 Co-authored-by: bigcat88 --- comfy_api_nodes/nodes_bytedance.py | 6 +- comfy_api_nodes/nodes_gemini.py | 5 +- comfy_api_nodes/nodes_recraft.py | 11 +-- comfy_api_nodes/util/__init__.py | 2 + comfy_api_nodes/util/conversions.py | 19 ++++- .../gemini_response_images_test.py | 57 +++++++++++++ .../image_conversions_test.py | 80 +++++++++++++++++++ 7 files changed, 170 insertions(+), 10 deletions(-) create mode 100644 tests-unit/comfy_api_nodes_test/gemini_response_images_test.py create mode 100644 tests-unit/comfy_api_nodes_test/image_conversions_test.py diff --git a/comfy_api_nodes/nodes_bytedance.py b/comfy_api_nodes/nodes_bytedance.py index 265f94d13..6040075f8 100644 --- a/comfy_api_nodes/nodes_bytedance.py +++ b/comfy_api_nodes/nodes_bytedance.py @@ -56,6 +56,8 @@ from comfy_api_nodes.util import ( ApiEndpoint, audio_bytes_to_audio_input, audio_input_to_mp3, + bytesio_to_image_tensor, + download_url_as_bytesio, download_url_to_image_tensor, download_url_to_video_output, downscale_image_tensor_by_max_side, @@ -1315,7 +1317,9 @@ class ByteDanceSeedreamLayerSeparationNode(IO.ComfyNode): left, top, rect_w, rect_h = spec["left"], spec["top"], spec["rect_w"], spec["rect_h"] async with semaphore: try: - rgba = (await download_url_to_image_tensor(str(item["url"])))[0] + # the layer math below needs the alpha channel, and ByteDance encodes + # alpha-less images as plain RGB (the base plate is one), so force RGBA + rgba = bytesio_to_image_tensor(await download_url_as_bytesio(str(item["url"])), mode="RGBA")[0] except ProcessingInterrupted: raise except Exception as exc: diff --git a/comfy_api_nodes/nodes_gemini.py b/comfy_api_nodes/nodes_gemini.py index 131590751..8343633e0 100644 --- a/comfy_api_nodes/nodes_gemini.py +++ b/comfy_api_nodes/nodes_gemini.py @@ -43,6 +43,7 @@ from comfy_api_nodes.util import ( download_url_to_image_tensor, download_url_to_video_output, get_number_of_images, + pad_images_to_common_channels, sync_op, tensor_to_base64_string, upload_audio_to_comfyapi, @@ -233,8 +234,8 @@ async def get_image_from_response(response: GeminiGenerateContentResponse, thoug "Try rephrasing your prompt or changing the response modality to 'IMAGE+TEXT' " "to see the model's reasoning." ) - return torch.zeros((1, 1024, 1024, 4)) - return torch.cat(image_tensors, dim=0) + return torch.zeros((1, 1024, 1024, 3)) + return torch.cat(pad_images_to_common_channels(image_tensors), dim=0) def get_text_from_interaction(interaction: GeminiInteraction) -> str: diff --git a/comfy_api_nodes/nodes_recraft.py b/comfy_api_nodes/nodes_recraft.py index 2605b9021..9f1823426 100644 --- a/comfy_api_nodes/nodes_recraft.py +++ b/comfy_api_nodes/nodes_recraft.py @@ -27,6 +27,7 @@ from comfy_api_nodes.util import ( ApiEndpoint, bytesio_to_image_tensor, download_url_as_bytesio, + pad_images_to_common_channels, resize_mask_to_image, sync_op, tensor_to_bytesio, @@ -621,7 +622,7 @@ class RecraftImageToImageNode(IO.ComfyNode): images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0)) pbar.update(1) - return IO.NodeOutput(torch.cat(images, dim=0)) + return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0)) class RecraftImageInpaintingNode(IO.ComfyNode): @@ -723,7 +724,7 @@ class RecraftImageInpaintingNode(IO.ComfyNode): images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0)) pbar.update(1) - return IO.NodeOutput(torch.cat(images, dim=0)) + return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0)) class RecraftTextToVectorNode(IO.ComfyNode): @@ -954,7 +955,7 @@ class RecraftReplaceBackgroundNode(IO.ComfyNode): images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0)) pbar.update(1) - return IO.NodeOutput(torch.cat(images, dim=0)) + return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0)) class RecraftRemoveBackgroundNode(IO.ComfyNode): @@ -995,7 +996,7 @@ class RecraftRemoveBackgroundNode(IO.ComfyNode): image=image[i], path="/proxy/recraft/images/removeBackground", ) - images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0)) + images.append(torch.cat([bytesio_to_image_tensor(x, mode="RGBA") for x in sub_bytes], dim=0)) pbar.update(1) images_tensor = torch.cat(images, dim=0) @@ -1047,7 +1048,7 @@ class RecraftCrispUpscaleNode(IO.ComfyNode): images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0)) pbar.update(1) - return IO.NodeOutput(torch.cat(images, dim=0)) + return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0)) class RecraftCreativeUpscaleNode(RecraftCrispUpscaleNode): diff --git a/comfy_api_nodes/util/__init__.py b/comfy_api_nodes/util/__init__.py index 1fb6b96cf..2bb4a1b04 100644 --- a/comfy_api_nodes/util/__init__.py +++ b/comfy_api_nodes/util/__init__.py @@ -18,6 +18,7 @@ from .conversions import ( downscale_image_tensor_by_max_side, downscale_video_to_max_pixels, image_tensor_pair_to_batch, + pad_images_to_common_channels, pil_to_bytesio, resize_mask_to_image, tensor_to_base64_string, @@ -92,6 +93,7 @@ __all__ = [ "downscale_image_tensor_by_max_side", "downscale_video_to_max_pixels", "image_tensor_pair_to_batch", + "pad_images_to_common_channels", "pil_to_bytesio", "resize_mask_to_image", "tensor_to_base64_string", diff --git a/comfy_api_nodes/util/conversions.py b/comfy_api_nodes/util/conversions.py index f46cac3f8..eb81447a0 100644 --- a/comfy_api_nodes/util/conversions.py +++ b/comfy_api_nodes/util/conversions.py @@ -16,12 +16,14 @@ from comfy_api.latest import Input, InputImpl, Types from ._helpers import mimetype_to_extension -def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch.Tensor: +def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str | None = None) -> torch.Tensor: """Converts image data from BytesIO to a torch.Tensor. Args: image_bytesio: BytesIO object containing the image data. - mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA"). + mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA"). Defaults + to RGBA when the decoded image carries transparency and RGB when it + does not, so an API that returns no alpha does not get an opaque one. Returns: A torch.Tensor representing the image (1, H, W, C). @@ -31,6 +33,8 @@ def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch ValueError: If the specified mode is invalid. """ image = Image.open(image_bytesio) + if mode is None: + mode = "RGBA" if "A" in image.getbands() or "transparency" in image.info else "RGB" image = image.convert(mode) image_array = np.array(image).astype(np.float32) / 255.0 return torch.from_numpy(image_array).unsqueeze(0) @@ -53,6 +57,17 @@ def image_tensor_pair_to_batch(image1: torch.Tensor, image2: torch.Tensor) -> to return torch.cat((image1, image2), dim=0) +def pad_images_to_common_channels(images: list[torch.Tensor]) -> list[torch.Tensor]: + """Pads [B, H, W, C] image tensors with opaque alpha so they all share the largest channel count.""" + channels = max(image.shape[-1] for image in images) + return [ + torch.nn.functional.pad(image, (0, channels - image.shape[-1]), value=1.0) + if image.shape[-1] < channels + else image + for image in images + ] + + def tensor_to_bytesio( image: torch.Tensor, *, diff --git a/tests-unit/comfy_api_nodes_test/gemini_response_images_test.py b/tests-unit/comfy_api_nodes_test/gemini_response_images_test.py new file mode 100644 index 000000000..7b4708b43 --- /dev/null +++ b/tests-unit/comfy_api_nodes_test/gemini_response_images_test.py @@ -0,0 +1,57 @@ +import asyncio +import base64 +from io import BytesIO + +import torch +from PIL import Image + +from comfy.cli_args import args + +if not torch.cuda.is_available(): + args.cpu = True + +from comfy_api_nodes.apis.gemini import ( # noqa: E402 + GeminiCandidate, + GeminiContent, + GeminiGenerateContentResponse, + GeminiInlineData, + GeminiPart, +) +from comfy_api_nodes.nodes_gemini import get_image_from_response # noqa: E402 + + +def image_part(mode, color): + buffer = BytesIO() + Image.new(mode, (4, 4), color).save(buffer, format="PNG") + return GeminiPart( + inlineData=GeminiInlineData( + data=base64.b64encode(buffer.getvalue()).decode(), + mimeType="image/png", + ) + ) + + +def response(*parts): + return GeminiGenerateContentResponse( + candidates=[GeminiCandidate(content=GeminiContent(parts=list(parts), role="model"))] + ) + + +def test_rgb_only_response_stays_three_channels(): + out = asyncio.run(get_image_from_response(response(image_part("RGB", (10, 20, 30))))) + assert out.shape == (1, 4, 4, 3) + + +def test_mixed_rgb_and_rgba_parts_are_padded_to_the_same_width(): + out = asyncio.run( + get_image_from_response( + response( + image_part("RGB", (10, 20, 30)), + image_part("RGBA", (10, 20, 30, 0)), + ) + ) + ) + assert out.shape == (2, 4, 4, 4) + # the part that had no alpha is padded opaque, the transparent one is preserved + assert out[0, ..., 3].min() == 1.0 + assert out[1, ..., 3].max() == 0.0 diff --git a/tests-unit/comfy_api_nodes_test/image_conversions_test.py b/tests-unit/comfy_api_nodes_test/image_conversions_test.py new file mode 100644 index 000000000..3caff2818 --- /dev/null +++ b/tests-unit/comfy_api_nodes_test/image_conversions_test.py @@ -0,0 +1,80 @@ +from io import BytesIO + +import pytest +import torch +from PIL import Image + +from comfy.cli_args import args + +if not torch.cuda.is_available(): + args.cpu = True + +from comfy_api_nodes.util.conversions import bytesio_to_image_tensor, pad_images_to_common_channels # noqa: E402 + + +def encode(image: Image.Image, image_format: str = "PNG") -> BytesIO: + buffer = BytesIO() + image.save(buffer, format=image_format) + buffer.seek(0) + return buffer + + +def test_rgb_png_stays_three_channels(): + tensor = bytesio_to_image_tensor(encode(Image.new("RGB", (4, 4), (10, 20, 30)))) + assert tensor.shape == (1, 4, 4, 3) + + +def test_jpeg_stays_three_channels(): + tensor = bytesio_to_image_tensor(encode(Image.new("RGB", (4, 4), (10, 20, 30)), "JPEG")) + assert tensor.shape == (1, 4, 4, 3) + + +def test_grayscale_is_expanded_to_rgb(): + tensor = bytesio_to_image_tensor(encode(Image.new("L", (4, 4), 128))) + assert tensor.shape == (1, 4, 4, 3) + + +def test_rgba_png_keeps_its_alpha(): + tensor = bytesio_to_image_tensor(encode(Image.new("RGBA", (4, 4), (10, 20, 30, 0)))) + assert tensor.shape == (1, 4, 4, 4) + assert tensor[..., 3].max() == 0.0 + + +def test_palette_png_with_transparency_keeps_its_alpha(): + image = Image.new("P", (4, 4), 1) + image.putpalette([0, 0, 0, 255, 255, 255]) + image.info["transparency"] = 0 + image.putpixel((0, 0), 0) + tensor = bytesio_to_image_tensor(encode(image)) + assert tensor.shape == (1, 4, 4, 4) + assert tensor[0, 0, 0, 3] == 0.0 + assert tensor[0, 1, 1, 3] == 1.0 + + +@pytest.mark.parametrize("mode,channels", [("RGB", 3), ("RGBA", 4)]) +def test_explicit_mode_is_respected(mode, channels): + tensor = bytesio_to_image_tensor(encode(Image.new("RGBA", (4, 4), (10, 20, 30, 128))), mode=mode) + assert tensor.shape == (1, 4, 4, channels) + + +def test_pad_mixed_channels_concatenates(): + rgb = torch.rand(1, 4, 4, 3) + rgba = torch.rand(2, 4, 4, 4) + padded = pad_images_to_common_channels([rgb, rgba]) + result = torch.cat(padded, dim=0) + assert result.shape == (3, 4, 4, 4) + + +def test_pad_adds_opaque_alpha_and_keeps_rgb_values(): + rgb = torch.rand(1, 4, 4, 3) + rgba = torch.rand(1, 4, 4, 4) + padded_rgb, padded_rgba = pad_images_to_common_channels([rgb, rgba]) + assert torch.equal(padded_rgb[..., :3], rgb) + assert padded_rgb[..., 3].min() == 1.0 + assert padded_rgba is rgba + + +def test_pad_leaves_homogeneous_channels_unchanged(): + images = [torch.rand(1, 4, 4, 3), torch.rand(2, 4, 4, 3)] + padded = pad_images_to_common_channels(images) + assert all(p is i for p, i in zip(padded, images))