Make LoadImage and LoadImageMask return masks in batch format.

2023-09-26 02:56:40 -04:00 · 2023-09-26 02:56:40 -04:00 · 9546a798fb
parent 1d36dfb9fe
commit 9546a798fb
1 changed files with 2 additions and 2 deletions
--- a/nodes.py
+++ b/nodes.py
@ -1369,7 +1369,7 @@ class LoadImage:
            mask = 1. - torch.from_numpy(mask)
        else:
            mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
-        return (image, mask)
+        return (image, mask.unsqueeze(0))

    @classmethod
    def IS_CHANGED(s, image):
@ -1416,7 +1416,7 @@ class LoadImageMask:
                mask = 1. - mask
        else:
            mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
-        return (mask,)
+        return (mask.unsqueeze(0),)

    @classmethod
    def IS_CHANGED(s, image, channel):