Commit 9546a798 authored by comfyanonymous's avatar comfyanonymous
Browse files

Make LoadImage and LoadImageMask return masks in batch format.

parent 1d36dfb9
...@@ -1369,7 +1369,7 @@ class LoadImage: ...@@ -1369,7 +1369,7 @@ class LoadImage:
mask = 1. - torch.from_numpy(mask) mask = 1. - torch.from_numpy(mask)
else: else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
return (image, mask) return (image, mask.unsqueeze(0))
@classmethod @classmethod
def IS_CHANGED(s, image): def IS_CHANGED(s, image):
...@@ -1416,7 +1416,7 @@ class LoadImageMask: ...@@ -1416,7 +1416,7 @@ class LoadImageMask:
mask = 1. - mask mask = 1. - mask
else: else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
return (mask,) return (mask.unsqueeze(0),)
@classmethod @classmethod
def IS_CHANGED(s, image, channel): def IS_CHANGED(s, image, channel):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment