skyadmin/cog-webui-sd

1import numpy as np2from tqdm import trange3 4import modules.scripts as scripts5import gradio as gr6 7from modules import processing, shared, sd_samplers, images8from modules.processing import Processed9from modules.sd_samplers import samplers10from modules.shared import opts, cmd_opts, state11from modules import deepbooru12 13 14class Script(scripts.Script):15    def title(self):16        return "Loopback"17 18    def show(self, is_img2img):19        return is_img2img20 21    def ui(self, is_img2img):        22        loops = gr.Slider(minimum=1, maximum=32, step=1, label='Loops', value=4, elem_id=self.elem_id("loops"))23        denoising_strength_change_factor = gr.Slider(minimum=0.9, maximum=1.1, step=0.01, label='Denoising strength change factor', value=1, elem_id=self.elem_id("denoising_strength_change_factor"))24        append_interrogation = gr.Dropdown(label="Append interrogated prompt at each iteration", choices=["None", "CLIP", "DeepBooru"], value="None")25 26        return [loops, denoising_strength_change_factor, append_interrogation]27 28    def run(self, p, loops, denoising_strength_change_factor, append_interrogation):29        processing.fix_seed(p)30        batch_count = p.n_iter31        p.extra_generation_params = {32            "Denoising strength change factor": denoising_strength_change_factor,33        }34 35        p.batch_size = 136        p.n_iter = 137 38        output_images, info = None, None39        initial_seed = None40        initial_info = None41 42        grids = []43        all_images = []44        original_init_image = p.init_images45        original_prompt = p.prompt46        state.job_count = loops * batch_count47 48        initial_color_corrections = [processing.setup_color_correction(p.init_images[0])]49 50        for n in range(batch_count):51            history = []52 53            # Reset to original init image at the start of each batch54            p.init_images = original_init_image55 56            for i in range(loops):57                p.n_iter = 158                p.batch_size = 159                p.do_not_save_grid = True60 61                if opts.img2img_color_correction:62                    p.color_corrections = initial_color_corrections63 64                if append_interrogation != "None":65                    p.prompt = original_prompt + ", " if original_prompt != "" else ""66                    if append_interrogation == "CLIP":67                        p.prompt += shared.interrogator.interrogate(p.init_images[0])68                    elif append_interrogation == "DeepBooru":69                        p.prompt += deepbooru.model.tag(p.init_images[0])70 71                state.job = f"Iteration {i + 1}/{loops}, batch {n + 1}/{batch_count}"72 73                processed = processing.process_images(p)74 75                if initial_seed is None:76                    initial_seed = processed.seed77                    initial_info = processed.info78 79                init_img = processed.images[0]80 81                p.init_images = [init_img]82                p.seed = processed.seed + 183                p.denoising_strength = min(max(p.denoising_strength * denoising_strength_change_factor, 0.1), 1)84                history.append(processed.images[0])85 86            grid = images.image_grid(history, rows=1)87            if opts.grid_save:88                images.save_image(grid, p.outpath_grids, "grid", initial_seed, p.prompt, opts.grid_format, info=info, short_filename=not opts.grid_extended_filename, grid=True, p=p)89 90            grids.append(grid)91            all_images += history92 93        if opts.return_grid:94            all_images = grids + all_images95 96        processed = Processed(p, all_images, initial_seed, initial_info)97 98        return processed99