From cac32fe37c95ac23d9c2762392ef7fd179ed9c74 Mon Sep 17 00:00:00 2001
From: Acly <aclysia@gmail.com>
Date: Sun, 21 Jul 2024 18:31:42 +0200
Subject: [PATCH] Move image channel permutation to separate functions

---
 nodes.py | 46 +++++++++++++++++++++++++++-------------------
 1 file changed, 27 insertions(+), 19 deletions(-)

diff --git a/nodes.py b/nodes.py
index bf03f45..39edb5d 100644
--- a/nodes.py
+++ b/nodes.py
@@ -118,6 +118,22 @@ class CropImage:
         return (out,)
 
 
+def to_bchw(image: torch.Tensor):
+    if image.ndim == 3:
+        image = image.unsqueeze(0)
+    return image.movedim(-1, 1)
+
+
+def to_bhwc(image: torch.Tensor):
+    return image.movedim(1, -1)
+
+
+def mask_batch(mask: torch.Tensor):
+    if mask.ndim == 2:
+        mask = mask.unsqueeze(0)
+    return mask
+
+
 class ApplyMaskToImage:
     @classmethod
     def INPUT_TYPES(cls):
@@ -133,29 +149,21 @@ class ApplyMaskToImage:
     FUNCTION = "apply_mask"
 
     def apply_mask(self, image: torch.Tensor, mask: torch.Tensor):
-        # Move the channel to the second dimension for processing
-        out = image.movedim(-1, 1)
-
-        # Check if the images are RGB, and if so, add an alpha channel initialized to 1
+        out = to_bchw(image)
         if out.shape[1] == 3:  # Assuming RGB images
             out = torch.cat([out, torch.ones_like(out[:, :1, :, :])], dim=1)
-
-        # Ensure masks are unsqueezed to match the alpha channel dimension if needed
-        if mask.ndim == 2:
-            mask = mask.unsqueeze(0)  # Add a batch dimension to masks
-        # For single mask, expand it to match size of image batch size.
-        if mask.shape[0] == 1:
-            mask = mask.repeat(out.shape[0], 1, 1)
+        mask = mask_batch(mask)
 
         assert mask.ndim == 3, f"Mask should have shape [B, H, W]. {mask.shape}"
-        assert out.ndim == 4, f"Image should have shsape [B, C, H, W]. {out.shape}"
-        assert out.shape[-2:] == mask.shape[-2:], f"{out.shape[-2:]} != {mask.shape[-2:]}"
-        assert out.shape[0] == mask.shape[0], f"{out.shape[0]} != {mask.shape[0]}"
+        assert out.ndim == 4, f"Image should have shape [B, C, H, W]. {out.shape}"
+        assert (
+            out.shape[-2:] == mask.shape[-2:]
+        ), f"Image size {out.shape[-2:]} must match mask size {mask.shape[-2:]}"
+        is_mask_batch = mask.shape[0] == out.shape[0]
+
         # Apply each mask in the batch to its corresponding image's alpha channel
         for i in range(out.shape[0]):
-            out[i, 3, :, :] = mask[i]
+            alpha = mask[i] if is_mask_batch else mask[0]
+            out[i, 3, :, :] = alpha
 
-        # Move the channel back to its original dimension
-        out = out.movedim(1, -1)
-
-        return (out,)
+        return (to_bhwc(out),)