get better masks

app.py

@@ -112,8 +112,10 @@ def run_segmentation(image, object_to_segment, device):
         outputs.pred_masks.cpu(),
         inputs["original_sizes"].cpu(),
         inputs["reshaped_input_sizes"].cpu(),
-    )
-
+    )[0]
+    # Merge the masks
+    masks = torch.max(masks[:, 0, ...], dim=0, keepdim=False).values
+    
     delete_model(seg_model)
     return masks
 
@@ -133,10 +135,10 @@ def run_inpainting(image, replaced_caption, masks, device):
     output = pipeline(
         prompt=prompt,
         image=image,
-        mask_image=Image.fromarray(masks[0][0][0, :, :].numpy()),
+        mask_image=Image.fromarray(masks.numpy()),
         negative_prompt=negative_prompt,
         guidance_scale=7.5,
-        strength=0.6,
+        strength=1.0,
     ).images[0]
 
     delete_model(pipeline)
@@ -174,7 +176,7 @@ def run_open_gen_fill(image, edit_prompt):
         replace_with,
         caption,
         replaced_caption,
-        Image.fromarray(masks[0][0][0, :, :].numpy()),
+        Image.fromarray(masks.numpy()),
         output,
     )