Merge pull request #20 from TashaSkyUp/patch-2

fix for incorrect gathering of images after sampling.
2024-03-02 01:17:19 +02:00
parent dd37991f40 233bc89e67
commit 6dcf0f29e0
1 changed files with 4 additions and 3 deletions
@@ -242,10 +242,11 @@ class SUPIR_Upscale:
        if not keep_model_loaded:
            self.model = None
            mm.soft_empty_cache()
-        if len(out[0]) == 3:
-            out_stacked = torch.stack(out, dim=0).cpu().to(torch.float32).permute(0, 2, 3, 1)
-        else:
+
+        if len(out[0].shape) == 4:
            out_stacked = torch.cat(out, dim=0).cpu().to(torch.float32).permute(0, 2, 3, 1)
+        else:
+            out_stacked = torch.stack(out, dim=0).cpu().to(torch.float32).permute(0, 2, 3, 1)

        return (out_stacked,)