Fix mask output

Masks were returned as single image, which wouldn't work with AnimateDiff.
2023-09-24 18:29:57 +03:00
parent 2666ef7aeb
commit ea33b45313
1 changed files with 1 additions and 1 deletions
@@ -562,7 +562,7 @@ class LoadImagesFromDirectory:
        if len(images) == 0:
            raise FileNotFoundError(f"No images could be loaded from directory '{directory}'.")

-        return (torch.cat(images, dim=0), torch.cat(masks, dim=0), image_count)
+        return (torch.cat(images, dim=0), torch.stack(masks, dim=0), image_count)