diff --git a/lama_cleaner/model/base.py b/lama_cleaner/model/base.py index b61ba8f..726a2be 100644 --- a/lama_cleaner/model/base.py +++ b/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] - result = result * (mask / 255) + image[:, :, ::-1] * (1 - (mask / 255)) + if config.sd_prevent_unmasked_area: + mask = mask[:, :, np.newaxis] + result = result * (mask / 255) + image[:, :, ::-1] * (1 - (mask / 255)) return result def forward_post_process(self, result, image, mask, config): diff --git a/lama_cleaner/model/sd.py b/lama_cleaner/model/sd.py index d56139a..f19b96c 100644 --- a/lama_cleaner/model/sd.py +++ b/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") diff --git a/lama_cleaner/schema.py b/lama_cleaner/schema.py index d384a08..8b7e172 100644 --- a/lama_cleaner/schema.py +++ b/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" diff --git a/lama_cleaner/tests/test_sd_model.py b/lama_cleaner/tests/test_sd_model.py index 9b80886..50135a0 100644 --- a/lama_cleaner/tests/test_sd_model.py +++ b/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,