add sd_prevent_unmasked_area

lama_cleaner/model/base.py

@@ -67,8 +67,9 @@ class InpaintModel:
 
         result, image, mask = self.forward_post_process(result, image, mask, config)
 
-        mask = mask[:, :, np.newaxis]
+        if config.sd_prevent_unmasked_area:
-        result = result * (mask / 255) + image[:, :, ::-1] * (1 - (mask / 255))
+            mask = mask[:, :, np.newaxis]
+            result = result * (mask / 255) + image[:, :, ::-1] * (1 - (mask / 255))
         return result
 
     def forward_post_process(self, result, image, mask, config):

lama_cleaner/model/sd.py

@@ -151,6 +151,7 @@ class SD(DiffusionInpaintModel):
             height=img_h,
             width=img_w,
             generator=torch.manual_seed(config.sd_seed),
+            callback_steps=1,
         ).images[0]
 
         output = (output * 255).round().astype("uint8")

lama_cleaner/schema.py

@@ -103,6 +103,9 @@ class Config(BaseModel):
     # lcm-lora
     sd_lcm_lora: bool = False
 
+    # preserving the unmasked area at the expense of some more unnatural transitions between the masked and unmasked areas.
+    sd_prevent_unmasked_area: bool = True
+
     # Configs for opencv inpainting
     # opencv document https://docs.opencv.org/4.6.0/d7/d8b/group__photo__inpaint.html#gga8002a65f5a3328fbf15df81b842d3c3ca05e763003a805e6c11c673a9f4ba7d07
     cv2_flag: str = "INPAINT_NS"

lama_cleaner/tests/test_sd_model.py

@@ -95,17 +95,20 @@ def test_runway_sd_1_5(sd_device, strategy, sampler, cpu_textencoder, disable_ns
     )
 
 
-@pytest.mark.parametrize("sd_device", ["cuda"])
+@pytest.mark.parametrize("sd_device", ["mps"])
 @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
 @pytest.mark.parametrize("sampler", [SDSampler.ddim])
-def test_runway_sd_1_5_negative_prompt(sd_device, strategy, sampler):
+@pytest.mark.parametrize("sd_prevent_unmasked_area", [False, True])
+def test_runway_sd_1_5_negative_prompt(
+    sd_device, strategy, sampler, sd_prevent_unmasked_area
+):
     def callback(i, t, latents):
         pass
 
     if sd_device == "cuda" and not torch.cuda.is_available():
         return
 
-    sd_steps = 50 if sd_device == "cuda" else 1
+    sd_steps = 50 if sd_device == "cuda" else 20
     model = ModelManager(
         name="sd1.5",
         device=torch.device(sd_device),
@@ -122,6 +125,7 @@ def test_runway_sd_1_5_negative_prompt(sd_device, strategy, sampler):
         negative_prompt="orange, yellow, small",
         sd_sampler=sampler,
         sd_match_histograms=True,
+        sd_prevent_unmasked_area=sd_prevent_unmasked_area,
     )
 
     name = f"{sampler}_negative_prompt"
@@ -129,7 +133,7 @@ def test_runway_sd_1_5_negative_prompt(sd_device, strategy, sampler):
     assert_equal(
         model,
         cfg,
-        f"runway_sd_{strategy.capitalize()}_{name}.png",
+        f"runway_sd_{strategy.capitalize()}_{name}_prevent_unmasked_area_{sd_prevent_unmasked_area}.png",
         img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
         mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
         fx=1,