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

Author: lstein

ldm/dream/image_util.py

@@ -50,6 +50,8 @@ def resize(self,width=None,height=None) -> Image:
         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 @@ def process_image(image,seed):
         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,8 +346,11 @@ def process_image(image,seed):
                                 f'Error running RealESRGAN - Your image was not upscaled.\n{e}'
                             )
                         if image_callback is not None:
-                            image_callback(image, seed, upscaled=True)
-                        else: # no callback passed, so we simply replace old image with rescaled one
+                            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
 
         except KeyboardInterrupt:
@@ -436,7 +432,7 @@ def _img2img(
         width,
         height,
         strength,
-        callback, # Currently not implemented for img2img
+        callback,  # Currently not implemented for img2img
     ):
         """
         An infinite iterator of images from the prompt and the initial image
@@ -445,13 +441,13 @@ def _img2img(
         # 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:
             sampler = self.sampler
 
-        init_image = self._load_img(init_img,width,height).to(self.device)
+        init_image = self._load_img(init_img, width, height).to(self.device)
         with precision_scope(self.device.type):
             init_latent = self.model.get_first_stage_encoding(
                 self.model.encode_first_stage(init_image)
@@ -514,7 +510,8 @@ def _sample_to_image(self, samples):
         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,8 +542,9 @@ def load_model(self):
                 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(traceback.format_exc(),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
 
             self._set_sampler()
@@ -606,10 +604,26 @@ def _load_img(self, path, width, height):
         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}')
 
-        image = InitImageResizer(image).resize(width,height)
-        print(f'resized input image to size {image.width}x{image.height}')
+        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)
+                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)
@@ -633,7 +647,7 @@ def _split_weighted_subprompts(text):
                 prompt = text[:idx]
                 remaining -= idx
                 # remove from main text
-                text = text[idx + 1 :]
+                text = text[idx + 1:]
                 # find value for weight
                 if ' ' in text:
                     idx = text.index(' ')   # first occurence
@@ -651,7 +665,7 @@ def _split_weighted_subprompts(text):
                     weight = 1.0
                 # remove from main text
                 remaining -= idx
-                text = text[idx + 1 :]
+                text = text[idx + 1:]
                 # append the sub-prompt and its weight
                 prompts.append(prompt)
                 weights.append(weight)
@@ -662,9 +676,9 @@ def _split_weighted_subprompts(text):
                     weights.append(1.0)
                 remaining = 0
         return prompts, weights
-        
-    # shows how the prompt is tokenized 
-    # usually tokens have '</w>' to indicate end-of-word, 
+
+    # shows how the prompt is tokenized
+    # usually tokens have '</w>' to indicate end-of-word,
     # but for readability it has been replaced with ' '
     def _log_tokenization(self, text):
         if not self.log_tokenization:
@@ -674,15 +688,31 @@ def _log_tokenization(self, text):
         discarded = ""
         usedTokens = 0
         totalTokens = len(tokens)
-        for i in range(0,totalTokens):                
-            token = tokens[i].replace('</w>',' ')
+        for i in range(0, totalTokens):
+            token = tokens[i].replace('</w>', ' ')
             # alternate color
             s = (usedTokens % 6) + 1
             if i < self.model.cond_stage_model.max_length:
                 tokenized = tokenized + f"\x1b[0;3{s};40m{token}"
                 usedTokens += 1
-            else: # over max token length
+            else:  # over max token length
                 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