From 07d8738a4f8a657f1c8d83634dd08d75b0d8a0b4 Mon Sep 17 00:00:00 2001
From: Tasha Upchurch <TashaSkyUp@gmail.com>
Date: Thu, 3 Oct 2024 18:16:49 -0600
Subject: [PATCH] possible fix for handling listed vaedecode and SVD

---
 image.py | 41 ++++++++++++++++++++++++++++++++++-------
 1 file changed, 34 insertions(+), 7 deletions(-)

diff --git a/image.py b/image.py
index 68fcae1..7d1db29 100644
--- a/image.py
+++ b/image.py
@@ -1938,16 +1938,43 @@ class ListVAEDecode:
         num_images = len(results)
 
         # also need to specify to make the final tensor on the CPU
-        final_tensor = torch.zeros((num_images, width, height, channels), device="cpu")
+
+        final_bsize = 0
+        # find the bsize
+        # sometimes a result will actually be shape (>1, W, H, C)
+        # sometimes it will be  (W, H, C)
+        # sometimes it will be (1, W, H, C)
+
         for i in range(num_images):
-            results_to_use = results[i]
-            if len(results_to_use.shape) == 3:
-                results_to_use = results_to_use.unsqueeze(0)
+            if len(results[i].shape) == 3:
+                results[i] = results[i].unsqueeze(0)
+                final_bsize += 1
+            elif len(results[i].shape) == 4:
+                if results[i].shape[0] == 1:
+                    results[i] = results[i]
+                    final_bsize += 1
+                else:
+                    final_bsize += results[i].shape[0]
 
-            print("FT:", final_tensor[i, :, :, :].shape)
-            print("RTU:", results_to_use.shape)
+        # assign
+        final_tensor = torch.zeros((final_bsize, width, height, channels), device="cpu")
+        cur_idx = 0
 
-            final_tensor[i, :, :, :] = results_to_use
+        while True:
+
+            if len(results[i].shape) == 4:
+                if results[i].shape[0] == 1:
+                    final_tensor[cur_idx] = results[i]
+                    cur_idx += 1
+                else:
+                    for j in range(results[i].shape[0]):
+                        res = results[i][j]
+                        if len(res.shape) == 3:
+                            res = res.unsqueeze(0)
+                        final_tensor[cur_idx] = res
+                        cur_idx += 1
+            if cur_idx >= final_bsize:
+                break
 
         return (results, final_tensor,)