huggingface-projects/stable-diffusion-multiplayer

1from PIL import Image2from PIL import ImageFilter3import cv24import numpy as np5import scipy6import scipy.signal7from scipy.spatial import cKDTree8 9import os10from perlin2d import *11 12patch_match_compiled = True13 14from PyPatchMatch import patch_match15 16 17def edge_pad(img, mask, mode=1):18    if mode == 0:19        nmask = mask.copy()20        nmask[nmask > 0] = 121        res0 = 1 - nmask22        res1 = nmask23        p0 = np.stack(res0.nonzero(), axis=0).transpose()24        p1 = np.stack(res1.nonzero(), axis=0).transpose()25        min_dists, min_dist_idx = cKDTree(p1).query(p0, 1)26        loc = p1[min_dist_idx]27        for (a, b), (c, d) in zip(p0, loc):28            img[a, b] = img[c, d]29    elif mode == 1:30        record = {}31        kernel = [[1] * 3 for _ in range(3)]32        nmask = mask.copy()33        nmask[nmask > 0] = 134        res = scipy.signal.convolve2d(35            nmask, kernel, mode="same", boundary="fill", fillvalue=136        )37        res[nmask < 1] = 038        res[res == 9] = 039        res[res > 0] = 140        ylst, xlst = res.nonzero()41        queue = [(y, x) for y, x in zip(ylst, xlst)]42        # bfs here43        cnt = res.astype(np.float32)44        acc = img.astype(np.float32)45        step = 146        h = acc.shape[0]47        w = acc.shape[1]48        offset = [(1, 0), (-1, 0), (0, 1), (0, -1)]49        while queue:50            target = []51            for y, x in queue:52                val = acc[y][x]53                for yo, xo in offset:54                    yn = y + yo55                    xn = x + xo56                    if 0 <= yn < h and 0 <= xn < w and nmask[yn][xn] < 1:57                        if record.get((yn, xn), step) == step:58                            acc[yn][xn] = acc[yn][xn] * cnt[yn][xn] + val59                            cnt[yn][xn] += 160                            acc[yn][xn] /= cnt[yn][xn]61                            if (yn, xn) not in record:62                                record[(yn, xn)] = step63                                target.append((yn, xn))64            step += 165            queue = target66        img = acc.astype(np.uint8)67    else:68        nmask = mask.copy()69        ylst, xlst = nmask.nonzero()70        yt, xt = ylst.min(), xlst.min()71        yb, xb = ylst.max(), xlst.max()72        content = img[yt : yb + 1, xt : xb + 1]73        img = np.pad(74            content,75            ((yt, mask.shape[0] - yb - 1), (xt, mask.shape[1] - xb - 1), (0, 0)),76            mode="edge",77        )78    return img, mask79 80 81def perlin_noise(img, mask):82    lin = np.linspace(0, 5, mask.shape[0], endpoint=False)83    x, y = np.meshgrid(lin, lin)84    avg = img.mean(axis=0).mean(axis=0)85    # noise=[((perlin(x, y)+1)*128+avg[i]).astype(np.uint8) for i in range(3)]86    noise = [((perlin(x, y) + 1) * 0.5 * 255).astype(np.uint8) for i in range(3)]87    noise = np.stack(noise, axis=-1)88    # mask=skimage.measure.block_reduce(mask,(8,8),np.min)89    # mask=mask.repeat(8, axis=0).repeat(8, axis=1)90    # mask_image=Image.fromarray(mask)91    # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 4))92    # mask=np.array(mask_image)93    nmask = mask.copy()94    # nmask=nmask/255.095    nmask[mask > 0] = 196    img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise97    # img=img.astype(np.uint8)98    return img, mask99 100 101def gaussian_noise(img, mask):102    noise = np.random.randn(mask.shape[0], mask.shape[1], 3)103    noise = (noise + 1) / 2 * 255104    noise = noise.astype(np.uint8)105    nmask = mask.copy()106    nmask[mask > 0] = 1107    img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise108    return img, mask109 110 111def cv2_telea(img, mask):112    ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_TELEA)113    return ret, mask114 115 116def cv2_ns(img, mask):117    ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_NS)118    return ret, mask119 120 121def patch_match_func(img, mask):122    ret = patch_match.inpaint(img, mask=255 - mask, patch_size=3)123    return ret, mask124 125 126def mean_fill(img, mask):127    avg = img.mean(axis=0).mean(axis=0)128    img[mask < 1] = avg129    return img, mask130 131 132functbl = {133    "gaussian": gaussian_noise,134    "perlin": perlin_noise,135    "edge_pad": edge_pad,136    "patchmatch": patch_match_func if (os.name != "nt" and patch_match_compiled) else edge_pad,137    "cv2_ns": cv2_ns,138    "cv2_telea": cv2_telea,139    "mean_fill": mean_fill,140}141