Spaces:
Runtime error
Runtime error
| import torch | |
| import numpy as np | |
| from PIL import Image, ImageFilter | |
| from modules.util import resample_image, set_image_shape_ceil, get_image_shape_ceil | |
| from modules.upscaler import perform_upscale | |
| import cv2 | |
| inpaint_head_model = None | |
| class InpaintHead(torch.nn.Module): | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.head = torch.nn.Parameter(torch.empty(size=(320, 5, 3, 3), device='cpu')) | |
| def __call__(self, x): | |
| x = torch.nn.functional.pad(x, (1, 1, 1, 1), "replicate") | |
| return torch.nn.functional.conv2d(input=x, weight=self.head) | |
| current_task = None | |
| def box_blur(x, k): | |
| x = Image.fromarray(x) | |
| x = x.filter(ImageFilter.BoxBlur(k)) | |
| return np.array(x) | |
| def max_filter_opencv(x, ksize=3): | |
| # Use OpenCV maximum filter | |
| # Make sure the input type is int16 | |
| return cv2.dilate(x, np.ones((ksize, ksize), dtype=np.int16)) | |
| def morphological_open(x): | |
| # Convert array to int16 type via threshold operation | |
| x_int16 = np.zeros_like(x, dtype=np.int16) | |
| x_int16[x > 127] = 256 | |
| for i in range(32): | |
| # Use int16 type to avoid overflow | |
| maxed = max_filter_opencv(x_int16, ksize=3) - 8 | |
| x_int16 = np.maximum(maxed, x_int16) | |
| # Clip negative values to 0 and convert back to uint8 type | |
| x_uint8 = np.clip(x_int16, 0, 255).astype(np.uint8) | |
| return x_uint8 | |
| def up255(x, t=0): | |
| y = np.zeros_like(x).astype(np.uint8) | |
| y[x > t] = 255 | |
| return y | |
| def imsave(x, path): | |
| x = Image.fromarray(x) | |
| x.save(path) | |
| def regulate_abcd(x, a, b, c, d): | |
| H, W = x.shape[:2] | |
| if a < 0: | |
| a = 0 | |
| if a > H: | |
| a = H | |
| if b < 0: | |
| b = 0 | |
| if b > H: | |
| b = H | |
| if c < 0: | |
| c = 0 | |
| if c > W: | |
| c = W | |
| if d < 0: | |
| d = 0 | |
| if d > W: | |
| d = W | |
| return int(a), int(b), int(c), int(d) | |
| def compute_initial_abcd(x): | |
| indices = np.where(x) | |
| a = np.min(indices[0]) | |
| b = np.max(indices[0]) | |
| c = np.min(indices[1]) | |
| d = np.max(indices[1]) | |
| abp = (b + a) // 2 | |
| abm = (b - a) // 2 | |
| cdp = (d + c) // 2 | |
| cdm = (d - c) // 2 | |
| l = int(max(abm, cdm) * 1.15) | |
| a = abp - l | |
| b = abp + l + 1 | |
| c = cdp - l | |
| d = cdp + l + 1 | |
| a, b, c, d = regulate_abcd(x, a, b, c, d) | |
| return a, b, c, d | |
| def solve_abcd(x, a, b, c, d, k): | |
| k = float(k) | |
| assert 0.0 <= k <= 1.0 | |
| H, W = x.shape[:2] | |
| if k == 1.0: | |
| return 0, H, 0, W | |
| while True: | |
| if b - a >= H * k and d - c >= W * k: | |
| break | |
| add_h = (b - a) < (d - c) | |
| add_w = not add_h | |
| if b - a == H: | |
| add_w = True | |
| if d - c == W: | |
| add_h = True | |
| if add_h: | |
| a -= 1 | |
| b += 1 | |
| if add_w: | |
| c -= 1 | |
| d += 1 | |
| a, b, c, d = regulate_abcd(x, a, b, c, d) | |
| return a, b, c, d | |
| def fooocus_fill(image, mask): | |
| current_image = image.copy() | |
| raw_image = image.copy() | |
| area = np.where(mask < 127) | |
| store = raw_image[area] | |
| for k, repeats in [(512, 2), (256, 2), (128, 4), (64, 4), (33, 8), (15, 8), (5, 16), (3, 16)]: | |
| for _ in range(repeats): | |
| current_image = box_blur(current_image, k) | |
| current_image[area] = store | |
| return current_image | |
| class InpaintWorker: | |
| def __init__(self, image, mask, use_fill=True, k=0.618): | |
| a, b, c, d = compute_initial_abcd(mask > 0) | |
| a, b, c, d = solve_abcd(mask, a, b, c, d, k=k) | |
| # interested area | |
| self.interested_area = (a, b, c, d) | |
| self.interested_mask = mask[a:b, c:d] | |
| self.interested_image = image[a:b, c:d] | |
| # super resolution | |
| if get_image_shape_ceil(self.interested_image) < 1024: | |
| self.interested_image = perform_upscale(self.interested_image) | |
| # resize to make images ready for diffusion | |
| self.interested_image = set_image_shape_ceil(self.interested_image, 1024) | |
| self.interested_fill = self.interested_image.copy() | |
| H, W, C = self.interested_image.shape | |
| # process mask | |
| self.interested_mask = up255(resample_image(self.interested_mask, W, H), t=127) | |
| # compute filling | |
| if use_fill: | |
| self.interested_fill = fooocus_fill(self.interested_image, self.interested_mask) | |
| # soft pixels | |
| self.mask = morphological_open(mask) | |
| self.image = image | |
| # ending | |
| self.latent = None | |
| self.latent_after_swap = None | |
| self.swapped = False | |
| self.latent_mask = None | |
| self.inpaint_head_feature = None | |
| return | |
| def load_latent(self, latent_fill, latent_mask, latent_swap=None): | |
| self.latent = latent_fill | |
| self.latent_mask = latent_mask | |
| self.latent_after_swap = latent_swap | |
| return | |
| def patch(self, inpaint_head_model_path, inpaint_latent, inpaint_latent_mask, model): | |
| global inpaint_head_model | |
| if inpaint_head_model is None: | |
| inpaint_head_model = InpaintHead() | |
| sd = torch.load(inpaint_head_model_path, map_location='cpu') | |
| inpaint_head_model.load_state_dict(sd) | |
| feed = torch.cat([ | |
| inpaint_latent_mask, | |
| model.model.process_latent_in(inpaint_latent) | |
| ], dim=1) | |
| inpaint_head_model.to(device=feed.device, dtype=feed.dtype) | |
| inpaint_head_feature = inpaint_head_model(feed) | |
| def input_block_patch(h, transformer_options): | |
| if transformer_options["block"][1] == 0: | |
| h = h + inpaint_head_feature.to(h) | |
| return h | |
| m = model.clone() | |
| m.set_model_input_block_patch(input_block_patch) | |
| return m | |
| def swap(self): | |
| if self.swapped: | |
| return | |
| if self.latent is None: | |
| return | |
| if self.latent_after_swap is None: | |
| return | |
| self.latent, self.latent_after_swap = self.latent_after_swap, self.latent | |
| self.swapped = True | |
| return | |
| def unswap(self): | |
| if not self.swapped: | |
| return | |
| if self.latent is None: | |
| return | |
| if self.latent_after_swap is None: | |
| return | |
| self.latent, self.latent_after_swap = self.latent_after_swap, self.latent | |
| self.swapped = False | |
| return | |
| def color_correction(self, img): | |
| fg = img.astype(np.float32) | |
| bg = self.image.copy().astype(np.float32) | |
| w = self.mask[:, :, None].astype(np.float32) / 255.0 | |
| y = fg * w + bg * (1 - w) | |
| return y.clip(0, 255).astype(np.uint8) | |
| def post_process(self, img): | |
| a, b, c, d = self.interested_area | |
| content = resample_image(img, d - c, b - a) | |
| result = self.image.copy() | |
| result[a:b, c:d] = content | |
| result = self.color_correction(result) | |
| return result | |
| def visualize_mask_processing(self): | |
| return [self.interested_fill, self.interested_mask, self.interested_image] | |