kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

add sd_prevent_unmasked_area

Browse Source
This commit is contained in:
Qing 2023-11-15 09:10:13 +08:00
parent f13bfb9085
commit 3e7a553273
4 changed files with 15 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
lama_cleaner/model/base.py
View File

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

1
lama_cleaner/model/sd.py
View File

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

3
lama_cleaner/schema.py
View File

@ -103,6 +103,9 @@ class Config(BaseModel):
# lcm-lora # lcm-lora
sd_lcm_lora: bool = False 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 # Configs for opencv inpainting
# opencv document https://docs.opencv.org/4.6.0/d7/d8b/group__photo__inpaint.html#gga8002a65f5a3328fbf15df81b842d3c3ca05e763003a805e6c11c673a9f4ba7d07 # opencv document https://docs.opencv.org/4.6.0/d7/d8b/group__photo__inpaint.html#gga8002a65f5a3328fbf15df81b842d3c3ca05e763003a805e6c11c673a9f4ba7d07
cv2_flag: str = "INPAINT_NS" cv2_flag: str = "INPAINT_NS"

12
lama_cleaner/tests/test_sd_model.py
View File

@ -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("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim]) @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): def callback(i, t, latents):
pass pass
if sd_device == "cuda" and not torch.cuda.is_available(): if sd_device == "cuda" and not torch.cuda.is_available():
return return
sd_steps = 50 if sd_device == "cuda" else 1 sd_steps = 50 if sd_device == "cuda" else 20
model = ModelManager( model = ModelManager(
name="sd1.5", name="sd1.5",
device=torch.device(sd_device), 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", negative_prompt="orange, yellow, small",
sd_sampler=sampler, sd_sampler=sampler,
sd_match_histograms=True, sd_match_histograms=True,
sd_prevent_unmasked_area=sd_prevent_unmasked_area,
) )
name = f"{sampler}_negative_prompt" name = f"{sampler}_negative_prompt"
@ -129,7 +133,7 @@ def test_runway_sd_1_5_negative_prompt(sd_device, strategy, sampler):
assert_equal( assert_equal(
model, model,
cfg, 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", img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png", mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1, fx=1,