Add an alpha mask output to ImageUpscaleWithModel

Upscale models take input_channels channels, so an RGBA image was handed
straight to a 3 channel model and failed. Split the alpha off, upscale it
to match the result and return it as a MASK using the LoadImage polarity.

comfy_extras/nodes_upscale_model.py

@@ -63,6 +63,7 @@ class ImageUpscaleWithModel(io.ComfyNode):
             ],
             outputs=[
                 io.Image.Output(),
+                io.Mask.Output("alpha", tooltip="The input's alpha channel, upscaled to match (1 = transparent, LoadImage convention). Wire it to Join Image with Alpha to put the transparency back; add Invert Mask first for ImageCompositeMasked. All zeros when the input had no alpha."),
             ],
         )
 
@@ -70,6 +71,11 @@ class ImageUpscaleWithModel(io.ComfyNode):
     def execute(cls, upscale_model, image) -> io.NodeOutput:
         device = upscale_model.patcher.load_device
 
+        alpha = None
+        if image.shape[-1] > upscale_model.input_channels:
+            alpha = image[..., upscale_model.input_channels:upscale_model.input_channels + 1]
+            image = image[..., :upscale_model.input_channels]
+
         memory_required = (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
         memory_required += image.nelement() * image.element_size()
         model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required, force_full_load=True)
@@ -95,7 +101,13 @@ class ImageUpscaleWithModel(io.ComfyNode):
                     raise e
 
         s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0).to(comfy.model_management.intermediate_dtype())
-        return io.NodeOutput(s)
+
+        if alpha is None:
+            mask = torch.zeros(s.shape[:-1], device=s.device, dtype=s.dtype)
+        else:
+            alpha = comfy.utils.common_upscale(alpha.to(s).movedim(-1, -3), s.shape[2], s.shape[1], "bilinear", "disabled")
+            mask = 1.0 - alpha.movedim(-3, -1)[..., 0].clamp(min=0.0, max=1.0)
+        return io.NodeOutput(s, mask)
 
     upscale = execute  # TODO: remove
 

tests-unit/comfy_extras_test/upscale_model_alpha_test.py

@@ -0,0 +1,65 @@
+import torch
+
+from comfy.cli_args import args
+
+if not torch.cuda.is_available():
+    args.cpu = True
+
+import comfy.model_management  # noqa: E402
+from comfy_extras.nodes_compositing import JoinImageWithAlpha  # noqa: E402
+from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel  # noqa: E402
+
+
+class StubUpscaleModel:
+    """Stands in for a spandrel ImageModelDescriptor that only accepts RGB."""
+
+    scale = 2
+    input_channels = 3
+
+    class patcher:
+        load_device = torch.device("cpu")
+
+    def __call__(self, image):
+        assert image.shape[1] == self.input_channels, "model was handed a non-RGB tensor"
+        return torch.nn.functional.interpolate(image, scale_factor=self.scale, mode="nearest")
+
+
+def rgba_image():
+    """16x16 RGBA where the top half is transparent and the bottom half opaque."""
+    image = torch.zeros(1, 16, 16, 4)
+    image[..., :3] = 0.5
+    image[0, 8:, :, 3] = 1.0
+    return image
+
+
+def upscale(image, monkeypatch):
+    monkeypatch.setattr(comfy.model_management, "load_models_gpu", lambda *args, **kwargs: None)
+    return ImageUpscaleWithModel.execute(StubUpscaleModel(), image).result
+
+
+def test_rgba_input_does_not_reach_the_model_and_alpha_is_upscaled(monkeypatch):
+    image, mask = upscale(rgba_image(), monkeypatch)
+
+    assert image.shape == (1, 32, 32, 3)
+    assert mask.shape == (1, 32, 32)
+    # inverted convention: transparent -> 1, opaque -> 0
+    assert mask[0, :14].min() > 0.9
+    assert mask[0, 18:].max() < 0.1
+
+
+def test_rgb_input_reports_a_fully_opaque_mask(monkeypatch):
+    image, mask = upscale(torch.zeros(1, 16, 16, 3), monkeypatch)
+
+    assert image.shape == (1, 32, 32, 3)
+    assert mask.shape == (1, 32, 32)
+    assert mask.max() == 0.0
+
+
+def test_mask_round_trips_through_join_image_with_alpha(monkeypatch):
+    image, mask = upscale(rgba_image(), monkeypatch)
+
+    rgba = JoinImageWithAlpha.execute(image, mask).result[0]
+
+    assert rgba.shape == (1, 32, 32, 4)
+    assert rgba[0, :14, :, 3].max() < 0.1
+    assert rgba[0, 18:, :, 3].min() > 0.9