added combine_masks

app.py

@@ -12,7 +12,7 @@ import argparse
 # Load configuration and models
 config = OmegaConf.load("config/inference_config.yaml")
 sd_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
-    "botp/stable-diffusion-v1-5-inpainting", torch_dtype=torch.float32
+    "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float32
 )
 clipaway = CLIPAway(
     sd_pipe=sd_pipeline, 
@@ -26,11 +26,19 @@ clipaway = CLIPAway(
 )
 
 def dilate_mask(mask, kernel_size=5, iterations=5):
-    mask = mask.convert("L").resize((512, 512), Image.NEAREST)
+    mask = mask.convert("L")
     kernel = np.ones((kernel_size, kernel_size), np.uint8)
     mask = cv2.dilate(np.array(mask), kernel, iterations=iterations)
     return Image.fromarray(mask)
 
+def combine_masks(uploaded_mask, sketched_mask):
+    if uploaded_mask is not None:
+        return uploaded_mask
+    elif sketched_mask is not None:
+        return sketched_mask
+    else:
+        raise ValueError("Please provide a mask")
+
 def remove_obj(image, uploaded_mask, seed):
     image_pil, sketched_mask = image["image"], image["mask"]
     mask = dilate_mask(combine_masks(uploaded_mask, sketched_mask))