From 09e41e8f769457b9ffd56901aae13991cdc6eeb8 Mon Sep 17 00:00:00 2001
From: Kyle Schouviller <kyle0654@hotmail.com>
Date: Sat, 5 Nov 2022 14:34:52 -0700
Subject: [PATCH] Add inpaint size options to inpaint at a larger size than the
 actual inpaint image, then scale back down for recombination

---
 ldm/generate.py                 |  7 ++++++-
 ldm/invoke/generator/inpaint.py | 18 +++++++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/ldm/generate.py b/ldm/generate.py
index 832dab40f8..a437d3baf4 100644
--- a/ldm/generate.py
+++ b/ldm/generate.py
@@ -299,6 +299,9 @@ class Generate:
             upscale          = None,
             # this is specific to inpainting and causes more extreme inpainting
             inpaint_replace  = 0.0,
+            # This controls the size at which inpaint occurs (scaled up for inpaint, then back down for the result)
+            inpaint_width    = None,
+            inpaint_height   = None,
             # This will help match inpainted areas to the original image more smoothly
             mask_blur_radius: int = 8,
             # Set this True to handle KeyboardInterrupt internally
@@ -490,7 +493,9 @@ class Generate:
                 seam_strength = seam_strength,
                 seam_steps = seam_steps,
                 tile_size = tile_size,
-                force_outpaint = force_outpaint
+                force_outpaint = force_outpaint,
+                inpaint_width  = inpaint_width,
+                inpaint_height = inpaint_height
             )
 
             if init_color:
diff --git a/ldm/invoke/generator/inpaint.py b/ldm/invoke/generator/inpaint.py
index 2ba45ed804..f3dd060cfb 100644
--- a/ldm/invoke/generator/inpaint.py
+++ b/ldm/invoke/generator/inpaint.py
@@ -150,7 +150,10 @@ class Inpaint(Img2Img):
                        seam_steps: int = 10,
                        tile_size: int = 32,
                        step_callback=None,
-                       inpaint_replace=False, **kwargs):
+                       inpaint_replace=False,
+                       inpaint_width=None,
+                       inpaint_height=None,
+                       **kwargs):
         """
         Returns a function returning an image derived from the prompt and
         the initial image + mask.  Return value depends on the seed at
@@ -168,11 +171,20 @@ class Inpaint(Img2Img):
             )
             init_filled.paste(init_image, (0,0), init_image.split()[-1])
 
+            # Resize if requested for inpainting
+            if inpaint_width and inpaint_height:
+                init_filled = init_filled.resize((inpaint_width, inpaint_height))
+
             # Create init tensor
             init_image = self._image_to_tensor(init_filled.convert('RGB'))
 
         if isinstance(mask_image, PIL.Image.Image):
             self.pil_mask = mask_image
+
+            # Resize if requested for inpainting
+            if inpaint_width and inpaint_height:
+                mask_image = mask_image.resize((inpaint_width, inpaint_height))
+
             mask_image = mask_image.resize(
                 (
                     mask_image.width // downsampling,
@@ -241,6 +253,10 @@ class Inpaint(Img2Img):
 
             result = self.sample_to_image(samples)
 
+            # Resize if necessary
+            if inpaint_width and inpaint_height:
+                result = result.resize(self.pil_image.size)
+
             # Seam paint if this is our first pass (seam_size set to 0 during seam painting)
             if seam_size > 0:
                 result = self.seam_paint(