V1.5 增加批量处理 + 输出MASK

BRIA_RMBG.py

@@ -60,12 +60,16 @@ class BRIA_RMBG_Zho:
             }
         }
 
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("image",)
+    RETURN_TYPES = ("IMAGE", "MASK", )
+    RETURN_NAMES = ("image", "mask", )
     FUNCTION = "remove_background"
     CATEGORY = "🧹BRIA RMBG"
   
     def remove_background(self, rmbgmodel, image):
+        processed_images = []
+        processed_masks = []
+
+        for image in image:
             orig_image = tensor2pil(image)
             w,h = orig_image.size
             image = resize_image(orig_image)
@@ -86,8 +90,18 @@ class BRIA_RMBG_Zho:
             pil_im = Image.fromarray(np.squeeze(im_array))
             new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
             new_im.paste(orig_image, mask=pil_im)
-        new_im = pil2tensor(new_im)
-        return (new_im,)
+
+            new_im_tensor = pil2tensor(new_im)  # 将PIL图像转换为Tensor
+            pil_im_tensor = pil2tensor(pil_im)  # 同上
+
+            processed_images.append(new_im_tensor)
+            processed_masks.append(pil_im_tensor)
+
+        new_ims = torch.cat(processed_images, dim=0)
+        new_masks = torch.cat(processed_masks, dim=0)
+
+        return new_ims, new_masks
+        
 
 NODE_CLASS_MAPPINGS = {
     "BRIA_RMBG_ModelLoader_Zho": BRIA_RMBG_ModelLoader_Zho,