fix: 🐛 batch support

melMass · Oct 9, 2023 · f1ff9fc · f1ff9fc
1 parent c1d42de
commit f1ff9fc
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 23 deletions.
diff --git a/nodes/image_processing.py b/nodes/image_processing.py
@@ -13,6 +13,7 @@
 from skimage.filters import gaussian
 from skimage.util import compare_images
 
+from ..log import log
 from ..utils import pil2tensor, tensor2np, tensor2pil
 
 # try:
@@ -318,22 +319,26 @@ def INPUT_TYPES(cls):
     FUNCTION = "render_mask"
 
     def render_mask(self, mask, color, background):
-        mask = tensor2np(mask)
-        mask = Image.fromarray(mask).convert("L")
+        masks = tensor2np(mask)
+        images = []
+        for m in masks:
+            _mask = Image.fromarray(m).convert("L")
 
-        image = Image.new("RGBA", mask.size, color=color)
-        # apply the mask
-        image = Image.composite(
-            image, Image.new("RGBA", mask.size, color=background), mask
-        )
+            log.debug(f"Converted mask to PIL Image format, size: {_mask.size}")
+
+            image = Image.new("RGBA", _mask.size, color=color)
+            # apply the mask
+            image = Image.composite(
+                image, Image.new("RGBA", _mask.size, color=background), _mask
+            )
 
-        # image = ImageChops.multiply(image, mask)
-        # apply over background
-        # image = Image.alpha_composite(Image.new("RGBA", image.size, color=background), image)
+            # image = ImageChops.multiply(image, mask)
+            # apply over background
+            # image = Image.alpha_composite(Image.new("RGBA", image.size, color=background), image)
 
-        image = pil2tensor(image.convert("RGB"))
+            images.append(image.convert("RGB"))
 
-        return (image,)
+        return (pil2tensor(images),)
 
 
 class ColoredImage:
@@ -381,19 +386,13 @@ def INPUT_TYPES(cls):
 
     CATEGORY = "mtb/image"
     RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("RGBA",)
     FUNCTION = "premultiply"
 
     def premultiply(self, image, mask, invert):
         images = tensor2pil(image)
-        if invert:
-            masks = tensor2pil(mask)  # .convert("L")
-        else:
-            masks = tensor2pil(1.0 - mask)
-
-        single = False
-        if len(mask) == 1:
-            single = True
-
+        masks = tensor2pil(mask) if invert else tensor2pil(1.0 - mask)
+        single = len(mask) == 1
         masks = [x.convert("L") for x in masks]
 
         out = []

diff --git a/nodes/mask.py b/nodes/mask.py
@@ -1,7 +1,8 @@
+import comfy.utils
+from PIL import Image
 from rembg import remove
+
 from ..utils import pil2tensor, tensor2pil
-from PIL import Image
-import comfy.utils
 
 
 class ImageRemoveBackgroundRembg:
@@ -91,6 +92,8 @@ def remove_background(
 
             image_on_bg.paste(img_rm, mask=mask)
 
+            image_on_bg = image_on_bg.convert("RGB")
+
             out_img.append(img_rm)
             out_mask.append(mask)
             out_img_on_bg.append(image_on_bg)