From ea33b4531324f03523249f593a8cc874869652e5 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 24 Sep 2023 18:29:57 +0300
Subject: [PATCH] Fix mask output

Masks were returned as single image, which wouldn't work with AnimateDiff.
---
 nodes.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 79a2e64..1a40aec 100644
--- a/nodes.py
+++ b/nodes.py
@@ -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)