22

app.py

@@ -141,8 +141,14 @@ def inference(
     control_image = processor(control_image, to_pil=True)
     control_image = control_image.resize((width, height))
 
+    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
+
     print("DEBUG: control_image_torch")
-    control_image_torch = get_image_latent(control_image, sample_size=[height, width])[:, :, 0]
+    control_image_torch = get_image_latent(control_image, sample_size=sample_size)[:, :, 0]
 
     # generation
     if randomize_seed: seed = random.randint(0, MAX_SEED)