diff --git a/lama_cleaner/model/manga.py b/lama_cleaner/model/manga.py index 577883a..cae4766 100644 --- a/lama_cleaner/model/manga.py +++ b/lama_cleaner/model/manga.py @@ -1,4 +1,5 @@ import os +import random import cv2 import numpy as np @@ -80,6 +81,7 @@ class Manga(InpaintModel): def init_model(self, device, **kwargs): self.inpaintor_model = load_jit_model(MANGA_INPAINTOR_MODEL_URL, device) self.line_model = load_jit_model(MANGA_LINE_MODEL_URL, device) + self.seed = 42 @staticmethod def is_downloaded() -> bool: @@ -95,10 +97,17 @@ class Manga(InpaintModel): mask: [H, W, 1] return: BGR IMAGE """ + seed = self.seed + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + gray_img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) gray_img = torch.from_numpy(gray_img[np.newaxis, np.newaxis, :, :].astype(np.float32)).to(self.device) start = time.time() lines = self.line_model(gray_img) + torch.cuda.empty_cache() lines = torch.clamp(lines, 0, 255) logger.info(f"erika_model time: {time.time() - start}") @@ -106,12 +115,13 @@ class Manga(InpaintModel): mask = mask.permute(0, 3, 1, 2) mask = torch.where(mask > 0.5, 1.0, 0.0) noise = torch.randn_like(mask) + ones = torch.ones_like(mask) gray_img = gray_img / 255 * 2 - 1.0 lines = lines / 255 * 2 - 1.0 start = time.time() - inpainted_image = self.inpaintor_model(gray_img, lines, mask, noise) + inpainted_image = self.inpaintor_model(gray_img, lines, mask, noise, ones) logger.info(f"image_inpaintor_model time: {time.time() - start}") cur_res = inpainted_image[0].permute(1, 2, 0).detach().cpu().numpy() diff --git a/setup.py b/setup.py index e75458f..d4ac6ca 100644 --- a/setup.py +++ b/setup.py @@ -21,7 +21,7 @@ def load_requirements(): # https://setuptools.readthedocs.io/en/latest/setuptools.html#including-data-files setuptools.setup( name="lama-cleaner", - version="0.27.0", + version="0.27.1", author="PanicByte", author_email="cwq1913@gmail.com", description="Image inpainting tool powered by SOTA AI Model",