From 233bc89e6730bc25b226ae6bd30c2d7ee0c1fa84 Mon Sep 17 00:00:00 2001 From: Tasha Upchurch Date: Fri, 1 Mar 2024 18:09:36 -0500 Subject: [PATCH] fix for incorrect gathering of images after sampling. --- nodes.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/nodes.py b/nodes.py index 2584993..fa85d73 100644 --- a/nodes.py +++ b/nodes.py @@ -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,)