From 92f0182fd6db66f65cb94bc278fc8e88ced71fe2 Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Wed, 13 Mar 2024 10:32:58 +0900
Subject: [PATCH] fix: Upscaler (SEGS) - compensate resolution

---
 modules/impact/config.py        |  2 +-
 modules/impact/segs_upscaler.py | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/modules/impact/config.py b/modules/impact/config.py
index 364f667..f26dc3f 100644
--- a/modules/impact/config.py
+++ b/modules/impact/config.py
@@ -2,7 +2,7 @@ import configparser
 import os
 
 
-version_code = [4, 83, 4]
+version_code = [4, 83, 5]
 version = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '')
 
 dependency_version = 20
diff --git a/modules/impact/segs_upscaler.py b/modules/impact/segs_upscaler.py
index 165ce77..4d44754 100644
--- a/modules/impact/segs_upscaler.py
+++ b/modules/impact/segs_upscaler.py
@@ -80,6 +80,16 @@ def upscaler(image, upscale_model, rescale_factor, resampling_method, supersampl
 def img2img_segs(image, model, clip, vae, seed, steps, cfg, sampler_name, scheduler,
                  positive, negative, denoise, noise_mask, control_net_wrapper=None,
                  inpaint_model=False, noise_mask_feather=0):
+
+    original_image_size = image.shape[1:3]
+
+    # Match to original image size
+    if original_image_size[0] % 8 > 0 or original_image_size[1] % 8 > 0:
+        scale = 8/min(original_image_size[0], original_image_size[1]) + 1
+        w = int(original_image_size[1] * scale)
+        h = int(original_image_size[0] * scale)
+        image = tensor_resize(image, w, h)
+
     if noise_mask is not None:
         noise_mask = tensor_gaussian_blur_mask(noise_mask, noise_mask_feather)
         noise_mask = noise_mask.squeeze(3)
@@ -106,6 +116,10 @@ def img2img_segs(image, model, clip, vae, seed, steps, cfg, sampler_name, schedu
     # prevent mixing of device
     refined_image = refined_image.cpu()
 
+    # Match to original image size
+    if refined_image.shape[1:3] != original_image_size:
+        refined_image = tensor_resize(refined_image, original_image_size[1], original_image_size[0])
+
     # don't convert to latent - latent break image
     # preserving pil is much better
     return refined_image