22

app.py

@@ -117,26 +117,26 @@ def inference(
 
     # process image
     print("DEBUG: process image")
-    if edit_dict is None:
-        print("Error: edit_dict is empty.")
+    if edit_dict is None or mask_image is None:
+        print("Error: edit_dict or mask_image is empty.")
         return None
     
-    width, height = edit_dict['background'].size
-
-    print("DEBUG: control_image_torch")
+    # rescale to prevent OOM
+    inpaint_image = edit_dict['background']
+    inpaint_image, width, height = image_utils.rescale_image(inpaint_image, 1, 8)
     sample_size = [height, width]
-
-    if mask_image is not None:
-        mask_image = get_image_latent(mask_image, sample_size=sample_size)[:, :1, 0]
-    else:
-        mask_image = torch.ones([1, 1, sample_size[0], sample_size[1]]) * 255
     
-    inpaint_image = edit_dict['background']
+    print("DEBUG: control_image_torch")
     if inpaint_image is not None:
         inpaint_image = get_image_latent(inpaint_image, sample_size=sample_size)[:, :, 0]
     else:
         inpaint_image = torch.zeros([1, 3, sample_size[0], sample_size[1]])
-    
+
+    if mask_image is not None:
+        mask_image, w, h = image_utils.rescale_image(mask_image, 1, 8)
+        mask_image = get_image_latent(mask_image, sample_size=sample_size)[:, :1, 0]
+    else:
+        mask_image = torch.ones([1, 1, sample_size[0], sample_size[1]]) * 255
 
     # generation
     if randomize_seed: seed = random.randint(0, MAX_SEED)