cymic/Waifu_Diffusion_Webui

1from collections import namedtuple2from copy import copy3from itertools import permutations, chain4import random5import csv6from io import StringIO7from PIL import Image8import numpy as np9 10import modules.scripts as scripts11import gradio as gr12 13from modules import images14from modules.processing import process_images, Processed15from modules.shared import opts, cmd_opts, state16import modules.shared as shared17import modules.sd_samplers18import modules.sd_models19import re20 21 22def apply_field(field):23    def fun(p, x, xs):24        setattr(p, field, x)25 26    return fun27 28 29def apply_prompt(p, x, xs):30    p.prompt = p.prompt.replace(xs[0], x)31    p.negative_prompt = p.negative_prompt.replace(xs[0], x)32 33 34def apply_order(p, x, xs):35    token_order = []36 37    # Initally grab the tokens from the prompt, so they can be replaced in order of earliest seen38    for token in x:39        token_order.append((p.prompt.find(token), token))40 41    token_order.sort(key=lambda t: t[0])42 43    prompt_parts = []44 45    # Split the prompt up, taking out the tokens46    for _, token in token_order:47        n = p.prompt.find(token)48        prompt_parts.append(p.prompt[0:n])49        p.prompt = p.prompt[n + len(token):]50 51    # Rebuild the prompt with the tokens in the order we want52    prompt_tmp = ""53    for idx, part in enumerate(prompt_parts):54        prompt_tmp += part55        prompt_tmp += x[idx]56    p.prompt = prompt_tmp + p.prompt57    58 59samplers_dict = {}60for i, sampler in enumerate(modules.sd_samplers.samplers):61    samplers_dict[sampler.name.lower()] = i62    for alias in sampler.aliases:63        samplers_dict[alias.lower()] = i64 65 66def apply_sampler(p, x, xs):67    sampler_index = samplers_dict.get(x.lower(), None)68    if sampler_index is None:69        raise RuntimeError(f"Unknown sampler: {x}")70 71    p.sampler_index = sampler_index72 73 74def apply_checkpoint(p, x, xs):75    info = modules.sd_models.get_closet_checkpoint_match(x)76    assert info is not None, f'Checkpoint for {x} not found'77    modules.sd_models.reload_model_weights(shared.sd_model, info)78 79 80def apply_hypernetwork(p, x, xs):81    hn = shared.hypernetworks.get(x, None)82    opts.data["sd_hypernetwork"] = hn.name if hn is not None else 'None'83 84 85def format_value_add_label(p, opt, x):86    if type(x) == float:87        x = round(x, 8)88 89    return f"{opt.label}: {x}"90 91 92def format_value(p, opt, x):93    if type(x) == float:94        x = round(x, 8)95    return x96 97 98def format_value_join_list(p, opt, x):99    return ", ".join(x)100 101 102def do_nothing(p, x, xs):103    pass104 105 106def format_nothing(p, opt, x):107    return ""108 109 110def str_permutations(x):111    """dummy function for specifying it in AxisOption's type when you want to get a list of permutations"""112    return x113 114 115AxisOption = namedtuple("AxisOption", ["label", "type", "apply", "format_value"])116AxisOptionImg2Img = namedtuple("AxisOptionImg2Img", ["label", "type", "apply", "format_value"])117 118 119axis_options = [120    AxisOption("Nothing", str, do_nothing, format_nothing),121    AxisOption("Seed", int, apply_field("seed"), format_value_add_label),122    AxisOption("Var. seed", int, apply_field("subseed"), format_value_add_label),123    AxisOption("Var. strength", float, apply_field("subseed_strength"), format_value_add_label),124    AxisOption("Steps", int, apply_field("steps"), format_value_add_label),125    AxisOption("CFG Scale", float, apply_field("cfg_scale"), format_value_add_label),126    AxisOption("Prompt S/R", str, apply_prompt, format_value),127    AxisOption("Prompt order", str_permutations, apply_order, format_value_join_list),128    AxisOption("Sampler", str, apply_sampler, format_value),129    AxisOption("Checkpoint name", str, apply_checkpoint, format_value),130    AxisOption("Hypernetwork", str, apply_hypernetwork, format_value),131    AxisOption("Sigma Churn", float, apply_field("s_churn"), format_value_add_label),132    AxisOption("Sigma min", float, apply_field("s_tmin"), format_value_add_label),133    AxisOption("Sigma max", float, apply_field("s_tmax"), format_value_add_label),134    AxisOption("Sigma noise", float, apply_field("s_noise"), format_value_add_label),135    AxisOption("Eta", float, apply_field("eta"), format_value_add_label),136    AxisOptionImg2Img("Denoising", float, apply_field("denoising_strength"), format_value_add_label),  # as it is now all AxisOptionImg2Img items must go after AxisOption ones137]138 139 140def draw_xy_grid(p, xs, ys, x_labels, y_labels, cell, draw_legend):141    res = []142 143    ver_texts = [[images.GridAnnotation(y)] for y in y_labels]144    hor_texts = [[images.GridAnnotation(x)] for x in x_labels]145 146    first_pocessed = None147 148    state.job_count = len(xs) * len(ys) * p.n_iter149 150    for iy, y in enumerate(ys):151        for ix, x in enumerate(xs):152            state.job = f"{ix + iy * len(xs) + 1} out of {len(xs) * len(ys)}"153 154            processed = cell(x, y)155            if first_pocessed is None:156                first_pocessed = processed157 158            try:159              res.append(processed.images[0])160            except:161              res.append(Image.new(res[0].mode, res[0].size))162 163    grid = images.image_grid(res, rows=len(ys))164    if draw_legend:165        grid = images.draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts)166 167    first_pocessed.images = [grid]168 169    return first_pocessed170 171 172re_range = re.compile(r"\s*([+-]?\s*\d+)\s*-\s*([+-]?\s*\d+)(?:\s*\(([+-]\d+)\s*\))?\s*")173re_range_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+(?:.\d*)?)(?:\s*\(([+-]\d+(?:.\d*)?)\s*\))?\s*")174 175re_range_count = re.compile(r"\s*([+-]?\s*\d+)\s*-\s*([+-]?\s*\d+)(?:\s*\[(\d+)\s*\])?\s*")176re_range_count_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+(?:.\d*)?)(?:\s*\[(\d+(?:.\d*)?)\s*\])?\s*")177 178class Script(scripts.Script):179    def title(self):180        return "X/Y plot"181 182    def ui(self, is_img2img):183        current_axis_options = [x for x in axis_options if type(x) == AxisOption or type(x) == AxisOptionImg2Img and is_img2img]184 185        with gr.Row():186            x_type = gr.Dropdown(label="X type", choices=[x.label for x in current_axis_options], value=current_axis_options[1].label, visible=False, type="index", elem_id="x_type")187            x_values = gr.Textbox(label="X values", visible=False, lines=1)188 189        with gr.Row():190            y_type = gr.Dropdown(label="Y type", choices=[x.label for x in current_axis_options], value=current_axis_options[4].label, visible=False, type="index", elem_id="y_type")191            y_values = gr.Textbox(label="Y values", visible=False, lines=1)192        193        draw_legend = gr.Checkbox(label='Draw legend', value=True)194        no_fixed_seeds = gr.Checkbox(label='Keep -1 for seeds', value=False)195 196        return [x_type, x_values, y_type, y_values, draw_legend, no_fixed_seeds]197 198    def run(self, p, x_type, x_values, y_type, y_values, draw_legend, no_fixed_seeds):199        modules.processing.fix_seed(p)200        p.batch_size = 1201 202        initial_hn = opts.sd_hypernetwork203 204        def process_axis(opt, vals):205            if opt.label == 'Nothing':206                return [0]207 208            valslist = [x.strip() for x in chain.from_iterable(csv.reader(StringIO(vals)))]209 210            if opt.type == int:211                valslist_ext = []212 213                for val in valslist:214                    m = re_range.fullmatch(val)215                    mc = re_range_count.fullmatch(val)216                    if m is not None:217 218                        start = int(m.group(1))219                        end = int(m.group(2))+1220                        step = int(m.group(3)) if m.group(3) is not None else 1221 222                        valslist_ext += list(range(start, end, step))223                    elif mc is not None:224                        start = int(mc.group(1))225                        end   = int(mc.group(2))226                        num   = int(mc.group(3)) if mc.group(3) is not None else 1227                        228                        valslist_ext += [int(x) for x in np.linspace(start=start, stop=end, num=num).tolist()]229                    else:230                        valslist_ext.append(val)231 232                valslist = valslist_ext233            elif opt.type == float:234                valslist_ext = []235 236                for val in valslist:237                    m = re_range_float.fullmatch(val)238                    mc = re_range_count_float.fullmatch(val)239                    if m is not None:240                        start = float(m.group(1))241                        end = float(m.group(2))242                        step = float(m.group(3)) if m.group(3) is not None else 1243 244                        valslist_ext += np.arange(start, end + step, step).tolist()245                    elif mc is not None:246                        start = float(mc.group(1))247                        end   = float(mc.group(2))248                        num   = int(mc.group(3)) if mc.group(3) is not None else 1249                        250                        valslist_ext += np.linspace(start=start, stop=end, num=num).tolist()251                    else:252                        valslist_ext.append(val)253 254                valslist = valslist_ext255            elif opt.type == str_permutations:256                valslist = list(permutations(valslist))257 258            valslist = [opt.type(x) for x in valslist]259 260            return valslist261 262        x_opt = axis_options[x_type]263        xs = process_axis(x_opt, x_values)264 265        y_opt = axis_options[y_type]266        ys = process_axis(y_opt, y_values)267 268        def fix_axis_seeds(axis_opt, axis_list):269            if axis_opt.label == 'Seed':270                return [int(random.randrange(4294967294)) if val is None or val == '' or val == -1 else val for val in axis_list]271            else:272                return axis_list273 274        if not no_fixed_seeds:275            xs = fix_axis_seeds(x_opt, xs)276            ys = fix_axis_seeds(y_opt, ys)277 278        if x_opt.label == 'Steps':279            total_steps = sum(xs) * len(ys)280        elif y_opt.label == 'Steps':281            total_steps = sum(ys) * len(xs)282        else:283            total_steps = p.steps * len(xs) * len(ys)284 285        print(f"X/Y plot will create {len(xs) * len(ys) * p.n_iter} images on a {len(xs)}x{len(ys)} grid. (Total steps to process: {total_steps * p.n_iter})")286        shared.total_tqdm.updateTotal(total_steps * p.n_iter)287 288        def cell(x, y):289            pc = copy(p)290            x_opt.apply(pc, x, xs)291            y_opt.apply(pc, y, ys)292 293            return process_images(pc)294 295        processed = draw_xy_grid(296            p,297            xs=xs,298            ys=ys,299            x_labels=[x_opt.format_value(p, x_opt, x) for x in xs],300            y_labels=[y_opt.format_value(p, y_opt, y) for y in ys],301            cell=cell,302            draw_legend=draw_legend303        )304 305        if opts.grid_save:306            images.save_image(processed.images[0], p.outpath_grids, "xy_grid", prompt=p.prompt, seed=processed.seed, grid=True, p=p)307 308        # restore checkpoint in case it was changed by axes309        modules.sd_models.reload_model_weights(shared.sd_model)310 311        opts.data["sd_hypernetwork"] = initial_hn312 313        return processed314