[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 <bigcat88@icloud.com>
Co-authored-by: bigcat88 <bigcat88@icloud.com>

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:

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:

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):

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",

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,
     *,

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

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))