skyadmin/cog-webui-sd

1import math2 3import modules.scripts as scripts4import gradio as gr5from PIL import Image, ImageDraw6 7from modules import images, processing, devices8from modules.processing import Processed, process_images9from modules.shared import opts, cmd_opts, state10 11 12class Script(scripts.Script):13    def title(self):14        return "Poor man's outpainting"15 16    def show(self, is_img2img):17        return is_img2img18 19    def ui(self, is_img2img):20        if not is_img2img:21            return None22        23        pixels = gr.Slider(label="Pixels to expand", minimum=8, maximum=256, step=8, value=128, elem_id=self.elem_id("pixels"))24        mask_blur = gr.Slider(label='Mask blur', minimum=0, maximum=64, step=1, value=4, elem_id=self.elem_id("mask_blur"))25        inpainting_fill = gr.Radio(label='Masked content', choices=['fill', 'original', 'latent noise', 'latent nothing'], value='fill', type="index", elem_id=self.elem_id("inpainting_fill"))26        direction = gr.CheckboxGroup(label="Outpainting direction", choices=['left', 'right', 'up', 'down'], value=['left', 'right', 'up', 'down'], elem_id=self.elem_id("direction"))27 28        return [pixels, mask_blur, inpainting_fill, direction]29 30    def run(self, p, pixels, mask_blur, inpainting_fill, direction):31        initial_seed = None32        initial_info = None33 34        p.mask_blur = mask_blur * 235        p.inpainting_fill = inpainting_fill36        p.inpaint_full_res = False37 38        left = pixels if "left" in direction else 039        right = pixels if "right" in direction else 040        up = pixels if "up" in direction else 041        down = pixels if "down" in direction else 042 43        init_img = p.init_images[0]44        target_w = math.ceil((init_img.width + left + right) / 64) * 6445        target_h = math.ceil((init_img.height + up + down) / 64) * 6446 47        if left > 0:48            left = left * (target_w - init_img.width) // (left + right)49        if right > 0:50            right = target_w - init_img.width - left51 52        if up > 0:53            up = up * (target_h - init_img.height) // (up + down)54 55        if down > 0:56            down = target_h - init_img.height - up57 58        img = Image.new("RGB", (target_w, target_h))59        img.paste(init_img, (left, up))60 61        mask = Image.new("L", (img.width, img.height), "white")62        draw = ImageDraw.Draw(mask)63        draw.rectangle((64            left + (mask_blur * 2 if left > 0 else 0),65            up + (mask_blur * 2 if up > 0 else 0),66            mask.width - right - (mask_blur * 2 if right > 0 else 0),67            mask.height - down - (mask_blur * 2 if down > 0 else 0)68        ), fill="black")69 70        latent_mask = Image.new("L", (img.width, img.height), "white")71        latent_draw = ImageDraw.Draw(latent_mask)72        latent_draw.rectangle((73             left + (mask_blur//2 if left > 0 else 0),74             up + (mask_blur//2 if up > 0 else 0),75             mask.width - right - (mask_blur//2 if right > 0 else 0),76             mask.height - down - (mask_blur//2 if down > 0 else 0)77        ), fill="black")78 79        devices.torch_gc()80 81        grid = images.split_grid(img, tile_w=p.width, tile_h=p.height, overlap=pixels)82        grid_mask = images.split_grid(mask, tile_w=p.width, tile_h=p.height, overlap=pixels)83        grid_latent_mask = images.split_grid(latent_mask, tile_w=p.width, tile_h=p.height, overlap=pixels)84 85        p.n_iter = 186        p.batch_size = 187        p.do_not_save_grid = True88        p.do_not_save_samples = True89 90        work = []91        work_mask = []92        work_latent_mask = []93        work_results = []94 95        for (y, h, row), (_, _, row_mask), (_, _, row_latent_mask) in zip(grid.tiles, grid_mask.tiles, grid_latent_mask.tiles):96            for tiledata, tiledata_mask, tiledata_latent_mask in zip(row, row_mask, row_latent_mask):97                x, w = tiledata[0:2]98 99                if x >= left and x+w <= img.width - right and y >= up and y+h <= img.height - down:100                    continue101 102                work.append(tiledata[2])103                work_mask.append(tiledata_mask[2])104                work_latent_mask.append(tiledata_latent_mask[2])105 106        batch_count = len(work)107        print(f"Poor man's outpainting will process a total of {len(work)} images tiled as {len(grid.tiles[0][2])}x{len(grid.tiles)}.")108 109        state.job_count = batch_count110 111        for i in range(batch_count):112            p.init_images = [work[i]]113            p.image_mask = work_mask[i]114            p.latent_mask = work_latent_mask[i]115 116            state.job = f"Batch {i + 1} out of {batch_count}"117            processed = process_images(p)118 119            if initial_seed is None:120                initial_seed = processed.seed121                initial_info = processed.info122 123            p.seed = processed.seed + 1124            work_results += processed.images125 126 127        image_index = 0128        for y, h, row in grid.tiles:129            for tiledata in row:130                x, w = tiledata[0:2]131 132                if x >= left and x+w <= img.width - right and y >= up and y+h <= img.height - down:133                    continue134 135                tiledata[2] = work_results[image_index] if image_index < len(work_results) else Image.new("RGB", (p.width, p.height))136                image_index += 1137 138        combined_image = images.combine_grid(grid)139 140        if opts.samples_save:141            images.save_image(combined_image, p.outpath_samples, "", initial_seed, p.prompt, opts.grid_format, info=initial_info, p=p)142 143        processed = Processed(p, [combined_image], initial_seed, initial_info)144 145        return processed146 147