Update WAS_Image_Threshold to accept batch images

lcolok · Sep 8, 2024 · 2a5be77 · 2a5be77
1 parent ee2e31a
commit 2a5be77
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py
@@ -5596,7 +5596,10 @@ def INPUT_TYPES(cls):
     CATEGORY = "WAS Suite/Image/Process"
 
     def image_threshold(self, image, threshold=0.5):
-        return (pil2tensor(self.apply_threshold(tensor2pil(image), threshold)), )
+        images = []
+        for img in image:
+            images.append(pil2tensor(self.apply_threshold(tensor2pil(img), threshold)))
+        return (torch.cat(images, dim=0), )
 
     def apply_threshold(self, input_image, threshold=0.5):
         # Convert the input image to grayscale