From 205170e1e57d96970935bf6061a7e9eb90f6eb13 Mon Sep 17 00:00:00 2001 From: Qing Date: Fri, 27 Jan 2023 20:59:22 +0800 Subject: [PATCH] add DiffusionInpaintModel --- lama_cleaner/make_gif.py | 29 --------------- lama_cleaner/model/base.py | 39 ++++++++++++++++++++ lama_cleaner/model/paint_by_example.py | 51 ++------------------------ lama_cleaner/model/sd.py | 50 ++----------------------- lama_cleaner/model/utils.py | 8 ++++ 5 files changed, 55 insertions(+), 122 deletions(-) diff --git a/lama_cleaner/make_gif.py b/lama_cleaner/make_gif.py index 342a1af..12a213f 100644 --- a/lama_cleaner/make_gif.py +++ b/lama_cleaner/make_gif.py @@ -123,32 +123,3 @@ def make_compare_gif( loop=0 ) return img_byte_arr.getvalue() - - -if __name__ == '__main__': - imgs = [ - ( - '/Users/qing/code/github/lama-cleaner/assets/unwant_person.jpg', - '/Users/qing/code/github/lama-cleaner/assets/unwant_person_clean.jpg' - ), - # ( - # '/Users/qing/code/github/lama-cleaner/assets/old_photo.jpg', - # '/Users/qing/code/github/lama-cleaner/assets/old_photo_clean.jpg' - # ), - # ( - # '/Users/qing/code/github/lama-cleaner/assets/unwant_object.jpg', - # '/Users/qing/code/github/lama-cleaner/assets/unwant_object_clean.jpg' - # ), - # ( - # '/Users/qing/code/github/lama-cleaner/assets/unwant_text.jpg', - # '/Users/qing/code/github/lama-cleaner/assets/unwant_text_clean.jpg' - # ), - # ( - # '/Users/qing/code/github/lama-cleaner/assets/watermark.jpg', - # '/Users/qing/code/github/lama-cleaner/assets/watermark_cleanup.jpg' - # ), - ] - for src_p, clean_p in imgs: - img_bytes = make_compare_gif(Image.open(src_p), Image.open(clean_p), max_side_length=600) - with open(Path(src_p).with_suffix('.gif'), 'wb') as f: - f.write(img_bytes) diff --git a/lama_cleaner/model/base.py b/lama_cleaner/model/base.py index 05d27f1..60eccb4 100644 --- a/lama_cleaner/model/base.py +++ b/lama_cleaner/model/base.py @@ -245,3 +245,42 @@ class InpaintModel: crop_img, crop_mask, [l, t, r, b] = self._crop_box(image, mask, box, config) return self._pad_forward(crop_img, crop_mask, config), [l, t, r, b] + + +class DiffusionInpaintModel(InpaintModel): + @torch.no_grad() + def __call__(self, image, mask, config: Config): + """ + images: [H, W, C] RGB, not normalized + masks: [H, W] + return: BGR IMAGE + """ + # boxes = boxes_from_mask(mask) + if config.use_croper: + crop_img, crop_mask, (l, t, r, b) = self._apply_cropper(image, mask, config) + crop_image = self._scaled_pad_forward(crop_img, crop_mask, config) + inpaint_result = image[:, :, ::-1] + inpaint_result[t:b, l:r, :] = crop_image + else: + inpaint_result = self._scaled_pad_forward(image, mask, config) + + return inpaint_result + + def _scaled_pad_forward(self, image, mask, config: Config): + longer_side_length = int(config.sd_scale * max(image.shape[:2])) + origin_size = image.shape[:2] + downsize_image = resize_max_size(image, size_limit=longer_side_length) + downsize_mask = resize_max_size(mask, size_limit=longer_side_length) + logger.info( + f"Resize image to do sd inpainting: {image.shape} -> {downsize_image.shape}" + ) + inpaint_result = self._pad_forward(downsize_image, downsize_mask, config) + # only paste masked area result + inpaint_result = cv2.resize( + inpaint_result, + (origin_size[1], origin_size[0]), + interpolation=cv2.INTER_CUBIC, + ) + original_pixel_indices = mask < 127 + inpaint_result[original_pixel_indices] = image[:, :, ::-1][original_pixel_indices] + return inpaint_result diff --git a/lama_cleaner/model/paint_by_example.py b/lama_cleaner/model/paint_by_example.py index b946a08..29ffa2a 100644 --- a/lama_cleaner/model/paint_by_example.py +++ b/lama_cleaner/model/paint_by_example.py @@ -1,19 +1,16 @@ -import random - import PIL import PIL.Image import cv2 -import numpy as np import torch from diffusers import DiffusionPipeline from loguru import logger -from lama_cleaner.helper import resize_max_size -from lama_cleaner.model.base import InpaintModel +from lama_cleaner.model.base import DiffusionInpaintModel +from lama_cleaner.model.utils import set_seed from lama_cleaner.schema import Config -class PaintByExample(InpaintModel): +class PaintByExample(DiffusionInpaintModel): pad_mod = 8 min_size = 512 @@ -53,11 +50,7 @@ class PaintByExample(InpaintModel): mask: [H, W, 1] 255 means area to repaint return: BGR IMAGE """ - seed = config.paint_by_example_seed - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) + set_seed(config.paint_by_example_seed) output = self.model( image=PIL.Image.fromarray(image), @@ -71,42 +64,6 @@ class PaintByExample(InpaintModel): output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR) return output - def _scaled_pad_forward(self, image, mask, config: Config): - longer_side_length = int(config.sd_scale * max(image.shape[:2])) - origin_size = image.shape[:2] - downsize_image = resize_max_size(image, size_limit=longer_side_length) - downsize_mask = resize_max_size(mask, size_limit=longer_side_length) - logger.info( - f"Resize image to do paint_by_example: {image.shape} -> {downsize_image.shape}" - ) - inpaint_result = self._pad_forward(downsize_image, downsize_mask, config) - # only paste masked area result - inpaint_result = cv2.resize( - inpaint_result, - (origin_size[1], origin_size[0]), - interpolation=cv2.INTER_CUBIC, - ) - original_pixel_indices = mask < 127 - inpaint_result[original_pixel_indices] = image[:, :, ::-1][original_pixel_indices] - return inpaint_result - - @torch.no_grad() - def __call__(self, image, mask, config: Config): - """ - images: [H, W, C] RGB, not normalized - masks: [H, W] - return: BGR IMAGE - """ - if config.use_croper: - crop_img, crop_mask, (l, t, r, b) = self._apply_cropper(image, mask, config) - crop_image = self._scaled_pad_forward(crop_img, crop_mask, config) - inpaint_result = image[:, :, ::-1] - inpaint_result[t:b, l:r, :] = crop_image - else: - inpaint_result = self._scaled_pad_forward(image, mask, config) - - return inpaint_result - def forward_post_process(self, result, image, mask, config): if config.paint_by_example_match_histograms: result = self._match_histograms(result, image[:, :, ::-1], mask) diff --git a/lama_cleaner/model/sd.py b/lama_cleaner/model/sd.py index 08585bf..96d0879 100644 --- a/lama_cleaner/model/sd.py +++ b/lama_cleaner/model/sd.py @@ -8,9 +8,8 @@ from diffusers import PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler, EulerD EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler from loguru import logger -from lama_cleaner.helper import resize_max_size -from lama_cleaner.model.base import InpaintModel -from lama_cleaner.model.utils import torch_gc +from lama_cleaner.model.base import DiffusionInpaintModel +from lama_cleaner.model.utils import torch_gc, set_seed from lama_cleaner.schema import Config, SDSampler @@ -28,7 +27,7 @@ class CPUTextEncoderWrapper: return [self.text_encoder(x.to(self.text_encoder.device), **kwargs)[0].to(input_device).to(self.torch_dtype)] -class SD(InpaintModel): +class SD(DiffusionInpaintModel): pad_mod = 8 min_size = 512 @@ -73,25 +72,6 @@ class SD(InpaintModel): self.callback = kwargs.pop("callback", None) - def _scaled_pad_forward(self, image, mask, config: Config): - longer_side_length = int(config.sd_scale * max(image.shape[:2])) - origin_size = image.shape[:2] - downsize_image = resize_max_size(image, size_limit=longer_side_length) - downsize_mask = resize_max_size(mask, size_limit=longer_side_length) - logger.info( - f"Resize image to do sd inpainting: {image.shape} -> {downsize_image.shape}" - ) - inpaint_result = self._pad_forward(downsize_image, downsize_mask, config) - # only paste masked area result - inpaint_result = cv2.resize( - inpaint_result, - (origin_size[1], origin_size[0]), - interpolation=cv2.INTER_CUBIC, - ) - original_pixel_indices = mask < 127 - inpaint_result[original_pixel_indices] = image[:, :, ::-1][original_pixel_indices] - return inpaint_result - def forward(self, image, mask, config: Config): """Input image and output image have same size image: [H, W, C] RGB @@ -118,11 +98,7 @@ class SD(InpaintModel): self.model.scheduler = scheduler - seed = config.sd_seed - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) + set_seed(config.sd_seed) if config.sd_mask_blur != 0: k = 2 * config.sd_mask_blur + 1 @@ -147,24 +123,6 @@ class SD(InpaintModel): output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR) return output - @torch.no_grad() - def __call__(self, image, mask, config: Config): - """ - images: [H, W, C] RGB, not normalized - masks: [H, W] - return: BGR IMAGE - """ - # boxes = boxes_from_mask(mask) - if config.use_croper: - crop_img, crop_mask, (l, t, r, b) = self._apply_cropper(image, mask, config) - crop_image = self._scaled_pad_forward(crop_img, crop_mask, config) - inpaint_result = image[:, :, ::-1] - inpaint_result[t:b, l:r, :] = crop_image - else: - inpaint_result = self._scaled_pad_forward(image, mask, config) - - return inpaint_result - def forward_post_process(self, result, image, mask, config): if config.sd_match_histograms: result = self._match_histograms(result, image[:, :, ::-1], mask) diff --git a/lama_cleaner/model/utils.py b/lama_cleaner/model/utils.py index 7c6ead4..52961b1 100644 --- a/lama_cleaner/model/utils.py +++ b/lama_cleaner/model/utils.py @@ -1,4 +1,5 @@ import math +import random from typing import Any import torch @@ -713,3 +714,10 @@ def torch_gc(): if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() + + +def set_seed(seed: int): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed)