skyadmin/cog-webui-sd

1import math2 3import numpy as np4import skimage5 6import modules.scripts as scripts7import gradio as gr8from PIL import Image, ImageDraw9 10from modules import images, processing, devices11from modules.processing import Processed, process_images12from modules.shared import opts, cmd_opts, state13 14 15# this function is taken from https://github.com/parlance-zz/g-diffuser-bot16def get_matched_noise(_np_src_image, np_mask_rgb, noise_q=1, color_variation=0.05):17    # helper fft routines that keep ortho normalization and auto-shift before and after fft18    def _fft2(data):19        if data.ndim > 2:  # has channels20            out_fft = np.zeros((data.shape[0], data.shape[1], data.shape[2]), dtype=np.complex128)21            for c in range(data.shape[2]):22                c_data = data[:, :, c]23                out_fft[:, :, c] = np.fft.fft2(np.fft.fftshift(c_data), norm="ortho")24                out_fft[:, :, c] = np.fft.ifftshift(out_fft[:, :, c])25        else:  # one channel26            out_fft = np.zeros((data.shape[0], data.shape[1]), dtype=np.complex128)27            out_fft[:, :] = np.fft.fft2(np.fft.fftshift(data), norm="ortho")28            out_fft[:, :] = np.fft.ifftshift(out_fft[:, :])29 30        return out_fft31 32    def _ifft2(data):33        if data.ndim > 2:  # has channels34            out_ifft = np.zeros((data.shape[0], data.shape[1], data.shape[2]), dtype=np.complex128)35            for c in range(data.shape[2]):36                c_data = data[:, :, c]37                out_ifft[:, :, c] = np.fft.ifft2(np.fft.fftshift(c_data), norm="ortho")38                out_ifft[:, :, c] = np.fft.ifftshift(out_ifft[:, :, c])39        else:  # one channel40            out_ifft = np.zeros((data.shape[0], data.shape[1]), dtype=np.complex128)41            out_ifft[:, :] = np.fft.ifft2(np.fft.fftshift(data), norm="ortho")42            out_ifft[:, :] = np.fft.ifftshift(out_ifft[:, :])43 44        return out_ifft45 46    def _get_gaussian_window(width, height, std=3.14, mode=0):47        window_scale_x = float(width / min(width, height))48        window_scale_y = float(height / min(width, height))49 50        window = np.zeros((width, height))51        x = (np.arange(width) / width * 2. - 1.) * window_scale_x52        for y in range(height):53            fy = (y / height * 2. - 1.) * window_scale_y54            if mode == 0:55                window[:, y] = np.exp(-(x ** 2 + fy ** 2) * std)56            else:57                window[:, y] = (1 / ((x ** 2 + 1.) * (fy ** 2 + 1.))) ** (std / 3.14)  # hey wait a minute that's not gaussian58 59        return window60 61    def _get_masked_window_rgb(np_mask_grey, hardness=1.):62        np_mask_rgb = np.zeros((np_mask_grey.shape[0], np_mask_grey.shape[1], 3))63        if hardness != 1.:64            hardened = np_mask_grey[:] ** hardness65        else:66            hardened = np_mask_grey[:]67        for c in range(3):68            np_mask_rgb[:, :, c] = hardened[:]69        return np_mask_rgb70 71    width = _np_src_image.shape[0]72    height = _np_src_image.shape[1]73    num_channels = _np_src_image.shape[2]74 75    np_src_image = _np_src_image[:] * (1. - np_mask_rgb)76    np_mask_grey = (np.sum(np_mask_rgb, axis=2) / 3.)77    img_mask = np_mask_grey > 1e-678    ref_mask = np_mask_grey < 1e-379 80    windowed_image = _np_src_image * (1. - _get_masked_window_rgb(np_mask_grey))81    windowed_image /= np.max(windowed_image)82    windowed_image += np.average(_np_src_image) * np_mask_rgb  # / (1.-np.average(np_mask_rgb))  # rather than leave the masked area black, we get better results from fft by filling the average unmasked color83 84    src_fft = _fft2(windowed_image)  # get feature statistics from masked src img85    src_dist = np.absolute(src_fft)86    src_phase = src_fft / src_dist87 88    # create a generator with a static seed to make outpainting deterministic / only follow global seed89    rng = np.random.default_rng(0)90 91    noise_window = _get_gaussian_window(width, height, mode=1)  # start with simple gaussian noise92    noise_rgb = rng.random((width, height, num_channels))93    noise_grey = (np.sum(noise_rgb, axis=2) / 3.)94    noise_rgb *= color_variation  # the colorfulness of the starting noise is blended to greyscale with a parameter95    for c in range(num_channels):96        noise_rgb[:, :, c] += (1. - color_variation) * noise_grey97 98    noise_fft = _fft2(noise_rgb)99    for c in range(num_channels):100        noise_fft[:, :, c] *= noise_window101    noise_rgb = np.real(_ifft2(noise_fft))102    shaped_noise_fft = _fft2(noise_rgb)103    shaped_noise_fft[:, :, :] = np.absolute(shaped_noise_fft[:, :, :]) ** 2 * (src_dist ** noise_q) * src_phase  # perform the actual shaping104 105    brightness_variation = 0.  # color_variation # todo: temporarily tieing brightness variation to color variation for now106    contrast_adjusted_np_src = _np_src_image[:] * (brightness_variation + 1.) - brightness_variation * 2.107 108    # scikit-image is used for histogram matching, very convenient!109    shaped_noise = np.real(_ifft2(shaped_noise_fft))110    shaped_noise -= np.min(shaped_noise)111    shaped_noise /= np.max(shaped_noise)112    shaped_noise[img_mask, :] = skimage.exposure.match_histograms(shaped_noise[img_mask, :] ** 1., contrast_adjusted_np_src[ref_mask, :], channel_axis=1)113    shaped_noise = _np_src_image[:] * (1. - np_mask_rgb) + shaped_noise * np_mask_rgb114 115    matched_noise = shaped_noise[:]116 117    return np.clip(matched_noise, 0., 1.)118 119 120 121class Script(scripts.Script):122    def title(self):123        return "Outpainting mk2"124 125    def show(self, is_img2img):126        return is_img2img127 128    def ui(self, is_img2img):129        if not is_img2img:130            return None131 132        info = gr.HTML("<p style=\"margin-bottom:0.75em\">Recommended settings: Sampling Steps: 80-100, Sampler: Euler a, Denoising strength: 0.8</p>")133 134        pixels = gr.Slider(label="Pixels to expand", minimum=8, maximum=256, step=8, value=128, elem_id=self.elem_id("pixels"))135        mask_blur = gr.Slider(label='Mask blur', minimum=0, maximum=64, step=1, value=8, elem_id=self.elem_id("mask_blur"))136        direction = gr.CheckboxGroup(label="Outpainting direction", choices=['left', 'right', 'up', 'down'], value=['left', 'right', 'up', 'down'], elem_id=self.elem_id("direction"))137        noise_q = gr.Slider(label="Fall-off exponent (lower=higher detail)", minimum=0.0, maximum=4.0, step=0.01, value=1.0, elem_id=self.elem_id("noise_q"))138        color_variation = gr.Slider(label="Color variation", minimum=0.0, maximum=1.0, step=0.01, value=0.05, elem_id=self.elem_id("color_variation"))139 140        return [info, pixels, mask_blur, direction, noise_q, color_variation]141 142    def run(self, p, _, pixels, mask_blur, direction, noise_q, color_variation):143        initial_seed_and_info = [None, None]144 145        process_width = p.width146        process_height = p.height147 148        p.mask_blur = mask_blur*4149        p.inpaint_full_res = False150        p.inpainting_fill = 1151        p.do_not_save_samples = True152        p.do_not_save_grid = True153 154        left = pixels if "left" in direction else 0155        right = pixels if "right" in direction else 0156        up = pixels if "up" in direction else 0157        down = pixels if "down" in direction else 0158 159        init_img = p.init_images[0]160        target_w = math.ceil((init_img.width + left + right) / 64) * 64161        target_h = math.ceil((init_img.height + up + down) / 64) * 64162 163        if left > 0:164            left = left * (target_w - init_img.width) // (left + right)165 166        if right > 0:167            right = target_w - init_img.width - left168 169        if up > 0:170            up = up * (target_h - init_img.height) // (up + down)171 172        if down > 0:173            down = target_h - init_img.height - up174 175        def expand(init, count, expand_pixels, is_left=False, is_right=False, is_top=False, is_bottom=False):176            is_horiz = is_left or is_right177            is_vert = is_top or is_bottom178            pixels_horiz = expand_pixels if is_horiz else 0179            pixels_vert = expand_pixels if is_vert else 0180 181            images_to_process = []182            output_images = []183            for n in range(count):184                res_w = init[n].width + pixels_horiz185                res_h = init[n].height + pixels_vert186                process_res_w = math.ceil(res_w / 64) * 64187                process_res_h = math.ceil(res_h / 64) * 64188 189                img = Image.new("RGB", (process_res_w, process_res_h))190                img.paste(init[n], (pixels_horiz if is_left else 0, pixels_vert if is_top else 0))191                mask = Image.new("RGB", (process_res_w, process_res_h), "white")192                draw = ImageDraw.Draw(mask)193                draw.rectangle((194                    expand_pixels + mask_blur if is_left else 0,195                    expand_pixels + mask_blur if is_top else 0,196                    mask.width - expand_pixels - mask_blur if is_right else res_w,197                    mask.height - expand_pixels - mask_blur if is_bottom else res_h,198                ), fill="black")199 200                np_image = (np.asarray(img) / 255.0).astype(np.float64)201                np_mask = (np.asarray(mask) / 255.0).astype(np.float64)202                noised = get_matched_noise(np_image, np_mask, noise_q, color_variation)203                output_images.append(Image.fromarray(np.clip(noised * 255., 0., 255.).astype(np.uint8), mode="RGB"))204 205                target_width = min(process_width, init[n].width + pixels_horiz) if is_horiz else img.width206                target_height = min(process_height, init[n].height + pixels_vert) if is_vert else img.height207                p.width = target_width if is_horiz else img.width208                p.height = target_height if is_vert else img.height209 210                crop_region = (211                    0 if is_left else output_images[n].width - target_width,212                    0 if is_top else output_images[n].height - target_height,213                    target_width if is_left else output_images[n].width,214                    target_height if is_top else output_images[n].height,215                )216                mask = mask.crop(crop_region)217                p.image_mask = mask218 219                image_to_process = output_images[n].crop(crop_region)220                images_to_process.append(image_to_process)221 222            p.init_images = images_to_process223 224            latent_mask = Image.new("RGB", (p.width, p.height), "white")225            draw = ImageDraw.Draw(latent_mask)226            draw.rectangle((227                expand_pixels + mask_blur * 2 if is_left else 0,228                expand_pixels + mask_blur * 2 if is_top else 0,229                mask.width - expand_pixels - mask_blur * 2 if is_right else res_w,230                mask.height - expand_pixels - mask_blur * 2 if is_bottom else res_h,231            ), fill="black")232            p.latent_mask = latent_mask233 234            proc = process_images(p)235 236            if initial_seed_and_info[0] is None:237                initial_seed_and_info[0] = proc.seed238                initial_seed_and_info[1] = proc.info239 240            for n in range(count):241                output_images[n].paste(proc.images[n], (0 if is_left else output_images[n].width - proc.images[n].width, 0 if is_top else output_images[n].height - proc.images[n].height))242                output_images[n] = output_images[n].crop((0, 0, res_w, res_h))243 244            return output_images245 246        batch_count = p.n_iter247        batch_size = p.batch_size248        p.n_iter = 1249        state.job_count = batch_count * ((1 if left > 0 else 0) + (1 if right > 0 else 0) + (1 if up > 0 else 0) + (1 if down > 0 else 0))250        all_processed_images = []251 252        for i in range(batch_count):253            imgs = [init_img] * batch_size254            state.job = f"Batch {i + 1} out of {batch_count}"255 256            if left > 0:257                imgs = expand(imgs, batch_size, left, is_left=True)258            if right > 0:259                imgs = expand(imgs, batch_size, right, is_right=True)260            if up > 0:261                imgs = expand(imgs, batch_size, up, is_top=True)262            if down > 0:263                imgs = expand(imgs, batch_size, down, is_bottom=True)264 265            all_processed_images += imgs266 267        all_images = all_processed_images268 269        combined_grid_image = images.image_grid(all_processed_images)270        unwanted_grid_because_of_img_count = len(all_processed_images) < 2 and opts.grid_only_if_multiple271        if opts.return_grid and not unwanted_grid_because_of_img_count:272            all_images = [combined_grid_image] + all_processed_images273 274        res = Processed(p, all_images, initial_seed_and_info[0], initial_seed_and_info[1])275 276        if opts.samples_save:277            for img in all_processed_images:278                images.save_image(img, p.outpath_samples, "", res.seed, p.prompt, opts.grid_format, info=res.info, p=p)279 280        if opts.grid_save and not unwanted_grid_because_of_img_count:281            images.save_image(combined_grid_image, p.outpath_grids, "grid", res.seed, p.prompt, opts.grid_format, info=res.info, short_filename=not opts.grid_extended_filename, grid=True, p=p)282 283        return res284