Fix outpainting params (huggingface#1089)

nod-ai · Feb 24, 2023 · bae208b · bae208b
1 parent b6c14ad
commit bae208b
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/apps/stable_diffusion/scripts/outpaint.py b/apps/stable_diffusion/scripts/outpaint.py
@@ -157,8 +157,8 @@ def outpaint_inf(
             prompt,
             negative_prompt,
             image,
-            args.pixels,
-            args.mask_blur,
+            pixels,
+            mask_blur,
             left,
             right,
             top,

diff --git a/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py b/apps/stable_diffusion/src/pipelines/pipeline_shark_stable_diffusion_outpaint.py
@@ -69,11 +69,13 @@ def prepare_latents(
         latents = latents * self.scheduler.init_noise_sigma
         return latents
 
-    def prepare_mask_and_masked_image(self, image, mask, mask_blur):
+    def prepare_mask_and_masked_image(
+        self, image, mask, mask_blur, width, height
+    ):
         if mask_blur > 0:
             mask = mask.filter(ImageFilter.GaussianBlur(mask_blur))
-        image = image.resize((512, 512))
-        mask = mask.resize((512, 512))
+        image = image.resize((width, height))
+        mask = mask.resize((width, height))
 
         # preprocess image
         if isinstance(image, (Image.Image, np.ndarray)):
@@ -478,7 +480,7 @@ def expand(
 
             # Preprocess mask and image
             mask, masked_image = self.prepare_mask_and_masked_image(
-                image_to_process, mask_to_process, mask_blur
+                image_to_process, mask_to_process, mask_blur, width, height
             )
 
             # Prepare mask latent variables

diff --git a/apps/stable_diffusion/src/utils/stable_args.py b/apps/stable_diffusion/src/utils/stable_args.py
@@ -61,13 +61,15 @@ def path_expand(s):
     "--height",
     type=int,
     default=512,
+    choices=range(384, 768, 8),
     help="the height of the output image.",
 )
 
 p.add_argument(
     "--width",
     type=int,
     default=512,
+    choices=range(384, 768, 8),
     help="the width of the output image.",
 )