22

app.py

@@ -94,8 +94,8 @@ pipe.to("cuda", torch.bfloat16)
 print("pipe ready.")
 
 def prepare(edit_dict):
-    return edit_dict['background']
-    # return edit_dict['layers'][0]
+    # return edit_dict['background']
+    return edit_dict['layers'][0]
 
 @spaces.GPU
 def inference(
@@ -148,7 +148,8 @@ def inference(
     # control_image = control_image.resize((width, height))
     # control_image_torch = get_image_latent(control_image, sample_size=sample_size)[:, :, 0]
 
-    mask_image = edit_dict['layers'][0]
+    # mask_image = edit_dict['layers'][0]
+    mask_image = edit_dict['composite']
     if mask_image is not None:
         mask_image = get_image_latent(mask_image, sample_size=sample_size)[:, :1, 0]
     else: