From c01e040f5538fc3bceeff1b79d5d591a75aa838e Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 2 Mar 2024 03:00:17 +0200
Subject: [PATCH] Resize rounded images back to original

---
 nodes.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 311d3aa..9595c66 100644
--- a/nodes.py
+++ b/nodes.py
@@ -7,6 +7,7 @@ import comfy.utils
 import comfy.model_management as mm
 import folder_paths
 from nodes import ImageScaleBy
+from nodes import ImageScale
 import torch.cuda
 from .sgm.util import instantiate_from_config
 from .SUPIR.util import convert_dtype, load_state_dict
@@ -236,12 +237,18 @@ class SUPIR_Upscale:
             self.model = None
             mm.soft_empty_cache()
 
+        original_height, original_width = H, W  
+        processed_height = samples.size(2)
+        target_width = int(processed_height * (original_width / original_height))
+
         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,)
+        final_image, = ImageScale.upscale(self, out_stacked, "lanczos", target_width, processed_height, crop="disabled")
+
+        return (final_image,)
 
 
 NODE_CLASS_MAPPINGS = {