Merge branch 'resolution-checker' of https://github.com/blessedcoolant/stable-diffusion into main

ldm/dream/image_util.py

@@ -50,6 +50,8 @@ class InitImageResizer():
         new_image = Image.new('RGB',(width,height))
         new_image.paste(resized_image,((width-rw)//2,(height-rh)//2))
 
+        print(f'>> Resized image size to {width}x{height}')
+
         return new_image
 
 def make_grid(image_list, rows=None, cols=None):

ldm/simplet2i.py

@@ -266,16 +266,9 @@ class T2I:
         assert (
             0.0 <= strength <= 1.0
         ), 'can only work with strength in [0.0, 1.0]'
-        w, h = map(
-            lambda x: x - x % 64, (width, height)
-        )  # resize to integer multiple of 64
 
-        if h != height or w != width:
-            print(
-                f'Height and width must be multiples of 64. Resizing to {h}x{w}.'
-            )
-            height = h
-            width  = w
+        if not(width == self.width and height == self.height):
+            width, height, _ = self._resolution_check(width, height, log=True)
 
         scope = autocast if self.precision == 'autocast' else nullcontext
 
@@ -353,6 +346,9 @@ class T2I:
                                 f'Error running RealESRGAN - Your image was not upscaled.\n{e}'
                             )
                         if image_callback is not None:
+                            if save_original:
+                                image_callback(image, seed)
+                            else:
                                 image_callback(image, seed, upscaled=True)
                         else:  # no callback passed, so we simply replace old image with rescaled one
                             result[0] = image
@@ -445,7 +441,7 @@ class T2I:
         # PLMS sampler not supported yet, so ignore previous sampler
         if self.sampler_name != 'ddim':
             print(
-                f"sampler '{self.sampler_name}' is not yet supported. Using DDM sampler"
+                f"sampler '{self.sampler_name}' is not yet supported. Using DDIM sampler"
             )
             sampler = DDIMSampler(self.model, device=self.device)
         else:
@@ -514,7 +510,8 @@ class T2I:
         x_samples = self.model.decode_first_stage(samples)
         x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
         if len(x_samples) != 1:
-            raise Exception(f'expected to get a single image, but got {len(x_samples)}')
+            raise Exception(
+                f'expected to get a single image, but got {len(x_samples)}')
         x_sample = 255.0 * rearrange(
             x_samples[0].cpu().numpy(), 'c h w -> h w c'
         )
@@ -545,7 +542,8 @@ class T2I:
                 self.model.cond_stage_model.device = self.device
             except AttributeError:
                 import traceback
-                print('Error loading model. Only the CUDA backend is supported',file=sys.stderr)
+                print(
+                    'Error loading model. Only the CUDA backend is supported', file=sys.stderr)
                 print(traceback.format_exc(), file=sys.stderr)
                 raise SystemExit
 
@@ -606,10 +604,26 @@ class T2I:
         print(f'image path = {path}, cwd = {os.getcwd()}')
         with Image.open(path) as img:
             image = img.convert('RGB')
-        print(f'loaded input image of size {image.width}x{image.height} from {path}')
+        print(
+            f'loaded input image of size {image.width}x{image.height} from {path}')
 
+        from ldm.dream.image_util import InitImageResizer
+        if width == self.width and height == self.height:
+            new_image_width, new_image_height, resize_needed = self._resolution_check(
+                image.width, image.height)
+        else:
+            if height == self.height:
+                new_image_width, new_image_height, resize_needed = self._resolution_check(
+                    width, image.height)
+            if width == self.width:
+                new_image_width, new_image_height, resize_needed = self._resolution_check(
+                    image.width, height)
+            else:
                 image = InitImageResizer(image).resize(width, height)
-        print(f'resized input image to size {image.width}x{image.height}')
+                resize_needed=False
+        if resize_needed:
+            image = InitImageResizer(image).resize(
+                new_image_width, new_image_height)
 
         image = np.array(image).astype(np.float32) / 255.0
         image = image[None].transpose(0, 3, 1, 2)
@@ -685,4 +699,20 @@ class T2I:
                 discarded = discarded + f"\x1b[0;3{s};40m{token}"
         print(f"\nTokens ({usedTokens}):\n{tokenized}\x1b[0m")
         if discarded != "":
-            print(f"Tokens Discarded ({totalTokens-usedTokens}):\n{discarded}\x1b[0m")
+            print(
+                f"Tokens Discarded ({totalTokens-usedTokens}):\n{discarded}\x1b[0m")
+
+    def _resolution_check(self, width, height, log=False):
+        resize_needed = False
+        w, h = map(
+            lambda x: x - x % 64, (width, height)
+        )  # resize to integer multiple of 64
+        if h != height or w != width:
+            if log:
+                print(
+                    f'>> Provided width and height must be multiples of 64. Auto-resizing to {w}x{h}'
+                )
+            height = h
+            width = w
+            resize_needed = True
+        return width, height, resize_needed