fix(api): match mask and image size before adding noise (#373)

ssube · May 4, 2023 · 23aaf65 · 23aaf65
1 parent d66bf9e
commit 23aaf65
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/api/onnx_web/chain/upscale_outpaint.py b/api/onnx_web/chain/upscale_outpaint.py
@@ -50,6 +50,11 @@ def upscale_outpaint(
         # if no mask was provided, keep the full source image
         stage_mask = Image.new("RGB", source.size, "black")
 
+    # masks start as 512x512, resize to cover the source, then trim the extra
+    mask_max = max(source.width, source.height)
+    stage_mask = ImageOps.contain(stage_mask, (mask_max, mask_max))
+    stage_mask = stage_mask.crop((0, 0, source.width, source.height))
+
     source, stage_mask, noise, full_size = expand_image(
         source,
         stage_mask,
@@ -59,11 +64,6 @@ def upscale_outpaint(
         mask_filter=mask_filter,
     )
 
-    # masks start as 512x512, resize to cover the source, then trim the extra
-    mask_max = max(source.width, source.height)
-    stage_mask = ImageOps.contain(stage_mask, (mask_max, mask_max))
-    stage_mask = stage_mask.crop((0, 0, source.width, source.height))
-
     full_latents = get_latents_from_seed(params.seed, Size(*full_size))
 
     draw_mask = ImageDraw.Draw(stage_mask)